From 49e709eafb2df5b060e9a8c95b5a809379257c53 Mon Sep 17 00:00:00 2001
From: ishandhanani <82981111+ishandhanani@users.noreply.github.com>
Date: Wed, 22 Jul 2026 15:18:16 -0700
Subject: [PATCH 01/43] (cherry-pick) feat: add native gRPC sidecar module
 launcher (#31076) (#32074)

Signed-off-by: Ishan Dhanani <ishandhanani@gmail.com>
Signed-off-by: Connor Carpenter <connorc@nvidia.com>
Co-authored-by: Connor Carpenter <connorc@nvidia.com>
---
 python/sglang/srt/entrypoints/http_server.py  |  10 ++
 python/sglang/srt/entrypoints/sidecar.py      | 130 +++++++++++++++++
 python/sglang/srt/server_args.py              |  33 +++++
 .../unit/server_args/test_server_args.py      | 134 ++++++++++++++++++
 4 files changed, 307 insertions(+)
 create mode 100644 python/sglang/srt/entrypoints/sidecar.py

diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py
index 9a1aa72c560d..51c85e63023f 100644
--- a/python/sglang/srt/entrypoints/http_server.py
+++ b/python/sglang/srt/entrypoints/http_server.py
@@ -265,6 +265,7 @@ async def init_multi_tokenizer() -> ServerArgs:
 @asynccontextmanager
 async def lifespan(fast_api_app: FastAPI):
     grpc_handle = None
+    sidecar = None
     warmup_thread = None
     if getattr(fast_api_app, "is_single_tokenizer_mode", False):
         server_args = fast_api_app.server_args
@@ -397,6 +398,10 @@ async def lifespan(fast_api_app: FastAPI):
                 template_manager=_global_state.template_manager,
                 scheduler_info=_global_state.scheduler_info,
             )
+            if server_args.sidecar is not None:
+                from sglang.srt.entrypoints.sidecar import start_sidecar
+
+                sidecar = start_sidecar(server_args)
 
         # Execute the general warmup
         warmup_thread = threading.Thread(
@@ -408,6 +413,11 @@ async def lifespan(fast_api_app: FastAPI):
         # Start the HTTP server
         yield
     finally:
+        if sidecar is not None:
+            try:
+                sidecar.stop()
+            except Exception:
+                logger.exception("Failed to stop sidecar")
         _shutdown_native_grpc_server(grpc_handle)
         if tool_server is not None and hasattr(tool_server, "aclose"):
             await tool_server.aclose()
diff --git a/python/sglang/srt/entrypoints/sidecar.py b/python/sglang/srt/entrypoints/sidecar.py
new file mode 100644
index 000000000000..ddf58b81e33a
--- /dev/null
+++ b/python/sglang/srt/entrypoints/sidecar.py
@@ -0,0 +1,130 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Lifecycle management for an optional local native gRPC sidecar."""
+
+import argparse
+import importlib
+import logging
+import multiprocessing as mp
+import os
+
+from sglang.srt.utils.common import kill_itself_when_parent_died, kill_process_tree
+from sglang.srt.utils.network import NetworkAddress
+from sglang.srt.utils.watchdog import SubprocessWatchdog
+
+logger = logging.getLogger(__name__)
+
+SGLANG_GRPC_ENDPOINT_ENV = "SGLANG_GRPC_ENDPOINT"
+_DEFAULT_SIDECAR_SHUTDOWN_TIMEOUT = 45.0
+
+
+def _loopback_host(host: str) -> str:
+    if not host or host == "0.0.0.0":
+        return "127.0.0.1"
+    if host in ("::", "[::]"):
+        return "::1"
+    return host
+
+
+def build_sidecar_endpoint(server_args) -> str:
+    return NetworkAddress(
+        _loopback_host(server_args.host), server_args.grpc_port
+    ).to_url()
+
+
+def _parse_sidecar_args(args: list[str] | None) -> tuple[list[str], float]:
+    parser = argparse.ArgumentParser(add_help=False, allow_abbrev=False)
+    parser.add_argument(
+        "--sidecar-shutdown-timeout",
+        type=float,
+        default=_DEFAULT_SIDECAR_SHUTDOWN_TIMEOUT,
+    )
+    parsed, provider_args = parser.parse_known_args(args or [])
+    if parsed.sidecar_shutdown_timeout <= 0:
+        raise ValueError("--sidecar-shutdown-timeout must be greater than 0.")
+    return provider_args, parsed.sidecar_shutdown_timeout
+
+
+def _run_sidecar(module_name: str, args: list[str], endpoint: str) -> None:
+    kill_itself_when_parent_died()
+    os.environ[SGLANG_GRPC_ENDPOINT_ENV] = endpoint
+    try:
+        main = getattr(importlib.import_module(module_name), "main")
+    except (AttributeError, ImportError) as e:
+        raise RuntimeError(
+            f"--sidecar requires importable module {module_name!r} "
+            "with a main(argv) function."
+        ) from e
+
+    if not callable(main):
+        raise RuntimeError(
+            f"--sidecar requires module {module_name!r} to expose "
+            "a callable main(argv)."
+        )
+
+    main(args)
+
+
+class Sidecar:
+    def __init__(
+        self,
+        proc,
+        module_name: str,
+        shutdown_timeout: float,
+    ):
+        self.proc = proc
+        self.module_name = module_name
+        self.shutdown_timeout = shutdown_timeout
+        self._watchdog = SubprocessWatchdog(
+            processes=[proc], process_names=[module_name]
+        )
+
+    def start(self) -> None:
+        self.proc.start()
+        self._watchdog.start()
+        logger.info(
+            "Sidecar module %s started pid=%s",
+            self.module_name,
+            self.proc.pid,
+        )
+
+    def stop(self) -> None:
+        self._watchdog.stop()
+        if self.proc.is_alive():
+            self.proc.terminate()
+            self.proc.join(timeout=self.shutdown_timeout)
+        else:
+            self.proc.join(timeout=0)
+
+        if self.proc.is_alive():
+            logger.warning("Sidecar module did not terminate; killing process tree")
+            kill_process_tree(self.proc.pid, wait_timeout=self.shutdown_timeout)
+
+
+def start_sidecar(server_args) -> Sidecar:
+    module_name = server_args.sidecar
+    assert module_name is not None
+    sidecar_args, shutdown_timeout = _parse_sidecar_args(server_args.sidecar_args)
+    endpoint = build_sidecar_endpoint(server_args)
+    proc = mp.get_context("spawn").Process(
+        name=f"sglang_sidecar_{module_name}",
+        target=_run_sidecar,
+        args=(module_name, sidecar_args, endpoint),
+    )
+    sidecar = Sidecar(
+        proc,
+        module_name,
+        shutdown_timeout=shutdown_timeout,
+    )
+    sidecar.start()
+    return sidecar
diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index be7028baf28f..838b925a34db 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -1070,6 +1070,22 @@ class ServerArgs:
         "default. In legacy --smg-grpc-mode this is the SMG server port and "
         "defaults to --port + 10000.",
     ] = None
+    sidecar: A[
+        Optional[str],
+        "Start a locally managed sidecar against the native gRPC server. "
+        "The selected module must expose main(argv) and read the resolved "
+        "native gRPC endpoint from SGLANG_GRPC_ENDPOINT. Requires --grpc-port "
+        "or SGLANG_GRPC_PORT.",
+    ] = None
+    sidecar_args: A[
+        Optional[List[str]],
+        Arg(
+            help="JSON array passed to the selected sidecar module's "
+            "main(argv) function. --sidecar-shutdown-timeout SECONDS is "
+            "consumed by SGLang.",
+            type_parser=json_list_type,
+        ),
+    ] = None
     skip_server_warmup: A[bool, "If set, skip warmup."] = False
     warmups: A[
         Optional[str],
@@ -3325,6 +3341,23 @@ def _handle_deprecated_args(self):
         # Native gRPC is incompatible with launch paths it doesn't wire into.
         # Legacy takes precedence over grpc_port, keeping re-runs idempotent.
         native_grpc = self.grpc_port is not None and not legacy_grpc
+        if self.sidecar_args is not None:
+            if self.sidecar is None:
+                raise ValueError("--sidecar-args requires --sidecar.")
+            if not isinstance(self.sidecar_args, list) or not all(
+                isinstance(arg, str) for arg in self.sidecar_args
+            ):
+                raise ValueError("--sidecar-args must be a JSON array of strings.")
+        if self.sidecar is not None:
+            if not self.sidecar.strip():
+                raise ValueError("--sidecar must not be empty.")
+            if legacy_grpc:
+                raise ValueError(
+                    "--sidecar requires SGLang's native gRPC server; "
+                    "it cannot be combined with --smg-grpc-mode/--grpc-mode."
+                )
+            if self.grpc_port is None:
+                raise ValueError("--sidecar requires --grpc-port or SGLANG_GRPC_PORT.")
         if native_grpc:
             if self.use_ray:
                 raise ValueError(
diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py
index 3887cbb7c03c..677e13b80511 100644
--- a/test/registered/unit/server_args/test_server_args.py
+++ b/test/registered/unit/server_args/test_server_args.py
@@ -9,6 +9,13 @@
 
 import sglang.srt.server_args as server_args_module
 from sglang.srt.arg_groups.speculative_hook import handle_speculative_decoding
+from sglang.srt.entrypoints.sidecar import (
+    SGLANG_GRPC_ENDPOINT_ENV,
+    Sidecar,
+    _run_sidecar,
+    build_sidecar_endpoint,
+    start_sidecar,
+)
 from sglang.srt.environ import envs
 from sglang.srt.layers.cp.base import is_cp_enabled, is_interleave
 from sglang.srt.model_executor.cuda_graph_config import (
@@ -1536,6 +1543,133 @@ def test_env_grpc_port_enables_native(self):
             sa._handle_deprecated_args()
         self.assertEqual(sa.grpc_port, 45000)
 
+    @staticmethod
+    def _sidecar_parser():
+        parser = server_args_module.argparse.ArgumentParser()
+        ServerArgs.add_cli_args(parser)
+        return parser
+
+    def test_sidecar_builds_loopback_grpc_endpoints(self):
+        self.assertEqual(
+            build_sidecar_endpoint(SimpleNamespace(host="0.0.0.0", grpc_port=50051)),
+            "http://127.0.0.1:50051",
+        )
+        self.assertEqual(
+            build_sidecar_endpoint(SimpleNamespace(host="::", grpc_port=50051)),
+            "http://[::1]:50051",
+        )
+        self.assertEqual(
+            build_sidecar_endpoint(SimpleNamespace(host="[::]", grpc_port=50051)),
+            "http://[::1]:50051",
+        )
+
+    def test_sidecar_args_parse_as_exact_json_argv(self):
+        argv = ["--flag", "value"]
+        parsed = self._sidecar_parser().parse_args(
+            ["--model-path", "dummy", "--sidecar-args", json.dumps(argv)]
+        )
+        self.assertEqual(parsed.sidecar_args, argv)
+
+    def test_start_sidecar_passes_endpoint_and_provider_argv_separately(self):
+        server_args = SimpleNamespace(
+            sidecar="example.sidecar",
+            sidecar_args=[
+                "--sidecar-shutdown-timeout",
+                "42",
+                "--grpc-connections",
+                "2",
+            ],
+            host="127.0.0.1",
+            grpc_port=50051,
+        )
+        with (
+            patch("sglang.srt.entrypoints.sidecar.mp.get_context") as get_context,
+            patch("sglang.srt.entrypoints.sidecar.Sidecar") as sidecar_class,
+        ):
+            start_sidecar(server_args)
+
+        process_kwargs = get_context.return_value.Process.call_args.kwargs
+        self.assertEqual(process_kwargs["name"], "sglang_sidecar_example.sidecar")
+        self.assertEqual(process_kwargs["target"], _run_sidecar)
+        self.assertEqual(
+            process_kwargs["args"],
+            (
+                "example.sidecar",
+                ["--grpc-connections", "2"],
+                "http://127.0.0.1:50051",
+            ),
+        )
+        sidecar_class.assert_called_once_with(
+            get_context.return_value.Process.return_value,
+            "example.sidecar",
+            shutdown_timeout=42.0,
+        )
+
+    def test_sidecar_requires_native_grpc(self):
+        sa = self._args(sidecar="example.sidecar")
+        with self.assertRaisesRegex(ValueError, "requires --grpc-port"):
+            sa._handle_deprecated_args()
+
+    def test_sidecar_rejects_legacy_grpc(self):
+        sa = self._args(sidecar="example.sidecar", smg_grpc_mode=True)
+        with self.assertRaisesRegex(ValueError, "native gRPC server"):
+            sa._handle_deprecated_args()
+
+    def test_sidecar_rejects_empty_value(self):
+        sa = self._args(sidecar="", grpc_port=50051)
+        with self.assertRaisesRegex(ValueError, "must not be empty"):
+            sa._handle_deprecated_args()
+
+    def test_sidecar_sets_endpoint_env_before_import_and_calls_main(self):
+        main = MagicMock()
+
+        def import_module(module_name):
+            self.assertEqual(module_name, "example.sidecar")
+            self.assertEqual(
+                os.environ[SGLANG_GRPC_ENDPOINT_ENV],
+                "http://127.0.0.1:50051",
+            )
+            self.assertEqual(os.environ["DYN_NAMESPACE"], "pluh")
+            return SimpleNamespace(main=main)
+
+        with (
+            patch.dict(
+                os.environ,
+                {
+                    SGLANG_GRPC_ENDPOINT_ENV: "http://stale.example:1",
+                    "DYN_NAMESPACE": "pluh",
+                },
+            ),
+            patch("sglang.srt.entrypoints.sidecar.kill_itself_when_parent_died"),
+            patch(
+                "sglang.srt.entrypoints.sidecar.importlib.import_module",
+                side_effect=import_module,
+            ),
+        ):
+            _run_sidecar(
+                "example.sidecar",
+                ["--provider-flag", "value"],
+                "http://127.0.0.1:50051",
+            )
+
+        main.assert_called_once_with(["--provider-flag", "value"])
+
+    def test_sidecar_stop_uses_configured_shutdown_timeout(self):
+        proc = MagicMock(pid=1234)
+        proc.is_alive.side_effect = [True, True]
+        sidecar = Sidecar(
+            proc,
+            "example.sidecar",
+            shutdown_timeout=42.0,
+        )
+
+        with patch("sglang.srt.entrypoints.sidecar.kill_process_tree") as kill_tree:
+            sidecar.stop()
+
+        proc.terminate.assert_called_once_with()
+        proc.join.assert_called_once_with(timeout=42.0)
+        kill_tree.assert_called_once_with(1234, wait_timeout=42.0)
+
     def test_legacy_smg_derives_grpc_port_from_http_port(self):
         sa = self._args(port=30000, smg_grpc_mode=True)
         sa._handle_deprecated_args()

From 6227ae88b94d9d76973eebfcb11364308b650561 Mon Sep 17 00:00:00 2001
From: Kangyan-Zhou <zky314343421@gmail.com>
Date: Thu, 23 Jul 2026 17:58:26 -0400
Subject: [PATCH 02/43] [Cherry-pick to release/v0.5.16] Fix nvfp4 online scale
 with pcg (#32246) (#32259)

Co-authored-by: Qiaolin Yu <liin1211@outlook.com>
---
 .../sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
index 72d44a474465..bc90aac8f6c9 100644
--- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
@@ -965,9 +965,15 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
         ):
             e4m3_max = 256.0
 
+        global_scale_inv = torch.full(
+            (1,),
+            1.0 / (e4m3_max * 6.0),
+            dtype=torch.float32,
+            device=hidden_states.device,
+        )
         hs_fp4_bytes, hs_sf_bytes, per_token_scale = nvfp4_quantize(
             hidden_states,
-            1.0 / (e4m3_max * 6.0),
+            global_scale_inv,
             sfLayout=SfLayout.layout_linear,
             per_token_activation=True,
             backend="cute-dsl",

From 6870be7448172b35594fd88bbcd2185c64035988 Mon Sep 17 00:00:00 2001
From: Kangyan-Zhou <zky314343421@gmail.com>
Date: Thu, 23 Jul 2026 18:04:40 -0400
Subject: [PATCH 03/43] [Cherry-pick to release/v0.5.16] [spec decoding] fix
 inkling multi layer mtp draft extend cuda graph (#32254) (#32260)

Co-authored-by: Qiaolin Yu <liin1211@outlook.com>
---
 .../sglang/srt/speculative/multi_layer_eagle_worker_v2.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
index 309efdc308f1..4f3720776161 100644
--- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
+++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
@@ -673,8 +673,12 @@ def _draft_extend_for_decode(
         # Batch 2: Draft extend
         draft_extend_input = EagleDraftExtendInput(
             hidden_states=batch_result.logits_output.hidden_states,
-            # Actual width: the multi-layer chain fills num_steps + 1 rows/req.
-            num_tokens_per_req=self.speculative_num_steps + 1,
+            # Actual width: the multi-layer chain fills num_steps + 1 rows/req,
+            # plus the boundary-KV front rows when the widened window is active
+            # (must match the capture width in the draft-extend graph runner).
+            num_tokens_per_req=self.speculative_num_steps
+            + 1
+            + self.draft_extend_num_front_tokens,
             num_tokens_for_logprob_per_req=1,
             num_front_tokens=self.draft_extend_num_front_tokens,
         )

From 7505dc627f70c7f66ab82a658e39da4cd6fc241f Mon Sep 17 00:00:00 2001
From: Qiaolin Yu <liin1211@outlook.com>
Date: Thu, 23 Jul 2026 23:15:43 -0700
Subject: [PATCH 04/43] [Cherry-pick to release/v0.5.16] Fix dynamo recompile
 limit in allreduce and bf16 gemm (#32239) (#32292)

---
 .../sglang/srt/distributed/parallel_state.py  | 114 ++++++++++++++----
 python/sglang/srt/layers/communicator.py      |   7 +-
 .../sglang/srt/layers/quantization/unquant.py |  41 ++++++-
 python/sglang/srt/server_args.py              |  13 ++
 4 files changed, 141 insertions(+), 34 deletions(-)

diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py
index b1d26d32f723..990d1b98ccfb 100644
--- a/python/sglang/srt/distributed/parallel_state.py
+++ b/python/sglang/srt/distributed/parallel_state.py
@@ -626,7 +626,10 @@ def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
 
         In addition, PyTorch custom ops do not support mutation or returning
         a new tensor in the same op. So we need to figure out if the op is
-        in-place or out-of-place ahead of time.
+        in-place or out-of-place ahead of time — except under Dynamo tracing,
+        where the method selection would guard on the symbolic shape; there we
+        always emit the out-of-place op with method "auto" and resolve the
+        method at runtime inside the op.
         """
         # Bypass the function if we are using only 1 GPU.
         if self.world_size == 1:
@@ -656,6 +659,32 @@ def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
         if self.npu_communicator is not None and not self.npu_communicator.disabled:
             return self.npu_communicator.all_reduce(input_)
 
+        if torch.compiler.is_compiling():
+            # Byte-size thresholds in method selection (e.g. `_pick_algo` or
+            # `should_mscclpp_allreduce`) would guard on the symbolic token dim
+            # and recompile per shape; defer the selection to runtime inside
+            # the opaque custom op. Groups without any accelerated
+            # communicator keep the inplace split op so their collective
+            # stays outside captured graphs. The symmetric-memory in-place
+            # path below is deliberately bypassed under compile: its raw
+            # pynccl call is untraceable (hard error with fullgraph, graph
+            # break otherwise) and its in-place contract does not fit the
+            # outplace custom op.
+            if (
+                self.ca_comm is None
+                and self.qr_comm is None
+                and self.pymscclpp_comm is None
+                and self.torch_symm_mem_comm is None
+                and self.pynccl_comm is None
+            ):
+                inplace_all_reduce(input_, group_name=self.unique_name)
+                return input_
+            return outplace_all_reduce(
+                input_,
+                group_name=self.unique_name,
+                outplace_all_reduce_method="auto",
+            )
+
         should_use_pymscclpp_allreduce = (
             self.pymscclpp_comm is not None
             and self.pymscclpp_comm.should_mscclpp_allreduce(input_)
@@ -670,31 +699,10 @@ def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
                 self.pynccl_comm.all_reduce(input_)
                 return input_
 
-        outplace_all_reduce_method = None
-        if (
-            self.ca_comm is not None
-            and not self.ca_comm.disabled
-            and not should_use_pymscclpp_allreduce
-            and self.ca_comm.should_custom_ar(input_)
-        ):
-            outplace_all_reduce_method = "ca"
-        elif (
-            self.qr_comm is not None
-            and not self.qr_comm.disabled
-            and self.qr_comm.should_quick_allreduce(input_)
-        ):
-            outplace_all_reduce_method = "qr"
-        elif self.pymscclpp_comm is not None and should_use_pymscclpp_allreduce:
-            outplace_all_reduce_method = "pymscclpp"
-        elif (
-            self.torch_symm_mem_comm is not None
-            and not self.torch_symm_mem_comm.disabled
-            and self.torch_symm_mem_comm.should_torch_symm_mem_allreduce(input_)
-        ):
-            outplace_all_reduce_method = "torch_symm_mem"
-        elif is_in_tc_piecewise_cuda_graph() and self.pynccl_comm is not None:
-            # For piecewise cuda graph, we use pynccl outplace allreduce
-            outplace_all_reduce_method = "pynccl"
+        outplace_all_reduce_method = self._resolve_outplace_all_reduce_method(
+            input_=input_,
+            should_use_pymscclpp_allreduce=should_use_pymscclpp_allreduce,
+        )
         if outplace_all_reduce_method is not None:
             return outplace_all_reduce(
                 input_,
@@ -780,9 +788,63 @@ def fused_allreduce_rmsnorm(
         )
         return fused_outputs
 
+    def _resolve_outplace_all_reduce_method(
+        self,
+        input_: torch.Tensor,
+        should_use_pymscclpp_allreduce: Optional[bool] = None,
+    ) -> Optional[str]:
+        if should_use_pymscclpp_allreduce is None:
+            should_use_pymscclpp_allreduce = (
+                self.pymscclpp_comm is not None
+                and self.pymscclpp_comm.should_mscclpp_allreduce(input_)
+            )
+        if (
+            self.ca_comm is not None
+            and not self.ca_comm.disabled
+            and not should_use_pymscclpp_allreduce
+            and self.ca_comm.should_custom_ar(input_)
+        ):
+            return "ca"
+        if (
+            self.qr_comm is not None
+            and not self.qr_comm.disabled
+            and self.qr_comm.should_quick_allreduce(input_)
+        ):
+            return "qr"
+        if self.pymscclpp_comm is not None and should_use_pymscclpp_allreduce:
+            return "pymscclpp"
+        if (
+            self.torch_symm_mem_comm is not None
+            and not self.torch_symm_mem_comm.disabled
+            and self.torch_symm_mem_comm.should_torch_symm_mem_allreduce(input_)
+        ):
+            return "torch_symm_mem"
+        if is_in_tc_piecewise_cuda_graph() and self.pynccl_comm is not None:
+            # For piecewise cuda graph, we use pynccl outplace allreduce
+            return "pynccl"
+        return None
+
     def _all_reduce_out_place(
         self, input_: torch.Tensor, outplace_all_reduce_method: str
     ) -> torch.Tensor:
+        if outplace_all_reduce_method == "auto":
+            outplace_all_reduce_method = self._resolve_outplace_all_reduce_method(
+                input_
+            )
+            if outplace_all_reduce_method == "pymscclpp":
+                # pymscclpp reduces in place and returns its input; feed it a
+                # clone to honor the op's no-mutation contract.
+                input_ = input_.clone()
+            elif outplace_all_reduce_method is None:
+                # Force pynccl over the in-place fallback: it is graph-capture
+                # safe and NCCL is natively out-of-place, avoiding the clone
+                # the in-place fallback needs.
+                if self.pynccl_comm is not None:
+                    outplace_all_reduce_method = "pynccl"
+                else:
+                    out = input_.clone()
+                    self._all_reduce_in_place(out)
+                    return out
         ca_comm = self.ca_comm
         qr_comm = self.qr_comm
         pymscclpp_comm = self.pymscclpp_comm
diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py
index 3271b7c845c7..a23e6e4f17fc 100644
--- a/python/sglang/srt/layers/communicator.py
+++ b/python/sglang/srt/layers/communicator.py
@@ -167,11 +167,14 @@ def apply_flashinfer_allreduce_fusion(batch_size: int):
         # Ref: https://github.com/sgl-project/sglang/issues/17237
         (_is_sm90_supported or _is_sm100_supported)
         and _is_flashinfer_available
-        and batch_size > 0
-        and batch_size <= FUSE_ALLREDUCE_MAX_BATCH_SIZE
         and not is_dp_attention_enabled()
         and get_server_args().flashinfer_allreduce_fusion_backend is not None
         and not is_flashinfer_allreduce_unavailable()
+        # Symbolic size checks stay last: under Dynamo tracing they guard on
+        # the dynamic token dim, so statically-off configs must short-circuit
+        # before reaching them.
+        and batch_size > 0
+        and batch_size <= FUSE_ALLREDUCE_MAX_BATCH_SIZE
     )
 
 
diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py
index 5c77f22a7b61..d2774728c1dc 100644
--- a/python/sglang/srt/layers/quantization/unquant.py
+++ b/python/sglang/srt/layers/quantization/unquant.py
@@ -40,6 +40,7 @@
     use_intel_amx_backend,
     use_intel_xpu_backend,
 )
+from sglang.srt.utils.custom_op import register_custom_op
 
 if TYPE_CHECKING:
     from sglang.srt.layers.moe.token_dispatcher import (
@@ -104,6 +105,25 @@ def initialize_bf16_gemm_config(server_args: ServerArgs) -> None:
     _BF16_GEMM_BACKEND = backend
 
 
+def _bf16_gemm_dispatch_fake(
+    x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor]
+) -> torch.Tensor:
+    return x.new_empty((*x.shape[:-1], weight.shape[0]))
+
+
+@register_custom_op(fake_impl=_bf16_gemm_dispatch_fake)
+def bf16_gemm_dispatch(
+    x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor]
+) -> torch.Tensor:
+    if _use_cutedsl_bf16_gemm is not None and _use_cutedsl_bf16_gemm(
+        x.numel() // x.shape[-1], weight.shape[0], weight.shape[1]
+    ):
+        return _cutedsl_bf16_gemm(x.view(-1, x.shape[-1]), weight, bias).view(
+            *x.shape[:-1], -1
+        )
+    return F.linear(x, weight, bias)
+
+
 def get_bf16_gemm_backend() -> Bf16GemmBackend:
     global _BF16_GEMM_BACKEND
     if _BF16_GEMM_BACKEND is None:
@@ -207,15 +227,24 @@ def apply(
             and x.dtype == torch.bfloat16
             and layer.weight.dtype == torch.bfloat16
             and (bias is None or bias.dtype == torch.bfloat16)
-            and _use_cutedsl_bf16_gemm(
+        ):
+            if torch.compiler.is_compiling():
+                # The m-dependent kernel heuristic would guard on the symbolic
+                # token dim under Dynamo and recompile per shape bucket; the
+                # opaque op resolves it at runtime with concrete shapes,
+                # keeping the per-shape kernel choice.
+                return bf16_gemm_dispatch(x, layer.weight, bias)
+            if _use_cutedsl_bf16_gemm(
                 x.numel() // x.shape[-1],
                 layer.weight.shape[0],
                 layer.weight.shape[1],
-            )
-        ):
-            x_shapes = x.shape
-            output = _cutedsl_bf16_gemm(x.view(-1, x_shapes[-1]), layer.weight, bias)
-            return output.view(*x_shapes[:-1], -1)
+            ):
+                x_shapes = x.shape
+                output = _cutedsl_bf16_gemm(
+                    x.view(-1, x_shapes[-1]), layer.weight, bias
+                )
+                return output.view(*x_shapes[:-1], -1)
+            return F.linear(x, layer.weight, bias)
 
         return F.linear(x, layer.weight, bias)
 
diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index 838b925a34db..66b7f770f291 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -3755,6 +3755,19 @@ def _disable_tc_piecewise_cudagraph_if_incompatible(self):
                 "decode context parallel (dcp_size > 1)",
                 lambda: self.dcp_size > 1,
             ),
+            # TcPiecewise makes the trtllm_mla prefill fall back to the
+            # flashinfer-MLA implementation, which faults (illegal address)
+            # on an FP8 KV cache.
+            (
+                "MLA attention with FP8 KV cache",
+                lambda: self.kv_cache_dtype.startswith("fp8")
+                and (
+                    _resolved_view(self).attention_backend
+                    in ("trtllm_mla", "flashinfer_mla")
+                    or _resolved_view(self).prefill_attention_backend
+                    in ("trtllm_mla", "flashinfer_mla")
+                ),
+            ),
         ]
         for _name, predicate in rules:
             if predicate():

From fdebc938f7f4d16fe6b9f55dcd9a767cf0899ea1 Mon Sep 17 00:00:00 2001
From: Kangyan-Zhou <zky314343421@gmail.com>
Date: Fri, 24 Jul 2026 16:25:42 -0400
Subject: [PATCH 05/43] [Cherry-pick to release/v0.5.16] Fix stale
 flashinfer-MLA fallback poisoning spec verify capture (trtllm_mla +
 tc_piecewise) (#32288) (#32346)

Co-authored-by: Qiaolin Yu <liin1211@outlook.com>
---
 .../srt/layers/attention/trtllm_mla_backend.py      | 11 ++++++++++-
 python/sglang/srt/server_args.py                    | 13 -------------
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py
index 14fba1e25f52..9dc858dc91d1 100755
--- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py
+++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py
@@ -814,8 +814,17 @@ def forward_extend(
         llama_4_scaling: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
 
+        # The fallback belongs to genuine extend forwards only. Target-verify /
+        # draft-extend must never honor it: `forward_prefill_metadata` is a
+        # stale leftover from the last prefill there (eager init clears it,
+        # but decode-graph capture does not), and capturing verify through the
+        # flashinfer path binds the graph to prefill-planned wrapper buffers,
+        # which fault (illegal address) at replay.
         if (
-            self.forward_prefill_metadata is not None
+            forward_batch.forward_mode.is_extend()
+            and not forward_batch.forward_mode.is_target_verify()
+            and not forward_batch.forward_mode.is_draft_extend_v2()
+            and self.forward_prefill_metadata is not None
             and self.forward_prefill_metadata.fallback_to_flashinfer_impl
         ):
             return super().forward_extend(
diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index 66b7f770f291..838b925a34db 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -3755,19 +3755,6 @@ def _disable_tc_piecewise_cudagraph_if_incompatible(self):
                 "decode context parallel (dcp_size > 1)",
                 lambda: self.dcp_size > 1,
             ),
-            # TcPiecewise makes the trtllm_mla prefill fall back to the
-            # flashinfer-MLA implementation, which faults (illegal address)
-            # on an FP8 KV cache.
-            (
-                "MLA attention with FP8 KV cache",
-                lambda: self.kv_cache_dtype.startswith("fp8")
-                and (
-                    _resolved_view(self).attention_backend
-                    in ("trtllm_mla", "flashinfer_mla")
-                    or _resolved_view(self).prefill_attention_backend
-                    in ("trtllm_mla", "flashinfer_mla")
-                ),
-            ),
         ]
         for _name, predicate in rules:
             if predicate():

From 5338b11556635f556a462343a1b3e490c526af73 Mon Sep 17 00:00:00 2001
From: Yueming Yuan <yym022502@gmail.com>
Date: Fri, 24 Jul 2026 21:33:54 -0700
Subject: [PATCH 06/43] [1/27] [sglang-miles] True on-policy training support
 (FSDP2 + qwen_dense) (#18639)

Squashes the qwen_dense on-policy follow-up into the original FSDP2 commit:
the follow-up rewrote the norm_kwargs / rl_on_policy_target call sites this
commit introduced, so they cannot be applied independently.

Rebased onto v0.5.16: rl_on_policy_target predicates now go through the
sglang.srt.true_on_policy helpers, and every legacy get_global_server_args()
call site uses runtime_context.get_server_args() so v0.5.16's
test_legacy_global_ratchet baseline still holds.

Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>
---
 python/sglang/srt/arg_groups/overrides.py     |   14 +-
 .../srt/distributed/communication_op.py       |    9 +
 .../sglang/srt/distributed/parallel_state.py  |    5 +-
 python/sglang/srt/layers/activation.py        |    5 +-
 python/sglang/srt/layers/attention/vision.py  |    8 +-
 python/sglang/srt/layers/communicator.py      |   10 +
 python/sglang/srt/layers/layernorm.py         |   62 +-
 python/sglang/srt/layers/linear.py            |   10 +-
 python/sglang/srt/layers/logits_processor.py  |   13 +-
 .../moe/moe_runner/triton_utils/fused_moe.py  |    6 +-
 python/sglang/srt/layers/on_policy_utils.py   |    6 +
 .../srt/layers/rotary_embedding/base.py       |   14 +-
 .../srt/layers/rotary_embedding/mrope.py      |    3 +-
 python/sglang/srt/layers/sampler.py           |   25 +-
 .../srt/model_executor/forward_batch_info.py  |    5 +-
 .../sglang/srt/model_executor/model_runner.py |    5 +
 .../runner/decode_cuda_graph_runner.py        |   35 +-
 python/sglang/srt/models/qwen2.py             |   32 +-
 python/sglang/srt/models/qwen2_moe.py         |    6 +-
 python/sglang/srt/models/qwen3.py             |   62 +-
 python/sglang/srt/models/qwen3_moe.py         |   38 +-
 python/sglang/srt/models/qwen3_vl.py          |  172 +-
 python/sglang/srt/models/sdar.py              |   34 +-
 python/sglang/srt/models/sdar_moe.py          |   34 +-
 python/sglang/srt/models/step3p5.py           |    7 +-
 python/sglang/srt/models/utils.py             |    2 +
 .../multimodal/processors/base_processor.py   |    3 +-
 python/sglang/srt/server_args.py              |   45 +-
 .../sglang/srt/tp_invariant_ops/__init__.py   |   23 +
 .../srt/tp_invariant_ops/tp_invariant_ops.py  | 1941 +++++++++++++++++
 python/sglang/srt/true_on_policy/__init__.py  |   49 +
 python/sglang/srt/true_on_policy/config.py    |  150 ++
 python/sglang/srt/true_on_policy/contracts.py |  111 +
 python/sglang/srt/true_on_policy/schema.py    |   33 +
 test/manual/layers/test_layernorm.py          |   20 +
 .../core/test_dense_deterministic_math.py     |  338 +++
 test/registered/core/test_on_policy_wiring.py |  606 +++++
 test/registered/core/test_tp_invariant_ops.py |  866 ++++++++
 38 files changed, 4518 insertions(+), 289 deletions(-)
 create mode 100644 python/sglang/srt/layers/on_policy_utils.py
 create mode 100644 python/sglang/srt/tp_invariant_ops/__init__.py
 create mode 100644 python/sglang/srt/tp_invariant_ops/tp_invariant_ops.py
 create mode 100644 python/sglang/srt/true_on_policy/__init__.py
 create mode 100644 python/sglang/srt/true_on_policy/config.py
 create mode 100644 python/sglang/srt/true_on_policy/contracts.py
 create mode 100644 python/sglang/srt/true_on_policy/schema.py
 create mode 100644 test/registered/core/test_dense_deterministic_math.py
 create mode 100644 test/registered/core/test_on_policy_wiring.py
 create mode 100644 test/registered/core/test_tp_invariant_ops.py

diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py
index 5e9676993c15..6e606762b450 100644
--- a/python/sglang/srt/arg_groups/overrides.py
+++ b/python/sglang/srt/arg_groups/overrides.py
@@ -1939,7 +1939,12 @@ def _page_size_default(view: Any) -> dict:
 @register_post_process
 def _data_parallelism_defaults(view: Any) -> dict:
     if view.dp_size == 1 and view.ep_join_mode != "scale":
-        return {"enable_dp_attention": False, "enable_dp_lm_head": False}
+        overrides = {"enable_dp_attention": False}
+        # Keep the dp LM head when attention context parallelism is enabled,
+        # even without data parallelism (context-parallel LM head).
+        if not (view.enable_dp_lm_head and view.attn_cp_size > 1):
+            overrides["enable_dp_lm_head"] = False
+        return overrides
     return {}
 
 
@@ -1948,9 +1953,10 @@ def _dp_lm_head_validation(view: Any) -> dict:
     """Read-only validation pass: dp-attention is a prerequisite for the
     dp LM head. Reads the mid-resolution values through the view."""
     if view.enable_dp_lm_head:
-        assert (
-            view.enable_dp_attention
-        ), "Please enable dp attention when setting enable_dp_lm_head. "
+        assert view.enable_dp_attention or view.attn_cp_size > 1, (
+            "Please enable dp attention when setting enable_dp_lm_head, "
+            "unless attention context parallelism is enabled."
+        )
     return {}
 
 
diff --git a/python/sglang/srt/distributed/communication_op.py b/python/sglang/srt/distributed/communication_op.py
index 89f9986e4f68..674300d3b718 100644
--- a/python/sglang/srt/distributed/communication_op.py
+++ b/python/sglang/srt/distributed/communication_op.py
@@ -7,6 +7,9 @@
 import torch
 import torch.distributed
 
+from sglang.srt.tp_invariant_ops import tree_all_reduce_sum
+from sglang.srt.true_on_policy import should_use_tp_invariant_tree_all_reduce
+
 from .parallel_state import (
     get_attn_tp_group,
     get_moe_ep_group,
@@ -17,6 +20,8 @@
 
 def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
     """All-reduce the input tensor across model parallel group."""
+    if should_use_tp_invariant_tree_all_reduce():
+        return tree_all_reduce_sum(input_, device_group=get_tp_group().device_group)
     return get_tp_group().all_reduce(input_)
 
 
@@ -64,6 +69,10 @@ def broadcast_tensor_dict(
 
 def attention_tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
     """All-reduce the input tensor across attention parallel group."""
+    if should_use_tp_invariant_tree_all_reduce():
+        return tree_all_reduce_sum(
+            input_, device_group=get_attn_tp_group().device_group
+        )
     return get_attn_tp_group().all_reduce(input_)
 
 
diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py
index 990d1b98ccfb..58a5a6469e51 100644
--- a/python/sglang/srt/distributed/parallel_state.py
+++ b/python/sglang/srt/distributed/parallel_state.py
@@ -2607,7 +2607,10 @@ def get_dcp_rank():
 
 def get_tensor_model_parallel_rank():
     """Return my rank for the tensor model parallel group."""
-    return get_tp_group().rank_in_group
+    try:
+        return get_tp_group().rank_in_group
+    except Exception:
+        return 0
 
 
 # ATTN_TP
diff --git a/python/sglang/srt/layers/activation.py b/python/sglang/srt/layers/activation.py
index a7db9d27afdd..307c2249e323 100644
--- a/python/sglang/srt/layers/activation.py
+++ b/python/sglang/srt/layers/activation.py
@@ -33,7 +33,8 @@
     Phase,
     check_cuda_graph_backend,
 )
-from sglang.srt.runtime_context import get_parallel, get_server_args
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.true_on_policy import is_true_on_policy_enabled
 from sglang.srt.utils import (
     cpu_has_amx_support,
     get_bool_env_var,
@@ -89,7 +90,7 @@ def _(x):
 class SiluAndMul(MultiPlatformOp):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        if get_server_args().rl_on_policy_target is not None:
+        if is_true_on_policy_enabled():
             self._forward_method = self.forward_native
         elif _use_aiter and envs.SGLANG_OPT_USE_AITER_SILU_MUL.get():
             self._forward_method = self.forward_aiter
diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py
index f11ed272232b..72843493c1d6 100644
--- a/python/sglang/srt/layers/attention/vision.py
+++ b/python/sglang/srt/layers/attention/vision.py
@@ -16,6 +16,7 @@
 from sglang.srt.environ import envs
 from sglang.srt.models.utils import apply_qk_norm
 from sglang.srt.runtime_context import get_parallel
+from sglang.srt.true_on_policy import is_true_on_policy_enabled
 from sglang.srt.utils import (
     cpu_has_amx_support,
     get_bool_env_var,
@@ -1124,7 +1125,7 @@ def _init_qk_norm(
                 weight_dtype=torch.float32,
                 cast_x_before_out_mul=True,
             )
-            if get_server_args().rl_on_policy_target is not None
+            if is_true_on_policy_enabled()
             else {}
         )
         q_norm = RMSNorm(
@@ -1256,10 +1257,7 @@ def forward(
         if x.dim() == 2:
             x = x.unsqueeze(0)
         assert x.dim() == 3, x.shape
-        if (
-            get_server_args().rl_on_policy_target is not None
-            and position_embeddings is not None
-        ):
+        if is_true_on_policy_enabled() and position_embeddings is not None:
             assert isinstance(position_embeddings, tuple), (
                 "expected position_embeddings to be a tuple of two tensors,\n"
                 f"but got {type(position_embeddings)}, change if needed"
diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py
index a23e6e4f17fc..f60957e93d8a 100644
--- a/python/sglang/srt/layers/communicator.py
+++ b/python/sglang/srt/layers/communicator.py
@@ -74,6 +74,10 @@
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.runtime_context import get_forward, get_parallel, get_server_args
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+from sglang.srt.true_on_policy import (
+    should_disable_mlp_allreduce_fusion_for_on_policy,
+    should_disable_reduce_scatter_for_on_policy,
+)
 from sglang.srt.utils import (
     get_bool_env_var,
     is_cuda,
@@ -756,6 +760,9 @@ def postprocess_layer(
         )
 
     def should_use_reduce_scatter(self, forward_batch: ForwardBatch):
+        if should_disable_reduce_scatter_for_on_policy():
+            return False
+
         if not self.allow_reduce_scatter:
             return False
         if (
@@ -776,6 +783,9 @@ def should_use_reduce_scatter(self, forward_batch: ForwardBatch):
     def should_fuse_mlp_allreduce_with_next_layer(
         self, forward_batch: ForwardBatch
     ) -> bool:
+        if should_disable_mlp_allreduce_fusion_for_on_policy():
+            return False
+
         # When MOE_FULL is active (moe_cp allgather), fusion must be disabled because
         # the fusion path skips postprocess_layer which contains the moe_cp scatter.
         # Without scatter, hidden_states remain at MOE_FULL size while residual is at
diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py
index 8b0ca51417e7..7c5f535169ef 100644
--- a/python/sglang/srt/layers/layernorm.py
+++ b/python/sglang/srt/layers/layernorm.py
@@ -32,6 +32,10 @@
     check_cuda_graph_backend,
 )
 from sglang.srt.runtime_context import get_parallel, get_server_args
+from sglang.srt.true_on_policy import (
+    get_on_policy_rms_norm_kwargs,
+    is_true_on_policy_enabled,
+)
 from sglang.srt.utils import (
     cpu_has_amx_support,
     get_bool_env_var,
@@ -184,7 +188,7 @@ def _forward_with_allreduce_fusion(
 
         if world_size > 1:
             if post_residual_addition is not None:
-                residual = residual + post_residual_addition
+                x = x + post_residual_addition
 
             # Prefer AITER fused AR+RMSNorm when enabled on AMD.
             if _use_aiter:
@@ -220,13 +224,32 @@ def __init__(
         eps: float = 1e-6,
         var_hidden_size: Optional[int] = None,
         cast_x_before_out_mul: bool = False,
-        fp32_residual: bool = False,
+        fp32_residual: bool = True,
         has_weight: bool = True,
-        weight_dtype: Optional = None,
-        override_orig_dtype: Optional = None,
+        weight_dtype: Optional[torch.dtype] = None,
+        override_orig_dtype: Optional[torch.dtype] = None,
         x_pad_to_multiple: int = 0,
+        true_on_policy_weight_dtype: Optional[torch.dtype] = None,
+        true_on_policy_override_orig_dtype: Optional[torch.dtype] = None,
+        true_on_policy_fp32_residual: bool = False,
     ) -> None:
         super().__init__()
+        true_on_policy_kwargs = get_on_policy_rms_norm_kwargs(
+            weight_dtype=true_on_policy_weight_dtype,
+            override_orig_dtype=true_on_policy_override_orig_dtype,
+            fp32_residual=true_on_policy_fp32_residual,
+        )
+        if not cast_x_before_out_mul:
+            cast_x_before_out_mul = true_on_policy_kwargs.get(
+                "cast_x_before_out_mul", cast_x_before_out_mul
+            )
+        fp32_residual = true_on_policy_kwargs.get("fp32_residual", fp32_residual)
+        if weight_dtype is None:
+            weight_dtype = true_on_policy_kwargs.get("weight_dtype", weight_dtype)
+        if override_orig_dtype is None:
+            override_orig_dtype = true_on_policy_kwargs.get(
+                "override_orig_dtype", override_orig_dtype
+            )
         self.has_weight = has_weight
         self.cast_x_before_out_mul = cast_x_before_out_mul
         self.fp32_residual = fp32_residual
@@ -280,11 +303,17 @@ def forward_cuda(
             x = x.contiguous().reshape(-1, original_shape[-1])
         if self.variance_size_override is not None:
             return self.forward_native(x, residual, post_residual_addition)
+        if (
+            self.weight.dtype != x.dtype
+            or self.cast_x_before_out_mul
+            or self.override_orig_dtype is not None
+        ):
+            return self.forward_native(x, residual, post_residual_addition)
         if is_batch_invariant_mode_enabled():
             if (
                 residual is not None
                 or self.cast_x_before_out_mul
-                or get_server_args().rl_on_policy_target == "fsdp"
+                or is_true_on_policy_enabled()
             ):
                 return self.forward_native(x, residual, post_residual_addition)
             out = rms_norm_batch_invariant(
@@ -420,10 +449,10 @@ def forward_aiter(
                 self.x_pad_to_multiple,
             )
         if residual is not None:
-            residual_out = torch.empty_like(x)
-            output = torch.empty_like(x)
             if post_residual_addition is not None:
                 residual = residual + post_residual_addition
+            residual_out = torch.empty_like(x)
+            output = torch.empty_like(x)
             fused_add_rms_norm(
                 output,
                 x,
@@ -465,10 +494,10 @@ def forward_hip(
             # NOTE: Remove this if aiter kernel supports discontinuous input
             x = x.contiguous()
         if residual is not None:
-            out = torch.empty_like(x)
-            residual_out = torch.empty_like(x)
             if post_residual_addition is not None:
                 residual = residual + post_residual_addition
+            out = torch.empty_like(x)
+            residual_out = torch.empty_like(x)
             fused_add_rms_norm(
                 out, x, residual_out, residual, self.weight.data, self.variance_epsilon
             )
@@ -509,15 +538,18 @@ def forward_native(
         if not x.is_contiguous():
             x = x.contiguous()
         orig_dtype = self.override_orig_dtype or x.dtype
+
+        if residual is not None and not self.fp32_residual:
+            x = x + residual
+            if post_residual_addition is not None:
+                x = x + post_residual_addition
+            residual = x.clone()
         x = x.to(torch.float32)
-        if residual is not None:
+        if residual is not None and self.fp32_residual:
             x = x + residual.to(torch.float32)
             if post_residual_addition is not None:
                 x = x + post_residual_addition.to(torch.float32)
-            if self.fp32_residual:
-                residual = x.clone()
-            else:
-                residual = x.to(orig_dtype)
+            residual = x.to(orig_dtype)
 
         hidden_size = x.shape[-1]
         if hidden_size != self.hidden_size:
@@ -745,7 +777,7 @@ def forward_native(
         orig_dtype = x.dtype
         if residual is not None:
             if post_residual_addition is not None:
-                residual = residual + post_residual_addition
+                x = x + post_residual_addition
             x = x + residual
             residual = x
 
diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py
index 086b64badac0..f91ec3f765bd 100644
--- a/python/sglang/srt/layers/linear.py
+++ b/python/sglang/srt/layers/linear.py
@@ -40,6 +40,7 @@
 )
 from sglang.srt.layers.utils import pad_or_narrow_weight
 from sglang.srt.runtime_context import get_parallel, get_server_args
+from sglang.srt.true_on_policy import should_use_tp_invariant_row_linear
 from sglang.srt.utils import get_bool_env_var, is_cpu, is_hip, is_npu, set_weight_attrs
 
 if TYPE_CHECKING:
@@ -1578,7 +1579,14 @@ def forward(self, input_, skip_all_reduce=False, forward_batch=None):
                 get_tp_group(), disabled=not is_allocation_symmetric()
             )
         with symm_ctx:
-            output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_)
+            if should_use_tp_invariant_row_linear(input_parallel.shape[-1]):
+                output_parallel = torch.ops.tp_inv_ops.matmul_tp_inv(
+                    input_parallel, self.weight.t(), bias_
+                )
+            else:
+                output_parallel = self.quant_method.apply(
+                    self, input_parallel, bias=bias_
+                )
 
         # skip_all_reduce: explicit call-site override. Also honor
         # ForwardFlags (fuse_mlp_allreduce / mlp_reduce_scatter) published by
diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py
index cf55b97cf446..5ab5e27448e9 100644
--- a/python/sglang/srt/layers/logits_processor.py
+++ b/python/sglang/srt/layers/logits_processor.py
@@ -48,6 +48,7 @@
     ForwardMode,
 )
 from sglang.srt.runtime_context import get_parallel, get_server_args
+from sglang.srt.true_on_policy import should_force_bfloat16_lm_head
 from sglang.srt.utils.common import (
     is_cpu,
     is_npu,
@@ -745,6 +746,11 @@ def _compute_lm_head(
                 logits = torch.matmul(
                     hidden_states.to(torch.float32), lm_head.weight.to(torch.float32).T
                 )
+            elif should_force_bfloat16_lm_head(use_fp32_lm_head=self.use_fp32_lm_head):
+                logits = torch.matmul(
+                    hidden_states.to(torch.bfloat16),
+                    lm_head.weight.to(torch.bfloat16).T,
+                ).to(torch.bfloat16)
             elif use_intel_amx_backend(lm_head):
                 logits = torch.ops.sgl_kernel.weight_packed_linear(
                     hidden_states.to(lm_head.weight.dtype),
@@ -752,11 +758,6 @@ def _compute_lm_head(
                     None,  # bias
                     True,  # is_vnni
                 )
-            elif self.rl_on_policy_target is not None:
-                # Due to tie-weight, we may not be able to change lm_head's weight dtype
-                logits = torch.matmul(
-                    hidden_states.bfloat16(), lm_head.weight.T.bfloat16()
-                )
             else:
                 logits = torch.matmul(
                     hidden_states.to(lm_head.weight.dtype), lm_head.weight.T
@@ -848,6 +849,8 @@ def _copy_logits_to_buffer(
             assert logits_buffer.dtype == torch.float
             logits_buffer.copy_(logits)
             logits = logits_buffer
+        elif should_force_bfloat16_lm_head(use_fp32_lm_head=self.use_fp32_lm_head):
+            logits = logits[:, : self.vocab_size].to(torch.bfloat16)
         else:
             logits = logits.float()
         return logits
diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py
index 6b1f0112831c..ab6d27a2b16f 100644
--- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py
+++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py
@@ -95,7 +95,11 @@
 
 
 def _use_moe_sum_reduce_torch_compile(num_tokens: int) -> bool:
-    return num_tokens <= 32 and not is_batch_invariant_mode_enabled()
+    return (
+        num_tokens <= 32
+        and not is_batch_invariant_mode_enabled()
+        and not get_server_args().enable_deterministic_inference
+    )
 
 
 @register_custom_op(mutates_args=["hidden_states"])
diff --git a/python/sglang/srt/layers/on_policy_utils.py b/python/sglang/srt/layers/on_policy_utils.py
new file mode 100644
index 000000000000..234397b9b33c
--- /dev/null
+++ b/python/sglang/srt/layers/on_policy_utils.py
@@ -0,0 +1,6 @@
+"""Compatibility imports for SGLang true-on-policy helpers.
+
+New code should import from :mod:`sglang.srt.true_on_policy`.
+"""
+
+from sglang.srt.true_on_policy import *  # noqa: F403
diff --git a/python/sglang/srt/layers/rotary_embedding/base.py b/python/sglang/srt/layers/rotary_embedding/base.py
index 2f68e4965111..891e924094b5 100644
--- a/python/sglang/srt/layers/rotary_embedding/base.py
+++ b/python/sglang/srt/layers/rotary_embedding/base.py
@@ -11,7 +11,7 @@
 from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb
 from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.platforms import current_platform
-from sglang.srt.runtime_context import get_server_args
+from sglang.srt.true_on_policy import is_true_on_policy_enabled
 from sglang.srt.utils import (
     cpu_has_amx_support,
     get_bool_env_var,
@@ -127,12 +127,8 @@ def __init__(
         self._apply_rotary_emb_wrapped = apply_rotary_emb
 
         # XXX (MUSA): Implement sgl_kernel.rotary_embedding support for MUSA backend
-        if get_server_args().rl_on_policy_target is not None or _is_musa:
+        if is_true_on_policy_enabled() or _is_musa:
             self._forward_method = self.forward_native
-            self._apply_rotary_emb_wrapped = torch.compile(
-                dynamic=True,
-                disable=_is_npu,
-            )(apply_rotary_emb)
         self.position_cos, self.position_sin = None, None
 
     def _match_cos_sin_cache_dtype(self, query: torch.Tensor) -> None:
@@ -150,9 +146,7 @@ def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor:
         # use CPU to compute the cache and then move it to GPU. However, we
         # create the cache on GPU for faster initialization. This may cause
         # a slight numerical difference between the HF implementation and ours.
-        init_device = (
-            "cpu" if get_server_args().rl_on_policy_target is not None else None
-        )
+        init_device = "cpu" if is_true_on_policy_enabled() else None
         inv_freq = 1.0 / (
             base
             ** (
@@ -162,7 +156,7 @@ def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor:
                 / self.rotary_dim
             )
         )
-        if get_server_args().rl_on_policy_target is not None:
+        if is_true_on_policy_enabled():
             inv_freq = inv_freq.cuda()
         return inv_freq
 
diff --git a/python/sglang/srt/layers/rotary_embedding/mrope.py b/python/sglang/srt/layers/rotary_embedding/mrope.py
index 979b9741b689..bdfa907cab38 100644
--- a/python/sglang/srt/layers/rotary_embedding/mrope.py
+++ b/python/sglang/srt/layers/rotary_embedding/mrope.py
@@ -19,6 +19,7 @@
     yarn_linear_ramp_mask,
 )
 from sglang.srt.runtime_context import get_server_args
+from sglang.srt.true_on_policy import is_true_on_policy_enabled
 from sglang.srt.utils import (
     cpu_has_amx_support,
     is_cuda,
@@ -132,7 +133,7 @@ def __init__(
             self.register_buffer("axis_map", axis_map, persistent=False)
         else:
             self.axis_map = None
-        if get_server_args().rl_on_policy_target is not None:
+        if is_true_on_policy_enabled():
             self._forward_method = self.forward_native
 
     def get_cos_sin_with_position(self, positions):
diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py
index de587ec5f466..72bd25ae7b9a 100644
--- a/python/sglang/srt/layers/sampler.py
+++ b/python/sglang/srt/layers/sampler.py
@@ -18,6 +18,7 @@
 from sglang.srt.runtime_context import get_parallel, get_server_args
 from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
 from sglang.srt.sampling.sampling_params import TOP_K_ALL
+from sglang.srt.true_on_policy import resolve_true_on_policy_runtime_policy
 from sglang.srt.utils.async_probe import sanitize_nan_logits
 from sglang.srt.utils.common import (
     get_bool_env_var,
@@ -74,11 +75,11 @@ def __init__(self):
         if is_dp_attention_enabled():
             self.tp_sync_group = get_parallel().attn_tp_group.device_group
 
-        self.rl_on_policy_target = get_server_args().rl_on_policy_target
+        true_on_policy = resolve_true_on_policy_runtime_policy(get_server_args())
         # In RL on-policy mode, deterministic inference is automatically enabled.
         self.enable_deterministic = get_server_args().enable_deterministic_inference
         # In RL on-policy mode, we use log_softmax to compute logprobs to match the trainer.
-        self.use_log_softmax_logprob = self.rl_on_policy_target is not None
+        self.use_log_softmax_logprob = true_on_policy.enabled
         self.use_ascend_backend = get_server_args().sampling_backend == "ascend"
 
         self.output_logprob_processor = OutputLogprobProcessor()
@@ -148,17 +149,13 @@ def forward(
             if return_logprob and SGLANG_RETURN_ORIGINAL_LOGPROB:
                 original_logprobs = torch.log_softmax(logits, dim=-1)
 
+            # Post process logits
+            logits.div_(sampling_info.temperatures)
+
             # In RL on-policy mode, we use log_softmax to compute logprobs to match the trainer.
             logprobs_via_logsoftmax_kernel = None
-            if self.rl_on_policy_target is not None:
-                # TODO: use more inplace ops to save memory
-                logits_div_temperature = (
-                    logits.bfloat16().div(sampling_info.temperatures).bfloat16()
-                )
-                logprobs_via_logsoftmax_kernel = torch.log_softmax(
-                    logits_div_temperature, dim=-1
-                )
-                del logits_div_temperature
+            if self.use_log_softmax_logprob:
+                logprobs_via_logsoftmax_kernel = torch.log_softmax(logits, dim=-1)
 
             if self.use_ascend_backend:
                 # Ascend backend: sample from logits directly.
@@ -184,8 +181,6 @@ def forward(
                     logprobs = logprobs_via_logsoftmax_kernel
             else:
                 # Standard path: do softmax and sample from probs.
-                logits.div_(sampling_info.temperatures)
-
                 # In-place op to save memory
                 logits[:] = torch.softmax(logits, dim=-1)
                 probs = logits
@@ -458,13 +453,13 @@ def _forward_ascend_backend(
         """Handle the full Ascend backend sampling path.
 
         Ascend backend has fused kernels that handle softmax internally,
-        so we sample directly from temperature-scaled logits.
+        so we sample directly from temperature-scaled logits. Temperature
+        scaling is already applied by the caller before branch dispatch.
 
         Returns:
             A tuple of (batch_next_token_ids, logprobs). logprobs is None
             when return_logprob is False or SGLANG_RETURN_ORIGINAL_LOGPROB is set.
         """
-        logits.div_(sampling_info.temperatures)
         batch_next_token_ids = self._sample_from_logits(
             logits, sampling_info, simple_sampling_case, positions
         )
diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py
index d3c884fbf871..81abdb78102b 100644
--- a/python/sglang/srt/model_executor/forward_batch_info.py
+++ b/python/sglang/srt/model_executor/forward_batch_info.py
@@ -52,6 +52,7 @@
     ForwardBatchDeepSeekMHAMixin,
 )
 from sglang.srt.runtime_context import get_parallel, get_server_args
+from sglang.srt.true_on_policy import is_true_on_policy_enabled
 from sglang.srt.utils import (
     is_cuda,
     is_hip,
@@ -1096,7 +1097,7 @@ def _compute_mrope_positions(self, model_runner: ModelRunner, batch: ScheduleBat
             mm_input = batch.multimodal_inputs[batch_idx]
             if self.forward_mode.is_decode():
                 # 3 * N
-                if mm_input is None or rl_on_policy_target is not None:
+                if mm_input is None or is_true_on_policy_enabled():
                     mrope_positions_list[batch_idx] = torch.full(
                         (3, 1),
                         self.seq_lens_cpu[batch_idx] - 1,
@@ -1112,7 +1113,7 @@ def _compute_mrope_positions(self, model_runner: ModelRunner, batch: ScheduleBat
                     batch.extend_lens[batch_idx],
                     batch.prefix_lens[batch_idx],
                 )
-                if mm_input is None or rl_on_policy_target is not None:
+                if mm_input is None or is_true_on_policy_enabled():
                     # text only
                     mrope_positions = torch.tensor(
                         [
diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py
index 547494b48eb2..12988d3188ff 100644
--- a/python/sglang/srt/model_executor/model_runner.py
+++ b/python/sglang/srt/model_executor/model_runner.py
@@ -183,6 +183,7 @@
     get_global_experts_capturer,
     set_global_experts_capturer,
 )
+from sglang.srt.true_on_policy import is_tp_invariant_target
 from sglang.srt.utils import (
     cpu_has_amx_support,
     enable_show_time_cost,
@@ -695,6 +696,10 @@ def maybe_enable_batch_invariant_mode(self):
             from sglang.srt.batch_invariant_ops import enable_batch_invariant_mode
 
             enable_batch_invariant_mode()
+        if is_tp_invariant_target():
+            from sglang.srt.tp_invariant_ops import enable_tp_invariant_mode
+
+            enable_tp_invariant_mode()
 
     def get_pp_proxy_topk_size(self) -> Optional[int]:
         return misc_utils.resolve_pp_proxy_topk_size(
diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
index ec1275999273..bc3ca55c73d4 100644
--- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
+++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
@@ -93,6 +93,9 @@
 from sglang.srt.multiplex.pdmux_context import get_current_stream_idx, get_stream_groups
 from sglang.srt.runtime_context import get_flags, get_parallel
 from sglang.srt.speculative.ragged_verify import resolve_ragged_verify_layout
+from sglang.srt.true_on_policy import (
+    patch_prefill_only_deterministic_inference_for_cuda_graph,
+)
 from sglang.srt.utils import (
     empty_context,
     get_available_gpu_memory,
@@ -849,22 +852,32 @@ def capture(self) -> None:
         # Trigger CUDA graph capture for specific shapes.
         # Capture the large shapes first so that the smaller shapes
         # can reuse the memory pool allocated for the large shapes.
-        with freeze_gc(self.model_runner.server_args.enable_cudagraph_gc):
-            if not self.enable_pdmux:
-                with graph_capture() as graph_capture_context, profile_context as prof:
-                    self.stream = graph_capture_context.stream
-                    with self.backend.capture_session(self.stream):
-                        self._capture_one_stream()
-            else:
-                set_pdmux_status(False)
-                for i, sg in enumerate(self.stream_groups):
+        with patch_prefill_only_deterministic_inference_for_cuda_graph(
+            self.model_runner.server_args,
+            attn_backend=getattr(self.model_runner, "attn_backend", None),
+            dvr_target_verify_cuda_graph=getattr(
+                self.model_runner, "enable_dvr_target_verify_cuda_graph", False
+            ),
+        ):
+            with freeze_gc(self.model_runner.server_args.enable_cudagraph_gc):
+                if not self.enable_pdmux:
                     with (
-                        graph_capture(stream=sg[1]) as graph_capture_context,
+                        graph_capture() as graph_capture_context,
                         profile_context as prof,
                     ):
                         self.stream = graph_capture_context.stream
                         with self.backend.capture_session(self.stream):
-                            self._capture_one_stream(i)
+                            self._capture_one_stream()
+                else:
+                    set_pdmux_status(False)
+                    for i, sg in enumerate(self.stream_groups):
+                        with (
+                            graph_capture(stream=sg[1]) as graph_capture_context,
+                            profile_context as prof,
+                        ):
+                            self.stream = graph_capture_context.stream
+                            with self.backend.capture_session(self.stream):
+                                self._capture_one_stream(i)
 
         if self.enable_profile_cuda_graph:
             self._post_process_after_profile(prof)
diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py
index 98ae54aeb371..c79fa483106b 100644
--- a/python/sglang/srt/models/qwen2.py
+++ b/python/sglang/srt/models/qwen2.py
@@ -50,7 +50,12 @@
     kv_cache_scales_loader,
 )
 from sglang.srt.platforms import current_platform
-from sglang.srt.runtime_context import get_parallel, get_server_args
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.true_on_policy import (
+    get_on_policy_rms_norm_kwargs,
+    is_true_on_policy_enabled,
+    should_force_bfloat16_dense_tensor_math,
+)
 from sglang.srt.utils import add_prefix, make_layers
 from sglang.srt.utils.hf_transformers_utils import get_rope_config
 
@@ -96,9 +101,11 @@ def forward(
         x: torch.Tensor,
         forward_batch: ForwardBatch = None,
     ) -> torch.Tensor:
-        if get_server_args().rl_on_policy_target is not None:
-            x = x.bfloat16()
-
+        if (
+            should_force_bfloat16_dense_tensor_math()
+            or x.dtype != self.gate_up_proj.weight.dtype
+        ):
+            x = x.to(self.gate_up_proj.weight.dtype)
         gate_up, _ = self.gate_up_proj(x)
         x = self.act_fn(gate_up)
         x, _ = self.down_proj(x, forward_batch=forward_batch)
@@ -289,11 +296,7 @@ def __init__(
                 quant_config=quant_config,
                 use_attn_tp_group=is_dp_attention_enabled(),
                 prefix=add_prefix("embed_tokens", prefix),
-                params_dtype=(
-                    torch.float32
-                    if get_server_args().rl_on_policy_target is not None
-                    else None
-                ),
+                params_dtype=(torch.float32 if is_true_on_policy_enabled() else None),
             )
         else:
             self.embed_tokens = PPMissingLayer()
@@ -320,16 +323,7 @@ def __init__(
             prefix=add_prefix("layers", prefix),
         )
         if self.pp_group.is_last_rank:
-            norm_kwargs = (
-                dict(
-                    weight_dtype=torch.float32,
-                    cast_x_before_out_mul=True,
-                    override_orig_dtype=torch.float32,
-                    fp32_residual=True,
-                )
-                if get_server_args().rl_on_policy_target is not None
-                else {}
-            )
+            norm_kwargs = get_on_policy_rms_norm_kwargs()
             self.norm = RMSNorm(
                 config.hidden_size, eps=config.rms_norm_eps, **norm_kwargs
             )
diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py
index 72e8ddee474c..67fbd324e74a 100644
--- a/python/sglang/srt/models/qwen2_moe.py
+++ b/python/sglang/srt/models/qwen2_moe.py
@@ -92,6 +92,7 @@
 from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.runtime_context import get_forward, get_parallel, get_server_args
+from sglang.srt.true_on_policy import get_on_policy_rms_norm_kwargs
 from sglang.srt.utils import (
     add_prefix,
     cpu_has_amx_support,
@@ -866,7 +867,10 @@ def __init__(
             prefix=add_prefix("layers", prefix),
         )
         if self.pp_group.is_last_rank:
-            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+            norm_kwargs = get_on_policy_rms_norm_kwargs(fp32_residual=False)
+            self.norm = RMSNorm(
+                config.hidden_size, eps=config.rms_norm_eps, **norm_kwargs
+            )
         else:
             self.norm = PPMissingLayer(return_tuple=True)
 
diff --git a/python/sglang/srt/models/qwen3.py b/python/sglang/srt/models/qwen3.py
index daaca2276434..c9db1b8c39d9 100644
--- a/python/sglang/srt/models/qwen3.py
+++ b/python/sglang/srt/models/qwen3.py
@@ -34,6 +34,10 @@
 from sglang.srt.models.qwen2 import Qwen2Model
 from sglang.srt.models.utils import apply_qk_norm
 from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream
+from sglang.srt.true_on_policy import (
+    should_disable_fused_qk_norm_mrope,
+    should_force_bfloat16_dense_tensor_math,
+)
 from sglang.srt.utils import add_prefix, get_bool_env_var, is_cuda, is_hip, is_npu
 
 Qwen3Config = None
@@ -106,16 +110,16 @@ def __init__(
         self.max_position_embeddings = max_position_embeddings
         self.tp_rank = get_parallel().tp_rank
 
-        norm_kwargs = (
-            dict(
-                weight_dtype=torch.float32,
-                cast_x_before_out_mul=True,
-            )
-            if get_server_args().rl_on_policy_target is not None
-            else {}
+        self.q_norm = RMSNorm(
+            self.head_dim,
+            eps=rms_norm_eps,
+            true_on_policy_weight_dtype=torch.float32,
+        )
+        self.k_norm = RMSNorm(
+            self.head_dim,
+            eps=rms_norm_eps,
+            true_on_policy_weight_dtype=torch.float32,
         )
-        self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps, **norm_kwargs)
-        self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps, **norm_kwargs)
 
         self.qkv_proj = QKVParallelLinear(
             hidden_size,
@@ -271,14 +275,20 @@ def forward(
         hidden_states: torch.Tensor,
         forward_batch: ForwardBatch,
     ) -> torch.Tensor:
-        if get_server_args().rl_on_policy_target is not None:
-            hidden_states = hidden_states.bfloat16()
+        if (
+            should_force_bfloat16_dense_tensor_math()
+            or hidden_states.dtype != self.qkv_proj.weight.dtype
+        ):
+            # True-on-policy RMSNorm can produce fp32 activations while dense
+            # projections remain bf16, including during cuda-graph capture when
+            # the global on-policy flag is temporarily cleared.
+            hidden_states = hidden_states.to(self.qkv_proj.weight.dtype)
 
         save_kv_cache = True
         use_aiter_fused = (
             self.use_fused_qk_norm_mrope
             and forward_batch.forward_mode.is_decode()
-            and get_server_args().rl_on_policy_target is None
+            and not should_disable_fused_qk_norm_mrope()
         )
 
         if use_aiter_fused:
@@ -298,9 +308,9 @@ def forward(
                 forward_batch=forward_batch,
             )
 
-        if get_server_args().rl_on_policy_target is not None:
-            q = q.to(torch.bfloat16)
-            k = k.to(torch.bfloat16)
+        if should_force_bfloat16_dense_tensor_math() or q.dtype != v.dtype:
+            q = q.to(v.dtype)
+            k = k.to(v.dtype)
 
         attn_output = self.attn(q, k, v, forward_batch, save_kv_cache=save_kv_cache)
         output, _ = self.o_proj(attn_output)
@@ -355,21 +365,19 @@ def __init__(
             prefix=add_prefix("mlp", prefix),
         )
 
-        norm_kwargs = (
-            dict(
-                weight_dtype=torch.float32,
-                cast_x_before_out_mul=True,
-                override_orig_dtype=torch.float32,
-                fp32_residual=True,
-            )
-            if get_server_args().rl_on_policy_target is not None
-            else {}
-        )
         self.input_layernorm = RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps, **norm_kwargs
+            config.hidden_size,
+            eps=config.rms_norm_eps,
+            true_on_policy_weight_dtype=torch.float32,
+            true_on_policy_override_orig_dtype=torch.float32,
+            true_on_policy_fp32_residual=True,
         )
         self.post_attention_layernorm = RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps, **norm_kwargs
+            config.hidden_size,
+            eps=config.rms_norm_eps,
+            true_on_policy_weight_dtype=torch.float32,
+            true_on_policy_override_orig_dtype=torch.float32,
+            true_on_policy_fp32_residual=True,
         )
 
         self.layer_scatter_modes = LayerScatterModes.init_new(
diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py
index 9cee1d21d8d7..c277c11ffa6e 100644
--- a/python/sglang/srt/models/qwen3_moe.py
+++ b/python/sglang/srt/models/qwen3_moe.py
@@ -22,6 +22,7 @@
 from typing import Any, Dict, Iterable, List, Optional, Tuple, TypeVar
 
 import torch
+import torch.nn.functional as F
 from torch import nn
 from transformers import PretrainedConfig
 
@@ -48,7 +49,7 @@
 )
 from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class
 from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
-from sglang.srt.layers.moe.topk import TopK
+from sglang.srt.layers.moe.topk import StandardTopKOutput, TopK
 from sglang.srt.layers.moe.utils import (
     RoutingMethodType,
     filter_moe_weight_param_global_expert,
@@ -78,6 +79,11 @@
     get_server_args,
     get_stream,
 )
+from sglang.srt.true_on_policy import (
+    get_on_policy_rms_norm_kwargs,
+    is_true_on_policy_enabled,
+    should_disable_fused_qk_norm_mrope,
+)
 from sglang.srt.utils import (
     LazyValue,
     add_prefix,
@@ -320,7 +326,20 @@ def forward_normal(
 
         # router_logits: (num_tokens, n_experts)
         router_logits, _ = self.gate(hidden_states)
-        topk_output = self.topk(hidden_states, router_logits)
+        if is_true_on_policy_enabled():
+            routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
+            routing_weights, selected_experts = torch.topk(
+                routing_weights, self.top_k, dim=-1
+            )
+            routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+            routing_weights = routing_weights.to(hidden_states.dtype)
+            topk_output = StandardTopKOutput(
+                topk_weights=routing_weights,
+                topk_ids=selected_experts,
+                router_logits=router_logits,
+            )
+        else:
+            topk_output = self.topk(hidden_states, router_logits)
         final_hidden_states = self.experts(hidden_states, topk_output)
 
         if self.ep_size > 1 and not should_skip_post_experts_all_reduce(
@@ -508,7 +527,7 @@ def __init__(
         )
         self.compatible_with_fused_kv_buffer = (
             False if isinstance(self.rotary_emb, MRotaryEmbedding) else True
-        )
+        ) and not is_true_on_policy_enabled()
         self.compatible_with_fused_qk_norm_rope = not isinstance(
             self.rotary_emb, MRotaryEmbedding
         ) and self.head_dim in (64, 128, 256)
@@ -523,6 +542,7 @@ def __init__(
                 torch.bfloat16,
                 _yarn_factor != 1.0,
             )
+            and not should_disable_fused_qk_norm_mrope()
         )
         self._used_fused_qk_norm_rope_last_call = False
 
@@ -535,8 +555,9 @@ def __init__(
             prefix=add_prefix("attn", prefix),
         )
 
-        self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
-        self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        norm_kwargs = get_on_policy_rms_norm_kwargs(fp32_residual=False)
+        self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps, **norm_kwargs)
+        self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps, **norm_kwargs)
         self.alt_stream = alt_stream
 
     def op_prepare(self, state):
@@ -779,9 +800,12 @@ def __init__(
                 quant_config=quant_config,
                 prefix=add_prefix("mlp", prefix),
             )
-        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        norm_kwargs = get_on_policy_rms_norm_kwargs(fp32_residual=False)
+        self.input_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps, **norm_kwargs
+        )
         self.post_attention_layernorm = RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps
+            config.hidden_size, eps=config.rms_norm_eps, **norm_kwargs
         )
 
         self.layer_communicator = LayerCommunicator(
diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py
index 7036cd717a1c..e75eeb7d82c0 100644
--- a/python/sglang/srt/models/qwen3_vl.py
+++ b/python/sglang/srt/models/qwen3_vl.py
@@ -455,75 +455,73 @@ def rot_pos_emb(
 
         return cos_combined, sin_combined
 
-    def _get_interpolation_indices(self, dim_size: int) -> torch.Tensor:
-        """
-        Compute continuous interpolation indices for a single dimension.
+    def fast_pos_embed_interpolate(self, grid_thw):
+        grid_ts, grid_hs, grid_ws = grid_thw[:, 0], grid_thw[:, 1], grid_thw[:, 2]
+        num_grid_per_side = int(self.num_position_embeddings**0.5)
+        device = self.pos_embed.weight.device
 
-        Returns continuous indices.
-        """
-        if self.align_corners:
-            indices = np.linspace(
-                0, self.num_grid_per_side - 1, dim_size, dtype=np.float32
-            )
-        else:
-            indices = (np.arange(dim_size, dtype=np.float32) + 0.5) * (
-                self.num_grid_per_side / dim_size
-            ) - 0.5
-            indices = np.clip(indices, 0, self.num_grid_per_side - 1)
-        return indices
+        idx_list = [[] for _ in range(4)]
+        weight_list = [[] for _ in range(4)]
 
-    def _calculate_indices_and_weights(self, h_idxs, w_idxs):
-        """
-        Compute bilinear interpolation indices and weights.
+        for t, h, w in zip(grid_ts, grid_hs, grid_ws):
+            h_idxs = torch.linspace(0, num_grid_per_side - 1, h)
+            w_idxs = torch.linspace(0, num_grid_per_side - 1, w)
 
-        Returns tuple of (indices, weights), each as 4 numpy arrays for the 4 corner points.
-        """
-        h_f = np.floor(h_idxs).astype(np.int64)
-        h_c = np.clip(h_f + 1, 0, self.num_grid_per_side - 1)
-        dh = h_idxs - h_f
+            h_idxs_floor = h_idxs.int()
+            w_idxs_floor = w_idxs.int()
+            h_idxs_ceil = (h_idxs.int() + 1).clip(max=num_grid_per_side - 1)
+            w_idxs_ceil = (w_idxs.int() + 1).clip(max=num_grid_per_side - 1)
 
-        w_f = np.floor(w_idxs).astype(np.int64)
-        w_c = np.clip(w_f + 1, 0, self.num_grid_per_side - 1)
-        dw = w_idxs - w_f
+            dh = h_idxs - h_idxs_floor
+            dw = w_idxs - w_idxs_floor
 
-        side = self.num_grid_per_side
+            base_h = h_idxs_floor * num_grid_per_side
+            base_h_ceil = h_idxs_ceil * num_grid_per_side
 
-        indices = [
-            (h_f[:, None] * side + w_f).flatten(),
-            (h_f[:, None] * side + w_c).flatten(),
-            (h_c[:, None] * side + w_f).flatten(),
-            (h_c[:, None] * side + w_c).flatten(),
-        ]
-        weights = [
-            ((1 - dh)[:, None] * (1 - dw)).flatten(),
-            ((1 - dh)[:, None] * dw).flatten(),
-            (dh[:, None] * (1 - dw)).flatten(),
-            (dh[:, None] * dw).flatten(),
-        ]
-        return indices, weights
+            indices = [
+                (base_h[None].T + w_idxs_floor[None]).flatten(),
+                (base_h[None].T + w_idxs_ceil[None]).flatten(),
+                (base_h_ceil[None].T + w_idxs_floor[None]).flatten(),
+                (base_h_ceil[None].T + w_idxs_ceil[None]).flatten(),
+            ]
 
-    def _get_position_embedding(self, patch_pos_embeds, grid_ts, grid_hs, grid_ws):
-        """
-        Tile and reorganize position embeddings to align with the token sequence.
-        """
-        result_parts = []
-        merge_size = self.spatial_merge_size
+            weights = [
+                ((1 - dh)[None].T * (1 - dw)[None]).flatten(),
+                ((1 - dh)[None].T * dw[None]).flatten(),
+                (dh[None].T * (1 - dw)[None]).flatten(),
+                (dh[None].T * dw[None]).flatten(),
+            ]
 
-        for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws):
-            pos_embed = pos_embed.repeat(t, 1)
+            for i in range(4):
+                idx_list[i].extend(indices[i].tolist())
+                weight_list[i].extend(weights[i].tolist())
 
-            h_merge = h // merge_size
-            w_merge = w // merge_size
+        idx_tensor = torch.tensor(idx_list, dtype=torch.long, device=device)
+        weight_tensor = torch.tensor(
+            weight_list, dtype=self.pos_embed.weight.dtype, device=device
+        )
+        pos_embeds = self.pos_embed(idx_tensor).to(device) * weight_tensor[:, :, None]
+        patch_pos_embeds = pos_embeds[0] + pos_embeds[1] + pos_embeds[2] + pos_embeds[3]
+
+        patch_pos_embeds = patch_pos_embeds.split(
+            [h * w for h, w in zip(grid_hs, grid_ws)]
+        )
 
+        patch_pos_embeds_permute = []
+        merge_size = self.spatial_merge_size
+        for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws):
+            pos_embed = pos_embed.repeat(t, 1)
             pos_embed = (
-                pos_embed.view(t, h_merge, merge_size, w_merge, merge_size, -1)
+                pos_embed.view(
+                    t, h // merge_size, merge_size, w // merge_size, merge_size, -1
+                )
                 .permute(0, 1, 3, 2, 4, 5)
                 .flatten(0, 4)
             )
 
-            result_parts.append(pos_embed)
+            patch_pos_embeds_permute.append(pos_embed)
 
-        return torch.cat(result_parts, dim=0)
+        return torch.cat(patch_pos_embeds_permute)
 
     def _torch_interp_indices(
         self, dim_size: int, device: torch.device
@@ -756,61 +754,6 @@ def bucket_flashinfer_max_seqlen(self, real_max_seqlen: int) -> int:
             round_up(real_max_seqlen, FLASHINFER_MAX_SEQLEN_BUCKETS[-1]),
         )
 
-    def fast_pos_embed_interpolate(self, grid_thw):
-        """Interpolate position embeddings for (batch, 3) size input dimensions.
-
-        Performs bilinear interpolation on spatial dimensions (height, width) and replicates
-        along temporal dimension. The result is reorganized according to spatial_merge_size.
-
-        Args:
-            grid_thw: Tensor of shape [batch_size, 3] with (temporal, height, width) dimensions
-                     in patches for each sample.
-
-        Returns:
-            Interpolated position embeddings tensor.
-        """
-        grid_thw_cpu = grid_thw.cpu().numpy()
-
-        # transfer data to CPU before loop
-        temporal_dims = grid_thw_cpu[:, 0].tolist()
-        height_dims = grid_thw_cpu[:, 1].tolist()
-        width_dims = grid_thw_cpu[:, 2].tolist()
-
-        device = self.pos_embed.weight.device
-        dtype = self.pos_embed.weight.dtype
-
-        patches_size = [h * w for h, w in zip(height_dims, width_dims)]
-        total_patches = sum(patches_size)
-        all_indices_np = np.zeros((4, total_patches), dtype=np.int64)
-        all_weights_np = np.zeros((4, total_patches), dtype=np.float32)
-
-        current_idx = 0
-
-        # calculate indices and weights on CPU
-        for t, h, w in zip(temporal_dims, height_dims, width_dims):
-            h_idxs = self._get_interpolation_indices(h)
-            w_idxs = self._get_interpolation_indices(w)
-
-            indices, weights = self._calculate_indices_and_weights(h_idxs, w_idxs)
-
-            end_idx = current_idx + h * w
-            for i in range(4):
-                all_indices_np[i, current_idx:end_idx] = indices[i]
-                all_weights_np[i, current_idx:end_idx] = weights[i]
-            current_idx = end_idx
-
-        idx_tensor = torch.from_numpy(all_indices_np).to(device)
-        weight_tensor = torch.from_numpy(all_weights_np).to(dtype=dtype, device=device)
-
-        # calculate interpolation
-        pos_embeds = self.pos_embed(idx_tensor.view(-1))
-        pos_embeds = pos_embeds.view(4, total_patches, -1)
-        patch_pos_embeds = (pos_embeds * weight_tensor.unsqueeze(-1)).sum(dim=0)
-        patch_pos_embeds = patch_pos_embeds.split(patches_size)
-        return self._get_position_embedding(
-            patch_pos_embeds, temporal_dims, height_dims, width_dims
-        )
-
     def compute_flashinfer_batch_offsets_packed(
         self,
         token_cu_seqlens: np.ndarray,
@@ -1163,14 +1106,19 @@ def forward(
                     hidden_states + residual if residual is not None else hidden_states
                 )
 
+            deepstack_embeds = None
+            if input_deepstack_embeds is not None:
+                prev_layer_idx = layer_idx - 1
+                if prev_layer_idx in self.deepstack_embed_to_decoder_layer:
+                    sep = self.hidden_size * prev_layer_idx
+                    deepstack_embeds = input_deepstack_embeds[
+                        :, sep : sep + self.hidden_size
+                    ]
+
             # SGLang applies residual at the START of the next layer, not at the END like HuggingFace.
             # See: https://github.com/huggingface/transformers/blob/v5.0.0rc0/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py#L549
             # To match HF behavior, deepstack must be added AFTER residual: (hidden_states + residual) + deepstack
             # The order matters because addition with different tensors is not associative in practice.
-            # Deepstack for prev_layer is applied at the start of current layer via post_residual_addition.
-            deepstack_embeds = self.get_deepstack_embeds(
-                layer_idx - 1, input_deepstack_embeds
-            )
             hidden_states, residual = layer(
                 positions,
                 hidden_states,
diff --git a/python/sglang/srt/models/sdar.py b/python/sglang/srt/models/sdar.py
index 91946db54b4e..bedb500895cf 100644
--- a/python/sglang/srt/models/sdar.py
+++ b/python/sglang/srt/models/sdar.py
@@ -47,6 +47,10 @@
     get_server_args,
     get_stream,
 )
+from sglang.srt.true_on_policy import (
+    get_on_policy_rms_norm_kwargs,
+    should_force_bfloat16_dense_tensor_math,
+)
 from sglang.srt.utils import add_prefix, is_cuda, make_layers
 
 logger = logging.getLogger(__name__)
@@ -207,7 +211,7 @@ def forward(
         hidden_states: torch.Tensor,
         forward_batch: ForwardBatch,
     ):
-        if get_server_args().rl_on_policy_target is not None:
+        if should_force_bfloat16_dense_tensor_math():
             hidden_states = hidden_states.bfloat16()
 
         qkv, _ = self.qkv_proj(hidden_states)
@@ -235,7 +239,7 @@ def forward(
             ),
         )
 
-        if get_server_args().rl_on_policy_target is not None:
+        if should_force_bfloat16_dense_tensor_math():
             q = q.to(torch.bfloat16)
             k = k.to(torch.bfloat16)
 
@@ -263,15 +267,10 @@ def __init__(
         self.hidden_size = config.hidden_size
         self.layer_id = layer_id
 
-        norm_kwargs = (
-            dict(
-                weight_dtype=torch.float32,
-                cast_x_before_out_mul=True,
-                override_orig_dtype=torch.float32,
-                fp32_residual=True,
-            )
-            if get_server_args().rl_on_policy_target is not None
-            else {}
+        norm_kwargs = get_on_policy_rms_norm_kwargs(
+            weight_dtype=torch.float32,
+            override_orig_dtype=torch.float32,
+            fp32_residual=True,
         )
         self.input_layernorm = RMSNorm(
             self.hidden_size, eps=config.rms_norm_eps, **norm_kwargs
@@ -388,15 +387,10 @@ def __init__(
             prefix=add_prefix("layers", prefix),
         )
         if self.pp_group.is_last_rank:
-            norm_kwargs = (
-                dict(
-                    weight_dtype=torch.float32,
-                    cast_x_before_out_mul=True,
-                    override_orig_dtype=torch.float32,
-                    fp32_residual=True,
-                )
-                if get_server_args().rl_on_policy_target is not None
-                else {}
+            norm_kwargs = get_on_policy_rms_norm_kwargs(
+                weight_dtype=torch.float32,
+                override_orig_dtype=torch.float32,
+                fp32_residual=True,
             )
             self.norm = RMSNorm(self.embed_dim, eps=config.rms_norm_eps, **norm_kwargs)
         else:
diff --git a/python/sglang/srt/models/sdar_moe.py b/python/sglang/srt/models/sdar_moe.py
index b91db7be3359..ff411f6684b7 100644
--- a/python/sglang/srt/models/sdar_moe.py
+++ b/python/sglang/srt/models/sdar_moe.py
@@ -63,6 +63,10 @@
     get_server_args,
     get_stream,
 )
+from sglang.srt.true_on_policy import (
+    get_on_policy_rms_norm_kwargs,
+    should_force_bfloat16_dense_tensor_math,
+)
 from sglang.srt.utils import LazyValue, add_prefix, is_cuda, make_layers
 
 logger = logging.getLogger(__name__)
@@ -274,7 +278,7 @@ def forward(
         hidden_states: torch.Tensor,
         forward_batch: ForwardBatch,
     ) -> torch.Tensor:
-        if get_server_args().rl_on_policy_target is not None:
+        if should_force_bfloat16_dense_tensor_math():
             hidden_states = hidden_states.bfloat16()
 
         qkv, _ = self.qkv_proj(hidden_states)
@@ -302,7 +306,7 @@ def forward(
             ),
         )
 
-        if get_server_args().rl_on_policy_target is not None:
+        if should_force_bfloat16_dense_tensor_math():
             q = q.to(torch.bfloat16)
             k = k.to(torch.bfloat16)
 
@@ -331,15 +335,10 @@ def __init__(
         self.hidden_size = config.hidden_size
         self.layer_id = layer_id
 
-        norm_kwargs = (
-            dict(
-                weight_dtype=torch.float32,
-                cast_x_before_out_mul=True,
-                override_orig_dtype=torch.float32,
-                fp32_residual=True,
-            )
-            if get_server_args().rl_on_policy_target is not None
-            else {}
+        norm_kwargs = get_on_policy_rms_norm_kwargs(
+            weight_dtype=torch.float32,
+            override_orig_dtype=torch.float32,
+            fp32_residual=True,
         )
         self.input_layernorm = RMSNorm(
             self.hidden_size, eps=config.rms_norm_eps, **norm_kwargs
@@ -471,15 +470,10 @@ def __init__(
         )
 
         if self.pp_group.is_last_rank:
-            norm_kwargs = (
-                dict(
-                    weight_dtype=torch.float32,
-                    cast_x_before_out_mul=True,
-                    override_orig_dtype=torch.float32,
-                    fp32_residual=True,
-                )
-                if get_server_args().rl_on_policy_target is not None
-                else {}
+            norm_kwargs = get_on_policy_rms_norm_kwargs(
+                weight_dtype=torch.float32,
+                override_orig_dtype=torch.float32,
+                fp32_residual=True,
             )
             self.norm = RMSNorm(self.embed_dim, eps=config.rms_norm_eps, **norm_kwargs)
         else:
diff --git a/python/sglang/srt/models/step3p5.py b/python/sglang/srt/models/step3p5.py
index 117b795bd0fa..f0c271f63f00 100644
--- a/python/sglang/srt/models/step3p5.py
+++ b/python/sglang/srt/models/step3p5.py
@@ -52,6 +52,7 @@
     get_server_args,
     get_stream,
 )
+from sglang.srt.true_on_policy import is_true_on_policy_enabled
 from sglang.srt.utils import add_prefix, is_cuda, is_non_idle_and_non_empty, make_layers
 
 Step3p5Config = None
@@ -674,11 +675,7 @@ def __init__(
                 quant_config=quant_config,
                 enable_tp=not is_dp_attention_enabled(),
                 prefix=add_prefix("embed_tokens", prefix),
-                params_dtype=(
-                    torch.float32
-                    if get_server_args().rl_on_policy_target is not None
-                    else None
-                ),
+                params_dtype=torch.float32 if is_true_on_policy_enabled() else None,
             )
         else:
             self.embed_tokens = PPMissingLayer()
diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py
index 38b16e514888..b780e9a123e9 100644
--- a/python/sglang/srt/models/utils.py
+++ b/python/sglang/srt/models/utils.py
@@ -484,6 +484,8 @@ def apply_qk_norm(
         _is_cuda  # TODO(dark): have not tested on ROCm or other backends
         and allow_inplace  # TODO(dark): this can be relaxed if needed
         and (q_eps == k_eps)  # TODO(dark): this can also be relaxed
+        and q_norm.weight.dtype == q.dtype
+        and k_norm.weight.dtype == k.dtype
         and not envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get()
         and get_server_args().cuda_graph_config.prefill.tc_compiler
         != "inductor"  # let inductor fuse QK norm
diff --git a/python/sglang/srt/multimodal/processors/base_processor.py b/python/sglang/srt/multimodal/processors/base_processor.py
index ed82bb72429d..db5d782cac12 100644
--- a/python/sglang/srt/multimodal/processors/base_processor.py
+++ b/python/sglang/srt/multimodal/processors/base_processor.py
@@ -20,6 +20,7 @@
     MultimodalProcessorOutput,
 )
 from sglang.srt.runtime_context import get_server_args
+from sglang.srt.true_on_policy import is_true_on_policy_enabled
 from sglang.srt.utils import (
     envs,
     is_cpu,
@@ -459,7 +460,7 @@ def process_mm_data(
             and isinstance(processor.image_processor, BaseImageProcessor)
             and not self.disable_fast_image_processor
         ):
-            if _is_cpu or get_server_args().rl_on_policy_target is not None:
+            if _is_cpu or is_true_on_policy_enabled():
                 kwargs["device"] = "cpu"
             elif _is_xpu:
                 kwargs["device"] = "xpu"
diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index 838b925a34db..228cbb8136ed 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -60,6 +60,10 @@
 from sglang.srt.parser.reasoning_parser import ReasoningParser
 from sglang.srt.platforms import current_platform
 from sglang.srt.speculative.decoupled_spec_io import DecoupledSpecIpcConfig
+from sglang.srt.true_on_policy.contracts import (
+    resolve_true_on_policy_runtime_policy,
+    validate_true_on_policy_contract,
+)
 from sglang.srt.utils.common import (
     LORA_TARGET_ALL_MODULES,
     SUPPORTED_LORA_TARGET_MODULES,
@@ -310,7 +314,7 @@ def add_chunked_prefix_cache_attention_backend(backend_name):
 RADIX_EVICTION_POLICY_CHOICES = ["lru", "lfu", "slru", "priority"]
 RETRACTION_POLICY_CHOICES = ["length", "priority"]
 
-RL_ON_POLICY_TARGET_CHOICES = ["fsdp"]
+RL_ON_POLICY_TARGET_CHOICES = ["fsdp", "fsdp_tp"]
 
 LORA_BACKEND_CHOICES = ["triton", "csgmv", "ascend", "torch_native"]
 
@@ -2755,6 +2759,15 @@ class ServerArgs:
         bool,
         "Enable deterministic inference mode with batch invariant ops.",
     ] = False
+    enable_prefill_only_deterministic_inference: A[
+        bool,
+        "Enable prefill-only deterministic inference mode with batch invariant ops.",
+    ] = False
+    true_on_policy_contract: A[
+        Optional[str],
+        "Internal true-on-policy parity contract selected by the launcher. "
+        "Normal users should prefer the Miles true_on_policy switch.",
+    ] = None
     rl_on_policy_target: A[
         Optional[str],
         Arg(
@@ -6731,16 +6744,40 @@ def _handle_cache_compatibility(self):
             raise ValueError("--swa-full-tokens-ratio should be in range (0, 1.0].")
 
     def _handle_deterministic_inference(self):
+        validate_true_on_policy_contract(self)
+
+        if self.enable_prefill_only_deterministic_inference:
+            self.enable_deterministic_inference = True
+
         if self.rl_on_policy_target is not None:
             logger.warning(
-                "Enable deterministic inference because of rl_on_policy_target."
+                "Enable deterministic inference because of legacy rl_on_policy_target."
+            )
+            self.enable_deterministic_inference = True
+
+        if self.true_on_policy_contract is not None:
+            logger.warning(
+                "Enable deterministic inference because of true_on_policy_contract."
             )
             self.enable_deterministic_inference = True
 
             # For VLM
             envs.SGLANG_VLM_CACHE_SIZE_MB.set(0)
-            # TODO remove this environment variable as a whole
-            envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.set(True)
+
+            if (
+                resolve_true_on_policy_runtime_policy(
+                    self
+                ).disable_flashinfer_allreduce_fusion
+                and self.enable_flashinfer_allreduce_fusion
+            ):
+                self.enable_flashinfer_allreduce_fusion = False
+                logger.warning(
+                    "Disable flashinfer allreduce fusion because of "
+                    "true_on_policy_contract with TP rollout."
+                )
+
+        if self.enable_deterministic_inference:
+            envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.set("1")
 
         if self.enable_deterministic_inference:
             if self.enable_aiter_allreduce_fusion:
diff --git a/python/sglang/srt/tp_invariant_ops/__init__.py b/python/sglang/srt/tp_invariant_ops/__init__.py
new file mode 100644
index 000000000000..2a3f7a07f4a7
--- /dev/null
+++ b/python/sglang/srt/tp_invariant_ops/__init__.py
@@ -0,0 +1,23 @@
+from .tp_invariant_ops import (
+    disable_tp_invariant_mode,
+    enable_tp_invariant_mode,
+    is_tp_invariant_mode_enabled,
+    matmul_tp_inv,
+    matmul_tp_persistent,
+    moe_sum_tree_reduce,
+    set_tp_invariant_mode,
+    tree_all_reduce_sum,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "matmul_tp_persistent",
+    "matmul_tp_inv",
+    "moe_sum_tree_reduce",
+    "tree_all_reduce_sum",
+    "set_tp_invariant_mode",
+    "is_tp_invariant_mode_enabled",
+    "disable_tp_invariant_mode",
+    "enable_tp_invariant_mode",
+]
diff --git a/python/sglang/srt/tp_invariant_ops/tp_invariant_ops.py b/python/sglang/srt/tp_invariant_ops/tp_invariant_ops.py
new file mode 100644
index 000000000000..b465ff4a1977
--- /dev/null
+++ b/python/sglang/srt/tp_invariant_ops/tp_invariant_ops.py
@@ -0,0 +1,1941 @@
+import contextlib
+import math
+import os
+import sys
+from typing import Any, Callable, Dict
+
+import torch
+import torch.distributed as dist
+import triton
+import triton.language as tl
+
+# Triton's constexpr tree unrolling causes deep AST recursion in the JIT
+# compiler.  The two-level tree (v2) bounds compilation depth to
+# max(log2(SUBTREE), log2(E/SUBTREE)) instead of log2(E), but we still
+# need headroom for the per-level AST visitor overhead.
+if sys.getrecursionlimit() < 16384:
+    sys.setrecursionlimit(16384)
+
+
+def _matmul_launch_metadata(
+    grid: Callable[..., Any], kernel: Any, args: Dict[str, Any]
+) -> Dict[str, Any]:
+    ret = {}
+    m, n, k = args["M"], args["N"], args["K"]
+    ret["name"] = f"{kernel.name} [M={m}, N={n}, K={k}]"
+    if "tiles_per_update" in args:
+        ret["name"] = (
+            f"{kernel.name} [M={m}, N={n}, K={k}, tiles_per_update={args['tiles_per_update']:02}]"
+        )
+    if "c_ptr" in args:
+        bytes_per_elem = args["c_ptr"].element_size()
+    else:
+        bytes_per_elem = 1 if args["FP8_OUTPUT"] else 2
+    ret[f"flops{bytes_per_elem * 8}"] = 2.0 * m * n * k
+    ret["bytes"] = bytes_per_elem * (m * k + n * k + m * n)
+    return ret
+
+
+@triton.jit
+def _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS):
+    group_id = tile_id // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + (tile_id % group_size_m)
+    pid_n = (tile_id % num_pid_in_group) // group_size_m
+    return pid_m, pid_n
+
+
+def _get_tl_dtype(dtype):
+    if dtype == torch.float32:
+        return tl.float32
+    elif dtype == torch.float16:
+        return tl.float16
+    elif dtype == torch.bfloat16:
+        return tl.bfloat16
+
+
+# ---- kernel ----
+@triton.jit(launch_metadata=_matmul_launch_metadata)
+def matmul_kernel_tp_persistent(
+    A_ptr,
+    B_ptr,
+    C_ptr,
+    M,
+    N,
+    K,
+    stride_am,
+    stride_ak,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    NUM_SMS: tl.constexpr,
+    LEVEL_K: tl.constexpr,
+    TILE_K: tl.constexpr,
+    FIRST_LEVEL_BLOCK: tl.constexpr,
+    NEXT_POWER_OF_LEVEL: tl.constexpr,
+    NEXT_POWER_OF_REMAIN_LEVEL: tl.constexpr,
+    ACC_DTYPE: tl.constexpr,
+    OUT_DTYPE: tl.constexpr,
+    A_LARGE: tl.constexpr,
+    B_LARGE: tl.constexpr,
+    C_LARGE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_M)
+    num_pid_n = tl.cdiv(N, BLOCK_N)
+    num_tiles = num_pid_m * num_pid_n
+
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+
+    manual_acc = 3
+    acc1 = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE)
+    acc2 = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE)
+    acc3 = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE)
+
+    S = tl.zeros((NEXT_POWER_OF_REMAIN_LEVEL, BLOCK_M, BLOCK_N), dtype=ACC_DTYPE)
+
+    S_mask = tl.arange(0, NEXT_POWER_OF_REMAIN_LEVEL)[:, None, None]
+    level_ids = tl.arange(0, NEXT_POWER_OF_LEVEL)
+
+    base_offs_m = tl.arange(0, BLOCK_M)
+    base_offs_n = tl.arange(0, BLOCK_N)
+    offs_k = tl.arange(0, BLOCK_K)
+    offs_k = tl.max_contiguous(tl.multiple_of(offs_k, BLOCK_K), BLOCK_K)
+
+    for tile_id in tl.range(pid, num_tiles, NUM_SMS, flatten=False):
+        pid_m, pid_n = _compute_pid(
+            tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS
+        )
+
+        start_m = pid_m * BLOCK_M
+        start_n = pid_n * BLOCK_N
+
+        offs_am = start_m + base_offs_m
+        if A_LARGE:
+            offs_am = offs_am.to(tl.int64)
+        offs_am = tl.where(offs_am < M, offs_am, 0)
+
+        offs_bn = start_n + base_offs_n
+        if B_LARGE:
+            offs_bn = offs_bn.to(tl.int64)
+        offs_bn = tl.where(offs_bn < N, offs_bn, 0)
+
+        offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_M), BLOCK_M)
+        offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_N), BLOCK_N)
+
+        a_ptrs = A_ptr + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak
+        b_ptrs = B_ptr + offs_bn[None, :] * stride_bn + offs_k[:, None] * stride_bk
+
+        count = tl.zeros((NEXT_POWER_OF_LEVEL,), dtype=tl.int32)
+
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE)
+        for s_tile_idx in range(0, TILE_K):
+            k0 = s_tile_idx * BLOCK_K
+            a = tl.load(
+                a_ptrs,
+                mask=(offs_am[:, None] < M) & ((k0 + offs_k)[None, :] < K),
+                other=0.0,
+            )
+            b = tl.load(
+                b_ptrs,
+                mask=((k0 + offs_k)[:, None] < K) & (offs_bn[None, :] < N),
+                other=0.0,
+            )
+            a_ptrs += BLOCK_K * stride_ak
+            b_ptrs += BLOCK_K * stride_bk
+
+            acc = tl.dot(a, b).to(ACC_DTYPE)
+
+            break_flag = 0
+            for level in range(LEVEL_K):
+                if break_flag == 0:
+                    idx_mask = level_ids == level
+
+                    count_value_added = tl.sum(count * idx_mask) + 1
+
+                    table_value = FIRST_LEVEL_BLOCK if level == 0 else 2
+
+                    carry_over = (table_value == count_value_added).to(tl.int1)
+
+                    if count_value_added > 1:
+                        if level == 0:
+                            acc = acc1 + acc
+                        elif level == 1:
+                            acc = acc2 + acc
+                        elif level == 2:
+                            acc = acc3 + acc
+                        else:
+                            tmp_acc_mask = S_mask == (level - manual_acc)
+                            acc = (
+                                tl.sum(S * tmp_acc_mask, axis=0, dtype=ACC_DTYPE) + acc
+                            )
+
+                    count = tl.where(
+                        idx_mask, count_value_added * (1 - carry_over), count
+                    )
+                    if not carry_over:
+                        break_flag = 1
+                        if level == 0:
+                            acc1 = acc
+                        elif level == 1:
+                            acc2 = acc
+                        elif level == 2:
+                            acc3 = acc
+                        else:
+                            tmp_acc_mask = S_mask == (level - manual_acc)
+                            S = tl.where(tmp_acc_mask, acc[None, :, :], S)
+
+        c_ptr = C_ptr + (offs_am[:, None] * stride_cm + offs_bn[None, :] * stride_cn)
+        offs_cm = start_m + base_offs_m
+        offs_cn = start_n + base_offs_n
+        if C_LARGE:
+            offs_cm = offs_cm.to(tl.int64)
+            offs_cn = offs_cn.to(tl.int64)
+        offs_cm = tl.where(offs_cm < M, offs_cm, 0)
+        offs_cn = tl.where(offs_cn < N, offs_cn, 0)
+        offs_cm = tl.max_contiguous(tl.multiple_of(offs_cm, BLOCK_M), BLOCK_M)
+        offs_cn = tl.max_contiguous(tl.multiple_of(offs_cn, BLOCK_N), BLOCK_N)
+        mask_c = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
+        tl.store(c_ptr, acc.to(OUT_DTYPE), mask=mask_c)
+
+
+@triton.jit(launch_metadata=_matmul_launch_metadata)
+def matmul_kernel_tp_persistent_optim(
+    A_ptr,
+    B_ptr,
+    C_ptr,
+    M,
+    N,
+    K,
+    stride_am,
+    stride_ak,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    NUM_SMS: tl.constexpr,
+    LEVEL_K: tl.constexpr,
+    TILE_K: tl.constexpr,
+    FIRST_LEVEL_BLOCK: tl.constexpr,
+    NEXT_POWER_OF_LEVEL: tl.constexpr,
+    NEXT_POWER_OF_REMAIN_LEVEL: tl.constexpr,
+    ACC_DTYPE: tl.constexpr,
+    OUT_DTYPE: tl.constexpr,
+    A_LARGE: tl.constexpr,
+    B_LARGE: tl.constexpr,
+    C_LARGE: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_M)
+    num_pid_n = tl.cdiv(N, BLOCK_N)
+    num_tiles = num_pid_m * num_pid_n
+
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+
+    # Flat loop (same structure as original) with inlined level-0 handling.
+    # Level-0 uses scalar counter; tree merge (levels 1+) only on carry.
+    # Accumulators are conditionally allocated based on LEVEL_K (constexpr)
+    # to minimize register pressure and maximize occupancy.
+    acc1 = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE)
+    acc2 = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE)
+    acc3 = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE)
+    acc4 = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE)
+    acc5 = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE)
+
+    S = tl.zeros((NEXT_POWER_OF_REMAIN_LEVEL, BLOCK_M, BLOCK_N), dtype=ACC_DTYPE)
+    S_mask = tl.arange(0, NEXT_POWER_OF_REMAIN_LEVEL)[:, None, None]
+    level_ids = tl.arange(0, NEXT_POWER_OF_LEVEL)
+
+    offs_k = tl.arange(0, BLOCK_K)
+    offs_k = tl.max_contiguous(tl.multiple_of(offs_k, BLOCK_K), BLOCK_K)
+
+    for tile_id in tl.range(pid, num_tiles, NUM_SMS, flatten=False):
+        pid_m, pid_n = _compute_pid(
+            tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS
+        )
+
+        start_m = pid_m * BLOCK_M
+        start_n = pid_n * BLOCK_N
+
+        offs_am = start_m + tl.arange(0, BLOCK_M)
+        mask_m = offs_am < M
+        if A_LARGE:
+            offs_am = offs_am.to(tl.int64)
+        offs_am = tl.where(mask_m, offs_am, 0)
+
+        offs_bn = start_n + tl.arange(0, BLOCK_N)
+        mask_n = offs_bn < N
+        if B_LARGE:
+            offs_bn = offs_bn.to(tl.int64)
+        offs_bn = tl.where(mask_n, offs_bn, 0)
+
+        offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_M), BLOCK_M)
+        offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_N), BLOCK_N)
+        mask_m_bc = mask_m[:, None]
+        mask_n_bc = mask_n[None, :]
+
+        a_ptrs = A_ptr + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak
+        b_ptrs = B_ptr + offs_bn[None, :] * stride_bn + offs_k[:, None] * stride_bk
+
+        c0 = 0
+        c1 = 0
+        c2 = 0
+        c3 = 0
+        c4 = 0
+        count = tl.zeros((NEXT_POWER_OF_LEVEL,), dtype=tl.int32)
+
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_DTYPE)
+        for _ in range(0, TILE_K):
+            a = tl.load(a_ptrs, mask=mask_m_bc, other=0.0)
+            b = tl.load(b_ptrs, mask=mask_n_bc, other=0.0)
+            a_ptrs += BLOCK_K * stride_ak
+            b_ptrs += BLOCK_K * stride_bk
+
+            acc = tl.dot(a, b).to(ACC_DTYPE)
+
+            c0 += 1
+            if c0 > 1:
+                acc = acc1 + acc
+            if c0 < FIRST_LEVEL_BLOCK:
+                acc1 = acc
+            else:
+                c0 = 0
+                break_flag = 0
+                for level in range(1, LEVEL_K):
+                    if break_flag == 0:
+                        if level == 1:
+                            c1 += 1
+                            if c1 > 1:
+                                acc = acc2 + acc
+                            if c1 == 2:
+                                c1 = 0
+                            else:
+                                acc2 = acc
+                                break_flag = 1
+                        elif level == 2:
+                            c2 += 1
+                            if c2 > 1:
+                                acc = acc3 + acc
+                            if c2 == 2:
+                                c2 = 0
+                            else:
+                                acc3 = acc
+                                break_flag = 1
+                        elif level == 3:
+                            c3 += 1
+                            if c3 > 1:
+                                acc = acc4 + acc
+                            if c3 == 2:
+                                c3 = 0
+                            else:
+                                acc4 = acc
+                                break_flag = 1
+                        elif level == 4:
+                            c4 += 1
+                            if c4 > 1:
+                                acc = acc5 + acc
+                            if c4 == 2:
+                                c4 = 0
+                            else:
+                                acc5 = acc
+                                break_flag = 1
+                        else:
+                            idx_mask = level_ids == level
+                            count_value_added = tl.sum(count * idx_mask) + 1
+                            carry_over = (2 == count_value_added).to(tl.int1)
+                            if count_value_added > 1:
+                                tmp_acc_mask = S_mask == (level - 5)
+                                acc = (
+                                    tl.sum(S * tmp_acc_mask, axis=0, dtype=ACC_DTYPE)
+                                    + acc
+                                )
+                            count = tl.where(
+                                idx_mask, count_value_added * (1 - carry_over), count
+                            )
+                            if not carry_over:
+                                break_flag = 1
+                                tmp_acc_mask = S_mask == (level - 5)
+                                S = tl.where(tmp_acc_mask, acc[None, :, :], S)
+
+        offs_cm = start_m + tl.arange(0, BLOCK_M)
+        offs_cn = start_n + tl.arange(0, BLOCK_N)
+        if C_LARGE:
+            offs_cm = offs_cm.to(tl.int64)
+            offs_cn = offs_cn.to(tl.int64)
+        offs_cm = tl.where(mask_m, offs_cm, 0)
+        offs_cn = tl.where(mask_n, offs_cn, 0)
+        offs_cm = tl.max_contiguous(tl.multiple_of(offs_cm, BLOCK_M), BLOCK_M)
+        offs_cn = tl.max_contiguous(tl.multiple_of(offs_cn, BLOCK_N), BLOCK_N)
+        c_ptr = C_ptr + (offs_cm[:, None] * stride_cm + offs_cn[None, :] * stride_cn)
+        mask_c = mask_m_bc & mask_n_bc
+        tl.store(c_ptr, acc.to(OUT_DTYPE), mask=mask_c)
+
+
+def _matmul_tp_persistent_impl(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    bias: torch.Tensor = None,
+    fp32_accum: bool = False,
+    use_optim_kernel: bool = False,
+):
+    assert A.shape[-1] == B.shape[-2], "Dim doesn't match"
+
+    out_dtype = A.dtype
+    acc_dtype = torch.float32 if fp32_accum else A.dtype
+
+    NUM_SMS = torch.cuda.get_device_properties(A.device).multi_processor_count
+
+    # 1D launch kernel where each block gets its own program.
+    def grid(META):
+        return (
+            min(
+                NUM_SMS,
+                triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),
+            ),
+        )
+
+    base_configs = {
+        torch.bfloat16: {
+            "BLOCK_SIZE_M": 64,
+            "BLOCK_SIZE_N": 128,
+            "BLOCK_SIZE_K": 128,
+            "GROUP_SIZE_M": 8,
+            "num_stages": 2,
+            "num_warps": 8,
+        },
+        torch.float16: {
+            "BLOCK_SIZE_M": 64,
+            "BLOCK_SIZE_N": 128,
+            "BLOCK_SIZE_K": 128,
+            "GROUP_SIZE_M": 8,
+            "num_stages": 2,
+            "num_warps": 8,
+        },
+        torch.float32: {
+            "BLOCK_SIZE_M": 32,
+            "BLOCK_SIZE_N": 64,
+            "BLOCK_SIZE_K": 128,
+            "GROUP_SIZE_M": 8,
+            "num_stages": 2,
+            "num_warps": 8,
+        },
+    }
+
+    configs = base_configs
+    BLOCK_M = configs[out_dtype]["BLOCK_SIZE_M"]
+    BLOCK_N = configs[out_dtype]["BLOCK_SIZE_N"]
+    BLOCK_K = configs[out_dtype]["BLOCK_SIZE_K"]
+    GROUP_SIZE_M = configs[out_dtype]["GROUP_SIZE_M"]
+    num_stages = configs[out_dtype]["num_stages"]
+    num_warps = configs[out_dtype]["num_warps"]
+
+    M, K = A.shape
+    _, N = B.shape
+    assert (
+        K % BLOCK_K == 0
+    ), f"Dimension K should be divisible by BLOCK_K. Got K={K}, BLOCK_K={BLOCK_K}."
+    T = K // BLOCK_K
+    FIRST_LEVEL_BLOCK = T
+
+    if use_optim_kernel:
+        num_n_tiles = triton.cdiv(N, BLOCK_N)
+        total_tiles = triton.cdiv(M, BLOCK_M) * num_n_tiles
+
+        if total_tiles * 4 <= NUM_SMS:
+            while BLOCK_M > 16 and triton.cdiv(M, BLOCK_M) * num_n_tiles < NUM_SMS:
+                BLOCK_M //= 2
+        elif total_tiles * 2 <= NUM_SMS:
+            while BLOCK_M > 32 and triton.cdiv(M, BLOCK_M) * num_n_tiles < NUM_SMS:
+                BLOCK_M //= 2
+
+        if out_dtype in (torch.bfloat16, torch.float16):
+            num_warps = 4 if BLOCK_M <= 16 else 8
+            num_stages = 3 if K >= 1024 else 2
+
+    LEVEL_K = 1
+    while FIRST_LEVEL_BLOCK > 2 and FIRST_LEVEL_BLOCK % 2 == 0:
+        FIRST_LEVEL_BLOCK //= 2
+        LEVEL_K += 1
+
+    C = torch.empty((M, N), device=A.device, dtype=out_dtype)
+
+    # Original kernel manually handles levels 0-2 (3 accumulators);
+    # optim kernel handles levels 0-4 (5 register accumulators),
+    # so S tensor is smaller / unused for common LEVEL_K <= 5.
+    manual_acc = 5 if use_optim_kernel else 3
+
+    NEXT_POWER_OF_LEVEL = 2 ** math.ceil(math.log2(LEVEL_K))
+    NEXT_POWER_OF_REMAIN_LEVEL = (
+        2 ** math.ceil(math.log2(LEVEL_K - manual_acc)) if LEVEL_K > manual_acc else 1
+    )
+
+    kernel = (
+        matmul_kernel_tp_persistent_optim
+        if use_optim_kernel
+        else matmul_kernel_tp_persistent
+    )
+    kernel[grid](
+        A,
+        B,
+        C,
+        M,
+        N,
+        K,
+        *A.stride(),
+        *B.stride(),
+        *C.stride(),
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        BLOCK_K=BLOCK_K,
+        GROUP_SIZE_M=GROUP_SIZE_M,
+        NUM_SMS=NUM_SMS,
+        NEXT_POWER_OF_REMAIN_LEVEL=NEXT_POWER_OF_REMAIN_LEVEL,
+        LEVEL_K=LEVEL_K,
+        TILE_K=T,
+        FIRST_LEVEL_BLOCK=FIRST_LEVEL_BLOCK,
+        NEXT_POWER_OF_LEVEL=NEXT_POWER_OF_LEVEL,
+        ACC_DTYPE=_get_tl_dtype(acc_dtype),
+        OUT_DTYPE=_get_tl_dtype(out_dtype),
+        A_LARGE=A.numel() > 2**31,
+        B_LARGE=B.numel() > 2**31,
+        C_LARGE=C.numel() > 2**31,
+        num_warps=num_warps,
+        num_stages=num_stages,
+    )
+    if bias is not None:
+        C += bias
+    return C
+
+
+def matmul_tp_persistent(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    bias: torch.Tensor = None,
+    fp32_accum: bool = False,
+):
+    return _matmul_tp_persistent_impl(
+        A=A,
+        B=B,
+        bias=bias,
+        fp32_accum=fp32_accum,
+        use_optim_kernel=False,
+    )
+
+
+def matmul_tp_inv(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    bias: torch.Tensor = None,
+    fp32_accum: bool = False,
+):
+    return matmul_tp_persistent(A, B, bias=bias, fp32_accum=fp32_accum)
+
+
+def matmul_tp_persistent_optim(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    bias: torch.Tensor = None,
+    fp32_accum: bool = False,
+):
+    return _matmul_tp_persistent_impl(
+        A=A,
+        B=B,
+        bias=bias,
+        fp32_accum=fp32_accum,
+        use_optim_kernel=True,
+    )
+
+
+def tree_all_reduce_sum(x: torch.Tensor, device_group=None) -> torch.Tensor:
+    rank = dist.get_rank(device_group)
+    world_size = dist.get_world_size(device_group)
+
+    if world_size & (world_size - 1) != 0:
+        raise ValueError(
+            "world_size must be a power of 2 in order to use all_reduce_sum."
+        )
+
+    result = [torch.zeros_like(x) for _ in range(world_size)]
+    dist.all_gather(result, x, group=device_group)
+
+    for level in range(1, world_size.bit_length()):
+        for left in range(0, world_size, 1 << level):
+            right = left + (1 << (level - 1))
+            result[left] += result[right]
+
+    return result[0]
+
+
+def tree_all_reduce_sum_optim(x: torch.Tensor, device_group=None) -> torch.Tensor:
+    if not x.is_cuda:
+        raise ValueError("x must be a CUDA tensor.")
+    if not x.is_contiguous():
+        raise ValueError(
+            "x must be contiguous. Call x = x.contiguous() OUTSIDE graph capture."
+        )
+
+    world_size = dist.get_world_size(device_group)
+    if world_size & (world_size - 1) != 0:
+        raise ValueError("world_size must be a power of 2.")
+
+    # cache
+    if not hasattr(tree_all_reduce_sum, "_cache"):
+        tree_all_reduce_sum._cache = {}
+
+    key = (id(device_group), x.device.index, tuple(x.shape), x.dtype, world_size)
+    st = tree_all_reduce_sum._cache.get(key)
+    if st is None:
+        gather = torch.empty(
+            (world_size,) + tuple(x.shape), device=x.device, dtype=x.dtype
+        )
+        out = torch.empty_like(x)
+        st = tree_all_reduce_sum._cache[key] = (gather, out)
+
+    gather, out = st
+
+    # 1) all_gather into one contiguous buffer
+    dist.all_gather_into_tensor(gather, x, group=device_group)
+
+    # 2) deterministic tree pairing EXACTLY like your original:
+    # for level in range(1, bit_length):
+    #   for left in range(0, world_size, 1<<level):
+    #       right = left + (1<<(level-1))
+    #       gather[left] += gather[right]
+    #
+    # vectorized form: left indices are 0:step: , right are half:step:
+    levels = world_size.bit_length()
+    for level in range(1, levels):
+        step = 1 << level
+        half = step >> 1
+        # Views only (no alloc); one add_ kernel per level
+        gather[0:world_size:step].add_(gather[half:world_size:step])
+
+    out.copy_(gather[0])
+    tree_all_reduce_sum._cache.clear()
+    return out
+
+
+_tp_inv_MODE = False
+
+try:
+    def_lib = torch.library.Library("tp_inv_ops", "DEF")
+    def_lib.define("matmul_tp_inv(Tensor a, Tensor b, Tensor? bias=None) -> Tensor")
+except RuntimeError:
+    pass
+
+try:
+    impl = torch.library.Library("tp_inv_ops", "IMPL")
+
+    impl.impl("matmul_tp_inv", matmul_tp_persistent, "CUDA")
+except RuntimeError:
+    pass
+
+
+def is_tp_invariant_mode_enabled():
+    return _tp_inv_MODE
+
+
+def enable_tp_invariant_mode():
+    global _tp_inv_MODE
+
+    if _tp_inv_MODE:
+        return
+
+    _tp_inv_MODE = True
+
+
+def disable_tp_invariant_mode():
+    global _tp_inv_MODE
+
+    _tp_inv_MODE = False
+
+
+@contextlib.contextmanager
+def set_tp_invariant_mode(enabled=True):
+    global _tp_inv_MODE
+
+    old_state = _tp_inv_MODE
+
+    if enabled:
+        enable_tp_invariant_mode()
+    else:
+        disable_tp_invariant_mode()
+
+    try:
+        yield
+    finally:
+        _tp_inv_MODE = old_state
+
+
+def scatter_input_by_local_expert(
+    topk: torch.Tensor, input: torch.Tensor, E: int
+) -> torch.Tensor:
+    """
+    Args:
+        topk: [M, topk], long, -1 means remote expert
+        input: [M, topk, hidden_size], float
+        E: int, number of local experts (expert ids in [0, E))
+    Returns:
+        output: [M, E, hidden_size]
+    """
+    M, _, hidden_size = input.shape
+
+    # Mask out remote experts in output
+    valid = (topk != -1).unsqueeze(-1)  # [M, topk, 1]
+    output_masked = input * valid.to(input.dtype)  # [M, topk, hidden_size]
+
+    # Replace -1 with 0 for safe indexing (value doesn't matter because output is zero)
+    topk_index = topk.clamp(min=0)  # turns -1 into 0, leaves others unchanged
+
+    # Expand index to match output
+    index = topk_index.unsqueeze(-1).expand(
+        -1, -1, hidden_size
+    )  # [M, topk, hidden_size]
+
+    # Initialize result
+    output = torch.zeros(M, E, hidden_size, device=input.device, dtype=input.dtype)
+
+    # Scatter add
+    output.scatter_add_(1, index, output_masked)
+
+    return output
+
+
+@triton.jit
+def _load_expert_tile(
+    input_ptr,
+    input_base,
+    input_stride_1,
+    topk_ids_base,
+    topk_ids_stride_1,
+    zero_ptrs,
+    offs_dim,
+    mask,
+    mask_token,
+    e: tl.constexpr,
+    TOPK: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+):
+    """Load tile for expert e: search topk_ids for slot, redirect to zero_buf if remote."""
+    found_slot = tl.full([BLOCK_M], -1, dtype=tl.int32)
+    for k in range(TOPK):
+        kid = tl.load(
+            topk_ids_base + k * topk_ids_stride_1,
+            mask=mask_token,
+            other=-1,
+        ).to(tl.int32)
+        match = (kid == e) & (found_slot == -1)
+        found_slot = tl.where(match, k, found_slot)
+
+    is_valid = found_slot != -1
+    slot_safe = tl.maximum(found_slot, 0)
+    input_ptrs = input_base + slot_safe[:, None] * input_stride_1 + offs_dim[None, :]
+    load_ptrs = tl.where(is_valid[:, None], input_ptrs, zero_ptrs)
+    return tl.load(load_ptrs, mask=mask, other=0.0).to(tl.float32)
+
+
+@triton.jit
+def _tree_reduce_pair(
+    input_ptr,
+    input_base,
+    input_stride_1,
+    topk_ids_base,
+    topk_ids_stride_1,
+    zero_ptrs,
+    offs_dim,
+    mask,
+    mask_token,
+    start: tl.constexpr,
+    size: tl.constexpr,
+    TOPK: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+):
+    """
+    Recursively compute binary-tree sum over experts [start, start+size).
+    Returns a fp32 tile [BLOCK_M, BLOCK_DIM].
+
+    Tree structure (e.g. size=4, start=0):
+      _tree_reduce_pair(0, 4)
+        = _tree_reduce_pair(0, 2) + _tree_reduce_pair(2, 2)
+        = (_tree_reduce_pair(0,1) + _tree_reduce_pair(1,1))
+          + (_tree_reduce_pair(2,1) + _tree_reduce_pair(3,1))
+        = (load(e0) + load(e1)) + (load(e2) + load(e3))
+
+    Since size is constexpr and always a power of 2, Triton fully unrolls this
+    into a fixed sequence of loads and adds with no dynamic branching.
+    Max register depth = log2(E) tiles, e.g. E=64 -> 6 tiles.
+    """
+    if size == 1:
+        return _load_expert_tile(
+            input_ptr,
+            input_base,
+            input_stride_1,
+            topk_ids_base,
+            topk_ids_stride_1,
+            zero_ptrs,
+            offs_dim,
+            mask,
+            mask_token,
+            start,
+            TOPK,
+            BLOCK_M,
+        )
+    else:
+        half: tl.constexpr = size // 2
+        left = _tree_reduce_pair(
+            input_ptr,
+            input_base,
+            input_stride_1,
+            topk_ids_base,
+            topk_ids_stride_1,
+            zero_ptrs,
+            offs_dim,
+            mask,
+            mask_token,
+            start,
+            half,
+            TOPK,
+            BLOCK_M,
+        )
+        right = _tree_reduce_pair(
+            input_ptr,
+            input_base,
+            input_stride_1,
+            topk_ids_base,
+            topk_ids_stride_1,
+            zero_ptrs,
+            offs_dim,
+            mask,
+            mask_token,
+            start + half,
+            half,
+            TOPK,
+            BLOCK_M,
+        )
+        return left + right
+
+
+@triton.jit
+def _fused_tree_reduce_kernel(
+    # input: [M, topk, hidden_dim], contiguous
+    input_ptr,
+    input_stride_0,
+    input_stride_1,
+    # topk_ids: [M, topk], -1 means remote
+    topk_ids_ptr,
+    topk_ids_stride_0,
+    topk_ids_stride_1,
+    # zero_buf: [hidden_dim], all zeros
+    zero_buf_ptr,
+    # output: [M, hidden_dim]
+    output_ptr,
+    output_stride_0,
+    # scalars
+    token_num,
+    hidden_dim,
+    routed_scaling_factor,
+    # constexpr
+    E: tl.constexpr,
+    TOPK: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_DIM: tl.constexpr,
+):
+    """
+    Fused scatter + binary-tree reduce in a single kernel.
+    Zero extra memory: all intermediate results live in registers.
+
+    For each token block, loads expert tiles on-the-fly (with zero-buf redirect
+    for remote experts), and reduces them in binary-tree order using recursive
+    constexpr unrolling. The tree structure is:
+      result = tree_sum(0, E)
+      tree_sum(s, n) = tree_sum(s, n/2) + tree_sum(s+n/2, n/2)   if n > 1
+      tree_sum(s, 1) = load_expert(s)                              base case
+
+    Register pressure: log2(E) tiles of [BLOCK_M, BLOCK_DIM] fp32.
+    E.g. E=64, BLOCK_M=1, BLOCK_DIM=2048 -> 6 * 8KB = 48KB, well within limits.
+    """
+    input_stride_0 = tl.cast(input_stride_0, dtype=tl.int64)
+    input_stride_1 = tl.cast(input_stride_1, dtype=tl.int64)
+    topk_ids_stride_0 = tl.cast(topk_ids_stride_0, dtype=tl.int64)
+    topk_ids_stride_1 = tl.cast(topk_ids_stride_1, dtype=tl.int64)
+    output_stride_0 = tl.cast(output_stride_0, dtype=tl.int64)
+
+    token_block_id = tl.program_id(0)
+    dim_block_id = tl.program_id(1)
+
+    offs_token = token_block_id * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_dim = dim_block_id * BLOCK_DIM + tl.arange(0, BLOCK_DIM)
+
+    mask_token = offs_token < token_num
+    mask_dim = offs_dim < hidden_dim
+    mask = mask_token[:, None] & mask_dim[None, :]
+
+    zero_ptrs = zero_buf_ptr + offs_dim[None, :]
+    input_base = input_ptr + offs_token[:, None] * input_stride_0
+    topk_ids_base = topk_ids_ptr + offs_token * topk_ids_stride_0
+
+    # Binary tree reduce over all E experts, entirely in registers.
+    result = _tree_reduce_pair(
+        input_ptr,
+        input_base,
+        input_stride_1,
+        topk_ids_base,
+        topk_ids_stride_1,
+        zero_ptrs,
+        offs_dim,
+        mask,
+        mask_token,
+        0,
+        E,
+        TOPK,
+        BLOCK_M,
+    )
+
+    result *= routed_scaling_factor
+
+    store_ptrs = output_ptr + offs_token[:, None] * output_stride_0 + offs_dim[None, :]
+    tl.store(store_ptrs, result.to(input_ptr.dtype.element_ty), mask=mask)
+
+
+# Persistent zero buffer: only [hidden_dim] elements, negligible memory.
+# Allocated once on first call, never freed.
+_zero_buf_cache: torch.Tensor | None = None
+
+
+def moe_sum_tree_reduce_v1(
+    input: torch.Tensor,  # [M, topk, hidden_dim]
+    output: torch.Tensor,  # [M, hidden_dim]
+    curr_topk_ids: torch.Tensor,  # [M, topk], -1 means remote
+    routed_scaling_factor: float,
+    E: int,
+):
+    """
+    Fused MoE tree reduce: zero extra memory, CUDA Graph safe.
+
+    Single kernel: loads expert tiles on-the-fly with zero-buf pointer redirect
+    for remote experts (L1 cache hit), reduces in binary-tree order entirely
+    in registers. No scratch buffer needed.
+
+    Invariant guarantee: binary tree reduce order is fixed by expert id,
+    identical across all EP ranks regardless of which experts are local.
+
+    Memory overhead: only a single [hidden_dim] zero buffer (~14KB for H=7168 bf16).
+    """
+    assert input.is_contiguous()
+    assert output.is_contiguous()
+
+    token_num, topk, hidden_dim = input.shape
+    assert output.shape[0] == token_num and output.shape[1] == hidden_dim
+    assert (E & (E - 1)) == 0, f"E must be power of 2, got {E}"
+
+    # Fast path for the dominant K=8 case: avoids generic expert-tree loads/scans.
+    if topk == 8:
+        _moe_sum_tree_reduce_k8_fast_path(
+            input=input,
+            output=output,
+            curr_topk_ids=curr_topk_ids,
+            routed_scaling_factor=routed_scaling_factor,
+            E=E,
+        )
+        return
+
+    # Zero buffer: [hidden_dim], allocated once, reused forever
+    global _zero_buf_cache
+    if (
+        _zero_buf_cache is None
+        or _zero_buf_cache.device != input.device
+        or _zero_buf_cache.dtype != input.dtype
+        or _zero_buf_cache.numel() < hidden_dim
+    ):
+        _zero_buf_cache = torch.zeros(
+            hidden_dim, device=input.device, dtype=input.dtype
+        )
+    zero_buf = _zero_buf_cache
+
+    BLOCK_M = 1
+    BLOCK_DIM = 2048
+    num_warps = 16
+
+    grid = (
+        triton.cdiv(token_num, BLOCK_M),
+        triton.cdiv(hidden_dim, BLOCK_DIM),
+    )
+
+    _fused_tree_reduce_kernel[grid](
+        input,
+        input.stride(0),
+        input.stride(1),
+        curr_topk_ids,
+        curr_topk_ids.stride(0),
+        curr_topk_ids.stride(1),
+        zero_buf,
+        output,
+        output.stride(0),
+        token_num=token_num,
+        hidden_dim=hidden_dim,
+        routed_scaling_factor=routed_scaling_factor,
+        E=E,
+        TOPK=topk,
+        BLOCK_M=BLOCK_M,
+        BLOCK_DIM=BLOCK_DIM,
+        num_warps=num_warps,
+    )
+    return
+
+
+import math
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _moe_sum_tree_reduce_k8_fused_kernel_opt2d(
+    x_ptr,
+    ids_ptr,
+    out_ptr,
+    sx_m: tl.constexpr,
+    sx_k: tl.constexpr,
+    sx_h: tl.constexpr,  # x: [M,8,H]
+    sid_m: tl.constexpr,
+    sid_k: tl.constexpr,  # ids: [M,8]
+    so_m: tl.constexpr,
+    so_h: tl.constexpr,  # out: [M,H]
+    M,
+    H,  # runtime OK
+    E_LEVEL,  # runtime int (log2(E))
+    routed_scaling_factor,  # runtime scalar
+    BLOCK_M: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_h = tl.program_id(1)
+
+    m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
+
+    mask_m = m < M
+    mask_h = h < H
+    mask_mh = mask_m[:, None] & mask_h[None, :]
+
+    # ---- load ids (int32), -1 means remote ----
+    ids0 = tl.load(ids_ptr + m * sid_m + 0 * sid_k, mask=mask_m, other=-1).to(tl.int32)
+    ids1 = tl.load(ids_ptr + m * sid_m + 1 * sid_k, mask=mask_m, other=-1).to(tl.int32)
+    ids2 = tl.load(ids_ptr + m * sid_m + 2 * sid_k, mask=mask_m, other=-1).to(tl.int32)
+    ids3 = tl.load(ids_ptr + m * sid_m + 3 * sid_k, mask=mask_m, other=-1).to(tl.int32)
+    ids4 = tl.load(ids_ptr + m * sid_m + 4 * sid_k, mask=mask_m, other=-1).to(tl.int32)
+    ids5 = tl.load(ids_ptr + m * sid_m + 5 * sid_k, mask=mask_m, other=-1).to(tl.int32)
+    ids6 = tl.load(ids_ptr + m * sid_m + 6 * sid_k, mask=mask_m, other=-1).to(tl.int32)
+    ids7 = tl.load(ids_ptr + m * sid_m + 7 * sid_k, mask=mask_m, other=-1).to(tl.int32)
+
+    # h comes from tl.arange and can be annotated for alignment.
+    tl.multiple_of(h, 8)
+    tl.max_contiguous(h, BLOCK_H)
+
+    # m is also a tensor.
+    tl.multiple_of(m, BLOCK_M)
+
+    # ---- load values: remote handled by mask->other=0 (NO zero_tile / NO tl.where big tiles) ----
+    m0 = mask_mh & (ids0 != -1)[:, None]
+    m1 = mask_mh & (ids1 != -1)[:, None]
+    m2 = mask_mh & (ids2 != -1)[:, None]
+    m3 = mask_mh & (ids3 != -1)[:, None]
+    m4 = mask_mh & (ids4 != -1)[:, None]
+    m5 = mask_mh & (ids5 != -1)[:, None]
+    m6 = mask_mh & (ids6 != -1)[:, None]
+    m7 = mask_mh & (ids7 != -1)[:, None]
+
+    v0 = tl.load(
+        x_ptr + m[:, None] * sx_m + 0 * sx_k + h[None, :] * sx_h, mask=m0, other=0.0
+    )
+    v1 = tl.load(
+        x_ptr + m[:, None] * sx_m + 1 * sx_k + h[None, :] * sx_h, mask=m1, other=0.0
+    )
+    v2 = tl.load(
+        x_ptr + m[:, None] * sx_m + 2 * sx_k + h[None, :] * sx_h, mask=m2, other=0.0
+    )
+    v3 = tl.load(
+        x_ptr + m[:, None] * sx_m + 3 * sx_k + h[None, :] * sx_h, mask=m3, other=0.0
+    )
+    v4 = tl.load(
+        x_ptr + m[:, None] * sx_m + 4 * sx_k + h[None, :] * sx_h, mask=m4, other=0.0
+    )
+    v5 = tl.load(
+        x_ptr + m[:, None] * sx_m + 5 * sx_k + h[None, :] * sx_h, mask=m5, other=0.0
+    )
+    v6 = tl.load(
+        x_ptr + m[:, None] * sx_m + 6 * sx_k + h[None, :] * sx_h, mask=m6, other=0.0
+    )
+    v7 = tl.load(
+        x_ptr + m[:, None] * sx_m + 7 * sx_k + h[None, :] * sx_h, mask=m7, other=0.0
+    )
+
+    x_dtype = x_ptr.dtype.element_ty
+
+    # ---- deterministic dense-tree-equivalent reduce (same order concept as baseline) ----
+    # Remote entries have already been masked to 0.0 at load time (m0..m7).
+
+    for bit in tl.range(0, E_LEVEL):
+        bitmask = 1 << bit
+
+        # ========== lane0 as source ==========
+        cond = (ids0 != -1) & ((ids0 & bitmask) != 0)
+        target = ids0 ^ bitmask
+        mm1 = cond & (ids1 == target)
+        mm2 = cond & (ids2 == target)
+        mm3 = cond & (ids3 == target)
+        mm4 = cond & (ids4 == target)
+        mm5 = cond & (ids5 == target)
+        mm6 = cond & (ids6 == target)
+        mm7 = cond & (ids7 == target)
+        hit = mm1 | mm2 | mm3 | mm4 | mm5 | mm6 | mm7
+
+        src = v0.to(tl.float32)
+        v1 = tl.where(mm1[:, None], (v1.to(tl.float32) + src).to(x_dtype), v1)
+        v2 = tl.where(mm2[:, None], (v2.to(tl.float32) + src).to(x_dtype), v2)
+        v3 = tl.where(mm3[:, None], (v3.to(tl.float32) + src).to(x_dtype), v3)
+        v4 = tl.where(mm4[:, None], (v4.to(tl.float32) + src).to(x_dtype), v4)
+        v5 = tl.where(mm5[:, None], (v5.to(tl.float32) + src).to(x_dtype), v5)
+        v6 = tl.where(mm6[:, None], (v6.to(tl.float32) + src).to(x_dtype), v6)
+        v7 = tl.where(mm7[:, None], (v7.to(tl.float32) + src).to(x_dtype), v7)
+
+        v0 = tl.where(hit[:, None], 0.0, v0)
+        ids0 = tl.where(hit, -1, ids0)
+        ids0 = tl.where(cond & (~hit), target, ids0)
+
+        # ========== lane1 as source ==========
+        cond = (ids1 != -1) & ((ids1 & bitmask) != 0)
+        target = ids1 ^ bitmask
+        mm0 = cond & (ids0 == target)
+        mm2 = cond & (ids2 == target)
+        mm3 = cond & (ids3 == target)
+        mm4 = cond & (ids4 == target)
+        mm5 = cond & (ids5 == target)
+        mm6 = cond & (ids6 == target)
+        mm7 = cond & (ids7 == target)
+        hit = mm0 | mm2 | mm3 | mm4 | mm5 | mm6 | mm7
+
+        src = v1.to(tl.float32)
+        v0 = tl.where(mm0[:, None], (v0.to(tl.float32) + src).to(x_dtype), v0)
+        v2 = tl.where(mm2[:, None], (v2.to(tl.float32) + src).to(x_dtype), v2)
+        v3 = tl.where(mm3[:, None], (v3.to(tl.float32) + src).to(x_dtype), v3)
+        v4 = tl.where(mm4[:, None], (v4.to(tl.float32) + src).to(x_dtype), v4)
+        v5 = tl.where(mm5[:, None], (v5.to(tl.float32) + src).to(x_dtype), v5)
+        v6 = tl.where(mm6[:, None], (v6.to(tl.float32) + src).to(x_dtype), v6)
+        v7 = tl.where(mm7[:, None], (v7.to(tl.float32) + src).to(x_dtype), v7)
+
+        v1 = tl.where(hit[:, None], 0.0, v1)
+        ids1 = tl.where(hit, -1, ids1)
+        ids1 = tl.where(cond & (~hit), target, ids1)
+
+        # ========== lane2 as source ==========
+        cond = (ids2 != -1) & ((ids2 & bitmask) != 0)
+        target = ids2 ^ bitmask
+        mm0 = cond & (ids0 == target)
+        mm1 = cond & (ids1 == target)
+        mm3 = cond & (ids3 == target)
+        mm4 = cond & (ids4 == target)
+        mm5 = cond & (ids5 == target)
+        mm6 = cond & (ids6 == target)
+        mm7 = cond & (ids7 == target)
+        hit = mm0 | mm1 | mm3 | mm4 | mm5 | mm6 | mm7
+
+        src = v2.to(tl.float32)
+        v0 = tl.where(mm0[:, None], (v0.to(tl.float32) + src).to(x_dtype), v0)
+        v1 = tl.where(mm1[:, None], (v1.to(tl.float32) + src).to(x_dtype), v1)
+        v3 = tl.where(mm3[:, None], (v3.to(tl.float32) + src).to(x_dtype), v3)
+        v4 = tl.where(mm4[:, None], (v4.to(tl.float32) + src).to(x_dtype), v4)
+        v5 = tl.where(mm5[:, None], (v5.to(tl.float32) + src).to(x_dtype), v5)
+        v6 = tl.where(mm6[:, None], (v6.to(tl.float32) + src).to(x_dtype), v6)
+        v7 = tl.where(mm7[:, None], (v7.to(tl.float32) + src).to(x_dtype), v7)
+
+        v2 = tl.where(hit[:, None], 0.0, v2)
+        ids2 = tl.where(hit, -1, ids2)
+        ids2 = tl.where(cond & (~hit), target, ids2)
+
+        # ========== lane3 as source ==========
+        cond = (ids3 != -1) & ((ids3 & bitmask) != 0)
+        target = ids3 ^ bitmask
+        mm0 = cond & (ids0 == target)
+        mm1 = cond & (ids1 == target)
+        mm2 = cond & (ids2 == target)
+        mm4 = cond & (ids4 == target)
+        mm5 = cond & (ids5 == target)
+        mm6 = cond & (ids6 == target)
+        mm7 = cond & (ids7 == target)
+        hit = mm0 | mm1 | mm2 | mm4 | mm5 | mm6 | mm7
+
+        src = v3.to(tl.float32)
+        v0 = tl.where(mm0[:, None], (v0.to(tl.float32) + src).to(x_dtype), v0)
+        v1 = tl.where(mm1[:, None], (v1.to(tl.float32) + src).to(x_dtype), v1)
+        v2 = tl.where(mm2[:, None], (v2.to(tl.float32) + src).to(x_dtype), v2)
+        v4 = tl.where(mm4[:, None], (v4.to(tl.float32) + src).to(x_dtype), v4)
+        v5 = tl.where(mm5[:, None], (v5.to(tl.float32) + src).to(x_dtype), v5)
+        v6 = tl.where(mm6[:, None], (v6.to(tl.float32) + src).to(x_dtype), v6)
+        v7 = tl.where(mm7[:, None], (v7.to(tl.float32) + src).to(x_dtype), v7)
+
+        v3 = tl.where(hit[:, None], 0.0, v3)
+        ids3 = tl.where(hit, -1, ids3)
+        ids3 = tl.where(cond & (~hit), target, ids3)
+
+        # ========== lane4 as source ==========
+        cond = (ids4 != -1) & ((ids4 & bitmask) != 0)
+        target = ids4 ^ bitmask
+        mm0 = cond & (ids0 == target)
+        mm1 = cond & (ids1 == target)
+        mm2 = cond & (ids2 == target)
+        mm3 = cond & (ids3 == target)
+        mm5 = cond & (ids5 == target)
+        mm6 = cond & (ids6 == target)
+        mm7 = cond & (ids7 == target)
+        hit = mm0 | mm1 | mm2 | mm3 | mm5 | mm6 | mm7
+
+        src = v4.to(tl.float32)
+        v0 = tl.where(mm0[:, None], (v0.to(tl.float32) + src).to(x_dtype), v0)
+        v1 = tl.where(mm1[:, None], (v1.to(tl.float32) + src).to(x_dtype), v1)
+        v2 = tl.where(mm2[:, None], (v2.to(tl.float32) + src).to(x_dtype), v2)
+        v3 = tl.where(mm3[:, None], (v3.to(tl.float32) + src).to(x_dtype), v3)
+        v5 = tl.where(mm5[:, None], (v5.to(tl.float32) + src).to(x_dtype), v5)
+        v6 = tl.where(mm6[:, None], (v6.to(tl.float32) + src).to(x_dtype), v6)
+        v7 = tl.where(mm7[:, None], (v7.to(tl.float32) + src).to(x_dtype), v7)
+
+        v4 = tl.where(hit[:, None], 0.0, v4)
+        ids4 = tl.where(hit, -1, ids4)
+        ids4 = tl.where(cond & (~hit), target, ids4)
+
+        # ========== lane5 as source ==========
+        cond = (ids5 != -1) & ((ids5 & bitmask) != 0)
+        target = ids5 ^ bitmask
+        mm0 = cond & (ids0 == target)
+        mm1 = cond & (ids1 == target)
+        mm2 = cond & (ids2 == target)
+        mm3 = cond & (ids3 == target)
+        mm4 = cond & (ids4 == target)
+        mm6 = cond & (ids6 == target)
+        mm7 = cond & (ids7 == target)
+        hit = mm0 | mm1 | mm2 | mm3 | mm4 | mm6 | mm7
+
+        src = v5.to(tl.float32)
+        v0 = tl.where(mm0[:, None], (v0.to(tl.float32) + src).to(x_dtype), v0)
+        v1 = tl.where(mm1[:, None], (v1.to(tl.float32) + src).to(x_dtype), v1)
+        v2 = tl.where(mm2[:, None], (v2.to(tl.float32) + src).to(x_dtype), v2)
+        v3 = tl.where(mm3[:, None], (v3.to(tl.float32) + src).to(x_dtype), v3)
+        v4 = tl.where(mm4[:, None], (v4.to(tl.float32) + src).to(x_dtype), v4)
+        v6 = tl.where(mm6[:, None], (v6.to(tl.float32) + src).to(x_dtype), v6)
+        v7 = tl.where(mm7[:, None], (v7.to(tl.float32) + src).to(x_dtype), v7)
+
+        v5 = tl.where(hit[:, None], 0.0, v5)
+        ids5 = tl.where(hit, -1, ids5)
+        ids5 = tl.where(cond & (~hit), target, ids5)
+
+        # ========== lane6 as source ==========
+        cond = (ids6 != -1) & ((ids6 & bitmask) != 0)
+        target = ids6 ^ bitmask
+        mm0 = cond & (ids0 == target)
+        mm1 = cond & (ids1 == target)
+        mm2 = cond & (ids2 == target)
+        mm3 = cond & (ids3 == target)
+        mm4 = cond & (ids4 == target)
+        mm5 = cond & (ids5 == target)
+        mm7 = cond & (ids7 == target)
+        hit = mm0 | mm1 | mm2 | mm3 | mm4 | mm5 | mm7
+
+        src = v6.to(tl.float32)
+        v0 = tl.where(mm0[:, None], (v0.to(tl.float32) + src).to(x_dtype), v0)
+        v1 = tl.where(mm1[:, None], (v1.to(tl.float32) + src).to(x_dtype), v1)
+        v2 = tl.where(mm2[:, None], (v2.to(tl.float32) + src).to(x_dtype), v2)
+        v3 = tl.where(mm3[:, None], (v3.to(tl.float32) + src).to(x_dtype), v3)
+        v4 = tl.where(mm4[:, None], (v4.to(tl.float32) + src).to(x_dtype), v4)
+        v5 = tl.where(mm5[:, None], (v5.to(tl.float32) + src).to(x_dtype), v5)
+        v7 = tl.where(mm7[:, None], (v7.to(tl.float32) + src).to(x_dtype), v7)
+
+        v6 = tl.where(hit[:, None], 0.0, v6)
+        ids6 = tl.where(hit, -1, ids6)
+        ids6 = tl.where(cond & (~hit), target, ids6)
+
+        # ========== lane7 as source ==========
+        cond = (ids7 != -1) & ((ids7 & bitmask) != 0)
+        target = ids7 ^ bitmask
+        mm0 = cond & (ids0 == target)
+        mm1 = cond & (ids1 == target)
+        mm2 = cond & (ids2 == target)
+        mm3 = cond & (ids3 == target)
+        mm4 = cond & (ids4 == target)
+        mm5 = cond & (ids5 == target)
+        mm6 = cond & (ids6 == target)
+        hit = mm0 | mm1 | mm2 | mm3 | mm4 | mm5 | mm6
+
+        src = v7.to(tl.float32)
+        v0 = tl.where(mm0[:, None], (v0.to(tl.float32) + src).to(x_dtype), v0)
+        v1 = tl.where(mm1[:, None], (v1.to(tl.float32) + src).to(x_dtype), v1)
+        v2 = tl.where(mm2[:, None], (v2.to(tl.float32) + src).to(x_dtype), v2)
+        v3 = tl.where(mm3[:, None], (v3.to(tl.float32) + src).to(x_dtype), v3)
+        v4 = tl.where(mm4[:, None], (v4.to(tl.float32) + src).to(x_dtype), v4)
+        v5 = tl.where(mm5[:, None], (v5.to(tl.float32) + src).to(x_dtype), v5)
+        v6 = tl.where(mm6[:, None], (v6.to(tl.float32) + src).to(x_dtype), v6)
+
+        v7 = tl.where(hit[:, None], 0.0, v7)
+        ids7 = tl.where(hit, -1, ids7)
+        ids7 = tl.where(cond & (~hit), target, ids7)
+
+    # ---- final: since ids unique per token, bucket0 is unique; compute directly in fp32 (no out_tile chain) ----
+    acc = tl.zeros((BLOCK_M, BLOCK_H), dtype=tl.float32)
+    acc += tl.where((ids0 == 0)[:, None], v0.to(tl.float32), 0.0)
+    acc += tl.where((ids1 == 0)[:, None], v1.to(tl.float32), 0.0)
+    acc += tl.where((ids2 == 0)[:, None], v2.to(tl.float32), 0.0)
+    acc += tl.where((ids3 == 0)[:, None], v3.to(tl.float32), 0.0)
+    acc += tl.where((ids4 == 0)[:, None], v4.to(tl.float32), 0.0)
+    acc += tl.where((ids5 == 0)[:, None], v5.to(tl.float32), 0.0)
+    acc += tl.where((ids6 == 0)[:, None], v6.to(tl.float32), 0.0)
+    acc += tl.where((ids7 == 0)[:, None], v7.to(tl.float32), 0.0)
+
+    acc *= routed_scaling_factor
+
+    out_ptrs = out_ptr + m[:, None] * so_m + h[None, :] * so_h
+    tl.store(out_ptrs, acc.to(out_ptr.dtype.element_ty), mask=mask_mh)
+
+
+def _moe_sum_tree_reduce_k8_fast_path(
+    input: torch.Tensor,  # [M, 8, H]
+    output: torch.Tensor,  # [M, H]
+    curr_topk_ids: torch.Tensor,  # [M, 8], -1 means remote
+    routed_scaling_factor: float,
+    E: int,
+):
+    assert input.is_contiguous()
+    assert output.is_contiguous()
+    assert curr_topk_ids.is_contiguous()
+    M, K, H = input.shape
+    assert K == 8
+    assert output.shape == (M, H)
+    assert (E & (E - 1)) == 0
+
+    E_LEVEL = int(math.log2(E))
+
+    # K=8 specialization: use wider H tiles for better memory throughput.
+    if H >= 4096:
+        BLOCK_M = 8
+        BLOCK_H = 128
+        num_warps = 4
+    elif H >= 2048:
+        BLOCK_M = 8
+        BLOCK_H = 256
+        num_warps = 8
+    else:
+        BLOCK_M = 8
+        BLOCK_H = 128
+        num_warps = 4
+
+    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(H, BLOCK_H))
+    _moe_sum_tree_reduce_k8_fused_kernel_opt2d[grid](
+        input,
+        curr_topk_ids,
+        output,
+        input.stride(0),
+        input.stride(1),
+        input.stride(2),
+        curr_topk_ids.stride(0),
+        curr_topk_ids.stride(1),
+        output.stride(0),
+        output.stride(1),
+        M,
+        H,
+        E_LEVEL,
+        routed_scaling_factor,
+        BLOCK_M=BLOCK_M,
+        BLOCK_H=BLOCK_H,
+        num_warps=num_warps,
+    )
+    return output
+
+
+@triton.jit
+def _moe_sum_tree_reduce_topk16_sparse_kernel(
+    x_ptr,
+    ids_ptr,
+    out_ptr,
+    sx_m,
+    sx_k,
+    sx_h,  # x: [M, K, H]
+    sid_m,
+    sid_k,  # ids: [M, K]
+    so_m,
+    so_h,  # out: [M, H]
+    M,
+    K,
+    H,  # runtime
+    E_LEVEL,  # log2(E)
+    routed_scaling_factor,
+    BLOCK_H: tl.constexpr,
+    MAX_TOPK: tl.constexpr,  # fixed compile-time capacity, e.g. 16
+):
+    # One program handles one token + one hidden tile.
+    pid_m = tl.program_id(0)
+    pid_h = tl.program_id(1)
+
+    m = pid_m
+    h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
+
+    mask_m = m < M
+    mask_h = h < H
+    mask_mh = mask_m & mask_h
+
+    slot = tl.arange(0, MAX_TOPK)
+    in_k = slot < K
+
+    ids = tl.load(
+        ids_ptr + m * sid_m + slot * sid_k, mask=(mask_m & in_k), other=-1
+    ).to(tl.int32)
+    valid = (ids != -1) & in_k & mask_m
+
+    vals = tl.zeros((MAX_TOPK, BLOCK_H), dtype=tl.float32)
+    for s in range(MAX_TOPK):
+        vmask = mask_h & valid[s]
+        v = tl.load(
+            x_ptr + m * sx_m + s * sx_k + h * sx_h,
+            mask=vmask,
+            other=0.0,
+        ).to(tl.float32)
+        vals = tl.where((slot == s)[:, None], v[None, :], vals)
+
+    idx = tl.arange(0, MAX_TOPK)
+    for bit in tl.range(0, E_LEVEL):
+        bitmask = 1 << bit
+
+        for s in range(MAX_TOPK):
+            id_s = ids[s]
+            cond = (id_s != -1) & ((id_s & bitmask) != 0)
+            target = id_s ^ bitmask
+
+            match = (ids == target) & (idx != s)
+            # 1-based index; 0 means no match.
+            match_pos1 = tl.max(tl.where(match, idx + 1, 0))
+            has = match_pos1 > 0
+            j = match_pos1 - 1
+
+            src = vals[s, :]
+            hit_mask = idx == j
+            vals = vals + tl.where(hit_mask[:, None] & has, src[None, :], 0.0)
+            vals = tl.where((idx == s)[:, None] & has, 0.0, vals)
+            ids = tl.where((idx == s) & has, -1, ids)
+
+            # No partner found at this level: update bucket id only.
+            ids = tl.where((idx == s) & cond & (~has), target, ids)
+
+    acc = tl.zeros((BLOCK_H,), dtype=tl.float32)
+    for s in range(MAX_TOPK):
+        acc += tl.where(ids[s] == 0, vals[s, :], 0.0)
+
+    acc *= routed_scaling_factor
+    tl.store(
+        out_ptr + m * so_m + h * so_h, acc.to(out_ptr.dtype.element_ty), mask=mask_mh
+    )
+
+
+def moe_sum_tree_reduce_v1_topk_sparse16(
+    input: torch.Tensor,  # [M, K, H]
+    output: torch.Tensor,  # [M, H]
+    curr_topk_ids: torch.Tensor,  # [M, K], -1 means remote
+    routed_scaling_factor: float,
+    E: int,
+    max_topk: int = 16,
+):
+    """
+    Supplemental kernel path for topk != 8:
+    - optimized for small topk (<=16)
+    - keeps deterministic tree-equivalent reduction semantics
+    - does not replace existing v1 entrypoint automatically
+    """
+    assert input.is_contiguous()
+    assert output.is_contiguous()
+    assert curr_topk_ids.is_contiguous()
+    M, K, H = input.shape
+    assert output.shape == (M, H)
+    assert K <= max_topk, f"K={K} > max_topk={max_topk}"
+    assert (E & (E - 1)) == 0, "E must be power of 2"
+
+    E_LEVEL = int(math.log2(E))
+    BLOCK_H = 128 if H >= 4096 else 256
+    num_warps = 4 if BLOCK_H == 128 else 8
+
+    grid = (M, triton.cdiv(H, BLOCK_H))
+    _moe_sum_tree_reduce_topk16_sparse_kernel[grid](
+        input,
+        curr_topk_ids,
+        output,
+        input.stride(0),
+        input.stride(1),
+        input.stride(2),
+        curr_topk_ids.stride(0),
+        curr_topk_ids.stride(1),
+        output.stride(0),
+        output.stride(1),
+        M,
+        K,
+        H,
+        E_LEVEL,
+        routed_scaling_factor,
+        BLOCK_H=BLOCK_H,
+        MAX_TOPK=max_topk,
+        num_warps=num_warps,
+    )
+    return output
+
+
+def moe_sum_tree_reduce_v0(
+    input: torch.Tensor,  # [M, 8, H]
+    output: torch.Tensor,  # [M, H]
+    curr_topk_ids: torch.Tensor,  # [M, 8], -1 means remote
+    routed_scaling_factor: float,
+    E: int,
+):
+    # Backward-compatible alias.
+    return _moe_sum_tree_reduce_k8_fast_path(
+        input=input,
+        output=output,
+        curr_topk_ids=curr_topk_ids,
+        routed_scaling_factor=routed_scaling_factor,
+        E=E,
+    )
+
+
+@triton.jit
+def _moe_tree_reduce_sparse_k8_kernel(
+    x_ptr,
+    ids_ptr,
+    out_ptr,
+    sx_m: tl.constexpr,
+    sx_k: tl.constexpr,
+    sx_h: tl.constexpr,
+    sid_m: tl.constexpr,
+    sid_k: tl.constexpr,
+    so_m: tl.constexpr,
+    so_h: tl.constexpr,
+    M,
+    H,
+    routed_scaling_factor,
+    LOGE,  # runtime: tl.range does not unroll
+    CAST_MODE: tl.constexpr,  # 0: per-level bf16 round, 1: no intermediate cast
+    BLOCK_M: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_h = tl.program_id(1)
+
+    m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
+
+    tl.multiple_of(m, BLOCK_M)
+    tl.multiple_of(h, 8)
+    tl.max_contiguous(h, BLOCK_H)
+
+    mask_m = m < M
+    mask_h = h < H
+    mask_mh = mask_m[:, None] & mask_h[None, :]
+
+    x_ty = x_ptr.dtype.element_ty
+    NEG2 = -2  # sentinel: topk_ids cannot be -2
+
+    # ---- ids ----
+    ids0 = tl.load(ids_ptr + m * sid_m + 0 * sid_k, mask=mask_m, other=-1).to(tl.int32)
+    ids1 = tl.load(ids_ptr + m * sid_m + 1 * sid_k, mask=mask_m, other=-1).to(tl.int32)
+    ids2 = tl.load(ids_ptr + m * sid_m + 2 * sid_k, mask=mask_m, other=-1).to(tl.int32)
+    ids3 = tl.load(ids_ptr + m * sid_m + 3 * sid_k, mask=mask_m, other=-1).to(tl.int32)
+    ids4 = tl.load(ids_ptr + m * sid_m + 4 * sid_k, mask=mask_m, other=-1).to(tl.int32)
+    ids5 = tl.load(ids_ptr + m * sid_m + 5 * sid_k, mask=mask_m, other=-1).to(tl.int32)
+    ids6 = tl.load(ids_ptr + m * sid_m + 6 * sid_k, mask=mask_m, other=-1).to(tl.int32)
+    ids7 = tl.load(ids_ptr + m * sid_m + 7 * sid_k, mask=mask_m, other=-1).to(tl.int32)
+
+    # ---- vals: load -> fp32 ----
+    f0 = tl.load(
+        x_ptr + m[:, None] * sx_m + 0 * sx_k + h[None, :] * sx_h,
+        mask=mask_mh & (ids0 != -1)[:, None],
+        other=0.0,
+    ).to(tl.float32)
+    f1 = tl.load(
+        x_ptr + m[:, None] * sx_m + 1 * sx_k + h[None, :] * sx_h,
+        mask=mask_mh & (ids1 != -1)[:, None],
+        other=0.0,
+    ).to(tl.float32)
+    f2 = tl.load(
+        x_ptr + m[:, None] * sx_m + 2 * sx_k + h[None, :] * sx_h,
+        mask=mask_mh & (ids2 != -1)[:, None],
+        other=0.0,
+    ).to(tl.float32)
+    f3 = tl.load(
+        x_ptr + m[:, None] * sx_m + 3 * sx_k + h[None, :] * sx_h,
+        mask=mask_mh & (ids3 != -1)[:, None],
+        other=0.0,
+    ).to(tl.float32)
+    f4 = tl.load(
+        x_ptr + m[:, None] * sx_m + 4 * sx_k + h[None, :] * sx_h,
+        mask=mask_mh & (ids4 != -1)[:, None],
+        other=0.0,
+    ).to(tl.float32)
+    f5 = tl.load(
+        x_ptr + m[:, None] * sx_m + 5 * sx_k + h[None, :] * sx_h,
+        mask=mask_mh & (ids5 != -1)[:, None],
+        other=0.0,
+    ).to(tl.float32)
+    f6 = tl.load(
+        x_ptr + m[:, None] * sx_m + 6 * sx_k + h[None, :] * sx_h,
+        mask=mask_mh & (ids6 != -1)[:, None],
+        other=0.0,
+    ).to(tl.float32)
+    f7 = tl.load(
+        x_ptr + m[:, None] * sx_m + 7 * sx_k + h[None, :] * sx_h,
+        mask=mask_mh & (ids7 != -1)[:, None],
+        other=0.0,
+    ).to(tl.float32)
+
+    # ---- main: runtime loop ----
+    for bit in tl.range(0, LOGE):
+        bitmask = 1 << bit
+
+        # snapshot ids for matching (critical!)
+        a0, a1, a2, a3, a4, a5, a6, a7 = ids0, ids1, ids2, ids3, ids4, ids5, ids6, ids7
+
+        # ---------------- s0 ----------------
+        c0 = (a0 != -1) & ((a0 & bitmask) != 0)
+        t0 = tl.where(c0, a0 ^ bitmask, NEG2)
+        m01 = a1 == t0
+        m02 = a2 == t0
+        m03 = a3 == t0
+        m04 = a4 == t0
+        m05 = a5 == t0
+        m06 = a6 == t0
+        m07 = a7 == t0
+        hit0 = m01 | m02 | m03 | m04 | m05 | m06 | m07
+        src = f0
+        f1 = tl.where(m01[:, None], f1 + src, f1)
+        f2 = tl.where(m02[:, None], f2 + src, f2)
+        f3 = tl.where(m03[:, None], f3 + src, f3)
+        f4 = tl.where(m04[:, None], f4 + src, f4)
+        f5 = tl.where(m05[:, None], f5 + src, f5)
+        f6 = tl.where(m06[:, None], f6 + src, f6)
+        f7 = tl.where(m07[:, None], f7 + src, f7)
+        f0 = tl.where(hit0[:, None], 0.0, f0)
+        ids0 = tl.where(c0, tl.where(hit0, -1, t0), ids0)
+
+        # ---------------- s1 ----------------
+        c1 = (a1 != -1) & ((a1 & bitmask) != 0)
+        t1 = tl.where(c1, a1 ^ bitmask, NEG2)
+        m10 = a0 == t1
+        m12 = a2 == t1
+        m13 = a3 == t1
+        m14 = a4 == t1
+        m15 = a5 == t1
+        m16 = a6 == t1
+        m17 = a7 == t1
+        hit1 = m10 | m12 | m13 | m14 | m15 | m16 | m17
+        src = f1
+        f0 = tl.where(m10[:, None], f0 + src, f0)
+        f2 = tl.where(m12[:, None], f2 + src, f2)
+        f3 = tl.where(m13[:, None], f3 + src, f3)
+        f4 = tl.where(m14[:, None], f4 + src, f4)
+        f5 = tl.where(m15[:, None], f5 + src, f5)
+        f6 = tl.where(m16[:, None], f6 + src, f6)
+        f7 = tl.where(m17[:, None], f7 + src, f7)
+        f1 = tl.where(hit1[:, None], 0.0, f1)
+        ids1 = tl.where(c1, tl.where(hit1, -1, t1), ids1)
+
+        # ---------------- s2 ----------------
+        c2 = (a2 != -1) & ((a2 & bitmask) != 0)
+        t2 = tl.where(c2, a2 ^ bitmask, NEG2)
+        m20 = a0 == t2
+        m21 = a1 == t2
+        m23 = a3 == t2
+        m24 = a4 == t2
+        m25 = a5 == t2
+        m26 = a6 == t2
+        m27 = a7 == t2
+        hit2 = m20 | m21 | m23 | m24 | m25 | m26 | m27
+        src = f2
+        f0 = tl.where(m20[:, None], f0 + src, f0)
+        f1 = tl.where(m21[:, None], f1 + src, f1)
+        f3 = tl.where(m23[:, None], f3 + src, f3)
+        f4 = tl.where(m24[:, None], f4 + src, f4)
+        f5 = tl.where(m25[:, None], f5 + src, f5)
+        f6 = tl.where(m26[:, None], f6 + src, f6)
+        f7 = tl.where(m27[:, None], f7 + src, f7)
+        f2 = tl.where(hit2[:, None], 0.0, f2)
+        ids2 = tl.where(c2, tl.where(hit2, -1, t2), ids2)
+
+        # ---------------- s3 ----------------
+        c3 = (a3 != -1) & ((a3 & bitmask) != 0)
+        t3 = tl.where(c3, a3 ^ bitmask, NEG2)
+        m30 = a0 == t3
+        m31 = a1 == t3
+        m32 = a2 == t3
+        m34 = a4 == t3
+        m35 = a5 == t3
+        m36 = a6 == t3
+        m37 = a7 == t3
+        hit3 = m30 | m31 | m32 | m34 | m35 | m36 | m37
+        src = f3
+        f0 = tl.where(m30[:, None], f0 + src, f0)
+        f1 = tl.where(m31[:, None], f1 + src, f1)
+        f2 = tl.where(m32[:, None], f2 + src, f2)
+        f4 = tl.where(m34[:, None], f4 + src, f4)
+        f5 = tl.where(m35[:, None], f5 + src, f5)
+        f6 = tl.where(m36[:, None], f6 + src, f6)
+        f7 = tl.where(m37[:, None], f7 + src, f7)
+        f3 = tl.where(hit3[:, None], 0.0, f3)
+        ids3 = tl.where(c3, tl.where(hit3, -1, t3), ids3)
+
+        # ---------------- s4 ----------------
+        c4 = (a4 != -1) & ((a4 & bitmask) != 0)
+        t4 = tl.where(c4, a4 ^ bitmask, NEG2)
+        m40 = a0 == t4
+        m41 = a1 == t4
+        m42 = a2 == t4
+        m43 = a3 == t4
+        m45 = a5 == t4
+        m46 = a6 == t4
+        m47 = a7 == t4
+        hit4 = m40 | m41 | m42 | m43 | m45 | m46 | m47
+        src = f4
+        f0 = tl.where(m40[:, None], f0 + src, f0)
+        f1 = tl.where(m41[:, None], f1 + src, f1)
+        f2 = tl.where(m42[:, None], f2 + src, f2)
+        f3 = tl.where(m43[:, None], f3 + src, f3)
+        f5 = tl.where(m45[:, None], f5 + src, f5)
+        f6 = tl.where(m46[:, None], f6 + src, f6)
+        f7 = tl.where(m47[:, None], f7 + src, f7)
+        f4 = tl.where(hit4[:, None], 0.0, f4)
+        ids4 = tl.where(c4, tl.where(hit4, -1, t4), ids4)
+
+        # ---------------- s5 ----------------
+        c5 = (a5 != -1) & ((a5 & bitmask) != 0)
+        t5 = tl.where(c5, a5 ^ bitmask, NEG2)
+        m50 = a0 == t5
+        m51 = a1 == t5
+        m52 = a2 == t5
+        m53 = a3 == t5
+        m54 = a4 == t5
+        m56 = a6 == t5
+        m57 = a7 == t5
+        hit5 = m50 | m51 | m52 | m53 | m54 | m56 | m57
+        src = f5
+        f0 = tl.where(m50[:, None], f0 + src, f0)
+        f1 = tl.where(m51[:, None], f1 + src, f1)
+        f2 = tl.where(m52[:, None], f2 + src, f2)
+        f3 = tl.where(m53[:, None], f3 + src, f3)
+        f4 = tl.where(m54[:, None], f4 + src, f4)
+        f6 = tl.where(m56[:, None], f6 + src, f6)
+        f7 = tl.where(m57[:, None], f7 + src, f7)
+        f5 = tl.where(hit5[:, None], 0.0, f5)
+        ids5 = tl.where(c5, tl.where(hit5, -1, t5), ids5)
+
+        # ---------------- s6 ----------------
+        c6 = (a6 != -1) & ((a6 & bitmask) != 0)
+        t6 = tl.where(c6, a6 ^ bitmask, NEG2)
+        m60 = a0 == t6
+        m61 = a1 == t6
+        m62 = a2 == t6
+        m63 = a3 == t6
+        m64 = a4 == t6
+        m65 = a5 == t6
+        m67 = a7 == t6
+        hit6 = m60 | m61 | m62 | m63 | m64 | m65 | m67
+        src = f6
+        f0 = tl.where(m60[:, None], f0 + src, f0)
+        f1 = tl.where(m61[:, None], f1 + src, f1)
+        f2 = tl.where(m62[:, None], f2 + src, f2)
+        f3 = tl.where(m63[:, None], f3 + src, f3)
+        f4 = tl.where(m64[:, None], f4 + src, f4)
+        f5 = tl.where(m65[:, None], f5 + src, f5)
+        f7 = tl.where(m67[:, None], f7 + src, f7)
+        f6 = tl.where(hit6[:, None], 0.0, f6)
+        ids6 = tl.where(c6, tl.where(hit6, -1, t6), ids6)
+
+        # ---------------- s7 ----------------
+        c7 = (a7 != -1) & ((a7 & bitmask) != 0)
+        t7 = tl.where(c7, a7 ^ bitmask, NEG2)
+        m70 = a0 == t7
+        m71 = a1 == t7
+        m72 = a2 == t7
+        m73 = a3 == t7
+        m74 = a4 == t7
+        m75 = a5 == t7
+        m76 = a6 == t7
+        hit7 = m70 | m71 | m72 | m73 | m74 | m75 | m76
+        src = f7
+        f0 = tl.where(m70[:, None], f0 + src, f0)
+        f1 = tl.where(m71[:, None], f1 + src, f1)
+        f2 = tl.where(m72[:, None], f2 + src, f2)
+        f3 = tl.where(m73[:, None], f3 + src, f3)
+        f4 = tl.where(m74[:, None], f4 + src, f4)
+        f5 = tl.where(m75[:, None], f5 + src, f5)
+        f6 = tl.where(m76[:, None], f6 + src, f6)
+        f7 = tl.where(hit7[:, None], 0.0, f7)
+        ids7 = tl.where(c7, tl.where(hit7, -1, t7), ids7)
+
+        # ---- per-level bf16 rounding (optional) ----
+        if CAST_MODE == 0:
+            # emulate "store bf16 per level then reload": round once per level
+            f0 = f0.to(x_ty).to(tl.float32)
+            f1 = f1.to(x_ty).to(tl.float32)
+            f2 = f2.to(x_ty).to(tl.float32)
+            f3 = f3.to(x_ty).to(tl.float32)
+            f4 = f4.to(x_ty).to(tl.float32)
+            f5 = f5.to(x_ty).to(tl.float32)
+            f6 = f6.to(x_ty).to(tl.float32)
+            f7 = f7.to(x_ty).to(tl.float32)
+
+    # ---- gather root (id==0) ----
+    acc = tl.zeros((BLOCK_M, BLOCK_H), dtype=tl.float32)
+    acc += tl.where((ids0 == 0)[:, None], f0, 0.0)
+    acc += tl.where((ids1 == 0)[:, None], f1, 0.0)
+    acc += tl.where((ids2 == 0)[:, None], f2, 0.0)
+    acc += tl.where((ids3 == 0)[:, None], f3, 0.0)
+    acc += tl.where((ids4 == 0)[:, None], f4, 0.0)
+    acc += tl.where((ids5 == 0)[:, None], f5, 0.0)
+    acc += tl.where((ids6 == 0)[:, None], f6, 0.0)
+    acc += tl.where((ids7 == 0)[:, None], f7, 0.0)
+
+    acc *= routed_scaling_factor
+
+    out_ptrs = out_ptr + m[:, None] * so_m + h[None, :] * so_h
+    tl.store(out_ptrs, acc.to(out_ptr.dtype.element_ty), mask=mask_mh)
+
+
+def _launch_sparse_tree_k8_k10(
+    input: torch.Tensor,  # [M, K, H], bf16 contiguous
+    output: torch.Tensor,  # [M, H], bf16 contiguous
+    curr_topk_ids: torch.Tensor,  # [M, K], int32/int64 contiguous, -1 remote
+    routed_scaling_factor: float,
+    E: int,
+    cast_mode: int = 1,  # <-- you asked for "last cast" version, so default = 1
+):
+    assert (
+        input.is_contiguous()
+        and output.is_contiguous()
+        and curr_topk_ids.is_contiguous()
+    )
+    M, K, H = input.shape
+    assert output.shape == (M, H)
+    assert K in (8, 10)
+    assert (E & (E - 1)) == 0
+    LOGE = int(math.log2(E))
+
+    # stable params for CUDA Graph
+    BLOCK_M = 8
+    BLOCK_H = 256
+    num_warps = 8
+    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(H, BLOCK_H))
+
+    # NOTE: for strict CUDA Graph safety, avoid dtype conversion here.
+    # Please ensure curr_topk_ids is int32 upstream.
+    assert (
+        curr_topk_ids.dtype == torch.int32
+    ), "make ids int32 upstream for CUDA Graph stability"
+
+    if K == 8:
+        _moe_tree_reduce_sparse_k8_kernel[grid](
+            input,
+            curr_topk_ids,
+            output,
+            input.stride(0),
+            input.stride(1),
+            input.stride(2),
+            curr_topk_ids.stride(0),
+            curr_topk_ids.stride(1),
+            output.stride(0),
+            output.stride(1),
+            M,
+            H,
+            routed_scaling_factor,
+            LOGE=LOGE,
+            CAST_MODE=cast_mode,
+            BLOCK_M=BLOCK_M,
+            BLOCK_H=BLOCK_H,
+            num_warps=num_warps,
+        )
+
+
+def moe_sum_tree_reduce_v2(
+    input: torch.Tensor,
+    output: torch.Tensor,
+    curr_topk_ids: torch.Tensor,
+    routed_scaling_factor: float,
+    E: int,
+    *,
+    cast_mode: int = 0,  # default to "last cast" as requested
+):
+    assert (
+        input.is_contiguous()
+        and output.is_contiguous()
+        and curr_topk_ids.is_contiguous()
+    )
+    M, K, H = input.shape
+    if M == 0:
+        return output
+    if K in (8, 10):
+        _launch_sparse_tree_k8_k10(
+            input=input,
+            output=output,
+            curr_topk_ids=curr_topk_ids,
+            routed_scaling_factor=routed_scaling_factor,
+            E=E,
+            cast_mode=cast_mode,
+        )
+        return output
+
+    # fallback: keep your existing generic path here
+    return output
+
+
+def moe_sum_tree_reduce(
+    input: torch.Tensor,
+    output: torch.Tensor,
+    curr_topk_ids: torch.Tensor,
+    routed_scaling_factor: float,
+    E: int,
+):
+    curr_topk_ids = curr_topk_ids.to(torch.int32)
+    if os.environ.get("SGLANG_MOE_TREE_REDUCE_USE_V2", "0") == "1":
+        return moe_sum_tree_reduce_v2(
+            input=input,
+            output=output,
+            curr_topk_ids=curr_topk_ids,
+            routed_scaling_factor=routed_scaling_factor,
+            E=E,
+        )
+
+    if input.shape[1] == 8:
+        return moe_sum_tree_reduce_v0(
+            input=input,
+            output=output,
+            curr_topk_ids=curr_topk_ids,
+            routed_scaling_factor=routed_scaling_factor,
+            E=E,
+        )
+    return moe_sum_tree_reduce_v1(
+        input=input,
+        output=output,
+        curr_topk_ids=curr_topk_ids,
+        routed_scaling_factor=routed_scaling_factor,
+        E=E,
+    )
+
+
+# if os.getenv("MOE_SUM_OPTIM", "0") == "1":
+#     moe_sum_tree_reduce = moe_sum_tree_reduce_optim
+#     print("using optimized moe_sum_tree_reduce_optim")
+# else:
+#     print("using original moe_sum_tree_reduce_original")
+
+if os.getenv("TREE_ALL_REDUCE_OPTIM", "0") == "1":
+    tree_all_reduce_sum = tree_all_reduce_sum_optim
diff --git a/python/sglang/srt/true_on_policy/__init__.py b/python/sglang/srt/true_on_policy/__init__.py
new file mode 100644
index 000000000000..9fbaaeee48f3
--- /dev/null
+++ b/python/sglang/srt/true_on_policy/__init__.py
@@ -0,0 +1,49 @@
+"""True-on-policy runtime contract helpers."""
+
+from .config import (
+    ROW_LINEAR_INV_BLOCK_K,
+    get_on_policy_rms_norm_kwargs,
+    get_rl_on_policy_target,
+    is_tp_invariant_target,
+    is_true_on_policy_enabled,
+    patch_prefill_only_deterministic_inference_for_cuda_graph,
+    should_disable_flashinfer_allreduce_fusion,
+    should_disable_fused_qk_norm_mrope,
+    should_disable_mlp_allreduce_fusion_for_on_policy,
+    should_disable_reduce_scatter_for_on_policy,
+    should_force_bfloat16_dense_tensor_math,
+    should_force_bfloat16_lm_head,
+    should_use_tp_invariant_row_linear,
+    should_use_tp_invariant_tree_all_reduce,
+)
+from .contracts import (
+    QWEN3_DENSE_TRUE_ON_POLICY_V1,
+    SGLangTrueOnPolicyContract,
+    SGLangTrueOnPolicyRuntimePolicy,
+    get_true_on_policy_contract,
+    resolve_true_on_policy_runtime_policy,
+    validate_true_on_policy_contract,
+)
+
+__all__ = [
+    "QWEN3_DENSE_TRUE_ON_POLICY_V1",
+    "ROW_LINEAR_INV_BLOCK_K",
+    "SGLangTrueOnPolicyContract",
+    "SGLangTrueOnPolicyRuntimePolicy",
+    "get_true_on_policy_contract",
+    "get_on_policy_rms_norm_kwargs",
+    "get_rl_on_policy_target",
+    "is_tp_invariant_target",
+    "is_true_on_policy_enabled",
+    "patch_prefill_only_deterministic_inference_for_cuda_graph",
+    "resolve_true_on_policy_runtime_policy",
+    "validate_true_on_policy_contract",
+    "should_disable_flashinfer_allreduce_fusion",
+    "should_disable_fused_qk_norm_mrope",
+    "should_disable_mlp_allreduce_fusion_for_on_policy",
+    "should_disable_reduce_scatter_for_on_policy",
+    "should_force_bfloat16_dense_tensor_math",
+    "should_force_bfloat16_lm_head",
+    "should_use_tp_invariant_row_linear",
+    "should_use_tp_invariant_tree_all_reduce",
+]
diff --git a/python/sglang/srt/true_on_policy/config.py b/python/sglang/srt/true_on_policy/config.py
new file mode 100644
index 000000000000..82e156d6d4e8
--- /dev/null
+++ b/python/sglang/srt/true_on_policy/config.py
@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+import contextlib
+from typing import Any, Iterator, Optional
+
+import torch
+
+from sglang.srt.true_on_policy.contracts import resolve_true_on_policy_runtime_policy
+
+ROW_LINEAR_INV_BLOCK_K = 128
+
+
+def _get_global_server_args() -> Any:
+    from sglang.srt.runtime_context import get_server_args
+
+    return get_server_args()
+
+
+def get_rl_on_policy_target() -> Optional[str]:
+    return getattr(_get_global_server_args(), "rl_on_policy_target", None)
+
+
+def is_true_on_policy_enabled() -> bool:
+    return resolve_true_on_policy_runtime_policy(_get_global_server_args()).enabled
+
+
+def is_tp_invariant_target() -> bool:
+    return resolve_true_on_policy_runtime_policy(
+        _get_global_server_args()
+    ).tp_invariant_row_linear
+
+
+def should_disable_reduce_scatter_for_on_policy() -> bool:
+    return resolve_true_on_policy_runtime_policy(
+        _get_global_server_args()
+    ).disable_reduce_scatter
+
+
+def should_disable_mlp_allreduce_fusion_for_on_policy() -> bool:
+    return resolve_true_on_policy_runtime_policy(
+        _get_global_server_args()
+    ).disable_mlp_allreduce_fusion
+
+
+def should_disable_flashinfer_allreduce_fusion() -> bool:
+    return resolve_true_on_policy_runtime_policy(
+        _get_global_server_args()
+    ).disable_flashinfer_allreduce_fusion
+
+
+def should_force_bfloat16_dense_tensor_math() -> bool:
+    return resolve_true_on_policy_runtime_policy(
+        _get_global_server_args()
+    ).force_bfloat16_dense_tensor_math
+
+
+def should_force_bfloat16_lm_head(
+    *,
+    use_fp32_lm_head: bool = False,
+) -> bool:
+    return (
+        resolve_true_on_policy_runtime_policy(
+            _get_global_server_args()
+        ).force_bfloat16_lm_head
+        and not use_fp32_lm_head
+    )
+
+
+def should_disable_fused_qk_norm_mrope() -> bool:
+    return resolve_true_on_policy_runtime_policy(
+        _get_global_server_args()
+    ).disable_fused_qk_norm_mrope
+
+
+def get_on_policy_rms_norm_kwargs(
+    *,
+    weight_dtype: Optional[torch.dtype] = None,
+    override_orig_dtype: Optional[torch.dtype] = None,
+    fp32_residual: bool = False,
+) -> dict[str, Any]:
+    if not is_true_on_policy_enabled():
+        return {}
+
+    kwargs: dict[str, Any] = {
+        "cast_x_before_out_mul": True,
+        "fp32_residual": fp32_residual,
+    }
+    if weight_dtype is not None:
+        kwargs["weight_dtype"] = weight_dtype
+    if override_orig_dtype is not None:
+        kwargs["override_orig_dtype"] = override_orig_dtype
+    return kwargs
+
+
+def should_use_tp_invariant_row_linear(
+    k_size: int,
+    row_linear_enable_inv: Optional[bool] = None,
+) -> bool:
+    policy_enabled = resolve_true_on_policy_runtime_policy(
+        _get_global_server_args()
+    ).tp_invariant_row_linear
+    if row_linear_enable_inv is not None:
+        policy_enabled = policy_enabled and row_linear_enable_inv
+
+    return (
+        policy_enabled
+        and k_size >= ROW_LINEAR_INV_BLOCK_K
+        and k_size % ROW_LINEAR_INV_BLOCK_K == 0
+    )
+
+
+def should_use_tp_invariant_tree_all_reduce(
+    accl_binary_tree_enabled: Optional[bool] = None,
+) -> bool:
+    policy_enabled = resolve_true_on_policy_runtime_policy(
+        _get_global_server_args()
+    ).deterministic_tree_all_reduce
+    if accl_binary_tree_enabled is not None:
+        policy_enabled = policy_enabled and not accl_binary_tree_enabled
+
+    return policy_enabled
+
+
+@contextlib.contextmanager
+def patch_prefill_only_deterministic_inference_for_cuda_graph(
+    server_args: Any,
+    *,
+    attn_backend: Optional[Any] = None,
+    dvr_target_verify_cuda_graph: bool = False,
+) -> Iterator[bool]:
+    enabled = (
+        getattr(server_args, "enable_prefill_only_deterministic_inference", False)
+        and not dvr_target_verify_cuda_graph
+    )
+    if not enabled:
+        yield False
+        return
+
+    saved_num_splits = None
+    if attn_backend is not None and hasattr(attn_backend, "num_splits"):
+        saved_num_splits = attn_backend.num_splits
+
+    try:
+        if attn_backend is not None and hasattr(attn_backend, "num_splits"):
+            attn_backend.num_splits = 0
+
+        yield True
+    finally:
+        if attn_backend is not None and hasattr(attn_backend, "num_splits"):
+            attn_backend.num_splits = saved_num_splits
diff --git a/python/sglang/srt/true_on_policy/contracts.py b/python/sglang/srt/true_on_policy/contracts.py
new file mode 100644
index 000000000000..310778c842ce
--- /dev/null
+++ b/python/sglang/srt/true_on_policy/contracts.py
@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from sglang.srt.true_on_policy.schema import (
+    QWEN3_DENSE_TRUE_ON_POLICY_V1_SCHEMA,
+    TrueOnPolicyContractName,
+    TrueOnPolicyContractSchema,
+)
+
+QWEN3_DENSE_TRUE_ON_POLICY_V1 = QWEN3_DENSE_TRUE_ON_POLICY_V1_SCHEMA.name
+
+
+@dataclass(frozen=True)
+class SGLangTrueOnPolicyRuntimePolicy:
+    """SGLang-local behavior implied by a true-on-policy parity contract."""
+
+    contract_name: Optional[str]
+    enabled: bool
+    force_bfloat16_dense_tensor_math: bool
+    force_bfloat16_lm_head: bool
+    disable_reduce_scatter: bool
+    disable_mlp_allreduce_fusion: bool
+    disable_flashinfer_allreduce_fusion: bool
+    tp_invariant_row_linear: bool
+    deterministic_tree_all_reduce: bool
+    disable_fused_qk_norm_mrope: bool
+
+
+DEFAULT_RUNTIME_POLICY = SGLangTrueOnPolicyRuntimePolicy(
+    contract_name=None,
+    enabled=False,
+    force_bfloat16_dense_tensor_math=False,
+    force_bfloat16_lm_head=False,
+    disable_reduce_scatter=False,
+    disable_mlp_allreduce_fusion=False,
+    disable_flashinfer_allreduce_fusion=False,
+    tp_invariant_row_linear=False,
+    deterministic_tree_all_reduce=False,
+    disable_fused_qk_norm_mrope=False,
+)
+
+
+@dataclass(frozen=True)
+class SGLangTrueOnPolicyContract:
+    """SGLang-local adapter from a shared contract schema to runtime policy."""
+
+    schema: TrueOnPolicyContractSchema
+
+    @property
+    def name(self) -> TrueOnPolicyContractName:
+        return self.schema.name
+
+    def policy_for(self, server_args: Any) -> SGLangTrueOnPolicyRuntimePolicy:
+        uses_tp_invariant_rollout = getattr(server_args, "tp_size", 1) > 1
+        return SGLangTrueOnPolicyRuntimePolicy(
+            contract_name=self.name,
+            enabled=True,
+            force_bfloat16_dense_tensor_math=True,
+            force_bfloat16_lm_head=True,
+            disable_reduce_scatter=True,
+            disable_mlp_allreduce_fusion=True,
+            disable_flashinfer_allreduce_fusion=uses_tp_invariant_rollout,
+            tp_invariant_row_linear=uses_tp_invariant_rollout,
+            deterministic_tree_all_reduce=uses_tp_invariant_rollout,
+            disable_fused_qk_norm_mrope=True,
+        )
+
+
+QWEN3_DENSE_TRUE_ON_POLICY_CONTRACT = SGLangTrueOnPolicyContract(
+    schema=QWEN3_DENSE_TRUE_ON_POLICY_V1_SCHEMA,
+)
+
+
+_CONTRACT_BY_NAME = {
+    QWEN3_DENSE_TRUE_ON_POLICY_CONTRACT.name: QWEN3_DENSE_TRUE_ON_POLICY_CONTRACT,
+}
+
+
+def get_true_on_policy_contract(contract_name: str) -> SGLangTrueOnPolicyContract:
+    try:
+        return _CONTRACT_BY_NAME[contract_name]
+    except KeyError as exc:
+        supported = ", ".join(sorted(_CONTRACT_BY_NAME))
+        raise ValueError(
+            f"Unsupported SGLang true-on-policy contract {contract_name!r}. "
+            f"Supported contracts: {supported}"
+        ) from exc
+
+
+def _contract_name_for(server_args: Any) -> Optional[str]:
+    return getattr(server_args, "true_on_policy_contract", None)
+
+
+def validate_true_on_policy_contract(server_args: Any) -> None:
+    contract_name = getattr(server_args, "true_on_policy_contract", None)
+    if contract_name is None:
+        return
+    get_true_on_policy_contract(contract_name)
+
+
+def resolve_true_on_policy_runtime_policy(
+    server_args: Any,
+) -> SGLangTrueOnPolicyRuntimePolicy:
+    contract_name = _contract_name_for(server_args)
+    if contract_name is None:
+        return DEFAULT_RUNTIME_POLICY
+
+    validate_true_on_policy_contract(server_args)
+    return get_true_on_policy_contract(contract_name).policy_for(server_args)
diff --git a/python/sglang/srt/true_on_policy/schema.py b/python/sglang/srt/true_on_policy/schema.py
new file mode 100644
index 000000000000..0628573c6880
--- /dev/null
+++ b/python/sglang/srt/true_on_policy/schema.py
@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+TrueOnPolicyContractName = Literal["qwen3_dense_true_on_policy_v1"]
+ModelFamily = Literal["qwen3_dense", "qwen3_moe", "qwen3_next"]
+KernelContract = Literal["qwen3_dense_sglang_math"]
+LogprobContract = Literal["sglang_prefill"]
+
+
+@dataclass(frozen=True)
+class TrueOnPolicyContractSchema:
+    """Declarative cross-repo identity for a true-on-policy parity contract."""
+
+    name: TrueOnPolicyContractName
+    model_family: ModelFamily
+    required_kernel_contracts: tuple[KernelContract, ...]
+    logprob_contract: LogprobContract
+    sglang_attention_backend: str
+    fsdp_attention_implementation: str
+    disable_megatron_sequence_parallel: bool
+
+
+QWEN3_DENSE_TRUE_ON_POLICY_V1_SCHEMA = TrueOnPolicyContractSchema(
+    name="qwen3_dense_true_on_policy_v1",
+    model_family="qwen3_dense",
+    required_kernel_contracts=("qwen3_dense_sglang_math",),
+    logprob_contract="sglang_prefill",
+    sglang_attention_backend="fa3",
+    fsdp_attention_implementation="flash_attention_3",
+    disable_megatron_sequence_parallel=True,
+)
diff --git a/test/manual/layers/test_layernorm.py b/test/manual/layers/test_layernorm.py
index 299e5dcffaf4..8a0e68601ba6 100644
--- a/test/manual/layers/test_layernorm.py
+++ b/test/manual/layers/test_layernorm.py
@@ -56,6 +56,26 @@ def test_rms_norm(self):
             ):
                 self._run_rms_norm_test(*params)
 
+    def test_rms_norm_cuda_uses_native_for_fp32_weight(self):
+        torch.manual_seed(0)
+
+        hidden_size = 256
+        layer = RMSNorm(
+            hidden_size,
+            cast_x_before_out_mul=True,
+            weight_dtype=torch.float32,
+            override_orig_dtype=torch.float32,
+        )
+        layer.weight.data.normal_(mean=1.0, std=0.1)
+        x = torch.randn(17, hidden_size, dtype=torch.bfloat16) / hidden_size
+
+        with torch.inference_mode():
+            ref_out = layer.forward_native(x)
+            out = layer(x)
+
+        self.assertEqual(out.dtype, torch.float32)
+        self.assertTrue(torch.equal(out, ref_out))
+
 
 class TestGemmaRMSNorm(CustomTestCase):
     DTYPES = [torch.half, torch.bfloat16]
diff --git a/test/registered/core/test_dense_deterministic_math.py b/test/registered/core/test_dense_deterministic_math.py
new file mode 100644
index 000000000000..7c5c79ee52bd
--- /dev/null
+++ b/test/registered/core/test_dense_deterministic_math.py
@@ -0,0 +1,338 @@
+import json
+import os
+import subprocess
+import textwrap
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.true_on_policy import (
+    QWEN3_DENSE_TRUE_ON_POLICY_V1,
+    get_on_policy_rms_norm_kwargs,
+    should_force_bfloat16_dense_tensor_math,
+    should_force_bfloat16_lm_head,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=12, suite="stage-a-test-cpu")
+
+
+def _run_dense_math_script(script_body: str) -> dict[str, object]:
+    stubbed_imports = textwrap.dedent("""
+        import importlib.machinery
+        import json
+        import sys
+        import types
+        from pydantic import BaseModel
+
+        def install_openai_stubs():
+            openai_mod = types.ModuleType("openai")
+            openai_types_mod = types.ModuleType("openai.types")
+            openai_responses_mod = types.ModuleType("openai.types.responses")
+            openai_response_mod = types.ModuleType("openai.types.responses.response")
+            openai_tool_mod = types.ModuleType("openai.types.responses.tool")
+
+            openai_mod.__spec__ = importlib.machinery.ModuleSpec("openai", loader=None)
+            openai_types_mod.__spec__ = importlib.machinery.ModuleSpec("openai.types", loader=None)
+            openai_responses_mod.__spec__ = importlib.machinery.ModuleSpec(
+                "openai.types.responses", loader=None
+            )
+            openai_response_mod.__spec__ = importlib.machinery.ModuleSpec(
+                "openai.types.responses.response", loader=None
+            )
+            openai_tool_mod.__spec__ = importlib.machinery.ModuleSpec(
+                "openai.types.responses.tool", loader=None
+            )
+
+            for name in [
+                "ResponseFunctionToolCall",
+                "ResponseInputItemParam",
+                "ResponseOutputItem",
+                "ResponseOutputMessage",
+                "ResponseOutputText",
+                "ResponseReasoningItem",
+            ]:
+                setattr(openai_responses_mod, name, type(name, (BaseModel,), {}))
+
+            openai_response_mod.ToolChoice = type("ToolChoice", (BaseModel,), {})
+            openai_tool_mod.Tool = type("Tool", (BaseModel,), {})
+
+            sys.modules.setdefault("openai", openai_mod)
+            sys.modules.setdefault("openai.types", openai_types_mod)
+            sys.modules.setdefault("openai.types.responses", openai_responses_mod)
+            sys.modules.setdefault("openai.types.responses.response", openai_response_mod)
+            sys.modules.setdefault("openai.types.responses.tool", openai_tool_mod)
+
+        install_openai_stubs()
+
+        hf_utils_mod = types.ModuleType("sglang.srt.utils.hf_transformers_utils")
+        hf_utils_mod.__spec__ = importlib.machinery.ModuleSpec(
+            "sglang.srt.utils.hf_transformers_utils", loader=None
+        )
+        hf_utils_mod.check_gguf_file = lambda *args, **kwargs: False
+        hf_utils_mod.get_rope_config = lambda config: (
+            getattr(config, "rope_theta", 1000000),
+            getattr(config, "rope_scaling", None),
+        )
+        sys.modules.setdefault("sglang.srt.utils.hf_transformers_utils", hf_utils_mod)
+
+        gguf_mod = types.ModuleType("gguf")
+        gguf_mod.__spec__ = importlib.machinery.ModuleSpec("gguf", loader=None)
+        gguf_mod.GGMLQuantizationType = type(
+            "GGMLQuantizationType",
+            (),
+            {
+                "F32": 0,
+                "F16": 1,
+                "BF16": 2,
+                "Q4_0": 3,
+                "Q4_1": 4,
+                "Q5_0": 5,
+                "Q5_1": 6,
+                "Q8_0": 7,
+                "Q8_1": 8,
+                "Q2_K": 9,
+                "Q3_K": 10,
+                "Q4_K": 11,
+                "Q5_K": 12,
+                "Q6_K": 13,
+                "IQ1_S": 14,
+                "IQ1_M": 15,
+                "IQ2_XXS": 16,
+                "IQ2_XS": 17,
+                "IQ2_S": 18,
+                "IQ3_XXS": 19,
+                "IQ3_S": 20,
+                "IQ4_NL": 21,
+                "IQ4_XS": 22,
+            },
+        )
+        sys.modules.setdefault("gguf", gguf_mod)
+        """)
+
+    env = dict(os.environ)
+    pythonpath = env.get("PYTHONPATH")
+    repo_python = "python"
+    env["PYTHONPATH"] = (
+        f"{repo_python}{os.pathsep}{pythonpath}" if pythonpath else repo_python
+    )
+    script = f"{stubbed_imports}\n{script_body}"
+    completed = subprocess.run(
+        ["python", "-c", script],
+        check=True,
+        capture_output=True,
+        text=True,
+        env=env,
+    )
+    return json.loads(completed.stdout)
+
+
+class TestDenseOnPolicyHelpers(unittest.TestCase):
+    def test_default_dense_math_helpers_are_inactive(self):
+        server_args = SimpleNamespace(
+            true_on_policy_contract=None,
+            tp_size=1,
+        )
+
+        self.assertFalse(should_force_bfloat16_dense_tensor_math(server_args))
+        self.assertFalse(
+            should_force_bfloat16_lm_head(
+                server_args=server_args,
+                use_fp32_lm_head=False,
+            )
+        )
+        self.assertEqual(get_on_policy_rms_norm_kwargs(server_args), {})
+
+    def test_on_policy_dense_math_helpers_enable_bfloat16_and_rms_norm_kwargs(self):
+        server_args = SimpleNamespace(
+            true_on_policy_contract=QWEN3_DENSE_TRUE_ON_POLICY_V1,
+            tp_size=1,
+        )
+
+        kwargs = get_on_policy_rms_norm_kwargs(
+            server_args,
+            weight_dtype=torch.float32,
+            override_orig_dtype=torch.float32,
+            fp32_residual=True,
+        )
+
+        self.assertTrue(should_force_bfloat16_dense_tensor_math(server_args))
+        self.assertTrue(
+            should_force_bfloat16_lm_head(
+                server_args=server_args,
+                use_fp32_lm_head=False,
+            )
+        )
+        self.assertFalse(
+            should_force_bfloat16_lm_head(
+                server_args=server_args,
+                use_fp32_lm_head=True,
+            )
+        )
+        self.assertEqual(kwargs["weight_dtype"], torch.float32)
+        self.assertEqual(kwargs["override_orig_dtype"], torch.float32)
+        self.assertTrue(kwargs["cast_x_before_out_mul"])
+        self.assertTrue(kwargs["fp32_residual"])
+
+
+class TestDenseOnPolicyContracts(unittest.TestCase):
+    def test_qwen3_style_rms_norm_keeps_fp32_weight_output_and_residual(self):
+        result = _run_dense_math_script(textwrap.dedent("""
+                import json
+                from types import SimpleNamespace
+
+                import torch
+
+                from sglang.srt.layers.layernorm import RMSNorm
+                from sglang.srt.true_on_policy import get_on_policy_rms_norm_kwargs
+
+                from sglang.srt.true_on_policy import QWEN3_DENSE_TRUE_ON_POLICY_V1
+
+                server_args = SimpleNamespace(
+                    true_on_policy_contract=QWEN3_DENSE_TRUE_ON_POLICY_V1,
+                    tp_size=1,
+                )
+                norm = RMSNorm(
+                    4,
+                    eps=1e-6,
+                    **get_on_policy_rms_norm_kwargs(
+                        server_args,
+                        weight_dtype=torch.float32,
+                        override_orig_dtype=torch.float32,
+                        fp32_residual=True,
+                    ),
+                )
+                x = torch.randn(2, 4, dtype=torch.bfloat16)
+                residual = torch.randn(2, 4, dtype=torch.bfloat16)
+                out, residual_out = norm.forward_native(x, residual)
+                print(
+                    json.dumps(
+                        {
+                            "weight_dtype": str(norm.weight.dtype),
+                            "out_dtype": str(out.dtype),
+                            "residual_dtype": str(residual_out.dtype),
+                        }
+                    )
+                )
+                """))
+
+        self.assertEqual(result["weight_dtype"], "torch.float32")
+        self.assertEqual(result["out_dtype"], "torch.float32")
+        self.assertEqual(result["residual_dtype"], "torch.float32")
+
+    def test_rms_norm_can_self_configure_from_true_on_policy_role_hints(self):
+        result = _run_dense_math_script(textwrap.dedent("""
+                import json
+
+                import torch
+
+                from sglang.srt.layers.layernorm import RMSNorm
+                from sglang.srt.server_args import (
+                    ServerArgs,
+                    get_global_server_args,
+                    set_global_server_args_for_scheduler,
+                )
+                from sglang.srt.true_on_policy import QWEN3_DENSE_TRUE_ON_POLICY_V1
+
+                set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+                server_args = get_global_server_args()
+                server_args.true_on_policy_contract = QWEN3_DENSE_TRUE_ON_POLICY_V1
+                server_args.tp_size = 1
+                norm = RMSNorm(
+                    4,
+                    eps=1e-6,
+                    true_on_policy_weight_dtype=torch.float32,
+                    true_on_policy_override_orig_dtype=torch.float32,
+                    true_on_policy_fp32_residual=True,
+                )
+                print(
+                    json.dumps(
+                        {
+                            "weight_dtype": str(norm.weight.dtype),
+                            "cast_x_before_out_mul": norm.cast_x_before_out_mul,
+                            "fp32_residual": norm.fp32_residual,
+                            "override_orig_dtype": str(norm.override_orig_dtype),
+                        }
+                    )
+                )
+                """))
+
+        self.assertEqual(result["weight_dtype"], "torch.float32")
+        self.assertTrue(result["cast_x_before_out_mul"])
+        self.assertTrue(result["fp32_residual"])
+        self.assertEqual(result["override_orig_dtype"], "torch.float32")
+
+    def test_on_policy_lm_head_forces_bfloat16_matmul_inputs(self):
+        result = _run_dense_math_script(textwrap.dedent("""
+                import json
+                from types import SimpleNamespace
+                from unittest.mock import patch
+
+                import torch
+                import torch.nn as nn
+
+                from sglang.srt.layers.logits_processor import LogitsProcessor
+                from sglang.srt.server_args import (
+                    ServerArgs,
+                    get_global_server_args,
+                    set_global_server_args_for_scheduler,
+                )
+                from sglang.srt.true_on_policy import QWEN3_DENSE_TRUE_ON_POLICY_V1
+
+                class DummyMeta:
+                    gathered_buffer = None
+                    next_token_logits_buffer = None
+
+                    def compute_dp_attention_metadata(self):
+                        return None
+
+                class LMHeadStub(nn.Module):
+                    def __init__(self):
+                        super().__init__()
+                        self.weight = nn.Parameter(torch.randn(8, 4, dtype=torch.float32))
+
+                set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+                get_global_server_args().enable_dp_lm_head = False
+                get_global_server_args().enable_fp32_lm_head = False
+                get_global_server_args().true_on_policy_contract = QWEN3_DENSE_TRUE_ON_POLICY_V1
+                get_global_server_args().tp_size = 1
+
+                processor = LogitsProcessor(
+                    SimpleNamespace(vocab_size=8, final_logit_softcapping=None),
+                    skip_all_gather=True,
+                    logit_scale=None,
+                )
+                hidden_states = torch.randn(2, 4, dtype=torch.float32)
+                head = LMHeadStub()
+                captured = {}
+
+                original_matmul = torch.matmul
+
+                def probe_matmul(a, b, *args, **kwargs):
+                    if not captured:
+                        captured["a_dtype"] = str(a.dtype)
+                        captured["b_dtype"] = str(b.dtype)
+                    return original_matmul(a, b, *args, **kwargs)
+
+                with patch("torch.matmul", new=probe_matmul):
+                    logits = processor._get_logits(hidden_states, head, DummyMeta())
+
+                print(
+                    json.dumps(
+                        {
+                            "a_dtype": captured["a_dtype"],
+                            "b_dtype": captured["b_dtype"],
+                            "logits_dtype": str(logits.dtype),
+                        }
+                    )
+                )
+                """))
+
+        self.assertEqual(result["a_dtype"], "torch.bfloat16")
+        self.assertEqual(result["b_dtype"], "torch.bfloat16")
+        self.assertEqual(result["logits_dtype"], "torch.bfloat16")
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)
diff --git a/test/registered/core/test_on_policy_wiring.py b/test/registered/core/test_on_policy_wiring.py
new file mode 100644
index 000000000000..946f1922a9e0
--- /dev/null
+++ b/test/registered/core/test_on_policy_wiring.py
@@ -0,0 +1,606 @@
+import json
+import os
+import subprocess
+import textwrap
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from sglang.srt.true_on_policy import (
+    QWEN3_DENSE_TRUE_ON_POLICY_V1,
+    get_rl_on_policy_target,
+    get_true_on_policy_contract,
+    is_tp_invariant_target,
+    is_true_on_policy_enabled,
+    patch_prefill_only_deterministic_inference_for_cuda_graph,
+    resolve_true_on_policy_runtime_policy,
+    should_disable_flashinfer_allreduce_fusion,
+    should_disable_fused_qk_norm_mrope,
+    should_disable_mlp_allreduce_fusion_for_on_policy,
+    should_disable_reduce_scatter_for_on_policy,
+    should_use_tp_invariant_row_linear,
+    should_use_tp_invariant_tree_all_reduce,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=12, suite="stage-a-test-cpu")
+
+_PATCH_TARGET = "sglang.srt.server_args.get_global_server_args"
+
+
+def _run_server_args_script(argv: list[str]) -> dict[str, object]:
+    stubbed_imports = textwrap.dedent("""
+        import argparse
+        import importlib.machinery
+        import json
+        import sys
+        import types
+        from types import SimpleNamespace
+        from unittest.mock import patch
+
+        from pydantic import BaseModel
+
+        def install_openai_stubs():
+            openai_mod = types.ModuleType("openai")
+            openai_types_mod = types.ModuleType("openai.types")
+            openai_responses_mod = types.ModuleType("openai.types.responses")
+            openai_response_mod = types.ModuleType("openai.types.responses.response")
+            openai_tool_mod = types.ModuleType("openai.types.responses.tool")
+
+            openai_mod.__spec__ = importlib.machinery.ModuleSpec("openai", loader=None)
+            openai_types_mod.__spec__ = importlib.machinery.ModuleSpec("openai.types", loader=None)
+            openai_responses_mod.__spec__ = importlib.machinery.ModuleSpec(
+                "openai.types.responses", loader=None
+            )
+            openai_response_mod.__spec__ = importlib.machinery.ModuleSpec(
+                "openai.types.responses.response", loader=None
+            )
+            openai_tool_mod.__spec__ = importlib.machinery.ModuleSpec(
+                "openai.types.responses.tool", loader=None
+            )
+
+            for name in [
+                "ResponseFunctionToolCall",
+                "ResponseInputItemParam",
+                "ResponseOutputItem",
+                "ResponseOutputMessage",
+                "ResponseOutputText",
+                "ResponseReasoningItem",
+            ]:
+                setattr(openai_responses_mod, name, type(name, (BaseModel,), {}))
+
+            openai_response_mod.ToolChoice = type("ToolChoice", (BaseModel,), {})
+            openai_tool_mod.Tool = type("Tool", (BaseModel,), {})
+
+            sys.modules.setdefault("openai", openai_mod)
+            sys.modules.setdefault("openai.types", openai_types_mod)
+            sys.modules.setdefault("openai.types.responses", openai_responses_mod)
+            sys.modules.setdefault("openai.types.responses.response", openai_response_mod)
+            sys.modules.setdefault("openai.types.responses.tool", openai_tool_mod)
+
+        install_openai_stubs()
+
+        hf_utils_mod = types.ModuleType("sglang.srt.utils.hf_transformers_utils")
+        hf_utils_mod.__spec__ = importlib.machinery.ModuleSpec(
+            "sglang.srt.utils.hf_transformers_utils", loader=None
+        )
+        hf_utils_mod.check_gguf_file = lambda *args, **kwargs: False
+        sys.modules.setdefault("sglang.srt.utils.hf_transformers_utils", hf_utils_mod)
+
+        from sglang.srt.server_args import ServerArgs
+
+        def _mock_model_config():
+            return SimpleNamespace(
+                hf_config=SimpleNamespace(architectures=["Qwen2ForCausalLM"])
+            )
+
+        parser = argparse.ArgumentParser()
+        ServerArgs.add_cli_args(parser)
+        cli_args = parser.parse_args(ARGV)
+
+        with patch("sglang.srt.server_args.get_device", return_value="cuda"), patch.object(
+            ServerArgs, "get_model_config", return_value=_mock_model_config()
+        ):
+            server_args = ServerArgs.from_cli_args(cli_args)
+            server_args._handle_deterministic_inference()
+
+        print(
+            json.dumps(
+                {
+                    "enable_deterministic_inference": server_args.enable_deterministic_inference,
+                    "enable_prefill_only_deterministic_inference": server_args.enable_prefill_only_deterministic_inference,
+                    "enable_flashinfer_allreduce_fusion": server_args.enable_flashinfer_allreduce_fusion,
+                    "rl_on_policy_target": server_args.rl_on_policy_target,
+                    "true_on_policy_contract": server_args.true_on_policy_contract,
+                    "sampling_backend": server_args.sampling_backend,
+                }
+            )
+        )
+        """)
+
+    env = dict(os.environ)
+    pythonpath = env.get("PYTHONPATH")
+    repo_python = "python"
+    env["PYTHONPATH"] = (
+        f"{repo_python}{os.pathsep}{pythonpath}" if pythonpath else repo_python
+    )
+    script = f"ARGV = {argv!r}\n{stubbed_imports}"
+    completed = subprocess.run(
+        ["python", "-c", script],
+        check=True,
+        capture_output=True,
+        text=True,
+        env=env,
+    )
+    return json.loads(completed.stdout)
+
+
+class TestOnPolicyServerArgs(unittest.TestCase):
+    def test_cli_parses_prefill_only_deterministic_flag(self):
+        result = _run_server_args_script(
+            [
+                "--model-path",
+                "dummy",
+                "--attention-backend",
+                "triton",
+                "--enable-prefill-only-deterministic-inference",
+            ]
+        )
+
+        self.assertTrue(result["enable_prefill_only_deterministic_inference"])
+        self.assertTrue(result["enable_deterministic_inference"])
+        self.assertIsNone(result["rl_on_policy_target"])
+        self.assertEqual(result["sampling_backend"], "pytorch")
+
+    def test_cli_accepts_fsdp_and_fsdp_tp_targets(self):
+        fsdp_tp_result = _run_server_args_script(
+            [
+                "--model-path",
+                "dummy",
+                "--attention-backend",
+                "triton",
+                "--rl-on-policy-target",
+                "fsdp_tp",
+            ]
+        )
+        self.assertEqual(fsdp_tp_result["rl_on_policy_target"], "fsdp_tp")
+        self.assertIsNone(fsdp_tp_result["true_on_policy_contract"])
+        self.assertTrue(fsdp_tp_result["enable_deterministic_inference"])
+
+        fsdp_result = _run_server_args_script(
+            [
+                "--model-path",
+                "dummy",
+                "--attention-backend",
+                "triton",
+                "--rl-on-policy-target",
+                "fsdp",
+            ]
+        )
+        self.assertEqual(fsdp_result["rl_on_policy_target"], "fsdp")
+        self.assertIsNone(fsdp_result["true_on_policy_contract"])
+        self.assertTrue(fsdp_result["enable_deterministic_inference"])
+
+    def test_cli_accepts_explicit_true_on_policy_contract(self):
+        result = _run_server_args_script(
+            [
+                "--model-path",
+                "dummy",
+                "--attention-backend",
+                "triton",
+                "--true-on-policy-contract",
+                QWEN3_DENSE_TRUE_ON_POLICY_V1,
+            ]
+        )
+
+        self.assertIsNone(result["rl_on_policy_target"])
+        self.assertEqual(
+            result["true_on_policy_contract"], QWEN3_DENSE_TRUE_ON_POLICY_V1
+        )
+        self.assertTrue(result["enable_deterministic_inference"])
+
+    def test_contract_tp_rollout_disables_flashinfer_allreduce_fusion(self):
+        result = _run_server_args_script(
+            [
+                "--model-path",
+                "dummy",
+                "--attention-backend",
+                "triton",
+                "--tensor-parallel-size",
+                "2",
+                "--true-on-policy-contract",
+                QWEN3_DENSE_TRUE_ON_POLICY_V1,
+                "--enable-flashinfer-allreduce-fusion",
+            ]
+        )
+        self.assertFalse(result["enable_flashinfer_allreduce_fusion"])
+
+    def test_legacy_target_keeps_flashinfer_allreduce_fusion_available(self):
+        result = _run_server_args_script(
+            [
+                "--model-path",
+                "dummy",
+                "--attention-backend",
+                "triton",
+                "--rl-on-policy-target",
+                "fsdp_tp",
+                "--enable-flashinfer-allreduce-fusion",
+            ]
+        )
+        self.assertTrue(result["enable_flashinfer_allreduce_fusion"])
+
+
+def _mock_args(**kwargs):
+    defaults = dict(
+        rl_on_policy_target=None,
+        true_on_policy_contract=None,
+        tp_size=1,
+    )
+    defaults.update(kwargs)
+    return SimpleNamespace(**defaults)
+
+
+def _contract_args(*, tp_size: int = 1):
+    return _mock_args(
+        true_on_policy_contract=QWEN3_DENSE_TRUE_ON_POLICY_V1,
+        tp_size=tp_size,
+    )
+
+
+class TestDefaultPathUnchanged(unittest.TestCase):
+    """Default serving must not enter true-on-policy policy paths."""
+
+    def setUp(self):
+        self.default_args = _mock_args()
+
+    def test_default_args_no_on_policy(self):
+        with patch(_PATCH_TARGET, return_value=self.default_args):
+            self.assertIsNone(get_rl_on_policy_target())
+            self.assertFalse(is_true_on_policy_enabled())
+            self.assertFalse(is_tp_invariant_target())
+
+    def test_default_args_row_linear_uses_quant_method(self):
+        with patch(_PATCH_TARGET, return_value=self.default_args):
+            self.assertFalse(
+                should_use_tp_invariant_row_linear(
+                    256,
+                    row_linear_enable_inv=True,
+                )
+            )
+
+    def test_default_args_tree_allreduce_not_selected(self):
+        with patch(_PATCH_TARGET, return_value=self.default_args):
+            self.assertFalse(
+                should_use_tp_invariant_tree_all_reduce(
+                    accl_binary_tree_enabled=False,
+                )
+            )
+
+    def test_default_args_reduce_scatter_available(self):
+        with patch(_PATCH_TARGET, return_value=self.default_args):
+            self.assertFalse(should_disable_reduce_scatter_for_on_policy())
+
+    def test_default_args_mlp_fusion_available(self):
+        with patch(_PATCH_TARGET, return_value=self.default_args):
+            self.assertFalse(should_disable_mlp_allreduce_fusion_for_on_policy())
+
+    def test_default_args_flashinfer_fusion_available(self):
+        with patch(_PATCH_TARGET, return_value=self.default_args):
+            self.assertFalse(should_disable_flashinfer_allreduce_fusion())
+
+    def test_default_server_args_cli_no_on_policy_flags(self):
+        result = _run_server_args_script(
+            ["--model-path", "dummy", "--attention-backend", "triton"]
+        )
+        self.assertIsNone(result["rl_on_policy_target"])
+        self.assertFalse(result["enable_deterministic_inference"])
+        self.assertFalse(result["enable_prefill_only_deterministic_inference"])
+
+
+class TestOnPolicyHelpers(unittest.TestCase):
+    def test_tp_invariant_row_linear_selection(self):
+        with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)):
+            self.assertTrue(
+                should_use_tp_invariant_row_linear(
+                    256,
+                    row_linear_enable_inv=True,
+                )
+            )
+
+    def test_tp_invariant_row_linear_selection_is_contract_owned(self):
+        with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)):
+            with patch.dict(os.environ, {"ROW_LINEAR_ENABLE_INV": "0"}):
+                self.assertTrue(should_use_tp_invariant_row_linear(256))
+
+    def test_contract_resolver_ignores_legacy_target_without_contract(self):
+        policy = resolve_true_on_policy_runtime_policy(
+            _mock_args(rl_on_policy_target="fsdp_tp", tp_size=2)
+        )
+
+        self.assertIsNone(policy.contract_name)
+        self.assertFalse(policy.enabled)
+        self.assertFalse(policy.tp_invariant_row_linear)
+        self.assertFalse(policy.deterministic_tree_all_reduce)
+
+    def test_contract_resolver_accepts_explicit_qwen3_dense_contract(self):
+        args_tp1 = _contract_args(tp_size=1)
+        policy = resolve_true_on_policy_runtime_policy(args_tp1)
+
+        self.assertTrue(policy.enabled)
+        self.assertTrue(policy.force_bfloat16_dense_tensor_math)
+        self.assertFalse(policy.tp_invariant_row_linear)
+        with patch(_PATCH_TARGET, return_value=args_tp1):
+            self.assertFalse(
+                should_use_tp_invariant_row_linear(
+                    96,
+                    row_linear_enable_inv=True,
+                )
+            )
+            self.assertFalse(
+                should_use_tp_invariant_row_linear(
+                    256,
+                    row_linear_enable_inv=True,
+                )
+            )
+
+    def test_contract_object_owns_sglang_runtime_policy_values(self):
+        contract = get_true_on_policy_contract(QWEN3_DENSE_TRUE_ON_POLICY_V1)
+
+        policy = contract.policy_for(_contract_args(tp_size=2))
+
+        self.assertEqual(contract.schema.name, QWEN3_DENSE_TRUE_ON_POLICY_V1)
+        self.assertEqual(contract.schema.model_family, "qwen3_dense")
+        self.assertEqual(policy.contract_name, QWEN3_DENSE_TRUE_ON_POLICY_V1)
+        self.assertTrue(policy.enabled)
+        self.assertTrue(policy.force_bfloat16_dense_tensor_math)
+        self.assertTrue(policy.force_bfloat16_lm_head)
+        self.assertTrue(policy.disable_reduce_scatter)
+        self.assertTrue(policy.disable_mlp_allreduce_fusion)
+        self.assertTrue(policy.disable_flashinfer_allreduce_fusion)
+        self.assertTrue(policy.tp_invariant_row_linear)
+        self.assertTrue(policy.deterministic_tree_all_reduce)
+        self.assertTrue(policy.disable_fused_qk_norm_mrope)
+
+    def test_reduce_scatter_and_fusion_are_disabled_for_contract(self):
+        with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=1)):
+            self.assertTrue(should_disable_reduce_scatter_for_on_policy())
+        with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)):
+            self.assertTrue(should_disable_mlp_allreduce_fusion_for_on_policy())
+        with patch(_PATCH_TARGET, return_value=_mock_args()):
+            self.assertFalse(should_disable_reduce_scatter_for_on_policy())
+
+    def test_tree_all_reduce_selection_requires_tp_rollout_and_no_accl(self):
+        with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)):
+            self.assertTrue(
+                should_use_tp_invariant_tree_all_reduce(
+                    accl_binary_tree_enabled=False,
+                )
+            )
+            self.assertFalse(
+                should_use_tp_invariant_tree_all_reduce(
+                    accl_binary_tree_enabled=True,
+                )
+            )
+        with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=1)):
+            self.assertFalse(
+                should_use_tp_invariant_tree_all_reduce(
+                    accl_binary_tree_enabled=False,
+                )
+            )
+
+    def test_tree_all_reduce_selection_is_contract_owned(self):
+        with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)):
+            with patch.dict(os.environ, {"ACCL_BINARY_TREE_ENABLE": "1"}):
+                self.assertTrue(should_use_tp_invariant_tree_all_reduce())
+
+    def test_attention_handoff_tree_reduce_uses_attention_tp_group(self):
+        from sglang.srt.layers.communicator import (
+            CommunicateWithAllReduceAndLayerNormFn,
+        )
+
+        hidden_states = torch.ones(2, 4)
+        residual = torch.full((2, 4), 3.0)
+
+        class FakeNorm:
+            def __call__(self, x, residual):
+                return x + residual, residual
+
+        with (
+            patch(
+                "sglang.srt.layers.communicator.get_attn_tp_context",
+                return_value=SimpleNamespace(input_scattered=False),
+            ),
+            patch(
+                "sglang.srt.layers.communicator.apply_aiter_all_reduce_fusion",
+                return_value=False,
+            ),
+            patch(
+                "sglang.srt.layers.communicator.apply_flashinfer_allreduce_fusion",
+                return_value=False,
+            ),
+            patch(
+                "sglang.srt.layers.communicator.attention_tensor_model_parallel_all_reduce",
+                side_effect=lambda x: x + 10.0,
+            ) as attn_tree_reduce,
+        ):
+            output, output_residual = (
+                CommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual(
+                    hidden_states,
+                    residual,
+                    forward_batch=None,
+                    layernorm=FakeNorm(),
+                    context=SimpleNamespace(attn_dp_size=1, cache=None),
+                    residual_input_mode=None,
+                )
+            )
+
+        attn_tree_reduce.assert_called_once()
+        torch.testing.assert_close(output, hidden_states + 10.0 + residual)
+        torch.testing.assert_close(output_residual, residual)
+
+    def test_prefill_only_cuda_graph_patch_only_scopes_attention_splits(self):
+        server_args = SimpleNamespace(
+            enable_prefill_only_deterministic_inference=True,
+            enable_deterministic_inference=True,
+            enable_flashinfer_allreduce_fusion=False,
+            rl_on_policy_target="fsdp_tp",
+            true_on_policy_contract=QWEN3_DENSE_TRUE_ON_POLICY_V1,
+            disable_custom_all_reduce=True,
+        )
+        attn_backend = SimpleNamespace(num_splits=1)
+
+        with patch.dict(
+            os.environ,
+            {
+                "SGLANG_ENABLE_DETERMINISTIC_INFERENCE": "1",
+                "SGLANG_DISABLE_CUSTOM_ALL_REDUCE": "1",
+                "NCCL_ALGO": "allreduce:tree",
+            },
+            clear=False,
+        ):
+            with patch_prefill_only_deterministic_inference_for_cuda_graph(
+                server_args,
+                attn_backend=attn_backend,
+            ) as patched:
+                self.assertTrue(patched)
+                self.assertTrue(server_args.enable_deterministic_inference)
+                self.assertFalse(server_args.enable_flashinfer_allreduce_fusion)
+                self.assertEqual(server_args.rl_on_policy_target, "fsdp_tp")
+                self.assertEqual(
+                    server_args.true_on_policy_contract,
+                    QWEN3_DENSE_TRUE_ON_POLICY_V1,
+                )
+                self.assertEqual(attn_backend.num_splits, 0)
+                self.assertEqual(
+                    os.environ["SGLANG_ENABLE_DETERMINISTIC_INFERENCE"], "1"
+                )
+                self.assertEqual(os.environ["SGLANG_DISABLE_CUSTOM_ALL_REDUCE"], "1")
+                self.assertEqual(os.environ["NCCL_ALGO"], "allreduce:tree")
+
+            self.assertTrue(server_args.enable_deterministic_inference)
+            self.assertFalse(server_args.enable_flashinfer_allreduce_fusion)
+            self.assertEqual(server_args.rl_on_policy_target, "fsdp_tp")
+            self.assertEqual(
+                server_args.true_on_policy_contract,
+                QWEN3_DENSE_TRUE_ON_POLICY_V1,
+            )
+            self.assertTrue(server_args.disable_custom_all_reduce)
+            self.assertEqual(attn_backend.num_splits, 1)
+            self.assertEqual(os.environ["SGLANG_ENABLE_DETERMINISTIC_INFERENCE"], "1")
+            self.assertEqual(os.environ["SGLANG_DISABLE_CUSTOM_ALL_REDUCE"], "1")
+            self.assertEqual(os.environ["NCCL_ALGO"], "allreduce:tree")
+
+    def test_row_linear_k_alignment_edge_cases(self):
+        with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)):
+            self.assertFalse(
+                should_use_tp_invariant_row_linear(64, row_linear_enable_inv=True),
+            )
+            self.assertTrue(
+                should_use_tp_invariant_row_linear(128, row_linear_enable_inv=True),
+            )
+            self.assertFalse(
+                should_use_tp_invariant_row_linear(300, row_linear_enable_inv=True),
+            )
+            self.assertTrue(
+                should_use_tp_invariant_row_linear(3584, row_linear_enable_inv=True),
+            )
+
+    def test_row_linear_explicit_override_can_disable(self):
+        with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)):
+            self.assertFalse(
+                should_use_tp_invariant_row_linear(256, row_linear_enable_inv=False)
+            )
+
+    def test_flashinfer_allreduce_fusion_helpers(self):
+        with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)):
+            self.assertTrue(should_disable_flashinfer_allreduce_fusion())
+        with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=1)):
+            self.assertFalse(should_disable_flashinfer_allreduce_fusion())
+        with patch(_PATCH_TARGET, return_value=_mock_args()):
+            self.assertFalse(should_disable_flashinfer_allreduce_fusion())
+
+    def test_fused_qk_norm_mrope_helper_follows_true_on_policy_contract(self):
+        with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=1)):
+            self.assertTrue(should_disable_fused_qk_norm_mrope())
+        with patch(_PATCH_TARGET, return_value=_mock_args()):
+            self.assertFalse(should_disable_fused_qk_norm_mrope())
+
+    def test_get_rl_on_policy_target_returns_correct_value(self):
+        with patch(
+            _PATCH_TARGET, return_value=_mock_args(rl_on_policy_target="fsdp_tp")
+        ):
+            self.assertEqual(get_rl_on_policy_target(), "fsdp_tp")
+        with patch(_PATCH_TARGET, return_value=_mock_args(rl_on_policy_target="fsdp")):
+            self.assertEqual(get_rl_on_policy_target(), "fsdp")
+        with patch(_PATCH_TARGET, return_value=_mock_args()):
+            self.assertIsNone(get_rl_on_policy_target())
+
+    def test_is_true_on_policy_enabled_for_both_targets(self):
+        with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=1)):
+            self.assertTrue(is_true_on_policy_enabled())
+        with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)):
+            self.assertTrue(is_true_on_policy_enabled())
+        with patch(_PATCH_TARGET, return_value=_mock_args()):
+            self.assertFalse(is_true_on_policy_enabled())
+
+    def test_is_tp_invariant_target_only_fsdp_tp(self):
+        with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=2)):
+            self.assertTrue(is_tp_invariant_target())
+        with patch(_PATCH_TARGET, return_value=_contract_args(tp_size=1)):
+            self.assertFalse(is_tp_invariant_target())
+        with patch(_PATCH_TARGET, return_value=_mock_args()):
+            self.assertFalse(is_tp_invariant_target())
+
+    def test_cuda_graph_patch_noop_when_disabled(self):
+        server_args = SimpleNamespace(
+            enable_prefill_only_deterministic_inference=False,
+            enable_deterministic_inference=True,
+            rl_on_policy_target="fsdp_tp",
+        )
+        with patch_prefill_only_deterministic_inference_for_cuda_graph(
+            server_args,
+        ) as patched:
+            self.assertFalse(patched)
+            self.assertTrue(server_args.enable_deterministic_inference)
+            self.assertEqual(server_args.rl_on_policy_target, "fsdp_tp")
+
+    def test_cuda_graph_patch_noop_when_dvr_verify(self):
+        server_args = SimpleNamespace(
+            enable_prefill_only_deterministic_inference=True,
+            enable_deterministic_inference=True,
+            enable_flashinfer_allreduce_fusion=False,
+            rl_on_policy_target="fsdp_tp",
+            disable_custom_all_reduce=True,
+        )
+        with patch_prefill_only_deterministic_inference_for_cuda_graph(
+            server_args,
+            dvr_target_verify_cuda_graph=True,
+        ) as patched:
+            self.assertFalse(patched)
+            self.assertTrue(server_args.enable_deterministic_inference)
+            self.assertEqual(server_args.rl_on_policy_target, "fsdp_tp")
+
+    def test_tp_invariant_ops_import_is_available(self):
+        import sglang.srt.tp_invariant_ops as tp_invariant_ops
+
+        self.assertTrue(hasattr(tp_invariant_ops, "matmul_tp_inv"))
+
+    def test_legacy_on_policy_utils_import_matches_true_on_policy_namespace(self):
+        from sglang.srt import true_on_policy
+        from sglang.srt.layers import on_policy_utils as legacy
+
+        self.assertIs(
+            legacy.should_use_tp_invariant_row_linear,
+            true_on_policy.should_use_tp_invariant_row_linear,
+        )
+        self.assertIs(
+            legacy.patch_prefill_only_deterministic_inference_for_cuda_graph,
+            true_on_policy.patch_prefill_only_deterministic_inference_for_cuda_graph,
+        )
+        self.assertTrue(hasattr(torch.ops, "tp_inv_ops"))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/registered/core/test_tp_invariant_ops.py b/test/registered/core/test_tp_invariant_ops.py
new file mode 100644
index 000000000000..aaa4fab087c9
--- /dev/null
+++ b/test/registered/core/test_tp_invariant_ops.py
@@ -0,0 +1,866 @@
+"""Tests for TP-invariant kernels (PR1).
+
+TP-invariance property:
+    Given the same TP degree and the same input data, matmul_tp_persistent
+    plus tree_all_reduce_sum produces bitwise identical results across runs.
+    When K/BLOCK_K is divisible by tp_size AND each shard yields a power-of-two
+    block count, TP=1 and TP=N also agree (isomorphic tree structure).
+
+    For production K values (e.g. 3584, 5120) where block counts per shard are
+    not power-of-two, the tree structure differs between TP degrees. The
+    invariance guarantee is *determinism for a fixed TP degree*.
+
+All bitwise assertions use torch.equal, never approximate tolerances.
+"""
+
+import os
+import random
+import unittest
+
+import torch
+import torch.distributed as dist
+
+from sglang.srt.tp_invariant_ops import (
+    disable_tp_invariant_mode,
+    enable_tp_invariant_mode,
+    is_tp_invariant_mode_enabled,
+    matmul_tp_inv,
+    matmul_tp_persistent,
+    moe_sum_tree_reduce,
+    set_tp_invariant_mode,
+    tree_all_reduce_sum,
+)
+from sglang.srt.tp_invariant_ops.tp_invariant_ops import (
+    _MATMUL_K_BLOCK,
+    _fixed_tree_sum_tensors,
+    _is_power_of_two,
+)
+from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci
+
+register_cpu_ci(est_time=12, suite="stage-a-test-cpu")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="8-gpu-h200")
+
+BLOCK_K = _MATMUL_K_BLOCK  # 128
+
+
+def _simulate_tp_matmul(A, B, tp_size, **kwargs):
+    """Simulate what production TP does: each rank runs matmul_tp_persistent
+    on its K-shard, then tree_all_reduce_sum gathers and tree-sums."""
+    K = A.shape[1]
+    shard = K // tp_size
+    partials = []
+    for r in range(tp_size):
+        start = r * shard
+        end = start + shard
+        partials.append(
+            matmul_tp_persistent(A[:, start:end], B[start:end, :], **kwargs)
+        )
+    return _fixed_tree_sum_tensors(partials)
+
+
+# ---------------------------------------------------------------------------
+# Mode flag
+# ---------------------------------------------------------------------------
+class TestTPInvariantMode(unittest.TestCase):
+    def tearDown(self):
+        disable_tp_invariant_mode()
+
+    def test_mode_context_restores_previous_state(self):
+        disable_tp_invariant_mode()
+        self.assertFalse(is_tp_invariant_mode_enabled())
+
+        with set_tp_invariant_mode(True):
+            self.assertTrue(is_tp_invariant_mode_enabled())
+        self.assertFalse(is_tp_invariant_mode_enabled())
+
+        enable_tp_invariant_mode()
+        with set_tp_invariant_mode(False):
+            self.assertFalse(is_tp_invariant_mode_enabled())
+        self.assertTrue(is_tp_invariant_mode_enabled())
+
+    def test_enable_is_idempotent(self):
+        enable_tp_invariant_mode()
+        enable_tp_invariant_mode()
+        self.assertTrue(is_tp_invariant_mode_enabled())
+        disable_tp_invariant_mode()
+        self.assertFalse(is_tp_invariant_mode_enabled())
+
+    def test_context_restores_after_exception(self):
+        disable_tp_invariant_mode()
+        try:
+            with set_tp_invariant_mode(True):
+                raise RuntimeError("deliberate")
+        except RuntimeError:
+            pass
+        self.assertFalse(is_tp_invariant_mode_enabled())
+
+
+# ---------------------------------------------------------------------------
+# Reference ops: correctness
+# ---------------------------------------------------------------------------
+class TestTPInvariantReferenceOps(unittest.TestCase):
+    def test_fixed_tree_sum_order_is_stable(self):
+        values = [
+            torch.tensor([1.0e20], dtype=torch.float32),
+            torch.tensor([1.0], dtype=torch.float32),
+            torch.tensor([-1.0e20], dtype=torch.float32),
+            torch.tensor([3.0], dtype=torch.float32),
+        ]
+        tree_result = _fixed_tree_sum_tensors(values)
+        sequential_result = values[0] + values[1] + values[2] + values[3]
+
+        self.assertEqual(tree_result.item(), 0.0)
+        self.assertEqual(sequential_result.item(), 3.0)
+
+    def test_matmul_tp_persistent_matches_torch_matmul_fp32(self):
+        torch.manual_seed(0)
+        A = torch.randn(5, 257, dtype=torch.float32)
+        B = torch.randn(257, 7, dtype=torch.float32)
+        bias = torch.randn(7, dtype=torch.float32)
+
+        actual = matmul_tp_persistent(A, B, bias=bias)
+        expected = A @ B + bias
+        torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
+
+    def test_matmul_tp_persistent_bf16_approximate_to_torch(self):
+        """BF16 block-tree matmul vs native torch matmul. Different accumulation
+        order means results differ by up to a few BF16 ULPs; this test only
+        checks that the magnitude is in the same ballpark, not bitwise equality."""
+        torch.manual_seed(0)
+        A = torch.randn(6, 256, dtype=torch.bfloat16)
+        B = torch.randn(256, 10, dtype=torch.bfloat16)
+        bias = torch.randn(10, dtype=torch.bfloat16)
+
+        actual = matmul_tp_persistent(A, B, bias=bias)
+        expected = A @ B + bias
+        torch.testing.assert_close(actual, expected, rtol=1e-1, atol=1e-1)
+
+    def test_torch_custom_op_dispatches_to_matmul(self):
+        A = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+        B = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+
+        actual = torch.ops.tp_inv_ops.matmul_tp_inv(A, B)
+        expected = A @ B
+        torch.testing.assert_close(actual, expected)
+
+    def test_torch_custom_op_dispatches_with_bias(self):
+        torch.manual_seed(99)
+        A = torch.randn(4, 128, dtype=torch.float32)
+        B = torch.randn(128, 8, dtype=torch.float32)
+        bias = torch.randn(8, dtype=torch.float32)
+
+        actual = torch.ops.tp_inv_ops.matmul_tp_inv(A, B, bias)
+        expected = matmul_tp_persistent(A, B, bias=bias)
+        self.assertTrue(torch.equal(actual, expected))
+
+    def test_matmul_tp_inv_matches_persistent(self):
+        torch.manual_seed(11)
+        A = torch.randn(4, 256, dtype=torch.bfloat16)
+        B = torch.randn(256, 16, dtype=torch.bfloat16)
+        self.assertTrue(torch.equal(matmul_tp_inv(A, B), matmul_tp_persistent(A, B)))
+
+    def test_moe_sum_tree_reduce_matches_expert_order_reference(self):
+        input_tensor = torch.tensor(
+            [
+                [
+                    [1.0e20, 1.0],
+                    [1.0, 2.0],
+                    [-1.0e20, 4.0],
+                    [3.0, 8.0],
+                ],
+                [
+                    [5.0, 7.0],
+                    [11.0, 13.0],
+                    [17.0, 19.0],
+                    [23.0, 29.0],
+                ],
+            ],
+            dtype=torch.float32,
+        )
+        curr_topk_ids = torch.tensor(
+            [[0, 1, 2, 3], [3, -1, 1, 0]],
+            dtype=torch.int64,
+        )
+        output = torch.empty(2, 2, dtype=torch.float32)
+
+        moe_sum_tree_reduce(
+            input=input_tensor,
+            output=output,
+            curr_topk_ids=curr_topk_ids,
+            routed_scaling_factor=0.5,
+            E=4,
+        )
+
+        expected = torch.tensor([[0.0, 7.5], [22.5, 27.5]], dtype=torch.float32)
+        torch.testing.assert_close(output, expected)
+
+    def test_moe_sum_tree_reduce_rejects_non_power_of_two_expert_count(self):
+        with self.assertRaisesRegex(ValueError, "power of two"):
+            moe_sum_tree_reduce(
+                input=torch.zeros(1, 1, 2),
+                output=torch.zeros(1, 2),
+                curr_topk_ids=torch.zeros(1, 1, dtype=torch.int64),
+                routed_scaling_factor=1.0,
+                E=3,
+            )
+
+
+# ---------------------------------------------------------------------------
+# Matmul TP invariance: the core contract
+#
+# Two classes of tests:
+#   1. Cross-TP: TP=1 == TP=N, requires isomorphic tree (K/BLOCK_K power-of-two
+#      multiple of tp_size). Uses torch.equal.
+#   2. Determinism: same TP degree, two runs == bitwise identical. Works for
+#      any valid K alignment.
+# ---------------------------------------------------------------------------
+class TestTPInvarianceCrossTP(unittest.TestCase):
+    """TP=1 == TP=N when the binary tree structure is isomorphic."""
+
+    def test_fp32_cross_tp_k512(self):
+        torch.manual_seed(42)
+        A = torch.randn(8, 512, dtype=torch.float32)
+        B = torch.randn(512, 16, dtype=torch.float32)
+
+        result_tp1 = matmul_tp_persistent(A, B)
+        result_tp2 = _simulate_tp_matmul(A, B, tp_size=2)
+        result_tp4 = _simulate_tp_matmul(A, B, tp_size=4)
+
+        self.assertTrue(torch.equal(result_tp1, result_tp2))
+        self.assertTrue(torch.equal(result_tp1, result_tp4))
+
+    def test_bf16_cross_tp_k512(self):
+        torch.manual_seed(42)
+        A = torch.randn(8, 512, dtype=torch.bfloat16)
+        B = torch.randn(512, 16, dtype=torch.bfloat16)
+
+        result_tp1 = matmul_tp_persistent(A, B)
+        result_tp2 = _simulate_tp_matmul(A, B, tp_size=2)
+        result_tp4 = _simulate_tp_matmul(A, B, tp_size=4)
+
+        self.assertTrue(torch.equal(result_tp1, result_tp2))
+        self.assertTrue(torch.equal(result_tp1, result_tp4))
+
+    def test_bf16_cross_tp_k1024(self):
+        torch.manual_seed(7)
+        A = torch.randn(4, 1024, dtype=torch.bfloat16)
+        B = torch.randn(1024, 32, dtype=torch.bfloat16)
+
+        result_tp1 = matmul_tp_persistent(A, B)
+        result_tp2 = _simulate_tp_matmul(A, B, tp_size=2)
+        result_tp4 = _simulate_tp_matmul(A, B, tp_size=4)
+        result_tp8 = _simulate_tp_matmul(A, B, tp_size=8)
+
+        self.assertTrue(torch.equal(result_tp1, result_tp2))
+        self.assertTrue(torch.equal(result_tp1, result_tp4))
+        self.assertTrue(torch.equal(result_tp1, result_tp8))
+
+    def test_bf16_cross_tp_k2048_all_sizes(self):
+        """K=2048: 16 blocks. TP={1,2,4,8,16} all produce isomorphic trees."""
+        torch.manual_seed(102)
+        A = torch.randn(4, 2048, dtype=torch.bfloat16)
+        B = torch.randn(2048, 16, dtype=torch.bfloat16)
+
+        result_tp1 = matmul_tp_persistent(A, B)
+        for tp_size in [2, 4, 8, 16]:
+            result_tpN = _simulate_tp_matmul(A, B, tp_size=tp_size)
+            self.assertTrue(
+                torch.equal(result_tp1, result_tpN),
+                f"TP=1 != TP={tp_size} for K=2048",
+            )
+
+    def test_bf16_cross_tp_k4096(self):
+        """K=4096: 32 blocks."""
+        torch.manual_seed(103)
+        A = torch.randn(2, 4096, dtype=torch.bfloat16)
+        B = torch.randn(4096, 8, dtype=torch.bfloat16)
+
+        result_tp1 = matmul_tp_persistent(A, B)
+        for tp_size in [2, 4, 8]:
+            result_tpN = _simulate_tp_matmul(A, B, tp_size=tp_size)
+            self.assertTrue(
+                torch.equal(result_tp1, result_tpN),
+                f"TP=1 != TP={tp_size} for K=4096",
+            )
+
+    def test_fp16_cross_tp_k512(self):
+        torch.manual_seed(500)
+        A = torch.randn(4, 512, dtype=torch.float16)
+        B = torch.randn(512, 16, dtype=torch.float16)
+
+        result_tp1 = matmul_tp_persistent(A, B)
+        result_tp4 = _simulate_tp_matmul(A, B, tp_size=4)
+
+        self.assertTrue(torch.equal(result_tp1, result_tp4))
+
+    def test_torch_ops_dispatch_bf16_cross_tp(self):
+        """torch.ops.tp_inv_ops.matmul_tp_inv dispatch preserves cross-TP invariance."""
+        torch.manual_seed(400)
+        A = torch.randn(4, 512, dtype=torch.bfloat16)
+        B = torch.randn(512, 16, dtype=torch.bfloat16)
+
+        result_full = torch.ops.tp_inv_ops.matmul_tp_inv(A, B)
+
+        K = A.shape[1]
+        shard = K // 4
+        partials = []
+        for r in range(4):
+            s, e = r * shard, (r + 1) * shard
+            partials.append(torch.ops.tp_inv_ops.matmul_tp_inv(A[:, s:e], B[s:e, :]))
+        result_tp4 = _fixed_tree_sum_tensors(partials)
+
+        self.assertTrue(torch.equal(result_full, result_tp4))
+
+
+class TestTPInvarianceDeterminism(unittest.TestCase):
+    """Same TP degree, two runs -> bitwise identical. Works for all production K."""
+
+    def _assert_deterministic(self, K, tp_size, dtype=torch.bfloat16):
+        torch.manual_seed(42)
+        A = torch.randn(4, K, dtype=dtype)
+        B = torch.randn(K, 16, dtype=dtype)
+
+        result_a = _simulate_tp_matmul(A, B, tp_size=tp_size)
+        result_b = _simulate_tp_matmul(A, B, tp_size=tp_size)
+        self.assertTrue(
+            torch.equal(result_a, result_b),
+            f"K={K} TP={tp_size} dtype={dtype} not deterministic",
+        )
+
+    def test_bf16_k3584_tp2(self):
+        """Qwen3-4B hidden_size. 3584/2=1792 -> 14 blocks per shard."""
+        self._assert_deterministic(3584, tp_size=2)
+
+    def test_bf16_k3584_tp4(self):
+        self._assert_deterministic(3584, tp_size=4)
+
+    def test_bf16_k5120_tp4(self):
+        """Qwen3-30B hidden_size. 5120/4=1280 -> 10 blocks per shard."""
+        self._assert_deterministic(5120, tp_size=4)
+
+    def test_bf16_k5120_tp8(self):
+        self._assert_deterministic(5120, tp_size=8)
+
+    def test_bf16_k4096_tp8(self):
+        self._assert_deterministic(4096, tp_size=8)
+
+    def test_fp32_k3584_tp2(self):
+        self._assert_deterministic(3584, tp_size=2, dtype=torch.float32)
+
+    def test_fp32_accum_deterministic(self):
+        """fp32_accum=True is deterministic for a fixed TP degree."""
+        torch.manual_seed(200)
+        A = torch.randn(4, 512, dtype=torch.bfloat16)
+        B = torch.randn(512, 16, dtype=torch.bfloat16)
+
+        result_a = _simulate_tp_matmul(A, B, tp_size=2, fp32_accum=True)
+        result_b = _simulate_tp_matmul(A, B, tp_size=2, fp32_accum=True)
+        self.assertTrue(torch.equal(result_a, result_b))
+
+    def test_fp32_accum_output_dtype_is_input_dtype(self):
+        torch.manual_seed(300)
+        A = torch.randn(4, 512, dtype=torch.bfloat16)
+        B = torch.randn(512, 16, dtype=torch.bfloat16)
+
+        result = matmul_tp_persistent(A, B, fp32_accum=True)
+        self.assertEqual(result.dtype, torch.bfloat16)
+
+
+# ---------------------------------------------------------------------------
+# BFloat16 ops: approximate correctness vs torch
+# ---------------------------------------------------------------------------
+class TestBFloat16Ops(unittest.TestCase):
+    def test_moe_sum_tree_reduce_bf16(self):
+        input_tensor = torch.tensor(
+            [[[1.0, 2.0], [3.0, 4.0]]],
+            dtype=torch.bfloat16,
+        )
+        curr_topk_ids = torch.tensor([[0, 1]], dtype=torch.int64)
+        output = torch.empty(1, 2, dtype=torch.bfloat16)
+
+        moe_sum_tree_reduce(
+            input=input_tensor,
+            output=output,
+            curr_topk_ids=curr_topk_ids,
+            routed_scaling_factor=1.0,
+            E=2,
+        )
+
+        expected = torch.tensor([[4.0, 6.0]], dtype=torch.bfloat16)
+        torch.testing.assert_close(output, expected)
+
+
+# ---------------------------------------------------------------------------
+# MoE tree reduce: EP/slot invariance
+# ---------------------------------------------------------------------------
+class TestMoeReduceSlotInvariance(unittest.TestCase):
+    """moe_sum_tree_reduce must produce bitwise identical results regardless
+    of which topk slot an expert appears in. This is the EP invariance
+    property: different EP ranks may route the same experts to different
+    slot positions, but the tree-reduce result must be bitwise identical."""
+
+    def _run_moe_reduce(self, input_tensor, topk_ids, E, scaling=1.0):
+        output = torch.zeros(
+            input_tensor.shape[0], input_tensor.shape[2], dtype=input_tensor.dtype
+        )
+        moe_sum_tree_reduce(
+            input=input_tensor,
+            output=output,
+            curr_topk_ids=topk_ids,
+            routed_scaling_factor=scaling,
+            E=E,
+        )
+        return output
+
+    def _permute_slots(self, values, ids_src, ids_dst, topk):
+        """Rearrange values so that each expert's data moves to its new slot."""
+        tokens = values.shape[0]
+        result = torch.zeros_like(values)
+        for t in range(tokens):
+            for slot_dst in range(topk):
+                eid = ids_dst[t, slot_dst].item()
+                if eid == -1:
+                    continue
+                slot_src = (ids_src[t] == eid).nonzero(as_tuple=True)[0].item()
+                result[t, slot_dst] = values[t, slot_src]
+        return result
+
+    def test_slot_permutation_gives_identical_result_fp32(self):
+        torch.manual_seed(10)
+        H = 64
+        values = torch.randn(1, 4, H, dtype=torch.float32)
+
+        ids_a = torch.tensor([[0, 1, 2, 3]], dtype=torch.int64)
+        ids_b = torch.tensor([[2, 0, 3, 1]], dtype=torch.int64)
+
+        input_b = self._permute_slots(values, ids_a, ids_b, topk=4)
+
+        result_a = self._run_moe_reduce(values, ids_a, E=4)
+        result_b = self._run_moe_reduce(input_b, ids_b, E=4)
+        self.assertTrue(torch.equal(result_a, result_b))
+
+    def test_slot_permutation_gives_identical_result_bf16(self):
+        torch.manual_seed(20)
+        H = 128
+        values = torch.randn(2, 8, H, dtype=torch.bfloat16)
+
+        expert_ids = list(range(8))
+        rng = random.Random(42)
+        shuffled = expert_ids.copy()
+        rng.shuffle(shuffled)
+
+        ids_a = torch.tensor([expert_ids, expert_ids], dtype=torch.int64)
+        ids_b = torch.tensor([shuffled, shuffled], dtype=torch.int64)
+
+        input_b = self._permute_slots(values, ids_a, ids_b, topk=8)
+
+        result_a = self._run_moe_reduce(values, ids_a, E=8)
+        result_b = self._run_moe_reduce(input_b, ids_b, E=8)
+        self.assertTrue(torch.equal(result_a, result_b))
+
+    def test_slot_invariance_with_remote_experts(self):
+        """Remote experts (-1) in different positions must not change the result."""
+        torch.manual_seed(30)
+        H = 32
+        values_a = torch.randn(1, 4, H, dtype=torch.float32)
+
+        ids_a = torch.tensor([[0, -1, 2, -1]], dtype=torch.int64)
+        ids_b = torch.tensor([[-1, 2, -1, 0]], dtype=torch.int64)
+
+        values_b = torch.randn(1, 4, H, dtype=torch.float32)
+        for slot_b in range(4):
+            eid = ids_b[0, slot_b].item()
+            if eid == -1:
+                continue
+            slot_a = (ids_a[0] == eid).nonzero(as_tuple=True)[0].item()
+            values_b[0, slot_b] = values_a[0, slot_a]
+
+        result_a = self._run_moe_reduce(values_a, ids_a, E=4)
+        result_b = self._run_moe_reduce(values_b, ids_b, E=4)
+        self.assertTrue(torch.equal(result_a, result_b))
+
+    def test_moe_bf16_large_hidden_invariance(self):
+        """Production-scale hidden dim (H=7168) with E=64 in BF16."""
+        torch.manual_seed(40)
+        H = 7168
+        topk = 8
+        E = 64
+        tokens = 2
+
+        expert_ids = torch.zeros(tokens, topk, dtype=torch.int64)
+        for t in range(tokens):
+            chosen = torch.randperm(E)[:topk]
+            expert_ids[t] = chosen
+
+        values = torch.randn(tokens, topk, H, dtype=torch.bfloat16)
+
+        output_a = torch.zeros(tokens, H, dtype=torch.bfloat16)
+        moe_sum_tree_reduce(
+            input=values,
+            output=output_a,
+            curr_topk_ids=expert_ids,
+            routed_scaling_factor=0.25,
+            E=E,
+        )
+
+        perm = torch.randperm(topk)
+        values_b = values[:, perm, :]
+        ids_b = expert_ids[:, perm]
+
+        output_b = torch.zeros(tokens, H, dtype=torch.bfloat16)
+        moe_sum_tree_reduce(
+            input=values_b,
+            output=output_b,
+            curr_topk_ids=ids_b,
+            routed_scaling_factor=0.25,
+            E=E,
+        )
+
+        self.assertTrue(torch.equal(output_a, output_b))
+
+    def test_moe_deterministic_two_runs(self):
+        """Same input, two calls -> bitwise identical."""
+        torch.manual_seed(50)
+        values = torch.randn(4, 4, 256, dtype=torch.bfloat16)
+        ids = torch.tensor(
+            [[0, 1, 2, 3], [3, 2, 1, 0], [0, 0, 1, 1], [2, 3, 0, 1]],
+            dtype=torch.int64,
+        )
+
+        out_a = torch.zeros(4, 256, dtype=torch.bfloat16)
+        moe_sum_tree_reduce(
+            input=values,
+            output=out_a,
+            curr_topk_ids=ids,
+            routed_scaling_factor=0.5,
+            E=4,
+        )
+
+        out_b = torch.zeros(4, 256, dtype=torch.bfloat16)
+        moe_sum_tree_reduce(
+            input=values,
+            output=out_b,
+            curr_topk_ids=ids,
+            routed_scaling_factor=0.5,
+            E=4,
+        )
+
+        self.assertTrue(torch.equal(out_a, out_b))
+
+
+# ---------------------------------------------------------------------------
+# tree_all_reduce_sum: non-distributed
+# ---------------------------------------------------------------------------
+class TestTreeAllReduceNonDistributed(unittest.TestCase):
+    def test_returns_clone_when_dist_not_initialized(self):
+        if dist.is_initialized():
+            self.skipTest("dist already initialized")
+        x = torch.tensor([1.0, 2.0, 3.0])
+        result = tree_all_reduce_sum(x)
+        self.assertTrue(torch.equal(x, result))
+        self.assertFalse(x.data_ptr() == result.data_ptr())
+
+    def test_fixed_tree_sum_is_order_deterministic(self):
+        torch.manual_seed(50)
+        world_size = 8
+        shards = [torch.randn(16, dtype=torch.bfloat16) for _ in range(world_size)]
+
+        result_a = _fixed_tree_sum_tensors(shards)
+        result_b = _fixed_tree_sum_tensors(list(shards))
+        self.assertTrue(torch.equal(result_a, result_b))
+
+    def test_tree_sum_power_of_two_sizes(self):
+        for n in [1, 2, 4, 8, 16]:
+            shards = [torch.tensor([float(i + 1)]) for i in range(n)]
+            result = _fixed_tree_sum_tensors(shards)
+            self.assertAlmostEqual(result.item(), n * (n + 1) / 2, places=5)
+
+
+# ---------------------------------------------------------------------------
+# MoE reduce: order-matters proof
+# ---------------------------------------------------------------------------
+class TestMoeReduceOrderMatters(unittest.TestCase):
+    def test_expert_order_differs_from_slot_order(self):
+        input_tensor = torch.tensor(
+            [[[1.0e16, 0.0], [1.0, 0.0], [-1.0e16, 0.0], [2.0, 0.0]]],
+            dtype=torch.float32,
+        )
+        curr_topk_ids = torch.tensor([[2, 0, 3, 1]], dtype=torch.int64)
+        output = torch.empty(1, 2, dtype=torch.float32)
+
+        moe_sum_tree_reduce(
+            input=input_tensor,
+            output=output,
+            curr_topk_ids=curr_topk_ids,
+            routed_scaling_factor=1.0,
+            E=4,
+        )
+
+        slot_order_sum = (
+            input_tensor[0, 0]
+            + input_tensor[0, 1]
+            + input_tensor[0, 2]
+            + input_tensor[0, 3]
+        )
+        self.assertNotEqual(output[0, 0].item(), slot_order_sum[0].item())
+
+
+# ---------------------------------------------------------------------------
+# Edge cases
+# ---------------------------------------------------------------------------
+class TestEdgeCases(unittest.TestCase):
+    def test_fixed_tree_sum_single_tensor(self):
+        t = torch.tensor([5.0])
+        result = _fixed_tree_sum_tensors([t])
+        self.assertEqual(result.item(), 5.0)
+
+    def test_fixed_tree_sum_odd_count(self):
+        values = [torch.tensor([1.0]), torch.tensor([2.0]), torch.tensor([3.0])]
+        result = _fixed_tree_sum_tensors(values)
+        self.assertEqual(result.item(), 6.0)
+
+    def test_fixed_tree_sum_empty_raises(self):
+        with self.assertRaises(ValueError):
+            _fixed_tree_sum_tensors([])
+
+    def test_matmul_tp_persistent_k_less_than_block(self):
+        torch.manual_seed(0)
+        A = torch.randn(3, 64, dtype=torch.float32)
+        B = torch.randn(64, 5, dtype=torch.float32)
+        actual = matmul_tp_persistent(A, B)
+        expected = A @ B
+        torch.testing.assert_close(actual, expected)
+
+    def test_matmul_tp_persistent_k_equals_block(self):
+        torch.manual_seed(0)
+        A = torch.randn(3, 128, dtype=torch.float32)
+        B = torch.randn(128, 5, dtype=torch.float32)
+        actual = matmul_tp_persistent(A, B)
+        expected = A @ B
+        torch.testing.assert_close(actual, expected)
+
+    def test_matmul_tp_persistent_k_multi_block_non_aligned(self):
+        """K not divisible by BLOCK_K: reference handles remainder gracefully."""
+        torch.manual_seed(0)
+        A = torch.randn(3, 300, dtype=torch.float32)
+        B = torch.randn(300, 5, dtype=torch.float32)
+        actual = matmul_tp_persistent(A, B)
+        expected = A @ B
+        torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
+
+    def test_moe_sum_tree_reduce_single_expert(self):
+        input_tensor = torch.tensor([[[1.0, 2.0]]], dtype=torch.float32)
+        curr_topk_ids = torch.tensor([[0]], dtype=torch.int64)
+        output = torch.empty(1, 2, dtype=torch.float32)
+
+        moe_sum_tree_reduce(
+            input=input_tensor,
+            output=output,
+            curr_topk_ids=curr_topk_ids,
+            routed_scaling_factor=1.0,
+            E=1,
+        )
+        expected = torch.tensor([[1.0, 2.0]], dtype=torch.float32)
+        torch.testing.assert_close(output, expected)
+
+    def test_moe_sum_tree_reduce_single_topk(self):
+        input_tensor = torch.tensor(
+            [[[10.0, 20.0]], [[30.0, 40.0]]],
+            dtype=torch.float32,
+        )
+        curr_topk_ids = torch.tensor([[1], [0]], dtype=torch.int64)
+        output = torch.empty(2, 2, dtype=torch.float32)
+
+        moe_sum_tree_reduce(
+            input=input_tensor,
+            output=output,
+            curr_topk_ids=curr_topk_ids,
+            routed_scaling_factor=0.5,
+            E=2,
+        )
+        expected = torch.tensor([[5.0, 10.0], [15.0, 20.0]], dtype=torch.float32)
+        torch.testing.assert_close(output, expected)
+
+    def test_moe_sum_tree_reduce_all_remote(self):
+        input_tensor = torch.tensor(
+            [[[99.0, 99.0], [99.0, 99.0]]],
+            dtype=torch.float32,
+        )
+        curr_topk_ids = torch.tensor([[-1, -1]], dtype=torch.int64)
+        output = torch.empty(1, 2, dtype=torch.float32)
+
+        moe_sum_tree_reduce(
+            input=input_tensor,
+            output=output,
+            curr_topk_ids=curr_topk_ids,
+            routed_scaling_factor=1.0,
+            E=2,
+        )
+        expected = torch.zeros(1, 2, dtype=torch.float32)
+        torch.testing.assert_close(output, expected)
+
+    def test_moe_sum_tree_reduce_duplicate_expert_ids(self):
+        """Same expert in multiple slots: both contributions must be summed."""
+        input_tensor = torch.tensor(
+            [[[10.0, 20.0], [30.0, 40.0]]],
+            dtype=torch.float32,
+        )
+        curr_topk_ids = torch.tensor([[0, 0]], dtype=torch.int64)
+        output = torch.empty(1, 2, dtype=torch.float32)
+
+        moe_sum_tree_reduce(
+            input=input_tensor,
+            output=output,
+            curr_topk_ids=curr_topk_ids,
+            routed_scaling_factor=1.0,
+            E=2,
+        )
+        expected = torch.tensor([[40.0, 60.0]], dtype=torch.float32)
+        torch.testing.assert_close(output, expected)
+
+
+# ---------------------------------------------------------------------------
+# Input validation
+# ---------------------------------------------------------------------------
+class TestInputValidation(unittest.TestCase):
+    def test_matmul_rejects_non_2d_inputs(self):
+        with self.assertRaisesRegex(ValueError, "expected 2D"):
+            matmul_tp_persistent(torch.randn(2, 3, 4), torch.randn(4, 5))
+        with self.assertRaisesRegex(ValueError, "expected 2D"):
+            matmul_tp_persistent(torch.randn(2, 3), torch.randn(3))
+
+    def test_matmul_rejects_dimension_mismatch(self):
+        with self.assertRaisesRegex(ValueError, "dimension mismatch"):
+            matmul_tp_persistent(torch.randn(2, 3), torch.randn(4, 5))
+
+    def test_moe_rejects_wrong_input_dims(self):
+        with self.assertRaisesRegex(ValueError, "tokens, topk, hidden"):
+            moe_sum_tree_reduce(
+                input=torch.zeros(4, 8),
+                output=torch.zeros(4, 8),
+                curr_topk_ids=torch.zeros(4, 2, dtype=torch.int64),
+                routed_scaling_factor=1.0,
+                E=2,
+            )
+
+    def test_moe_rejects_wrong_topk_ids_dims(self):
+        with self.assertRaisesRegex(ValueError, "tokens, topk"):
+            moe_sum_tree_reduce(
+                input=torch.zeros(4, 2, 8),
+                output=torch.zeros(4, 8),
+                curr_topk_ids=torch.zeros(4, dtype=torch.int64),
+                routed_scaling_factor=1.0,
+                E=2,
+            )
+
+    def test_moe_rejects_shape_mismatch_between_input_and_ids(self):
+        with self.assertRaisesRegex(ValueError, "must match"):
+            moe_sum_tree_reduce(
+                input=torch.zeros(4, 2, 8),
+                output=torch.zeros(4, 8),
+                curr_topk_ids=torch.zeros(4, 3, dtype=torch.int64),
+                routed_scaling_factor=1.0,
+                E=2,
+            )
+
+    def test_moe_rejects_wrong_output_shape(self):
+        with self.assertRaisesRegex(ValueError, "output must have shape"):
+            moe_sum_tree_reduce(
+                input=torch.zeros(4, 2, 8),
+                output=torch.zeros(4, 4),
+                curr_topk_ids=torch.zeros(4, 2, dtype=torch.int64),
+                routed_scaling_factor=1.0,
+                E=2,
+            )
+
+    def test_is_power_of_two_helper(self):
+        self.assertTrue(_is_power_of_two(1))
+        self.assertTrue(_is_power_of_two(2))
+        self.assertTrue(_is_power_of_two(64))
+        self.assertFalse(_is_power_of_two(0))
+        self.assertFalse(_is_power_of_two(3))
+        self.assertFalse(_is_power_of_two(6))
+
+
+# ---------------------------------------------------------------------------
+# Distributed tree all-reduce (multi-GPU only)
+# ---------------------------------------------------------------------------
+class TestDistributedTreeAllReduce(unittest.TestCase):
+    @unittest.skipUnless(
+        int(os.environ.get("WORLD_SIZE", "1")) > 1,
+        "requires torchrun with WORLD_SIZE > 1",
+    )
+    def test_tree_all_reduce_sum_distributed(self):
+        own_pg = False
+        if not dist.is_initialized():
+            backend = "nccl" if torch.cuda.is_available() else "gloo"
+            dist.init_process_group(backend=backend)
+            own_pg = True
+
+        try:
+            world_size = dist.get_world_size()
+            if world_size & (world_size - 1) != 0:
+                self.skipTest("requires power-of-two world size")
+
+            local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+            if torch.cuda.is_available():
+                torch.cuda.set_device(local_rank)
+                device = torch.device("cuda", local_rank)
+            else:
+                device = torch.device("cpu")
+
+            rank = dist.get_rank()
+            value = torch.full((4,), float(rank + 1), device=device)
+            actual = tree_all_reduce_sum(value)
+            expected = torch.full(
+                (4,),
+                float(world_size * (world_size + 1) // 2),
+                device=device,
+            )
+
+            torch.testing.assert_close(actual, expected)
+            dist.barrier()
+        finally:
+            if own_pg:
+                dist.destroy_process_group()
+
+    @unittest.skipUnless(
+        int(os.environ.get("WORLD_SIZE", "1")) > 1,
+        "requires torchrun with WORLD_SIZE > 1",
+    )
+    def test_tree_all_reduce_bf16_bitwise_deterministic(self):
+        """Run tree all-reduce twice with same inputs, verify bitwise identical."""
+        own_pg = False
+        if not dist.is_initialized():
+            backend = "nccl" if torch.cuda.is_available() else "gloo"
+            dist.init_process_group(backend=backend)
+            own_pg = True
+
+        try:
+            world_size = dist.get_world_size()
+            if world_size & (world_size - 1) != 0:
+                self.skipTest("requires power-of-two world size")
+
+            local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+            if torch.cuda.is_available():
+                torch.cuda.set_device(local_rank)
+                device = torch.device("cuda", local_rank)
+            else:
+                device = torch.device("cpu")
+
+            rank = dist.get_rank()
+            torch.manual_seed(rank * 1000 + 777)
+            value = torch.randn(256, device=device, dtype=torch.bfloat16)
+
+            result_a = tree_all_reduce_sum(value)
+            result_b = tree_all_reduce_sum(value)
+
+            self.assertTrue(torch.equal(result_a, result_b))
+            dist.barrier()
+        finally:
+            if own_pg:
+                dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    unittest.main()

From 726efd8df8161a38992c454960e35813f28d78ee Mon Sep 17 00:00:00 2001
From: Yueming Yuan <yym022502@gmail.com>
Date: Fri, 24 Jul 2026 21:34:05 -0700
Subject: [PATCH 07/43] [2/27] [sglang-miles] R3 (Rollout Routing Replay)
 DeepEP and MTP support (#18642)

---
 python/sglang/srt/managers/scheduler.py          |  1 +
 .../scheduler_components/output_streamer.py      |  3 +++
 python/sglang/srt/model_executor/model_runner.py | 16 +++++++++++++---
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
index c46a8c596812..2599fde3f96c 100644
--- a/python/sglang/srt/managers/scheduler.py
+++ b/python/sglang/srt/managers/scheduler.py
@@ -1853,6 +1853,7 @@ def init_output_streamer(self) -> None:
             ps=self.ps,
             server_args=self.server_args,
             is_generation=self.is_generation,
+            is_multimodal_gen=getattr(self.model_config, "is_multimodal_gen", False),
             spec_algorithm=self.spec_algorithm,
             disaggregation_mode=self.disaggregation_mode,
             enable_hicache_storage=lambda: self.enable_hicache_storage,
diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py
index 278ccf42818a..e03d172e8643 100644
--- a/python/sglang/srt/managers/scheduler_components/output_streamer.py
+++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py
@@ -42,6 +42,7 @@ class SchedulerOutputStreamer:
     ps: ParallelState
     server_args: ServerArgs
     is_generation: bool
+    is_multimodal_gen: bool
     spec_algorithm: SpeculativeAlgorithm
     disaggregation_mode: DisaggregationMode
     enable_hicache_storage: Callable[[], bool]
@@ -160,6 +161,8 @@ def _stream_output_generation(
             self._maybe_log_time_stats(req=req)
 
         # Send to detokenizer
+        if self.is_multimodal_gen:
+            return
         payload = acc.to_payload(
             dp_rank=self.ps.dp_rank,
             is_idle_batch=is_idle_batch,
diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py
index 12988d3188ff..85ae1ca8ebda 100644
--- a/python/sglang/srt/model_executor/model_runner.py
+++ b/python/sglang/srt/model_executor/model_runner.py
@@ -770,7 +770,8 @@ def _init_post_memory_pool_components(self):
 
         self.maybe_init_hisparse_coordinator()
 
-        self.init_routed_experts_capturer()
+        if not self.is_draft_worker:
+            self.init_routed_experts_capturer()
         self.init_indexer_capturer()
 
         self.graph_shared_output = None
@@ -1296,6 +1297,15 @@ def forward(
         output.expert_distribution_metrics = recorder_outputs.get("metrics")
 
         no_copy_to_cpu = not self.server_args.disable_overlap_schedule
+        # In speculative decoding, num_tokens_per_bs > 1, so pass the actual
+        # number of tokens per DP rank in CUDA graph, not the batch size.
+        cuda_graph_num_tokens = None
+        if getattr(self.decode_cuda_graph_runner, "bs", None):
+            cuda_graph_num_tokens = (
+                self.decode_cuda_graph_runner.bs
+                * self.decode_cuda_graph_runner.num_tokens_per_bs
+            )
+
         if (
             not self.is_draft_worker
             and (experts_capturer := get_global_experts_capturer()) is not None
@@ -1303,7 +1313,7 @@ def forward(
             output.routed_experts_output = experts_capturer.on_forward_end(
                 forward_batch=forward_batch,
                 can_run_graph=output.can_run_graph,
-                cuda_graph_batch=getattr(self.decode_cuda_graph_runner, "bs", None),
+                cuda_graph_batch=cuda_graph_num_tokens,
                 no_copy_to_cpu=no_copy_to_cpu,
             )
 
@@ -1311,7 +1321,7 @@ def forward(
             output.indexer_topk_output = indexer_capturer.on_forward_end(
                 forward_batch=forward_batch,
                 can_run_graph=output.can_run_graph,
-                cuda_graph_batch=getattr(self.decode_cuda_graph_runner, "bs", None),
+                cuda_graph_batch=cuda_graph_num_tokens,
                 no_copy_to_cpu=no_copy_to_cpu,
             )
 

From 23b0989ef44c3495c2c27d5361353a8f5e232ce3 Mon Sep 17 00:00:00 2001
From: Yueming Yuan <yym022502@gmail.com>
Date: Fri, 24 Jul 2026 21:34:05 -0700
Subject: [PATCH 08/43] [3/27] [sglang-miles] PD disaggregation for RL (#18646)

---
 python/sglang/srt/managers/schedule_batch.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py
index a8cea65ec910..868b25bf882d 100755
--- a/python/sglang/srt/managers/schedule_batch.py
+++ b/python/sglang/srt/managers/schedule_batch.py
@@ -2613,8 +2613,10 @@ def retract_decode(
         while first_iter or (
             not self.check_decode_mem(selected_indices=sorted_indices)
         ):
-            if len(sorted_indices) == 1:
-                # Always keep at least one request
+            # We should allow all requests to be retracted in decode disaggregation mode
+            # because there can be prealloc prefill requests.
+            num_minimum_reqs = 0 if server_args.disaggregation_mode == "decode" else 1
+            if len(sorted_indices) == num_minimum_reqs:
                 break
 
             first_iter = False

From 3a013290bf87406a713e7b1f7b9772a8c87f41fd Mon Sep 17 00:00:00 2001
From: Yueming Yuan <yym022502@gmail.com>
Date: Fri, 24 Jul 2026 21:34:17 -0700
Subject: [PATCH 09/43] [4/27] [sglang-miles] MTP related fix (#18647)

---
 python/sglang/srt/server_args.py                            | 4 ++++
 .../sglang/srt/speculative/eagle_draft_cuda_graph_runner.py | 6 ++++--
 python/sglang/srt/speculative/eagle_worker_v2.py            | 6 +++++-
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index 228cbb8136ed..c44a6f04b61a 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -1576,6 +1576,10 @@ class ServerArgs:
         "Disable the decode-phase CUDA graph. Convenience for --cuda-graph-backend-decode=disabled.",
     ] = False
     disable_cuda_graph: A[bool, Arg(no_cli=True)] = False
+    disable_draft_cuda_graph: A[
+        bool,
+        "Disable cuda graph for draft model in speculative decoding.",
+    ] = False
     disable_cuda_graph_padding: A[
         bool,
         "Disable cuda graph when padding is needed. Still uses cuda graph when padding is not needed.",
diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
index 5f1f46894e3d..df88600f72da 100644
--- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
+++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
@@ -566,8 +566,10 @@ def execute(self, forward_batch: ForwardBatch):
             forward_batch.seq_lens,
             forward_batch.out_cache_loc,
             forward_batch.positions,
-            forward_batch.spec_info.topk_p,
-            forward_batch.spec_info.topk_index,
+            forward_batch.spec_info.topk_p.clamp(0, 1),
+            forward_batch.spec_info.topk_index.clamp(
+                0, self.model_runner.model_config.vocab_size - 1
+            ),
             forward_batch.req_pool_indices,
         ]
         if buffers.rids_int is not None and forward_batch.rids_int is not None:
diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py
index e8493012cbc0..f14399aa0fd2 100644
--- a/python/sglang/srt/speculative/eagle_worker_v2.py
+++ b/python/sglang/srt/speculative/eagle_worker_v2.py
@@ -376,7 +376,11 @@ def _capture_cuda_graphs(self):
         self.cuda_graph_runner = None
         self.cuda_graph_runner_for_draft_extend = None
 
-        if _is_cpu or check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED):
+        if (
+            _is_cpu
+            or check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED)
+            or self.server_args.disable_draft_cuda_graph
+        ):
             return
 
         if self.server_args.model_impl == "mindspore":

From 502c4d5e3fe86327fc976ee0b4da300a8aa15a87 Mon Sep 17 00:00:00 2001
From: Yueming Yuan <yym022502@gmail.com>
Date: Fri, 24 Jul 2026 21:34:24 -0700
Subject: [PATCH 10/43] [5/27] [sglang-miles] VLM training multimodal fallback
 fixes (#18781)

---
 python/sglang/srt/multimodal/processors/qwen_vl.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/sglang/srt/multimodal/processors/qwen_vl.py b/python/sglang/srt/multimodal/processors/qwen_vl.py
index afcc64436663..5e62893c4588 100644
--- a/python/sglang/srt/multimodal/processors/qwen_vl.py
+++ b/python/sglang/srt/multimodal/processors/qwen_vl.py
@@ -683,7 +683,7 @@ async def process_mm_data_async(
         **kwargs,
     ):
         entry_time = time.perf_counter()
-        base_output = await self.load_mm_data(
+        base_output = await self.legacy_load_mm_data(
             prompt=input_text,
             image_data=image_data,
             video_data=request_obj.video_data,

From 9c3ea38494182a77abf2bd3922a63ec72484f0e5 Mon Sep 17 00:00:00 2001
From: maocheng23 <maocheng@berkeley.edu>
Date: Fri, 24 Jul 2026 21:34:24 -0700
Subject: [PATCH 11/43] [6/27] [sglang-miles] Fix pause-aware weight update
 deadlocks (#22754, #22623)

---
 python/sglang/srt/layers/quantization/fp8.py     | 16 ++++++++++++----
 python/sglang/srt/managers/scheduler.py          |  1 +
 .../scheduler_components/idle_sleeper.py         |  6 ++++--
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py
index 33c2f16f9308..c145af57b364 100644
--- a/python/sglang/srt/layers/quantization/fp8.py
+++ b/python/sglang/srt/layers/quantization/fp8.py
@@ -489,11 +489,19 @@ def validate_block_quant_shapes(
                         f"{input_size_per_partition} is not divisible by "
                         f"weight quantization block_k = {block_k}."
                     )
-            # Required by column parallel or enabling merged weights
-            if (
+            # Required by column parallel or enabling merged weights.
+            is_tp_split = (
                 tp_size > 1 and output_size // output_size_per_partition == tp_size
-            ) or len(output_partition_sizes) > 1:
-                for output_partition_size in output_partition_sizes:
+            )
+            is_merged_gemm = len(output_partition_sizes) > 1
+            if is_tp_split or is_merged_gemm:
+                sizes_to_check = output_partition_sizes
+                if not is_tp_split and is_merged_gemm:
+                    # Match validate_fp8_block_shape: merged weights may have a
+                    # ragged final logical matrix, and scale tensors are already
+                    # allocated with ceil-divided block counts.
+                    sizes_to_check = output_partition_sizes[:-1]
+                for output_partition_size in sizes_to_check:
                     if output_partition_size % block_n != 0:
                         raise ValueError(
                             f"Weight output_partition_size = "
diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
index 2599fde3f96c..fbcfa485a229 100644
--- a/python/sglang/srt/managers/scheduler.py
+++ b/python/sglang/srt/managers/scheduler.py
@@ -648,6 +648,7 @@ def init_idle_sleeper(self) -> None:
                     self.ipc_channels.recv_from_tokenizer,
                     self.ipc_channels.recv_from_rpc,
                 ],
+                can_empty_cache=lambda: not self._engine_paused,
             )
         else:
             self.idle_sleeper = None
diff --git a/python/sglang/srt/managers/scheduler_components/idle_sleeper.py b/python/sglang/srt/managers/scheduler_components/idle_sleeper.py
index 38c0c7a2ebdf..4af8f5cbb491 100644
--- a/python/sglang/srt/managers/scheduler_components/idle_sleeper.py
+++ b/python/sglang/srt/managers/scheduler_components/idle_sleeper.py
@@ -17,9 +17,10 @@ class IdleSleeper:
     data that needs handling immediately.
     """
 
-    def __init__(self, sockets):
+    def __init__(self, sockets, can_empty_cache=None):
         self.poller = zmq.Poller()
         self.last_empty_time = real_time()
+        self.can_empty_cache = can_empty_cache
         for s in sockets:
             self.poller.register(s, zmq.POLLIN)
 
@@ -32,4 +33,5 @@ def maybe_sleep(self):
             and real_time() - self.last_empty_time > self.empty_cache_interval
         ):
             self.last_empty_time = real_time()
-            current_platform.empty_cache()
+            if self.can_empty_cache is None or self.can_empty_cache():
+                current_platform.empty_cache()

From 6005dccce425b4dba8db35db25ae84f1c097472c Mon Sep 17 00:00:00 2001
From: zyzshishui <zyzshishui@gmail.com>
Date: Fri, 24 Jul 2026 21:34:24 -0700
Subject: [PATCH 12/43] [7/27] [sglang-miles] R3 support on PD disaggregation
 mini_lb (#22916)

---
 python/sglang/srt/disaggregation/prefill.py   |  2 +
 .../bindings/python/pyproject.toml            |  1 +
 .../python/src/sglang_router/mini_lb.py       | 54 +++++++++++++++----
 3 files changed, 47 insertions(+), 10 deletions(-)

diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py
index 1f4c947fa3fa..f9aa8351b8ab 100644
--- a/python/sglang/srt/disaggregation/prefill.py
+++ b/python/sglang/srt/disaggregation/prefill.py
@@ -817,6 +817,8 @@ def process_disagg_prefill_inflight_queue(
                 # todo: set Transferring correctly in backend
                 undone_reqs.append(req)
             elif poll == KVPoll.Success:  # transfer done
+                if req.return_routed_experts:
+                    self.batch_result_processor._maybe_collect_routed_experts(req)
                 release_kv_cache(req, self.tree_cache)  # unlock the tree
                 if not isinstance(req.finished_reason, FINISH_ABORT):
                     req.finished_reason = FINISH_LENGTH(length=0)
diff --git a/sgl-model-gateway/bindings/python/pyproject.toml b/sgl-model-gateway/bindings/python/pyproject.toml
index c44b1b96abb4..98a9fbe3d9de 100644
--- a/sgl-model-gateway/bindings/python/pyproject.toml
+++ b/sgl-model-gateway/bindings/python/pyproject.toml
@@ -32,6 +32,7 @@ dependencies = [
     "setproctitle",
     "aiohttp",
     "orjson",
+    "pybase64",
     "uvicorn",
     "fastapi",
 ]
diff --git a/sgl-model-gateway/bindings/python/src/sglang_router/mini_lb.py b/sgl-model-gateway/bindings/python/src/sglang_router/mini_lb.py
index abdbd95233a0..a9d7a23fe86d 100644
--- a/sgl-model-gateway/bindings/python/src/sglang_router/mini_lb.py
+++ b/sgl-model-gateway/bindings/python/src/sglang_router/mini_lb.py
@@ -14,6 +14,7 @@
 
 import aiohttp
 import orjson
+import pybase64
 import uvicorn
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import ORJSONResponse, Response, StreamingResponse
@@ -34,6 +35,44 @@ def maybe_wrap_ipv6_address(address: str) -> str:
         return address
 
 
+def _merge_routed_experts(prefill: dict, decode: dict):
+    if "routed_experts" not in prefill or "routed_experts" not in decode:
+        return False
+
+    prefill_bytes = pybase64.b64decode(prefill["routed_experts"], validate=True)
+    decode_bytes = pybase64.b64decode(decode["routed_experts"], validate=True)
+    decode["routed_experts"] = pybase64.b64encode(
+        prefill_bytes + decode_bytes[len(prefill_bytes) :]
+    ).decode("utf-8")
+    return True
+
+
+def _merge_input_token_logprobs(prefill_meta: dict, decode_meta: dict):
+    if (
+        "input_token_logprobs" not in prefill_meta
+        or "input_token_logprobs" not in decode_meta
+    ):
+        return
+
+    decode_meta["input_token_logprobs"] = (
+        prefill_meta["input_token_logprobs"] + decode_meta["input_token_logprobs"]
+    )
+
+
+def _merge_prefill_json(prefill_json, decode_json):
+    if "meta_info" in prefill_json and "meta_info" in decode_json:
+        prefill_meta = prefill_json["meta_info"]
+        decode_meta = decode_json["meta_info"]
+        _merge_input_token_logprobs(prefill_meta, decode_meta)
+        _merge_routed_experts(prefill_meta, decode_meta)
+
+    if "sglext" not in prefill_json:
+        return
+
+    if "sglext" in decode_json:
+        _merge_routed_experts(prefill_json["sglext"], decode_json["sglext"])
+
+
 class MiniLoadBalancer:
     def __init__(
         self,
@@ -140,18 +179,13 @@ async def generate(
             # Wait for both responses to complete. Prefill should end first.
             prefill_response, decode_response = await asyncio.gather(*tasks)
 
-            if "return_logprob" in modified_request:
-
+            if (
+                "return_logprob" in modified_request
+                or "return_routed_experts" in modified_request
+            ):
                 prefill_json = await prefill_response.json()
                 ret_json = await decode_response.json()
-
-                # merge `meta_info.input_token_logprobs` from prefill to decode
-                if "meta_info" in ret_json:
-                    if "input_token_logprobs" in ret_json["meta_info"]:
-                        ret_json["meta_info"]["input_token_logprobs"] = (
-                            prefill_json["meta_info"]["input_token_logprobs"]
-                            + ret_json["meta_info"]["input_token_logprobs"]
-                        )
+                _merge_prefill_json(prefill_json, ret_json)
             else:
                 ret_json = await decode_response.json()
 

From 4014c32ec6a279dc4090862421b5ce7cba7eac2b Mon Sep 17 00:00:00 2001
From: Byron Hsu <byronhsu1230@gmail.com>
Date: Fri, 24 Jul 2026 21:34:47 -0700
Subject: [PATCH 13/43] [8/27] [sglang-miles] Improve PD pause handling
 (#23672, #23887)

---
 python/sglang/srt/managers/scheduler.py            | 14 ++++++++++++++
 .../managers/test_scheduler_pause_generation.py    |  1 +
 2 files changed, 15 insertions(+)

diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
index fbcfa485a229..e0b379dedd14 100644
--- a/python/sglang/srt/managers/scheduler.py
+++ b/python/sglang/srt/managers/scheduler.py
@@ -4166,6 +4166,20 @@ def _pause_engine(self) -> Tuple[List[Req], int]:
 
     def pause_generation(self, recv_req: PauseGenerationReqInput):
         assert recv_req.mode in ("in_place", "retract")
+        # PD disaggregation: `retract` has no decode-to-prefill rebootstrap path,
+        # so retracted decode-side requests cannot be re-prefilled under new
+        # weights. Fail fast before mutating any scheduler state to avoid
+        # leaving the engine in a half-paused / inconsistent state that later
+        # crashes inside radix-cache cleanup on flush_cache.
+        assert not (
+            recv_req.mode == "retract"
+            and self.disaggregation_mode != DisaggregationMode.NULL
+        ), (
+            "pause_generation(mode='retract') is not supported in PD "
+            "disaggregation mode yet. Decode-side retracted requests need "
+            "a rebootstrap path back to prefill."
+        )
+
         self._engine_paused = True
 
         if recv_req.mode == "in_place":
diff --git a/test/registered/unit/managers/test_scheduler_pause_generation.py b/test/registered/unit/managers/test_scheduler_pause_generation.py
index e9a24cef5d2e..384179fc9d41 100644
--- a/test/registered/unit/managers/test_scheduler_pause_generation.py
+++ b/test/registered/unit/managers/test_scheduler_pause_generation.py
@@ -34,6 +34,7 @@ def _new_scheduler(self) -> Scheduler:
         scheduler.last_batch = None
         scheduler.cur_batch_for_debug = None
         scheduler.chunked_req = None
+        scheduler.disaggregation_mode = DisaggregationMode.NULL
         scheduler.running_batch = MagicMock()
         scheduler.running_batch.reqs = []
         scheduler.running_batch.is_empty.return_value = True

From e08222b82e925dc74728425ecd15c00cab6e28d5 Mon Sep 17 00:00:00 2001
From: Jiajun Li <48857426+guapisolo@users.noreply.github.com>
Date: Fri, 24 Jul 2026 21:34:47 -0700
Subject: [PATCH 14/43] [9/27] [sglang-miles] Add KimiK2 raw tool call id
 parser (#25196)

---
 .../srt/entrypoints/openai/serving_chat.py    | 17 +++++++++----
 python/sglang/srt/function_call/core_types.py |  6 +++++
 .../srt/function_call/function_call_parser.py |  6 ++++-
 .../srt/function_call/kimik2_detector.py      | 24 +++++++++++++++++++
 4 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py
index 021d222834aa..142ac5d4909a 100644
--- a/python/sglang/srt/entrypoints/openai/serving_chat.py
+++ b/python/sglang/srt/entrypoints/openai/serving_chat.py
@@ -1663,11 +1663,7 @@ def _process_tool_call_id(
         history_tool_calls_cnt: int,
     ) -> str:
         """Process for generating a new and unique `tool_call_id`"""
-        if self.tool_call_parser != "kimi_k2":
-            # A simple uuid is sufficient for all models except for Kimi-K2.
-            tool_call_id = f"call_{uuid.uuid4().hex[:24]}"
-            return tool_call_id
-        else:
+        if self.tool_call_parser == "kimi_k2":
             # Align with Kimi-K2 format: functions.{name}:{index}
             # Kimi-K2 allows multiple tool_calls in one message; SGLang sets call_item.tool_index to the *local* position inside that message.
             # Therefore, the index must be corrected by using `history_tool_calls_cnt + call_item.tool_index` to ensure globally unique and properly ordered.
@@ -1676,6 +1672,17 @@ def _process_tool_call_id(
                 f"Process tool call idx, parser: {self.tool_call_parser}, tool_call_id: {tool_call_id}, history_cnt: {history_tool_calls_cnt}"
             )
             return tool_call_id
+        if self.tool_call_parser == "kimi_k2_raw_id":
+            # RL training needs the model-emitted tool_call_id round-tripped verbatim,
+            # so we skip the history-based renumbering above and return whatever the
+            # detector captured. Fall back to the canonical Kimi-K2 reconstruction
+            # (without history offset) if for any reason the detector did not record
+            # a raw id — the raw id field is best-effort but the format is stable.
+            if call_item.tool_call_id:
+                return call_item.tool_call_id
+            return f"functions.{call_item.name}:{call_item.tool_index}"
+        # A simple uuid is sufficient for all other models.
+        return f"call_{uuid.uuid4().hex[:24]}"
 
     def _process_tool_calls(
         self,
diff --git a/python/sglang/srt/function_call/core_types.py b/python/sglang/srt/function_call/core_types.py
index 1ea87df798c8..297dd2712ef3 100644
--- a/python/sglang/srt/function_call/core_types.py
+++ b/python/sglang/srt/function_call/core_types.py
@@ -10,6 +10,12 @@ class ToolCallItem(BaseModel):
     tool_index: int
     name: Optional[str] = None
     parameters: str  # JSON string
+    # The tool_call_id string emitted by the model, captured verbatim.
+    # Only populated by detectors whose downstream consumers need the exact
+    # model-emitted id (e.g. RL training trajectories). Existing detectors
+    # leave this as None and the serving layer falls back to its usual id
+    # generation strategy.
+    tool_call_id: Optional[str] = None
 
 
 class StreamingParseResult(BaseModel):
diff --git a/python/sglang/srt/function_call/function_call_parser.py b/python/sglang/srt/function_call/function_call_parser.py
index 0263d6e8477d..507d5c9f4c75 100644
--- a/python/sglang/srt/function_call/function_call_parser.py
+++ b/python/sglang/srt/function_call/function_call_parser.py
@@ -28,7 +28,10 @@
 from sglang.srt.function_call.hunyuan_detector import HunyuanDetector
 from sglang.srt.function_call.inkling_detector import InklingDetector
 from sglang.srt.function_call.internlm_detector import InternlmDetector
-from sglang.srt.function_call.kimik2_detector import KimiK2Detector
+from sglang.srt.function_call.kimik2_detector import (
+    KimiK2Detector,
+    KimiK2RawIdDetector,
+)
 from sglang.srt.function_call.lfm2_detector import Lfm2Detector
 from sglang.srt.function_call.llama32_detector import Llama32Detector
 from sglang.srt.function_call.mimo_detector import MiMoDetector
@@ -71,6 +74,7 @@ class FunctionCallParser:
         "glm47": Glm47MoeDetector,
         "gpt-oss": GptOssDetector,
         "kimi_k2": KimiK2Detector,
+        "kimi_k2_raw_id": KimiK2RawIdDetector,
         "lfm2": Lfm2Detector,
         "llama3": Llama32Detector,
         "mimo": MiMoDetector,
diff --git a/python/sglang/srt/function_call/kimik2_detector.py b/python/sglang/srt/function_call/kimik2_detector.py
index 11909fa6ffe5..417c2d3cf84a 100644
--- a/python/sglang/srt/function_call/kimik2_detector.py
+++ b/python/sglang/srt/function_call/kimik2_detector.py
@@ -197,6 +197,7 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult
                         tool_index=local_tool_index,
                         name=function_name,
                         parameters=function_args,
+                        tool_call_id=function_id,
                     )
                 )
                 local_tool_index += 1
@@ -328,6 +329,9 @@ def parse_streaming_increment(
                                 else None
                             ),
                             parameters=argument_diff,
+                            # Capture the model-emitted id on the name-carrying
+                            # delta so kimi_k2_raw_id can round-trip it (RL).
+                            tool_call_id=function_id if name_just_resolved else None,
                         )
                     )
                     if argument_diff:
@@ -467,3 +471,23 @@ def get_structural_tag(
 
     def get_structural_tag_name(self) -> str:
         return "kimi"
+
+
+class KimiK2RawIdDetector(KimiK2Detector):
+    """
+    Variant of KimiK2Detector that preserves the model-emitted tool_call_id verbatim.
+
+    The default kimi_k2 path renumbers ids via `history_tool_calls_cnt + tool_index`
+    in the serving layer so that multi-turn conversations get globally unique,
+    monotonically increasing ids (see PR #10600). That is the right behavior for
+    chat use cases.
+
+    RL training has the opposite requirement: the trajectory must round-trip the
+    exact tool_call_id the model produced (e.g. `functions.foo:5`), so that the
+    follow-up tool result turn references the same id the policy emitted. This
+    subclass exists purely as a marker so the serving layer can branch on the
+    parser name and use `ToolCallItem.tool_call_id` directly. Parsing logic
+    is identical to KimiK2Detector.
+    """
+
+    pass

From 496d8200315548a60d394dcf97c85f704fbaed98 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Fri, 24 Jul 2026 21:34:47 -0700
Subject: [PATCH 15/43] [10/27] Fix GLM4 MoE Lite shared expert TP flag

---
 python/sglang/srt/models/glm4_moe_lite.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/python/sglang/srt/models/glm4_moe_lite.py b/python/sglang/srt/models/glm4_moe_lite.py
index c650585a58b1..0ed6b272be85 100644
--- a/python/sglang/srt/models/glm4_moe_lite.py
+++ b/python/sglang/srt/models/glm4_moe_lite.py
@@ -251,6 +251,11 @@ def __init__(
         if config.n_shared_experts is not None and self.num_fused_shared_experts == 0:
             intermediate_size = config.moe_intermediate_size * config.n_shared_experts
             # disable tp for shared experts when enable deepep moe, or with fp4 allgather
+            _shared_expert_use_tp1 = (
+                get_moe_a2a_backend().is_deepep()
+                or get_moe_a2a_backend().is_mooncake()
+                or should_use_flashinfer_cutlass_moe_fp4_allgather()
+            )
             self.shared_experts = Glm4MoeLiteMLP(
                 hidden_size=config.hidden_size,
                 intermediate_size=intermediate_size,
@@ -258,14 +263,9 @@ def __init__(
                 quant_config=quant_config,
                 reduce_results=False,
                 prefix=add_prefix("shared_experts", prefix),
-                **(
-                    dict(tp_rank=0, tp_size=1)
-                    if get_moe_a2a_backend().is_deepep()
-                    or get_moe_a2a_backend().is_mooncake()
-                    or should_use_flashinfer_cutlass_moe_fp4_allgather()
-                    else {}
-                ),
+                **(dict(tp_rank=0, tp_size=1) if _shared_expert_use_tp1 else {}),
             )
+            self._shared_expert_tp1 = _shared_expert_use_tp1
             is_packed_weight = hasattr(
                 self.shared_experts.gate_up_proj.quant_method, "quant_config"
             )

From 448997ac3c1c497beef025610db83d457ea90e87 Mon Sep 17 00:00:00 2001
From: Nan Jiang <59716405+nanjiangwill@users.noreply.github.com>
Date: Fri, 24 Jul 2026 21:36:42 -0700
Subject: [PATCH 16/43] [11/27] [sglang-miles] MoE-LoRA: kimi 2.5/2.6, GLM-5.2
 FP8, cuda-graph/RL fixes (#25141, #29874, #31251)

Squashes the two follow-up PRs into the MoE-LoRA base commit: both rewrite the
mem_pool sharding helpers this commit introduces and cannot be applied
independently.

Rebased onto v0.5.16:
  - the shard probes keep v0.5.16's shared-MoE-over-full-TP carve-out at EP=1
    (`is_shared_moe_module`) and only fall through to the probed shard for
    non-MoE modules;
  - #29874's `free_lora` is dropped -- v0.5.16 already releases the slot on
    unload through `LoRAMemoryPool.remove_lora()`, which additionally zeroes the
    buffers for graph-captured replay, so `lora_manager.unload` needs no second
    call;
  - #31831's expected_checksums check is likewise already upstream in
    `tp_worker.load_lora_adapter_from_tensors`.

Co-authored-by: Ethan (Yusheng) Su <yushengsu@radixark.ai>
---
 .../sglang/kernels/ops/moe/virtual_experts.py | 154 +++++++----
 python/sglang/srt/configs/model_config.py     |   5 +-
 python/sglang/srt/entrypoints/engine.py       |  11 +-
 .../sglang/srt/lora/backend/base_backend.py   | 131 ++++++++-
 python/sglang/srt/lora/layers.py              |  78 ++++--
 python/sglang/srt/lora/lora_manager.py        |  69 ++++-
 python/sglang/srt/lora/lora_moe_runners.py    |   7 +-
 python/sglang/srt/lora/mem_pool.py            | 255 ++++++++++++++++--
 python/sglang/srt/managers/io_struct.py       |   2 +-
 .../srt/managers/tokenizer_control_mixin.py   |  16 +-
 python/sglang/srt/managers/tp_worker.py       |  16 +-
 .../srt/model_executor/forward_batch_info.py  |   8 +
 .../unit/lora/test_moe_lora_tail_stamp.py     | 134 +++++++++
 13 files changed, 759 insertions(+), 127 deletions(-)
 create mode 100644 test/registered/unit/lora/test_moe_lora_tail_stamp.py

diff --git a/python/sglang/kernels/ops/moe/virtual_experts.py b/python/sglang/kernels/ops/moe/virtual_experts.py
index 23fc8355b6dc..eb2fbaf109cd 100644
--- a/python/sglang/kernels/ops/moe/virtual_experts.py
+++ b/python/sglang/kernels/ops/moe/virtual_experts.py
@@ -515,7 +515,7 @@ def _merged_experts_fused_moe_lora_add_impl(
     output: torch.Tensor,
     hidden_states: torch.Tensor,
     lora_a: torch.Tensor,
-    lora_b: torch.Tensor,
+    lora_b: torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor, ...],
     topk_ids: torch.Tensor,
     topk_weights: torch.Tensor,
     token_lora_mapping: torch.Tensor,
@@ -524,13 +524,30 @@ def _merged_experts_fused_moe_lora_add_impl(
     experts_shared_outer_loras_b: bool,
     routing_cache: dict | None = None,
 ) -> None:
+    """Fused virtual-experts LoRA delta add.
+
+    ``lora_b`` accepts either a single tensor or a sequence of tensors stacked
+    along the output dim. Length-2 is the gate_up case where A has rank ``2*r``
+    (gate's A and up's A concatenated along rank) and each B has rank ``r``.
+    The shrink runs once over the full ``2*r`` rank; the expand runs once per
+    B, each reading its half of the intermediate and writing to its slice of
+    ``output``.
     """
-    1. Prepare virtual expert routing metadata from topk_ids + token_lora_mapping * num_experts.
-    2. Flatten LoRA weights from [max_loras, num_experts, ...] to [max_loras * num_experts, ...].
-    3. Run regular SGLang fused-MoE kernels for LoRA A and LoRA B.
-    4. Mask out tokens with token_lora_mapping == -1 on the add path.
-    """
+    lora_b_list: list[torch.Tensor] = (
+        list(lora_b) if isinstance(lora_b, (list, tuple)) else [lora_b]
+    )
+    n_b = len(lora_b_list)
+    assert n_b in (1, 2), f"lora_b must be length 1 or 2, got {n_b}"
+    b_rank = lora_b_list[0].shape[3]
+    for b in lora_b_list[1:]:
+        assert (
+            b.shape == lora_b_list[0].shape
+        ), f"all lora_b tensors must share shape; got {[tuple(t.shape) for t in lora_b_list]}"
+
     max_loras, _, max_lora_rank, _ = lora_a.shape
+    assert (
+        max_lora_rank == n_b * b_rank
+    ), f"lora_a rank {max_lora_rank} != n_b ({n_b}) * lora_b rank {b_rank}"
     input_top_k = 1 if hidden_states.shape[0] == topk_ids.numel() else topk_ids.shape[1]
 
     def _merge_lora_expert_weight(t: torch.Tensor) -> torch.Tensor:
@@ -614,16 +631,18 @@ def _get_routing(
             block_size=block_size,
             num_experts=virtual_num_experts,
         )
-        # _align_block_size uses a worst-case padded allocation. Trim the routing buffers
-        # to a tighter upper bound so we keep the real routed work but drop unused padding
-        num_tokens = topk_ids.numel()
-        max_nonempty = min(num_tokens, virtual_num_experts)
-        tight_padded = (
-            triton.cdiv(num_tokens + max_nonempty * (block_size - 1), block_size)
-            * block_size
-        )
-        sorted_token_ids = sorted_token_ids[:tight_padded]
-        expert_ids = expert_ids[: tight_padded // block_size]
+        # NOTE: do NOT trim sorted_token_ids / expert_ids to a tighter upper bound here.
+        # The downstream kernels (_moe_lora_shrink_splitk_kernel, fused_moe_kernel) read
+        # sorted_token_ids[pid_m*BLOCK : +BLOCK] and expert_ids[pid_m] WITHOUT a bounds mask
+        # for every block up to num_tokens_post_padded (a GPU-side count loaded at run time).
+        # num_tokens_post_padded comes from _align_block_size with `virtual_num_experts` buckets
+        # and can exceed a tighter `numel + min(numel,virtual_num_experts)*(block-1)` bound
+        # (most so for shared-outer, where virtual_num_experts = max_loras is small), so trimming
+        # made those unmasked reads land PAST the view. In eager mode the slack still lives inside
+        # the same _align_block_size allocation (garbage, masked out downstream) so it worked; under
+        # CUDA-graph capture/replay the graph mempool packs tensors tightly and that slack may belong
+        # to another pooled tensor / lie past a page -> cudaErrorIllegalInstruction during capture.
+        # Keep the full worst-case-allocated buffers so every unmasked read stays in-allocation.
         expert_ids = fused_sanitize_expert_ids(expert_ids, virtual_num_experts)
         result = (
             sorted_token_ids,
@@ -642,12 +661,22 @@ def _get_routing(
     )
 
     lora_a_virtual = _merge_lora_expert_weight(lora_a)
-    lora_b_virtual = _merge_lora_expert_weight(lora_b)
+    lora_b_virtuals = [_merge_lora_expert_weight(b) for b in lora_b_list]
     num_experts_a = lora_a.shape[1]
-    num_experts_b = lora_b.shape[1]
-
+    num_experts_b = lora_b_list[0].shape[1]
+    half_out = lora_b_list[0].shape[2]
+
+    # The kernels index token_lora_mapping / intermediate by token ids up to
+    # topk_ids.shape[0] (the DP-gathered token count under --enable-dp-attention). An
+    # under-sized mapping means unmasked OOB reads/writes that surface as a sticky,
+    # hard-to-attribute CUDA IMA — fail loudly on the host instead.
+    assert token_lora_mapping.shape[0] >= topk_ids.shape[0], (
+        f"token_lora_mapping covers {token_lora_mapping.shape[0]} tokens but the MoE runs on "
+        f"{topk_ids.shape[0]} (DP-gathered?) tokens; mapping was sized before the dp gather "
+        f"length was known (see get_gathered_moe_num_tokens)"
+    )
     intermediate = torch.zeros(
-        [token_lora_mapping.shape[0], topk_ids.shape[1], max_lora_rank],
+        [topk_ids.shape[0], topk_ids.shape[1], max_lora_rank],
         dtype=hidden_states.dtype,
         device=hidden_states.device,
     )
@@ -678,7 +707,7 @@ def _get_routing(
         a_stage_config,
     )
 
-    b_stage_config = _get_stage_config(lora_b_virtual, 1)
+    b_stage_config = _get_stage_config(lora_b_virtuals[0], 1)
     (
         sorted_token_ids,
         expert_ids,
@@ -692,33 +721,53 @@ def _get_routing(
         b_stage_config["BLOCK_SIZE_M"],
     )
 
-    invoke_fused_moe_kernel(
-        intermediate.view(-1, max_lora_rank),
-        lora_b_virtual,
-        None,
-        output,
-        None,
-        None,
-        None,
-        topk_weights,
-        topk_ids,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_padded,
-        mul_routed_weight,
-        1,
-        b_stage_config,
-        tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16,
-        False,
-        False,
-        False,
-        False,
-        False,
-        None,
-        fuse_add_to_output=True,
-        add_output_mask=token_lora_mask,
-        router_topk=topk_ids.shape[1],
-    )
+    # n_b expands. For len 1: K=b_rank covers full intermediate, write full output.
+    # For len 2 (gate_up): split intermediate along rank into [gate, up] halves
+    # (each contiguous, K=b_rank=r) and output along last dim into [gate, up]
+    # halves (each of width half_out). Each B in lora_b_virtuals is its own
+    # half's weight tensor, naturally K=b_rank.
+    for b_idx, b_virtual in enumerate(lora_b_virtuals):
+        if n_b == 1:
+            inter_arg = intermediate.view(-1, b_rank)
+            out_arg = output
+        else:
+            inter_arg = (
+                intermediate[..., b_idx * b_rank : (b_idx + 1) * b_rank]
+                .contiguous()
+                .view(-1, b_rank)
+            )
+            out_arg = output[
+                ..., b_idx * half_out : (b_idx + 1) * half_out
+            ].contiguous()
+        invoke_fused_moe_kernel(
+            inter_arg,
+            b_virtual,
+            None,
+            out_arg,
+            None,
+            None,
+            None,
+            topk_weights,
+            topk_ids,
+            sorted_token_ids,
+            expert_ids,
+            num_tokens_post_padded,
+            mul_routed_weight,
+            1,
+            b_stage_config,
+            tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16,
+            False,
+            False,
+            False,
+            False,
+            False,
+            None,
+            fuse_add_to_output=True,
+            add_output_mask=token_lora_mask,
+            router_topk=topk_ids.shape[1],
+        )
+        if n_b != 1:
+            output[..., b_idx * half_out : (b_idx + 1) * half_out].copy_(out_arg)
 
 
 def _merged_experts_fused_moe_lora_add_op(
@@ -761,7 +810,7 @@ def merged_experts_fused_moe_lora_add(
     output: torch.Tensor,
     hidden_states: torch.Tensor,
     lora_a: torch.Tensor,
-    lora_b: torch.Tensor,
+    lora_b: torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor, ...],
     topk_ids: torch.Tensor,
     topk_weights: torch.Tensor,
     token_lora_mapping: torch.Tensor,
@@ -770,7 +819,12 @@ def merged_experts_fused_moe_lora_add(
     experts_shared_outer_loras_b: bool,
     routing_cache: dict | None = None,
 ) -> None:
-    """Public API: wraps the registered op with routing_cache support."""
+    """Public API: wraps the registered op with routing_cache support.
+
+    ``lora_b`` accepts a sequence of length 2 for the gate_up case (each B
+    holds one half of the stacked output, rank ``r``, with A's rank ``2*r``);
+    a single tensor is used for the down case.
+    """
     _merged_experts_fused_moe_lora_add_impl(
         output,
         hidden_states,
diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py
index 020995054c2e..63b4cac0d6ef 100644
--- a/python/sglang/srt/configs/model_config.py
+++ b/python/sglang/srt/configs/model_config.py
@@ -207,7 +207,10 @@ def dsa_layer_skips_topk(config: PretrainedConfig, layer_id: int) -> bool:
 
 
 def get_dsa_index_n_heads(config: PretrainedConfig) -> int:
-    assert is_deepseek_dsa(config)
+    # Permit both DSA (V3.2-family) and V4: both carry the indexer (index_n_heads) and this must
+    # match get_dsa_index_head_dim's contract, else LoRA buffer init for indexer.wq_b /
+    # indexer.weights_proj on a V4 model asserts here while indexer.wk (which uses head_dim) succeeds.
+    assert is_deepseek_dsa(config) or is_deepseek_v4(config)
     return config.index_n_heads
 
 
diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py
index 5c158535e5da..1e0b27dc357b 100644
--- a/python/sglang/srt/entrypoints/engine.py
+++ b/python/sglang/srt/entrypoints/engine.py
@@ -1143,15 +1143,16 @@ def load_lora_adapter_from_tensors(
         load_format: Optional[str] = None,
     ):
         if load_format == "flattened_bucket":
-            serialized_tensors = tensors
+            serialized_named_tensors = list(tensors)
         else:
-            serialized_tensors = MultiprocessingSerializer.serialize(
-                tensors, output_str=True
-            )
+            serialized_named_tensors = [
+                MultiprocessingSerializer.serialize(tensors, output_str=True)
+                for _ in range(self.server_args.tp_size)
+            ]
         lora_req = LoadLoRAAdapterFromTensorsReqInput(
             lora_name=lora_name,
             config_dict=config_dict,
-            serialized_tensors=serialized_tensors,
+            serialized_named_tensors=serialized_named_tensors,
             load_format=load_format,
         )
         return self.loop.run_until_complete(
diff --git a/python/sglang/srt/lora/backend/base_backend.py b/python/sglang/srt/lora/backend/base_backend.py
index 16879b547435..45443b5da440 100644
--- a/python/sglang/srt/lora/backend/base_backend.py
+++ b/python/sglang/srt/lora/backend/base_backend.py
@@ -7,6 +7,38 @@
 from sglang.srt.lora.backend.lmhead_mixing import LoRABackendLmHeadMixing
 from sglang.srt.lora.utils import LoRABatchInfo, MoELoRABatchInfo
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.utils.common import ceil_align
+
+
+def get_gathered_moe_num_tokens(forward_batch: ForwardBatch, num_tokens: int) -> int:
+    """Token count the MoE-LoRA mapping must cover: gathered under --enable-dp-attention, else per-rank.
+
+    In the eager path prepare_lora_batch runs from ForwardBatch.init_new, BEFORE
+    prepare_mlp_sync_batch assigns forward_batch.global_dp_buffer_len (only the cuda-graph
+    capture path pre-sets it), so that field alone under-sizes the MoE-LoRA token mapping to the
+    per-rank length and the MoE-LoRA kernels index past it (sticky CUDA IMA). When it is unset,
+    derive an upper bound of the gathered length from global_num_tokens_cpu (assigned in
+    init_new before prepare_lora_batch runs), mirroring prepare_mlp_sync_batch's attn-tp/cp
+    alignment; max*n covers both SUM_LEN and MAX_LEN padding modes. Over-allocation is harmless:
+    the kernels index at most the actual gathered length.
+    """
+    if forward_batch.global_dp_buffer_len is not None:
+        return max(forward_batch.global_dp_buffer_len, num_tokens)
+    global_num_tokens = getattr(forward_batch, "global_num_tokens_cpu", None)
+    if not global_num_tokens:
+        return num_tokens
+    from sglang.srt.layers.dp_attention import get_attention_tp_size
+
+    # Local import: a module-level cp_utils import here is circular (see forward_batch_info).
+    from sglang.srt.layers.utils.cp_utils import get_cp_padding_align_size
+
+    attn_tp_size = get_attention_tp_size()
+    cp_align_size = get_cp_padding_align_size()
+    upper = max(
+        ceil_align(ceil_align(t, attn_tp_size), cp_align_size)
+        for t in global_num_tokens
+    ) * len(global_num_tokens)
+    return max(upper, num_tokens)
 
 
 class BaseLoRABackend(LoRABackendLmHeadMixing):
@@ -22,6 +54,7 @@ class BaseLoRABackend(LoRABackendLmHeadMixing):
     def __init__(self, max_loras_per_batch: int, device: torch.device):
         self.max_loras_per_batch = max_loras_per_batch
         self.device = device
+        self.batch_info = None
         self.init_lm_head_config()
         self._is_moe_lora = False
 
@@ -188,15 +221,30 @@ def init_cuda_graph_moe_buffers(
         """
         base = moe_layer.base_layer
         top_k = base.top_k
-        qinfo = moe_layer._quant_info
-        E, N, _ = qinfo.w13_weight.shape
-        hidden_dim = qinfo.w2_weight.shape[1]
-        device = qinfo.w13_weight.device
+        # Derive dims from the base FusedMoE rather than quant-specific tensors,
+        # so this works for any scheme (FP, WNA16, Marlin-packed, etc.).
+        E = base.num_local_experts
+        hidden_dim = base.hidden_size
+        N = 2 * base.intermediate_size_per_partition
+        device = next(base.parameters()).device
         dtype = compute_dtype
         num_experts = base.num_experts
 
+        # Under --enable-dp-attention the MoE runs on DP-GATHERED tokens (the global DP buffer of
+        # length up to max_bs * attn_dp_size), not the per-rank batch. The per-token LoRA routing
+        # buffers below are indexed by that gathered token count, so size them for the gathered
+        # maximum; otherwise the MoE-LoRA kernels read token_lora_mapping / sorted_token_ids past a
+        # per-rank-sized buffer -> cudaErrorIllegalInstruction during cuda-graph capture. Expert- and
+        # adapter-indexed buffers (cumsum_buffer, adapter_enabled, lora_ids) are unaffected.
+        from sglang.srt.layers.dp_attention import get_attention_dp_size
+
+        dp_size = max(1, get_attention_dp_size())
+        max_moe_tokens = max_bs * dp_size
+
         block_size_m = 64
-        max_num_tokens_padded = max_bs * top_k + num_experts * (block_size_m - 1)
+        max_num_tokens_padded = max_moe_tokens * top_k + num_experts * (
+            block_size_m - 1
+        )
         max_num_tokens_padded = (
             (max_num_tokens_padded + block_size_m - 1) // block_size_m
         ) * block_size_m
@@ -233,7 +281,7 @@ def init_cuda_graph_moe_buffers(
             # LongTensor.  weight_indices itself must stay int32 because the
             # CUDA moe_lora_align kernel casts it to int32_t*.
             "weight_indices_long": torch.zeros(
-                max_bs, dtype=torch.int64, device=device
+                max_moe_tokens, dtype=torch.int64, device=device
             ),
             "lora_ids": torch.arange(max_loras, dtype=torch.int32, device=device),
             "cumsum_buffer": torch.zeros(
@@ -242,14 +290,14 @@ def init_cuda_graph_moe_buffers(
                 device=device,
             ),
             "token_mask": torch.empty(
-                (max_loras * max_bs * top_k,),
+                (max_loras * max_moe_tokens * top_k,),
                 dtype=torch.int32,
                 device=device,
             ),
             "max_num_tokens_padded": max_num_tokens_padded,
             "max_num_m_blocks": max_num_m_blocks,
             "token_lora_mapping": torch.full(
-                (max_bs,), -1, dtype=torch.int32, device=device
+                (max_moe_tokens,), -1, dtype=torch.int32, device=device
             ),
         }
 
@@ -291,6 +339,19 @@ def _add_moe_lora_info(
             seg_indptr = batch_info.seg_indptr[: num_moe_segments + 1]
             req_to_lora = batch_info.weight_indices[:num_moe_segments]
 
+        # --enable-dp-attention all-gathers tokens into the MoE, so the MoE-LoRA kernels index
+        # token_lora_mapping by the GATHERED token count, not the per-rank num_tokens. Size the
+        # mapping to (an upper bound of) the gathered count so those reads stay in-bounds — see
+        # get_gathered_moe_num_tokens for why global_dp_buffer_len alone is NOT enough in the
+        # eager path (the per-rank segments still fill only [0, num_tokens); the tail stays -1).
+        moe_num_tokens = get_gathered_moe_num_tokens(forward_batch, num_tokens)
+        if batch_info.use_cuda_graph:
+            # Static capture buffers hold max_bs*dp tokens; a REAL replay's gathered length never
+            # exceeds that (the captured graph could not address it), so cap the upper bound at
+            # the buffer size. Batches whose gathered bound exceeds it are demoted to the eager
+            # prep path in LoRAManager.prepare_lora_batch before we get here.
+            moe_num_tokens = min(moe_num_tokens, token_lora_mapping.shape[0])
+
         adapter_enabled, token_lora_mapping = _compute_moe_lora_info(
             num_tokens,
             seg_indptr,
@@ -299,8 +360,34 @@ def _add_moe_lora_info(
             adapter_enabled,
             token_lora_mapping,
             max_len=max_len,
+            mapping_len=moe_num_tokens,
         )
 
+        # Tier-1 (colocate RL, exactly one adapter loaded): the DP-gathered tail
+        # [num_tokens, moe_num_tokens) covers OTHER dp ranks' tokens, which under colocate RL all
+        # use that single adapter. The per-rank fill above only wrote [0, num_tokens), leaving the
+        # tail -1 (adapter-disabled) -> cross-rank gathered tokens would miss the LoRA delta on
+        # this rank's local experts. The stamps below are ARMED BY LoRAManager.prepare_lora_batch
+        # (policy lives there; both are host ints, no GPU sync -> cuda-graph safe, written each
+        # batch in the eager prep path):
+        #   * _single_loaded_buffer_id: armed only when exactly one adapter is loaded AND this
+        #     rank's local requests actively use it -> stamp the tail with that adapter.
+        #   * _idle_rank_active_buffer_id: armed only on a true idle rank (num_tokens == 0) ->
+        #     _compute_moe_lora_info left the whole mapping -1 / adapter_enabled all-0, so stamp
+        #     the whole gathered buffer and enable the adapter.
+        # Multi-adapter batches and base-only local batches arm neither stamp and keep the -1
+        # tail: foreign tokens get base rather than a delta this rank cannot attribute.
+        if moe_num_tokens > num_tokens:
+            if num_tokens > 0:
+                single_bid = getattr(self, "_single_loaded_buffer_id", None)
+                if single_bid is not None:
+                    token_lora_mapping[num_tokens:moe_num_tokens].fill_(single_bid)
+            else:
+                idle_bid = getattr(self, "_idle_rank_active_buffer_id", None)
+                if idle_bid is not None:
+                    token_lora_mapping.fill_(idle_bid)
+                    adapter_enabled[idle_bid] = 1
+
         batch_info.moe_lora_info = MoELoRABatchInfo(
             seg_indptr=seg_indptr,
             req_to_lora=req_to_lora,
@@ -373,16 +460,28 @@ def _compute_moe_lora_info(
     adapter_enabled: torch.Tensor | None,
     token_lora_mapping: torch.Tensor | None,
     max_len: int,
+    mapping_len: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
+    # ``num_tokens`` is the PER-RANK fill count (segments cover this rank's tokens). ``mapping_len``
+    # is the length of the token_lora_mapping the MoE-LoRA kernels actually index -- under
+    # --enable-dp-attention the MoE runs on DP-GATHERED tokens (mapping_len = global_dp_buffer_len
+    # >= num_tokens), so the returned mapping must span the gathered count to keep those kernels
+    # in-bounds. The DP-gathered tail [num_tokens, mapping_len) defaults to -1 (adapter-disabled).
+    if mapping_len is None:
+        mapping_len = num_tokens
+    assert mapping_len >= num_tokens
     if token_lora_mapping is not None:
         assert (
-            num_tokens <= token_lora_mapping.shape[0]
-        ), "num_tokens must be less than or equal to the shape of token_lora_mapping"
-        token_lora_mapping = token_lora_mapping[:num_tokens]
+            mapping_len <= token_lora_mapping.shape[0]
+        ), "mapping_len must be less than or equal to the shape of token_lora_mapping"
+        token_lora_mapping = token_lora_mapping[:mapping_len]
     else:
         token_lora_mapping = torch.empty(
-            (num_tokens,), dtype=torch.int32, device=seg_indptr.device
+            (mapping_len,), dtype=torch.int32, device=seg_indptr.device
         )
+    if mapping_len > num_tokens:
+        # clean the gathered tail before the per-rank fill writes [0, num_tokens)
+        token_lora_mapping.fill_(-1)
 
     if adapter_enabled is not None:
         assert (
@@ -440,8 +539,12 @@ def _compute_moe_lora_info(
         torch.searchsorted(seg_indptr.to(torch.int32), token_positions, right=True) - 1
     )
 
-    token_lora_mapping = torch.index_select(
-        weight_indices.to(torch.int32), 0, req_indices, out=token_lora_mapping
+    # Fill only the per-rank prefix [0, num_tokens); the gathered tail keeps the -1 set above.
+    torch.index_select(
+        weight_indices.to(torch.int32),
+        0,
+        req_indices,
+        out=token_lora_mapping[:num_tokens],
     )
 
     return adapter_enabled, token_lora_mapping
diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py
index 33c685caa0b6..fb1a2d6c70a5 100644
--- a/python/sglang/srt/lora/layers.py
+++ b/python/sglang/srt/lora/layers.py
@@ -45,6 +45,15 @@ def __init__(
             self.weight = self.base_layer.weight
         if hasattr(self.base_layer, "bias") and self.base_layer.bias is not None:
             self.bias = self.base_layer.bias
+        if hasattr(self.base_layer, "reduce_results"):
+            self.reduce_results = self.base_layer.reduce_results
+        # Alias remaining base-layer parameters onto the wrapper so
+        # `named_parameters(remove_duplicate=True)` yields them at the outer
+        # path — weight loaders (e.g. FusedMoE's `w13_weight_packed`) lookup
+        # names without the `.base_layer.` segment.
+        for _name, _param in base_layer.named_parameters(recurse=False):
+            if not hasattr(self, _name):
+                setattr(self, _name, _param)
 
     def forward(self, x: torch.Tensor):
         return self.base_layer.forward(x)
@@ -211,8 +220,9 @@ def forward(self, input_: torch.Tensor):
         ):
             base_output = self.extra_token_embedding(input_, base_output)
 
-        # Apply LoRA if configured
-        if self.set_lora:
+        # Apply LoRA if configured. Skip if no batch_info (DP-attention idle
+        # forward): the base path is correct because no real tokens need LoRA.
+        if self.set_lora and self.lora_backend.batch_info is not None:
             # The backend's run_lora_a_embedding now handles both regular
             # and extra tokens efficiently with CUDA graph support
             base_output = self.apply_lora(base_output, input_, batch_info)
@@ -377,8 +387,8 @@ def forward(self, hidden_states: torch.Tensor):
             hidden_states, self.weight, bias=getattr(self.base_layer, "bias", None)
         )
 
-        # Apply LoRA if set
-        if self.set_lora:
+        # Apply LoRA if set. Skip in DP-attention idle forward (batch_info unset).
+        if self.set_lora and self.lora_backend.batch_info is not None:
             base_output = self.apply_lora(base_output, hidden_states)
 
         return base_output
@@ -467,7 +477,7 @@ def forward(self, input_: torch.Tensor):
             self.base_layer, input_, bias
         )
 
-        if self.set_lora:
+        if self.set_lora and self.lora_backend.batch_info is not None:
             output_parallel = self.apply_lora(output_parallel, input_)
 
         if self.base_layer.gather_output:
@@ -481,9 +491,13 @@ def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
         return A
 
     def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+        # See RowParallelLinearWithLoRA.slice_lora_a_weights for why base_layer.tp_rank
+        # is authoritative: DP-attention makes output_partition_sizes attn_tp-local while
+        # the caller passes global tp_rank.
+        local_tp_rank = getattr(self.base_layer, "tp_rank", tp_rank)
         shard_size = self.base_layer.output_partition_sizes[0]
-        start_idx = tp_rank * shard_size
-        end_idx = (tp_rank + 1) * shard_size
+        start_idx = local_tp_rank * shard_size
+        end_idx = (local_tp_rank + 1) * shard_size
         B = B[start_idx:end_idx, :]
         return B
 
@@ -577,12 +591,15 @@ def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
         return A
 
     def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
+        # base_layer.tp_rank is authoritative under DP-attention: the caller passes
+        # the global tp_rank but output_partition_sizes is attn_tp-local.
+        local_tp_rank = getattr(self.base_layer, "tp_rank", tp_rank)
         partition_sizes = self.base_layer.output_partition_sizes
         output_sizes = self.base_layer.output_sizes
         slices = []
         offset = 0
         for full_size, part_size in zip(output_sizes, partition_sizes):
-            start_idx = tp_rank * part_size
+            start_idx = local_tp_rank * part_size
             end_idx = start_idx + part_size
             slices.append(B[offset + start_idx : offset + end_idx, :])
             offset += full_size
@@ -686,11 +703,14 @@ def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int) -> torch.Tensor:
         q_proj_shard_size = base_layer.q_proj_shard_size
         kv_proj_shard_size = base_layer.kv_proj_shard_size
         num_kv_head_replicas = base_layer.num_kv_head_replicas
+        # See RowParallelLinearWithLoRA.slice_lora_a_weights for why base_layer.tp_rank
+        # is authoritative under DP-attention.
+        local_tp_rank = getattr(base_layer, "tp_rank", tp_rank)
 
-        q_start_idx = q_proj_shard_size * tp_rank
+        q_start_idx = q_proj_shard_size * local_tp_rank
         q_end_idx = q_start_idx + q_proj_shard_size
 
-        kv_shard_id = tp_rank // num_kv_head_replicas
+        kv_shard_id = local_tp_rank // num_kv_head_replicas
         kv_start_idx = kv_proj_shard_size * kv_shard_id
         kv_end_idx = kv_start_idx + kv_proj_shard_size
 
@@ -773,7 +793,9 @@ def forward(self, input_: torch.Tensor, skip_all_reduce=False, forward_batch=Non
             and not should_skip_mlp_all_reduce()
         )
 
-        if self.set_lora and should_reduce:
+        # LoRA skipped when batch_info is None (DP-attention idle forward).
+        have_batch_info = self.lora_backend.batch_info is not None
+        if self.set_lora and have_batch_info and should_reduce:
             lora_a_output = self.lora_backend.run_lora_a_sgemm(
                 input_parallel, self.A_buffer
             )
@@ -787,7 +809,7 @@ def forward(self, input_: torch.Tensor, skip_all_reduce=False, forward_batch=Non
                 base_output=output_,
             )
         else:
-            if self.set_lora:
+            if self.set_lora and have_batch_info:
                 output_parallel = self.apply_lora(output_parallel, input_parallel)
             if should_reduce:
                 output_ = tensor_model_parallel_all_reduce(output_parallel)
@@ -798,9 +820,15 @@ def forward(self, input_: torch.Tensor, skip_all_reduce=False, forward_batch=Non
         return output_, output_bias
 
     def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
+        # Use base_layer.tp_rank (not the argument) so the slicing rank matches
+        # the partition group the base layer was built on. For MLA o_proj under
+        # DP-attention, base_layer.tp_rank is attn_tp_rank while the caller
+        # passes the global tp_rank; input_size_per_partition is already
+        # attn_tp-sized, so using global tp_rank overshoots to empty.
+        local_tp_rank = getattr(self.base_layer, "tp_rank", tp_rank)
         shard_size = self.base_layer.input_size_per_partition
-        start_idx = tp_rank * shard_size
-        end_idx = (tp_rank + 1) * shard_size
+        start_idx = local_tp_rank * shard_size
+        end_idx = (local_tp_rank + 1) * shard_size
         A = A[:, start_idx:end_idx].contiguous()
         return A
 
@@ -886,7 +914,7 @@ def apply_lora(self, base_output: torch.Tensor, x: torch.Tensor) -> torch.Tensor
     def forward(self, x: torch.Tensor):
         bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None
         output = self.base_layer.quant_method.apply(self.base_layer, x, bias)
-        if self.set_lora:
+        if self.set_lora and self.lora_backend.batch_info is not None:
             output = self.apply_lora(output, x)
         output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
         return output, output_bias
@@ -922,6 +950,13 @@ def __init__(
 
         self.experts_shared_outer_loras: bool = False
         self.lora_use_virtual_experts: bool = False
+        # Forward for the model's own forward-path dispatch — the outer model
+        # reads several FusedMoE attributes (e.g. `self.experts.moe_runner_config`,
+        # `self.experts.dispatcher`, `self.experts.num_local_experts`,
+        # `self.experts.quant_method`) directly on the wrapper. Quant
+        # post-processing iterators skip LoRA wrappers via
+        # `isinstance(module, BaseLayerWithLoRA)` so the packed params on the
+        # inner FusedMoE get processed there, not here.
         self.quant_method = base_layer.quant_method
         self.moe_runner_config = base_layer.moe_runner_config
         self.dispatcher = base_layer.dispatcher
@@ -929,6 +964,8 @@ def __init__(
         self.should_fuse_routed_scaling_factor_in_topk = (
             base_layer.should_fuse_routed_scaling_factor_in_topk
         )
+        if hasattr(base_layer, "scheme"):
+            self.scheme = base_layer.scheme
 
         self.tp_size = getattr(base_layer, "moe_tp_size", 1)
         self.tp_rank = getattr(base_layer, "moe_tp_rank", 0)
@@ -953,6 +990,12 @@ def __init__(
             and base_layer.quant_method.runner is not None
         ):
             runner_backend = base_layer.quant_method.runner.runner_backend
+        elif (
+            hasattr(base_layer, "scheme")
+            and hasattr(base_layer.scheme, "runner")
+            and base_layer.scheme.runner is not None
+        ):
+            runner_backend = base_layer.scheme.runner.runner_backend
         else:
             runner_backend = MoeRunnerBackend.TRITON
 
@@ -1085,6 +1128,8 @@ def forward(self, hidden_states: torch.Tensor, topk_output: TopKOutput, **kwargs
         1. After gate_up projection, before activation
         2. After down projection, before final reduction
         """
+        if self.lora_backend.batch_info is None:
+            return self.base_layer.forward(hidden_states, topk_output, **kwargs)
 
         # Build LoRA info for this batch
         lora_info = self._get_lora_info()
@@ -1112,6 +1157,9 @@ def _forward_with_lora(
 
         # Use pre-computed quant info (doesn't change so not sure why we need to pass it in every time)
         quant_info = self._quant_info
+        quant_info.expert_map = getattr(
+            base_layer.dispatcher, "local_expert_mapping", None
+        )
 
         # ===== TO BE REFACTORED ====
         if self._lora_runner_backend.is_experimental_sgl_trtllm():
diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py
index 498cf4787aab..8aa8e453f99e 100644
--- a/python/sglang/srt/lora/lora_manager.py
+++ b/python/sglang/srt/lora/lora_manager.py
@@ -374,6 +374,22 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch):
             and bs <= self.max_bs_in_cuda_graph
             and forward_batch.forward_mode.is_cuda_graph()
         )
+        if use_cuda_graph and self.lora_backend.is_moe_lora:
+            # This flag is a HEURISTIC computed before the runner's real replay decision. Under
+            # --enable-dp-attention a batch that is graph-eligible on THIS rank (small local bs)
+            # can still have a DP-gathered token count exceeding the static MoE capture buffers
+            # (max_bs*dp) when other ranks carry more tokens — such a batch cannot replay the
+            # captured graph and runs eager, so prep the eager (freshly sized) buffers for it
+            # instead of an under-sized static mapping.
+            from sglang.srt.lora.backend.base_backend import get_gathered_moe_num_tokens
+
+            moe_cg_buffers = getattr(self.lora_backend, "moe_cg_buffers", None)
+            if (
+                moe_cg_buffers is not None
+                and get_gathered_moe_num_tokens(forward_batch, bs)
+                > moe_cg_buffers["token_lora_mapping"].shape[0]
+            ):
+                use_cuda_graph = False
 
         weight_indices = [0] * len(forward_batch.lora_ids)
         lora_ranks = [0] * self.max_loras_per_batch
@@ -386,6 +402,55 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch):
                 lora = self.loras[uid]
                 lora_ranks[weight_indices[i]] = lora.config.r
                 scalings[weight_indices[i]] = lora.scaling
+
+        local_active = any(lora_ranks[wi] > 0 for wi in weight_indices)
+
+        # MoE-expert LoRA under --enable-dp-attention: the MoE runs on DP-GATHERED tokens, so this
+        # rank's local experts also process OTHER ranks' tokens, whose adapter identity is not known
+        # host-side. Under Tier-1 (exactly one adapter loaded — the colocate RL case) the gathered
+        # tail can be attributed to that single adapter. Two mutually-exclusive stamps, both armed
+        # here (policy) and mechanically applied by the backend in _add_moe_lora_info:
+        #   * idle rank (forward_mode.is_idle(), zero local tokens): inject the adapter into
+        #     lora_ranks/scalings (nothing else carries them into the batch tensors) and record its
+        #     buffer id so the backend stamps the WHOLE gathered mapping and enables the adapter —
+        #     otherwise foreign tokens routed to this rank's experts silently lose the LoRA delta.
+        #   * active rank (local requests DO use the adapter): record the buffer id so the backend
+        #     stamps the gathered tail [num_tokens, moe_num_tokens) with it instead of copying an
+        #     arbitrary local token's slot. Base-only local batches (local_active False) and
+        #     multi-adapter batches arm NOTHING: the tail stays -1 (adapter-disabled), foreign
+        #     tokens get base — never a delta this rank cannot attribute.
+        # Buffer ids are host ints (no GPU sync -> cuda-graph safe). NOTE: gate on
+        # global_num_tokens_cpu too, not just global_dp_buffer_len — the eager path runs
+        # prepare_lora_batch from ForwardBatch.init_new BEFORE prepare_mlp_sync_batch assigns
+        # global_dp_buffer_len (only cuda-graph capture pre-sets it).
+        idle_rank_active_buffer_id = None
+        single_active_buffer_id = None
+        if self.lora_backend.is_moe_lora and (
+            getattr(forward_batch, "global_dp_buffer_len", None) is not None
+            or len(getattr(forward_batch, "global_num_tokens_cpu", None) or []) > 1
+        ):
+            loaded = [
+                (uid, bid)
+                for uid, bid in self.memory_pool.uid_to_buffer_id.items()
+                if uid is not None
+                and uid in self.loras
+                and self.loras[uid].config.r > 0
+            ]
+            if len(loaded) == 1:
+                uid, bid = loaded[0]
+                if forward_batch.forward_mode.is_idle():
+                    lora = self.loras[uid]
+                    lora_ranks[bid] = lora.config.r
+                    scalings[bid] = lora.scaling
+                    idle_rank_active_buffer_id = bid
+                elif local_active:
+                    single_active_buffer_id = bid
+
+        # Pass the stamps to the backend (read in _add_moe_lora_info); reset each batch so a
+        # previous batch's stamp never leaks into a later one.
+        self.lora_backend._idle_rank_active_buffer_id = idle_rank_active_buffer_id
+        self.lora_backend._single_loaded_buffer_id = single_active_buffer_id
+
         # Do in-place updates when CUDA graph is enabled and the batch forward mode
         # could use CUDA graph.
         self.lora_backend.prepare_lora_batch(
@@ -395,8 +460,8 @@ def prepare_lora_batch(self, forward_batch: ForwardBatch):
             scalings=scalings,
             use_cuda_graph=use_cuda_graph,
         )
-        self.lora_backend.batch_info.has_active_lora = any(
-            lora_ranks[wi] > 0 for wi in weight_indices
+        self.lora_backend.batch_info.has_active_lora = local_active or (
+            idle_rank_active_buffer_id is not None
         )
 
     def update_lora_info(self):
diff --git a/python/sglang/srt/lora/lora_moe_runners.py b/python/sglang/srt/lora/lora_moe_runners.py
index 927e429859f9..15a70066cd61 100644
--- a/python/sglang/srt/lora/lora_moe_runners.py
+++ b/python/sglang/srt/lora/lora_moe_runners.py
@@ -328,7 +328,6 @@ def _add_lora_gate_up_delta(
     r = lora_info.max_lora_rank
     gate_up_a = lora_info.gate_up_lora_a_weights
     gate_up_b = lora_info.gate_up_lora_b_weights
-
     if lora_info.experts_shared_outer_loras and not lora_info.lora_use_virtual_experts:
         gate_up_a = gate_up_a.expand(-1, lora_info.num_experts, -1, -1)
 
@@ -338,6 +337,8 @@ def _add_lora_gate_up_delta(
     if is_gated:
         inter_size = gate_up_b.shape[2] // 2
         lora_a_stacked = [gate_up_a[:, :, :r, :], gate_up_a[:, :, r : 2 * r, :]]
+        # B halves are also the tuple form the virtual-experts kernel wants
+        # (one shrink at K=2*r, two expands at K=r each).
         lora_b_stacked = [
             gate_up_b[:, :, :inter_size, :],
             gate_up_b[:, :, inter_size:, :],
@@ -351,7 +352,7 @@ def _add_lora_gate_up_delta(
             output=intermediate_cache,
             hidden_states=hidden_states,
             lora_a=gate_up_a,
-            lora_b=gate_up_b,
+            lora_b=tuple(lora_b_stacked) if is_gated else gate_up_b,
             topk_ids=topk_ids,
             topk_weights=topk_weights,
             token_lora_mapping=token_lora_mapping,
@@ -418,6 +419,8 @@ def _add_lora_down_delta(
     down_lora_a = lora_info.down_lora_a_weights
     down_lora_b = lora_info.down_lora_b_weights
     if lora_info.experts_shared_outer_loras and not lora_info.lora_use_virtual_experts:
+        # fused_moe_lora requires B's expert_dim to match A's; expand the
+        # shared B view.
         down_lora_b = down_lora_b.expand(-1, lora_info.num_experts, -1, -1)
 
     if lora_info.fully_sharded and lora_info.tp_size > 1:
diff --git a/python/sglang/srt/lora/mem_pool.py b/python/sglang/srt/lora/mem_pool.py
index 1a09e038597a..b099baab30f1 100644
--- a/python/sglang/srt/lora/mem_pool.py
+++ b/python/sglang/srt/lora/mem_pool.py
@@ -353,6 +353,157 @@ def _iter_local_expert_weights(
             f"Expected dict or 3D torch.Tensor, got {type(weights).__name__}."
         )
 
+    def _row_parallel_shard_tp(
+        self, module_name: str, base_model: torch.nn.Module, layer_idx: int
+    ) -> int:
+        """Shard count for a non-MoE row-parallel module's activation axis.
+
+        Probes the base module's ``input_size // input_size_per_partition`` so
+        the LoRA buffer matches the actual shard regardless of which TP group
+        owns it — covers DP-attention (``o_proj`` uses ``attn_tp_size``) and
+        shared-expert dense-vs-MoE per-layer-TP differences. Falls back to
+        ``self.tp_size``. Cached per ``(module_name, layer_idx)``.
+
+        MoE-internal names go through ``self.moe_tp_size`` upstream.
+        """
+        cache = getattr(self, "_row_parallel_tp_cache", None)
+        if cache is None:
+            cache = {}
+            setattr(self, "_row_parallel_tp_cache", cache)
+        key = (module_name, layer_idx)
+        if key in cache:
+            return cache[key]
+
+        layer_markers = (f".layers.{layer_idx}.", f"layers.{layer_idx}.")
+
+        def _probe(m):
+            in_size = getattr(m, "input_size", None)
+            per_part = getattr(m, "input_size_per_partition", None)
+            if in_size is not None and per_part is not None and per_part > 0:
+                return max(1, in_size // per_part)
+            inner = getattr(m, "base_layer", None)
+            if inner is not None and inner is not m:
+                return _probe(inner)
+            return None
+
+        suffix = f".{module_name}"
+        found = None
+        for _name, module in base_model.named_modules():
+            if not _name.endswith(suffix):
+                continue
+            if not any(marker in _name for marker in layer_markers):
+                continue
+            r = _probe(module)
+            if r is not None:
+                found = r
+                break
+
+        out = found if found is not None else self.tp_size
+        cache[key] = out
+        return out
+
+    def _column_parallel_shard_tp(
+        self, module_name: str, base_model: torch.nn.Module, layer_idx: int
+    ) -> int:
+        """Shard count for a non-MoE column-parallel module's output axis.
+
+        Probes the base module's ``output_size // output_size_per_partition``.
+        The input-axis probe used for row-parallel modules is poisoned here:
+        quantized linear methods (e.g. ``Fp8LinearMethod.create_weights``)
+        stamp ``input_size_per_partition == input_size`` on column-parallel
+        layers (whose input is never sharded), which reports a shard count of
+        1 and leaves the LoRA-B buffer at the full output dim while the base
+        stays TP-sharded -- ``set_lora_info`` then fails with "LoRA B output
+        dim != base partition prefix dim" (e.g. blockwise-fp8 GLM-5.2
+        ``shared_experts.gate_up_proj``). The output-axis ratio is
+        quantization-independent: ``output_size_per_partition`` is set in
+        ``ColumnParallelLinear.__init__`` before the quant method runs. Falls
+        back to ``self.tp_size``. Cached per ``(module_name, layer_idx)``.
+        """
+        cache = getattr(self, "_col_parallel_tp_cache", None)
+        if cache is None:
+            cache = {}
+            setattr(self, "_col_parallel_tp_cache", cache)
+        key = (module_name, layer_idx)
+        if key in cache:
+            return cache[key]
+
+        layer_markers = (f".layers.{layer_idx}.", f"layers.{layer_idx}.")
+
+        def _probe(m):
+            out_size = getattr(m, "output_size", None)
+            per_part = getattr(m, "output_size_per_partition", None)
+            if out_size is not None and per_part is not None and per_part > 0:
+                return max(1, out_size // per_part)
+            inner = getattr(m, "base_layer", None)
+            if inner is not None and inner is not m:
+                return _probe(inner)
+            return None
+
+        suffix = f".{module_name}"
+        found = None
+        for _name, module in base_model.named_modules():
+            if not _name.endswith(suffix):
+                continue
+            if not any(marker in _name for marker in layer_markers):
+                continue
+            r = _probe(module)
+            if r is not None:
+                found = r
+                break
+
+        out = found if found is not None else self.tp_size
+        cache[key] = out
+        return out
+
+    def _column_parallel_out_partition(
+        self, module_name: str, base_model: torch.nn.Module, layer_idx: int
+    ):
+        """Actual per-rank output dim of a non-MoE column-parallel base module.
+
+        Reads ``output_size_per_partition`` from the matching base linear -- the
+        ground truth that ``set_lora_info`` validates against. Critically handles
+        the dense MLP ``gate_up_proj`` that is fully REPLICATED under
+        ``--moe-dense-tp-size 1`` (``output_size_per_partition == output_size``),
+        where dividing ``get_hidden_dim``'s output by the global ``tp_size``
+        undersizes LoRA-B and raises "LoRA B output dim != base partition prefix
+        dim". Returns ``None`` if no base module is found. Cached per
+        ``(module_name, layer_idx)``.
+        """
+        cache = getattr(self, "_col_parallel_out_cache", None)
+        if cache is None:
+            cache = {}
+            setattr(self, "_col_parallel_out_cache", cache)
+        key = (module_name, layer_idx)
+        if key in cache:
+            return cache[key]
+
+        layer_markers = (f".layers.{layer_idx}.", f"layers.{layer_idx}.")
+
+        def _probe(m):
+            ops = getattr(m, "output_size_per_partition", None)
+            if ops is not None and ops > 0:
+                return ops
+            inner = getattr(m, "base_layer", None)
+            if inner is not None and inner is not m:
+                return _probe(inner)
+            return None
+
+        suffix = f".{module_name}"
+        found = None
+        for _name, module in base_model.named_modules():
+            if not _name.endswith(suffix):
+                continue
+            if not any(marker in _name for marker in layer_markers):
+                continue
+            r = _probe(module)
+            if r is not None:
+                found = r
+                break
+
+        cache[key] = found
+        return found
+
     def _get_standard_shape(
         self,
         module_name: str,
@@ -365,8 +516,12 @@ def _get_standard_shape(
             module_name, self.base_hf_config, base_model, layer_idx
         )
         c = get_stacked_multiply(module_name, base_model)
-        if self.tp_size > 1 and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES:
-            input_dim = divide(input_dim, self.tp_size)
+        # Non-MoE row-parallel modules: probe the actual shard size so o_proj /
+        # down_proj match attn_tp under DP-attention and the shared-experts
+        # dense-vs-MoE per-layer-TP differences.
+        row_tp = self._row_parallel_shard_tp(module_name, base_model, layer_idx)
+        if row_tp > 1 and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES:
+            input_dim = divide(input_dim, row_tp)
         return (self.max_loras_per_batch, max_lora_dim * c, input_dim)
 
     def get_lora_A_shape(
@@ -387,12 +542,17 @@ def get_lora_A_shape(
             module_name, self.base_hf_config, base_model, layer_idx
         )
         c = get_stacked_multiply(module_name, base_model)
-        # Routed MoE shards along moe_tp_size; shared MoE shards over full TP at EP=1.
+        # Routed MoE shards along moe_tp_size; shared MoE shards over full TP
+        # at EP=1. Non-MoE row-parallel modules use a probed shard that may be
+        # attn_tp under DP-attention.
         effective_tp_size = (
-            self.tp_size
-            if not self.is_moe_module(module_name)
-            or self.is_shared_moe_module(module_name)
-            else self.moe_tp_size
+            (
+                self.tp_size
+                if self.is_shared_moe_module(module_name)
+                else self.moe_tp_size
+            )
+            if self.is_moe_module(module_name)
+            else self._row_parallel_shard_tp(module_name, base_model, layer_idx)
         )
         if (
             effective_tp_size > 1
@@ -490,21 +650,48 @@ def get_lora_B_shape(
         _, output_dim = get_hidden_dim(
             module_name, self.base_hf_config, base_model, layer_idx
         )
-        # Same TP-vs-moe-TP sharding rule as get_lora_A_shape above.
+        # Routed MoE shards along moe_tp_size; shared MoE shards over full TP
+        # at EP=1. Non-MoE column-parallel modules probe the OUTPUT axis
+        # (quantization-independent). This used to call `_row_parallel_shard_tp`,
+        # which was a latent bug: its input-axis probe is meaningless for
+        # column-parallel layers (their input is never sharded) and only worked on
+        # bf16 by falling through to the `tp_size` fallback -- quantized linear
+        # methods stamp `input_size_per_partition == input_size`, so the probe
+        # reported the layer as unsharded and the B buffer was allocated at the
+        # full output dim (set_lora_info crash / oversized buffer).
         effective_tp_size = (
-            self.tp_size
-            if not self.is_moe_module(module_name)
-            or self.is_shared_moe_module(module_name)
-            else self.moe_tp_size
+            (
+                self.tp_size
+                if self.is_shared_moe_module(module_name)
+                else self.moe_tp_size
+            )
+            if self.is_moe_module(module_name)
+            else self._column_parallel_shard_tp(module_name, base_model, layer_idx)
         )
         if (
             effective_tp_size > 1
             and module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES
             and module_name not in REPLICATED_LINEAR_LORA_NAMES
         ):
-            output_dim = self._column_parallel_lora_b_per_rank_dim(
-                module_name, output_dim, effective_tp_size
+            # If the base column-parallel module is fully REPLICATED (its actual
+            # output_size_per_partition still equals the full output_dim -- e.g. the
+            # dense MLP gate_up under --moe-dense-tp-size 1), its output is NOT
+            # sharded, so keep LoRA-B at the full output dim. Dividing by the global
+            # tp_size here undersizes B and crashes set_lora_info ("LoRA B output dim
+            # != base partition prefix dim"). Non-MoE only; MoE shards by moe_tp_size.
+            probed_out = (
+                None
+                if self.is_moe_module(module_name)
+                else self._column_parallel_out_partition(
+                    module_name, base_model, layer_idx
+                )
             )
+            if probed_out is not None and probed_out == output_dim:
+                pass  # replicated base: keep full B output dim
+            else:
+                output_dim = self._column_parallel_lora_b_per_rank_dim(
+                    module_name, output_dim, effective_tp_size
+                )
 
         # Check if MoE module and return appropriate shape
         if self.is_moe_module(module_name):
@@ -963,19 +1150,45 @@ def load_lora_weight_tensor(
                         temp_B_buffer[target_module] = weights
                         temp_B_cache_keys[target_module] = name
                 elif expert_match:
-                    # Per-expert MoE weight — 2D tensors, one per expert
+                    # Per-expert MoE weight — 2D tensors, one per expert.
+                    # Init A and B INDEPENDENTLY (both buffer and cache_keys). Under
+                    # ``experts_shared_outer_loras`` one side of a projection is a
+                    # shared 3D Tensor (set by the dim()==3 branch below) and the
+                    # other is this per-expert dict: fc1 = shared A + per-expert B,
+                    # fc2 = the opposite. The old coupled init keyed every dict off
+                    # ``temp_A_buffer is None``, which either left the per-expert
+                    # side's cache_keys as None (-> TypeError at
+                    # ``[expert_id] = name``) or clobbered the shared side the 3D
+                    # branch already populated. So each side now guards its own
+                    # buffer + cache_keys and never touches the other side.
                     target_module = target_module + "_moe"
-                    if temp_A_buffer[target_module] is None:
-                        temp_A_buffer[target_module] = {}
-                        temp_B_buffer[target_module] = {}
-                        temp_A_cache_keys[target_module] = {}
-                        temp_B_cache_keys[target_module] = {}
-
                     expert_id = int(expert_match.group(1))
                     if "lora_A" in name:
+                        assert not isinstance(
+                            temp_A_buffer[target_module], torch.Tensor
+                        ), (
+                            f"{target_module} lora_A already holds a shared-outer 3D "
+                            f"tensor but also got per-expert weight '{name}'; a "
+                            f"projection side must use one layout, not both."
+                        )
+                        if temp_A_buffer[target_module] is None:
+                            temp_A_buffer[target_module] = {}
+                        if temp_A_cache_keys[target_module] is None:
+                            temp_A_cache_keys[target_module] = {}
                         temp_A_buffer[target_module][expert_id] = weights
                         temp_A_cache_keys[target_module][expert_id] = name
                     else:
+                        assert not isinstance(
+                            temp_B_buffer[target_module], torch.Tensor
+                        ), (
+                            f"{target_module} lora_B already holds a shared-outer 3D "
+                            f"tensor but also got per-expert weight '{name}'; a "
+                            f"projection side must use one layout, not both."
+                        )
+                        if temp_B_buffer[target_module] is None:
+                            temp_B_buffer[target_module] = {}
+                        if temp_B_cache_keys[target_module] is None:
+                            temp_B_cache_keys[target_module] = {}
                         temp_B_buffer[target_module][expert_id] = weights
                         temp_B_cache_keys[target_module][expert_id] = name
                 elif "experts" in name and weights.dim() == 3:
diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py
index fa359c548306..cc00f48ee568 100644
--- a/python/sglang/srt/managers/io_struct.py
+++ b/python/sglang/srt/managers/io_struct.py
@@ -2036,7 +2036,7 @@ class LoadLoRAAdapterFromTensorsReqInput(BaseReq, kw_only=True):
     # The PEFT adapter_config.json, already JSON — a tighter type would only add
     # decode strictness with no benefit.
     config_dict: Dict[str, Any]
-    serialized_tensors: str
+    serialized_named_tensors: Annotated[List[bytes], Base64Bytes()]
     pinned: bool = False
     added_tokens_config: Optional[Dict[str, int]] = None
     lora_id: Optional[str] = None
diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py
index 86b7b378f896..ad92b301f245 100644
--- a/python/sglang/srt/managers/tokenizer_control_mixin.py
+++ b/python/sglang/srt/managers/tokenizer_control_mixin.py
@@ -574,11 +574,9 @@ async def load_lora_adapter(
                     "LoRA is not enabled. Please set `--enable-lora` to enable LoRA."
                 )
 
-            # TODO (lifuhuang): Remove this after we verify that dynamic lora loading works
-            # with dp_size > 1.
             assert (
-                self.server_args.dp_size == 1
-            ), "dp_size must be 1 for dynamic lora loading"
+                self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
+            ), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading"
             logger.info(
                 "Start load Lora adapter. Lora name=%s, path=%s",
                 obj.lora_name,
@@ -653,8 +651,8 @@ async def load_lora_adapter_from_tensors(
                 )
 
             assert (
-                self.server_args.dp_size == 1
-            ), "dp_size must be 1 for dynamic lora loading"
+                self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
+            ), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading"
             logger.info(
                 "Start load Lora adapter from tensors. Lora name=%s",
                 obj.lora_name,
@@ -726,11 +724,9 @@ async def unload_lora_adapter(
                 obj.lora_name is not None
             ), "lora_name must be provided to unload LoRA adapter"
 
-            # TODO (lifuhuang): Remove this after we verify that dynamic lora loading works
-            # with dp_size > 1.
             assert (
-                self.server_args.dp_size == 1
-            ), "dp_size must be 1 for dynamic lora loading"
+                self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
+            ), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading"
             logger.info(
                 "Start unload Lora adapter. Lora name=%s",
                 obj.lora_name,
diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py
index 6d09e81332f0..525d1f973d9f 100644
--- a/python/sglang/srt/managers/tp_worker.py
+++ b/python/sglang/srt/managers/tp_worker.py
@@ -208,19 +208,23 @@ def unload_lora_adapter(self, recv_req: UnloadLoRAAdapterReqInput):
     def load_lora_adapter_from_tensors(
         self, recv_req: LoadLoRAAdapterFromTensorsReqInput
     ):
-        # The LoRA code handles TP sharding internally using slice_lora_a_weights
-        # and slice_lora_b_weights methods (see lora/layers.py:46-49, mem_pool.py:437-440).
+        # TP sharding for LoRA happens inside the lora module (see
+        # lora/layers.py:46-49 and mem_pool.py:437-440). Each TP rank
+        # deserializes its own producer's bytes — same convention as
+        # ``update_weights_from_tensor`` above. One producer per one
+        # consumer means the CUDA-IPC ref counter on the producer's
+        # bucket drops cleanly each cycle.
+        monkey_patch_torch_reductions()
+        serialized = recv_req.serialized_named_tensors[self.tp_rank]
         if recv_req.load_format == "flattened_bucket":
-            flattened_data = MultiprocessingSerializer.deserialize(
-                recv_req.serialized_tensors
-            )
+            flattened_data = MultiprocessingSerializer.deserialize(serialized)
             bucket = FlattenedTensorBucket(
                 flattened_tensor=flattened_data["flattened_tensor"],
                 metadata=flattened_data["metadata"],
             )
             tensors = dict(bucket.reconstruct_tensors())
         else:
-            tensors = MultiprocessingSerializer.deserialize(recv_req.serialized_tensors)
+            tensors = MultiprocessingSerializer.deserialize(serialized)
         if recv_req.expected_checksums is not None:
             import hashlib
 
diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py
index 81abdb78102b..78e9edc59250 100644
--- a/python/sglang/srt/model_executor/forward_batch_info.py
+++ b/python/sglang/srt/model_executor/forward_batch_info.py
@@ -809,6 +809,14 @@ def init_new(
 
         if ret.forward_mode.is_idle():
             ret.positions = torch.empty((0,), dtype=torch.int64, device=device)
+            # Under --enable-dp-attention an IDLE rank still runs its local experts over the
+            # DP-GATHERED tokens of the other ranks, so the (MoE-)LoRA batch info must be
+            # re-prepared for THIS batch: sized to the gathered length and stamped via the
+            # Tier-1 idle-rank path. Skipping it here leaves the backend serving the PREVIOUS
+            # batch's stale token_lora_mapping — undersized (OOB reads/writes in the MoE-LoRA
+            # kernels) and/or pointing foreign tokens at the wrong adapter.
+            if model_runner.server_args.enable_lora:
+                model_runner.lora_manager.prepare_lora_batch(ret)
             return ret
 
         # Override the positions with diffusion LLM or spec_info
diff --git a/test/registered/unit/lora/test_moe_lora_tail_stamp.py b/test/registered/unit/lora/test_moe_lora_tail_stamp.py
new file mode 100644
index 000000000000..b10d53389331
--- /dev/null
+++ b/test/registered/unit/lora/test_moe_lora_tail_stamp.py
@@ -0,0 +1,134 @@
+"""Unit tests for the MoE-LoRA DP-gathered tail stamping in _add_moe_lora_info.
+
+Under --enable-dp-attention the MoE runs on DP-GATHERED tokens, so the per-token
+LoRA mapping must cover [0, moe_num_tokens) while the per-rank segments only fill
+[0, num_tokens). These tests pin the Tier-1 (single loaded adapter) semantics of
+the gathered tail [num_tokens, moe_num_tokens):
+
+  * active rank (local requests use the single adapter, ``_single_loaded_buffer_id``
+    armed): tail stamped with that adapter's buffer id;
+  * true idle rank (num_tokens == 0, ``_idle_rank_active_buffer_id`` armed): the
+    whole mapping stamped and the adapter enabled;
+  * base-only local batches, multi-adapter batches, and stale idle stamps on
+    token-bearing batches: tail stays -1 (adapter-disabled) — foreign tokens are
+    never given a delta this rank cannot attribute.
+
+The backend object is stubbed (only the fields _add_moe_lora_info reads are
+populated) so the tests run hermetically without a server or dist groups; the
+gathered length is forced via forward_batch.global_dp_buffer_len.
+
+Usage:
+    python -m pytest test/registered/unit/lora/test_moe_lora_tail_stamp.py -v
+"""
+
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+
+register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-small")
+register_amd_ci(est_time=9, suite="stage-b-test-1-gpu-small-amd")
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.lora.backend.base_backend import BaseLoRABackend
+from sglang.srt.lora.utils import LoRABatchInfo
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+GATHERED = 16  # forced gathered length (> any local num_tokens used below)
+# buffer slots: 0 = the None/base uid slot (rank 0), 1/2 = adapters (rank 8), 3 = free
+RANKS = [0, 8, 8, 0]
+
+
+def _forward_batch(num_tokens: int):
+    mode = SimpleNamespace(
+        is_extend=lambda: False,
+        is_idle=lambda: num_tokens == 0,
+        is_cuda_graph=lambda: False,
+    )
+    return SimpleNamespace(
+        forward_mode=mode,
+        batch_size=num_tokens,  # decode: 1 token per seq
+        extend_seq_lens_cpu=None,
+        global_dp_buffer_len=GATHERED,
+        global_num_tokens_cpu=[num_tokens, GATHERED - num_tokens],
+    )
+
+
+def _batch_info(weight_indices, num_tokens):
+    bs = len(weight_indices)
+    return LoRABatchInfo(
+        use_cuda_graph=False,
+        bs=bs,
+        num_segments=bs,
+        seg_indptr=torch.arange(bs + 1, dtype=torch.int32, device=DEVICE),
+        weight_indices=torch.tensor(weight_indices, dtype=torch.int32, device=DEVICE),
+        lora_ranks=torch.tensor(RANKS, dtype=torch.int64, device=DEVICE),
+        scalings=torch.ones(len(RANKS), dtype=torch.float, device=DEVICE),
+        max_len=1,
+        seg_lens=torch.ones(bs, dtype=torch.int32, device=DEVICE),
+        permutation=None,
+        expected_tokens=num_tokens,
+    )
+
+
+def _run(weight_indices, num_tokens, idle_bid=None, single_bid=None):
+    stub = SimpleNamespace(
+        is_moe_lora=True,
+        _idle_rank_active_buffer_id=idle_bid,
+        _single_loaded_buffer_id=single_bid,
+    )
+    out = BaseLoRABackend._add_moe_lora_info(
+        stub, _forward_batch(num_tokens), _batch_info(weight_indices, num_tokens)
+    )
+    info = out.moe_lora_info
+    return info.token_lora_mapping.tolist(), info.adapter_enabled.tolist()
+
+
+class TestMoELoRATailStamp(unittest.TestCase):
+    def test_base_only_local_batch_keeps_disabled_tail(self):
+        # All local tokens on the base (None-uid) slot, nothing armed: the
+        # gathered tail must stay -1 and no adapter may be enabled.
+        mapping, enabled = _run([0, 0, 0, 0], num_tokens=4)
+        self.assertTrue(all(x == -1 for x in mapping[4:]), mapping)
+        self.assertEqual(enabled, [0, 0, 0, 0])
+
+    def test_stale_idle_stamp_ignored_on_token_bearing_batch(self):
+        # Defense in depth: an idle stamp must only be consumed when the rank is
+        # truly idle (num_tokens == 0), never on a batch with local tokens.
+        mapping, enabled = _run([0, 0, 0, 0], num_tokens=4, idle_bid=2)
+        self.assertTrue(all(x == -1 for x in mapping[4:]), mapping)
+        self.assertEqual(enabled[2], 0)
+
+    def test_multi_adapter_batch_keeps_disabled_tail(self):
+        # Two adapters used locally: the foreign tokens' adapter identity is
+        # unknowable host-side, so the tail must stay -1 (no mis-stamping with
+        # whatever the last local token happened to use).
+        mapping, _ = _run([1, 1, 2, 2], num_tokens=4)
+        self.assertTrue(all(x == -1 for x in mapping[4:]), mapping)
+
+    def test_idle_rank_stamps_whole_mapping_and_enables_adapter(self):
+        # True idle rank under Tier-1: the whole gathered mapping is stamped with
+        # the single loaded adapter and the adapter is enabled, so foreign tokens
+        # routed to this rank's experts get the LoRA delta.
+        mapping, enabled = _run([], num_tokens=0, idle_bid=2)
+        self.assertTrue(all(x == 2 for x in mapping), mapping)
+        self.assertEqual(enabled[2], 1)
+
+    def test_active_rank_stamps_tail_with_single_adapter(self):
+        # Local tokens actively use the single loaded adapter: the tail is
+        # stamped with its buffer id (not a copy of an arbitrary local slot).
+        mapping, enabled = _run([2, 2], num_tokens=2, single_bid=2)
+        self.assertTrue(all(x == 2 for x in mapping[2:]), mapping)
+        self.assertEqual(enabled[2], 1)
+
+    def test_single_adapter_loaded_but_local_base_keeps_disabled_tail(self):
+        # An adapter is loaded but this rank's local tokens are all base (the
+        # manager arms no stamp in this case): tail stays -1.
+        mapping, enabled = _run([0, 0], num_tokens=2)
+        self.assertTrue(all(x == -1 for x in mapping[2:]), mapping)
+        self.assertEqual(enabled, [0, 0, 0, 0])
+
+
+if __name__ == "__main__":
+    unittest.main()

From ed1b342e1c2e07011acd08fe3b3ea5d136d54c2e Mon Sep 17 00:00:00 2001
From: Yueming Yuan <yym022502@gmail.com>
Date: Fri, 24 Jul 2026 21:36:54 -0700
Subject: [PATCH 17/43] [12/27] [sglang-miles] DeepSeek V4 RL fixes (#27131,
 #27603, #27604, #27728)

---
 .../layers/deep_gemm_wrapper/compile_utils.py |  8 +++--
 python/sglang/srt/layers/linear.py            | 32 +++++++++++++++++++
 .../srt/layers/quantization/kv_cache.py       |  2 +-
 python/sglang/srt/models/deepseek_v4.py       |  4 ++-
 4 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py
index dec6bb512421..5a7128bc3f57 100644
--- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py
+++ b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py
@@ -35,10 +35,14 @@
 _IN_PRECOMPILE_STAGE = envs.SGLANG_IN_DEEPGEMM_PRECOMPILE_STAGE.get()
 _FAST_WARMUP = envs.SGLANG_JIT_DEEPGEMM_FAST_WARMUP.get()
 
-# Force redirect deep_gemm cache_dir
-os.environ["DG_JIT_CACHE_DIR"] = os.getenv(
+# Force redirect deep_gemm cache_dir.
+_dg_cache_dir = os.getenv(
     "SGLANG_DG_CACHE_DIR", os.path.join(os.path.expanduser("~"), ".cache", "deep_gemm")
 )
+if os.getenv("SGLANG_DG_CACHE_DIR_PER_PROCESS", "0").lower() in ("1", "true", "yes"):
+    _dg_cache_dir = os.path.join(_dg_cache_dir, f"pid_{os.getpid()}")
+os.makedirs(_dg_cache_dir, exist_ok=True)
+os.environ["DG_JIT_CACHE_DIR"] = _dg_cache_dir
 
 # Refer to https://github.com/deepseek-ai/DeepGEMM/commit/d75b218b7b8f4a5dd5406ac87905039ead3ae42f
 # NVRTC may have performance loss with some cases.
diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py
index f91ec3f765bd..b2385a95753e 100644
--- a/python/sglang/srt/layers/linear.py
+++ b/python/sglang/srt/layers/linear.py
@@ -1536,6 +1536,15 @@ def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor
             assert loaded_weight.numel() == 1
             loaded_weight = loaded_weight.reshape(1)
 
+        if (
+            isinstance(param, BlockQuantScaleParameter)
+            and getattr(param, "format_ue8m0", False)
+            and loaded_weight.dtype == torch.int32
+            and not self.use_presharded_weights
+        ):
+            self._load_ue8m0_packed_scale(param, loaded_weight)
+            return
+
         if isinstance(param, RowvLLMParameter):
             # This `BasevLLMParameter` is defined in sglang/srt/layers/parameter.py,
             # It supports additional parameters like tp_rank and use_presharded_weights.
@@ -1558,6 +1567,29 @@ def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor
                 # Fallback for parameters that don't accept additional args
                 param.load_row_parallel_weight(loaded_weight)
 
+    def _load_ue8m0_packed_scale(
+        self, param: BlockQuantScaleParameter, loaded_weight: torch.Tensor
+    ):
+        # TODO: This is hacky, better to have a more stable solution
+        # The loaded scale is a full-size DeepGEMM UE8M0 packed tensor (int32,
+        # mn-major, 4 scale columns per word; see transform_scale_ue8m0). The
+        # input-dim shard boundary sits at sub-word granularity whenever
+        # (k_local / block_k) % 4 != 0, so the packed form cannot be narrowed
+        # directly: unpack to fp32, narrow in scale-column units, then repack.
+        from sglang.srt.layers.quantization.fp8_utils import (
+            inverse_transform_scale_ue8m0,
+            transform_scale_ue8m0,
+        )
+
+        block_k = self.quant_method.quant_config.weight_block_size[1]
+        mn = param.data.shape[0]
+        sf_fp32 = inverse_transform_scale_ue8m0(loaded_weight, mn=mn)
+        shard_size = (self.input_size_per_partition + block_k - 1) // block_k
+        sf_local = sf_fp32.narrow(1, self.tp_rank * shard_size, shard_size).contiguous()
+        packed = transform_scale_ue8m0(sf_local, mn=mn)
+        assert param.data.shape == packed.shape, f"{param.data.shape=} {packed.shape=}"
+        param.data.copy_(packed)
+
     def forward(self, input_, skip_all_reduce=False, forward_batch=None):
         if self.input_is_parallel:
             input_parallel = input_
diff --git a/python/sglang/srt/layers/quantization/kv_cache.py b/python/sglang/srt/layers/quantization/kv_cache.py
index 105cb983938b..184c170447c5 100644
--- a/python/sglang/srt/layers/quantization/kv_cache.py
+++ b/python/sglang/srt/layers/quantization/kv_cache.py
@@ -56,7 +56,7 @@ def process_weights_after_loading(self, layer) -> None:
             if is_fp8_fnuz():
                 k_scale *= 2
                 v_scale *= 2
-        elif layer.k_scale < 0.0 and layer.v_scale < 0.0:
+        elif layer.k_scale <= 0.0 and layer.v_scale <= 0.0:
             # If no scales were loaded (both scales are invalid negative
             # values), use the default value of 1.0
             k_scale = 1.0
diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py
index 264806123933..c123d642cbaa 100644
--- a/python/sglang/srt/models/deepseek_v4.py
+++ b/python/sglang/srt/models/deepseek_v4.py
@@ -3055,7 +3055,9 @@ def auto_weight_loader(module):
         }
         if unloaded_params:
             logger.warning(
-                f"Some weights are not initialized from checkpoints: {unloaded_params}"
+                "Some weights are not initialized from checkpoints: "
+                f"count={len(unloaded_params)}. "
+                "Ignore this message for RL weight update."
             )
 
         self.post_load_weights(is_nextn=is_nextn, weight_names=weight_names)

From 976db7a6e3cd59eb3b974a5ce2a779ed8b68f615 Mon Sep 17 00:00:00 2001
From: Yisheng Gong <yishenggong9437@gmail.com>
Date: Fri, 24 Jul 2026 21:36:54 -0700
Subject: [PATCH 18/43] [13/27] [sglang-miles]
 load_lora_adapter_from_distributed API, with upsert (#27268, #31759, #30913)

Squashes the two follow-ups into the API commit: #31759 fixes this commit's
HTTP handler (the LoRAUpdateOutput was returned unserialized) and #30913
rewrites its registry-registration block to support in-place upsert, so
neither applies on its own.

Co-authored-by: Ethan (Yusheng) Su <yushengsu@radixark.ai>
Co-authored-by: Mathew Han <49226490+mathewjhan@users.noreply.github.com>
---
 python/sglang/srt/entrypoints/engine.py       |  29 +
 python/sglang/srt/entrypoints/http_server.py  |  13 +
 python/sglang/srt/lora/lora_manager.py        | 151 ++++-
 python/sglang/srt/lora/lora_registry.py       |  46 +-
 python/sglang/srt/managers/io_struct.py       |  26 +-
 python/sglang/srt/managers/scheduler.py       |  14 +
 .../srt/managers/tokenizer_control_mixin.py   | 110 ++++
 python/sglang/srt/managers/tp_worker.py       |  17 +
 .../sglang/srt/model_executor/model_runner.py |  75 ++-
 test/registered/unit/lora/test_lora_upsert.py | 537 ++++++++++++++++++
 10 files changed, 983 insertions(+), 35 deletions(-)
 create mode 100644 test/registered/unit/lora/test_lora_upsert.py

diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py
index 1e0b27dc357b..e85f3dc7a7eb 100644
--- a/python/sglang/srt/entrypoints/engine.py
+++ b/python/sglang/srt/entrypoints/engine.py
@@ -65,6 +65,7 @@
     GenerateReqInput,
     GetWeightsByNameReqInput,
     InitWeightsUpdateGroupReqInput,
+    LoadLoRAAdapterFromDistributedReqInput,
     LoadLoRAAdapterFromTensorsReqInput,
     LoadLoRAAdapterReqInput,
     MultimodalDataInputFormat,
@@ -1159,6 +1160,34 @@ def load_lora_adapter_from_tensors(
             self.tokenizer_manager.load_lora_adapter_from_tensors(lora_req, None)
         )
 
+    def load_lora_adapter_from_distributed(
+        self,
+        lora_name: str,
+        config_dict: Dict,
+        names: list[str],
+        dtypes: list[str],
+        shapes: list[list[int]],
+        group_name: str = "weight_update_group",
+        pinned: bool = False,
+        added_tokens_config: Optional[Dict] = None,
+    ):
+        """Load a new LoRA adapter whose weights are broadcast over
+        a process group. The weight-update group must already be
+        initialized via `init_weights_update_group`."""
+        lora_req = LoadLoRAAdapterFromDistributedReqInput(
+            lora_name=lora_name,
+            config_dict=config_dict,
+            names=names,
+            dtypes=dtypes,
+            shapes=shapes,
+            group_name=group_name,
+            pinned=pinned,
+            added_tokens_config=added_tokens_config,
+        )
+        return self.loop.run_until_complete(
+            self.tokenizer_manager.load_lora_adapter_from_distributed(lora_req, None)
+        )
+
     def load_lora_adapter(self, lora_name: str, lora_path: str, pinned: bool = False):
         """Load a new LoRA adapter without re-launching the engine."""
 
diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py
index 51c85e63023f..06a652673e29 100644
--- a/python/sglang/srt/entrypoints/http_server.py
+++ b/python/sglang/srt/entrypoints/http_server.py
@@ -124,6 +124,7 @@
     GetWeightsByNameReqInput,
     InitWeightsSendGroupForRemoteInstanceReqInput,
     InitWeightsUpdateGroupReqInput,
+    LoadLoRAAdapterFromDistributedReqInput,
     LoadLoRAAdapterFromTensorsReqInput,
     LoadLoRAAdapterReqInput,
     OpenSessionReqInput,
@@ -1496,6 +1497,18 @@ async def load_lora_adapter_from_tensors(
     return ORJSONResponse(msgspec_to_builtins(result), status_code=status_code)
 
 
+@app.api_route("/load_lora_adapter_from_distributed", methods=["POST"])
+async def load_lora_adapter_from_distributed(
+    obj: Annotated[LoadLoRAAdapterFromDistributedReqInput, Body()], request: Request
+):
+    """Load a new LoRA adapter broadcast over a process group without re-launching the server."""
+    result = await _global_state.tokenizer_manager.load_lora_adapter_from_distributed(
+        obj, request
+    )
+    status_code = HTTPStatus.OK if result.success else HTTPStatus.BAD_REQUEST
+    return ORJSONResponse(msgspec_to_builtins(result), status_code=status_code)
+
+
 @app.api_route("/unload_lora_adapter", methods=["POST"])
 @auth_level(AuthLevel.ADMIN_OPTIONAL)
 async def unload_lora_adapter(
diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py
index 8aa8e453f99e..f03db5093751 100644
--- a/python/sglang/srt/lora/lora_manager.py
+++ b/python/sglang/srt/lora/lora_manager.py
@@ -214,9 +214,19 @@ def _load_lora_adapter(self, lora_ref: LoRARef) -> LoRAUpdateOutput:
 
         return self.create_lora_update_result(success=True)
 
-    def validate_new_adapter(self, lora_config: LoRAConfig, lora_ref: LoRARef):
+    def validate_new_adapter(
+        self,
+        lora_config: LoRAConfig,
+        lora_ref: LoRARef,
+        is_update: bool = False,
+        old_ref: Optional[LoRARef] = None,
+    ):
         """
         Validate if an adapter can be loaded into the current LoRA memory pool and generate error if it is incompatible.
+
+        For an in-place refresh (``is_update``), pass the currently-loaded ref as
+        ``old_ref`` so checks that count loaded adapters exclude the adapter's
+        own current state.
         """
         if lora_config.lora_added_tokens_size > 0:
             raise ValueError(
@@ -228,18 +238,19 @@ def validate_new_adapter(self, lora_config: LoRAConfig, lora_ref: LoRARef):
                 f"Failed to load {lora_ref.lora_name} because LoRA serving currently doesn't support DoRA adapters"
             )
 
-        # Check if this LoRA adapter is already loaded
-        for existing_lora_ref in self.lora_refs.values():
-            if lora_ref.lora_name == existing_lora_ref.lora_name:
-                raise ValueError(
-                    f"Failed to load LoRA adapter {lora_ref.lora_name} because it is already loaded"
-                )
+        # Reject duplicates unless refreshing an existing adapter in place.
+        if not is_update:
+            for existing_lora_ref in self.lora_refs.values():
+                if lora_ref.lora_name == existing_lora_ref.lora_name:
+                    raise ValueError(
+                        f"Failed to load LoRA adapter {lora_ref.lora_name} because it is already loaded"
+                    )
 
-            if lora_ref.lora_path == existing_lora_ref.lora_path:
-                logger.warning(
-                    f"{lora_ref.lora_path} is already loaded with name: {existing_lora_ref.lora_name}, "
-                    f"but another copy is being loaded with name: {lora_ref.lora_name}"
-                )
+                if lora_ref.lora_path == existing_lora_ref.lora_path:
+                    logger.warning(
+                        f"{lora_ref.lora_path} is already loaded with name: {existing_lora_ref.lora_name}, "
+                        f"but another copy is being loaded with name: {lora_ref.lora_name}"
+                    )
 
         # Check if the LoRA adapter shape is compatible with the current LoRA memory pool configuration.
         memory_pool = getattr(self, "memory_pool", None)
@@ -252,7 +263,13 @@ def validate_new_adapter(self, lora_config: LoRAConfig, lora_ref: LoRARef):
             )
 
         # Ensure pinned LoRA adapters does not exceed maximal limit or cause starvation.
-        if lora_ref.pinned and self.num_pinned_loras >= self.max_loras_per_batch - 1:
+        # On an in-place refresh the adapter's own pin is already counted in
+        # num_pinned_loras; refreshing occupies no additional pinned slot, so
+        # exclude it or a pinned adapter at the limit could never be refreshed.
+        num_other_pinned_loras = self.num_pinned_loras
+        if is_update and old_ref is not None:
+            num_other_pinned_loras -= int(bool(old_ref.pinned))
+        if lora_ref.pinned and num_other_pinned_loras >= self.max_loras_per_batch - 1:
             raise ValueError(
                 f"Failed to load LoRA adapter {lora_ref.lora_name} as a pinned adapter. It is not allowed to pin all slots "
                 "in the LoRA memory pool to avoid starvation for unpinned adapters and base models. Please increase your "
@@ -785,16 +802,25 @@ def load_lora_weights_from_tensors(
         """
         Load the weights of a LoRA adapter from tensors to CPU memory.
         """
+        self.loras[lora_ref.lora_id] = self._create_lora_adapter_from_tensors(
+            lora_ref, self.configs[lora_ref.lora_id], tensors
+        )
+
+    def _create_lora_adapter_from_tensors(
+        self, lora_ref: LoRARef, config: LoRAConfig, tensors: Dict[str, torch.Tensor]
+    ) -> LoRAAdapter:
+        """Build and initialize a LoRAAdapter without touching served state,
+        so callers can stage it and commit only after every fallible step."""
         lora_adapter = LoRAAdapter(
             lora_ref.lora_id,
-            self.configs[lora_ref.lora_id],
+            config,
             self.base_hf_config,
             self.load_config,
             self.lora_backend,
             base_model=self.base_model,
         )
         lora_adapter.initialize_weights_from_tensors(tensors)
-        self.loras[lora_ref.lora_id] = lora_adapter
+        return lora_adapter
 
     def load_lora_adapter_from_tensors(
         self,
@@ -802,10 +828,11 @@ def load_lora_adapter_from_tensors(
         tensors: Dict[str, torch.Tensor],
         config_dict: Dict,
         added_tokens_config: Optional[Dict] = None,
+        upsert: bool = False,
     ) -> LoRAUpdateOutput:
         logger.info(f"LoRA adapter loading from tensors starts: {lora_ref}.")
         result = self._load_lora_adapter_from_tensors(
-            lora_ref, tensors, config_dict, added_tokens_config
+            lora_ref, tensors, config_dict, added_tokens_config, upsert=upsert
         )
         logger.info(f"LoRA adapter loading from tensors completes: {lora_ref}.")
         return result
@@ -816,36 +843,104 @@ def _load_lora_adapter_from_tensors(
         tensors: Dict[str, torch.Tensor],
         config_dict: Dict,
         added_tokens_config: Optional[Dict] = None,
+        upsert: bool = False,
     ) -> LoRAUpdateOutput:
         """
         Load a single LoRA adapter from tensors and config dict.
+
+        With ``upsert``, an already-loaded adapter has its weights refreshed in
+        place (reusing its lora_id); otherwise the adapter must not be loaded yet.
         """
         assert (
             lora_ref.lora_name is not None and lora_ref.lora_path is not None
         ), "LoRARef must have both lora_name and lora_path set for loading."
-        assert (
-            lora_ref.lora_id not in self.loras
-        ), f"LoRA adapter with ID {lora_ref.lora_id} is already loaded. This should have been verified before request is sent to the backend."
-
+        is_update = upsert and (lora_ref.lora_id in self.loras)
+        if not is_update:
+            assert (
+                lora_ref.lora_id not in self.loras
+            ), f"LoRA adapter with ID {lora_ref.lora_id} is already loaded. This should have been verified before request is sent to the backend."
+
+        uid = lora_ref.lora_id
+        old_config = self.configs.get(uid)
+        old_lora = self.loras.get(uid)
+        old_ref = self.lora_refs.get(uid)
+
+        # Stage every fallible step before mutating served state: a failed
+        # request must not leave a live adapter with new metadata over old
+        # weights (or leak unreachable entries on a fresh insert).
         try:
-            new_adapter = LoRAConfig.from_dict(
+            new_config = LoRAConfig.from_dict(
                 config_dict,
                 added_tokens_config,
                 base_vocab_size=self.base_hf_config.vocab_size,
             )
-            self.validate_new_adapter(new_adapter, lora_ref)
-            self.configs[lora_ref.lora_id] = new_adapter
-
-            self.load_lora_weights_from_tensors(lora_ref, tensors)
-
-            self.lora_refs[lora_ref.lora_id] = lora_ref
-            self.num_pinned_loras += int(lora_ref.pinned)
+            self.validate_new_adapter(
+                new_config, lora_ref, is_update=is_update, old_ref=old_ref
+            )
+            new_lora = self._create_lora_adapter_from_tensors(
+                lora_ref, new_config, tensors
+            )
         except Exception as e:
             return self.create_lora_update_result(
                 success=False,
                 error_message=str(e),
             )
 
+        self.configs[uid] = new_config
+        self.loras[uid] = new_lora
+
+        if (
+            is_update
+            and getattr(self, "memory_pool", None) is not None
+            and uid in self.memory_pool.uid_to_buffer_id
+        ):
+            buffer_id = self.memory_pool.uid_to_buffer_id[uid]
+            try:
+                # The served buffer slot is rewritten in place. Wait for
+                # already-launched kernels first: under overlap scheduling /
+                # CUDA-graph replay a forward pass can still be executing on
+                # another stream, and it must not read a half-rewritten slot.
+                if self.device.type == "cuda":
+                    torch.cuda.synchronize(self.device)
+                self.memory_pool.load_lora_weight_to_buffer(
+                    uid,
+                    buffer_id,
+                    new_lora,
+                    self.lora_modules,
+                    self.embed_tokens_module,
+                    self.lm_head_module,
+                )
+            except Exception as e:
+                # Roll back so the adapter keeps serving its previous weights
+                # instead of a torn half-old/half-new buffer.
+                self.configs[uid] = old_config
+                self.loras[uid] = old_lora
+                try:
+                    self.memory_pool.load_lora_weight_to_buffer(
+                        uid,
+                        buffer_id,
+                        old_lora,
+                        self.lora_modules,
+                        self.embed_tokens_module,
+                        self.lm_head_module,
+                    )
+                except Exception:
+                    logger.exception(
+                        f"Failed to restore previous weights for LoRA adapter "
+                        f"{lora_ref.lora_name} (buffer slot {buffer_id}) after a "
+                        "failed upsert; the served buffer may be corrupted."
+                    )
+                return self.create_lora_update_result(
+                    success=False,
+                    error_message=str(e),
+                )
+
+        self.lora_refs[uid] = lora_ref
+        # An upsert may change ``pinned``; track the delta against the replaced
+        # ref so the counter stays consistent with lora_refs.
+        old_pinned = int(bool(old_ref.pinned)) if is_update else 0
+        self.num_pinned_loras += int(bool(lora_ref.pinned)) - old_pinned
+
         return self.create_lora_update_result(success=True)
 
     def init_memory_pool(self):
diff --git a/python/sglang/srt/lora/lora_registry.py b/python/sglang/srt/lora/lora_registry.py
index 4e72efab56a4..8ee933a8c402 100644
--- a/python/sglang/srt/lora/lora_registry.py
+++ b/python/sglang/srt/lora/lora_registry.py
@@ -15,11 +15,11 @@
 
 import asyncio
 from collections import OrderedDict
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Tuple, Union
 from uuid import NAMESPACE_URL, uuid4, uuid5
 
 import msgspec
-from msgspec.structs import fields
+from msgspec.structs import fields, replace
 
 from sglang.srt.utils import ConcurrentCounter
 from sglang.srt.utils.aio_rwlock import RWLock
@@ -123,6 +123,48 @@ async def unregister(self, lora_name: str) -> str:
 
         return lora_ref.lora_id
 
+    async def get_lora_id(self, lora_name: str) -> Optional[str]:
+        """Return the ``lora_id`` of a registered adapter, or ``None``."""
+        async with self._registry_lock.reader_lock:
+            lora_ref = self._registry.get(lora_name, None)
+            return lora_ref.lora_id if lora_ref is not None else None
+
+    async def register_or_reuse(
+        self, lora_ref: LoRARef, upsert: bool = False
+    ) -> Tuple[LoRARef, bool]:
+        """Resolve which identity a load request should use.
+
+        Returns ``(ref, reused)``. With ``upsert`` and a same-name adapter
+        already registered, the returned ref adopts the existing ``lora_id``
+        (``reused=True``) so the backend refreshes that adapter in place;
+        otherwise ``lora_ref`` is returned unchanged (``reused=False``).
+        Nothing is registered here: the caller commits the resolved ref with
+        ``register`` / ``refresh`` once the backend load succeeded, keeping
+        failed loads invisible to the registry.
+        """
+        if not upsert:
+            return lora_ref, False
+        async with self._registry_lock.reader_lock:
+            existing = self._registry.get(lora_ref.lora_name, None)
+        if existing is None:
+            return lora_ref, False
+        return replace(lora_ref, lora_id=existing.lora_id), True
+
+    async def refresh(self, lora_ref: LoRARef):
+        """Replace a registered adapter's ref after a successful upsert.
+
+        Keeps the id (asserted) while adopting the new path/pinned metadata,
+        and counts as a use for LRU ordering.
+        """
+        async with self._registry_lock.writer_lock:
+            existing = self._registry.get(lora_ref.lora_name, None)
+            assert existing is not None and existing.lora_id == lora_ref.lora_id, (
+                f"refresh() must target a registered adapter with the same lora_id; "
+                f"got {lora_ref}, registered: {existing}"
+            )
+            self._registry[lora_ref.lora_name] = lora_ref
+            self._registry.move_to_end(lora_ref.lora_name)
+
     async def acquire(self, lora_name: Union[str, List[str]]) -> Union[str, List[str]]:
         """
         Queries registry for LoRA IDs based on LoRA names and start tracking the usage of the corresponding LoRA adapters
diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py
index cc00f48ee568..1fefc17c44d6 100644
--- a/python/sglang/srt/managers/io_struct.py
+++ b/python/sglang/srt/managers/io_struct.py
@@ -2041,6 +2041,8 @@ class LoadLoRAAdapterFromTensorsReqInput(BaseReq, kw_only=True):
     added_tokens_config: Optional[Dict[str, int]] = None
     lora_id: Optional[str] = None
     load_format: Optional[str] = None
+    # If already loaded, refresh weights in place instead of failing.
+    upsert: bool = False
     expected_checksums: Optional[Dict[str, str]] = None
 
     def to_ref(self) -> LoRARef:
@@ -2052,6 +2054,28 @@ def to_ref(self) -> LoRARef:
         )
 
 
+class LoadLoRAAdapterFromDistributedReqInput(BaseReq, kw_only=True):
+    lora_name: str
+    config_dict: Dict[str, Any]
+    names: List[str]
+    dtypes: List[str]
+    shapes: List[List[int]]
+    group_name: str = "weight_update_group"
+    pinned: bool = False
+    added_tokens_config: Optional[Dict[str, Any]] = None
+    lora_id: Optional[str] = None
+    # If already loaded, refresh weights in place instead of failing.
+    upsert: bool = False
+
+    def to_ref(self) -> LoRARef:
+        return LoRARef(
+            lora_id=self.lora_id,
+            lora_name=self.lora_name,
+            lora_path="__distributed__",
+            pinned=self.pinned,
+        )
+
+
 class LoRAUpdateOutput(BaseReq, kw_only=True):
     success: bool
     error_message: Optional[str] = None
@@ -2060,7 +2084,7 @@ class LoRAUpdateOutput(BaseReq, kw_only=True):
 
 LoadLoRAAdapterReqOutput = UnloadLoRAAdapterReqOutput = (
     LoadLoRAAdapterFromTensorsReqOutput
-) = LoRAUpdateOutput
+) = LoadLoRAAdapterFromDistributedReqOutput = LoRAUpdateOutput
 
 
 class BlockReqType(Enum):
diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
index e0b379dedd14..327a7ee414a4 100644
--- a/python/sglang/srt/managers/scheduler.py
+++ b/python/sglang/srt/managers/scheduler.py
@@ -114,6 +114,8 @@
     InitWeightsUpdateGroupReqInput,
     ListExternalCorporaReqInput,
     ListExternalCorporaReqOutput,
+    LoadLoRAAdapterFromDistributedReqInput,
+    LoadLoRAAdapterFromDistributedReqOutput,
     LoadLoRAAdapterFromTensorsReqInput,
     LoadLoRAAdapterFromTensorsReqOutput,
     LoadLoRAAdapterReqInput,
@@ -1409,6 +1411,10 @@ def init_request_dispatcher(self):
                     LoadLoRAAdapterFromTensorsReqInput,
                     self.load_lora_adapter_from_tensors,
                 ),
+                (
+                    LoadLoRAAdapterFromDistributedReqInput,
+                    self.load_lora_adapter_from_distributed,
+                ),
                 (UnloadLoRAAdapterReqInput, self.unload_lora_adapter),
                 (PauseGenerationReqInput, self.pause_generation),
                 (ContinueGenerationReqInput, self.continue_generation),
@@ -4383,6 +4389,14 @@ def load_lora_adapter_from_tensors(
         result = self.tp_worker.load_lora_adapter_from_tensors(recv_req)
         return result
 
+    def load_lora_adapter_from_distributed(
+        self, recv_req: LoadLoRAAdapterFromDistributedReqInput
+    ) -> LoadLoRAAdapterFromDistributedReqOutput:
+        """In-place loading a new lora adapter broadcast over a process group."""
+
+        result = self.tp_worker.load_lora_adapter_from_distributed(recv_req)
+        return result
+
     def unload_lora_adapter(
         self, recv_req: UnloadLoRAAdapterReqInput
     ) -> UnloadLoRAAdapterReqOutput:
diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py
index ad92b301f245..f5e3d5807e23 100644
--- a/python/sglang/srt/managers/tokenizer_control_mixin.py
+++ b/python/sglang/srt/managers/tokenizer_control_mixin.py
@@ -42,6 +42,8 @@
     InitWeightsUpdateGroupReqOutput,
     ListExternalCorporaReqInput,
     ListExternalCorporaReqOutput,
+    LoadLoRAAdapterFromDistributedReqInput,
+    LoadLoRAAdapterFromDistributedReqOutput,
     LoadLoRAAdapterFromTensorsReqInput,
     LoadLoRAAdapterFromTensorsReqOutput,
     LoadLoRAAdapterReqInput,
@@ -637,6 +639,24 @@ async def load_lora_adapter(
                 error_message=str(e),
             )
 
+    def _validate_lora_upsert_supported(
+        self: TokenizerManager,
+        obj: LoadLoRAAdapterFromDistributedReqInput,
+    ) -> None:
+        """Upsert resolves lora_name -> lora_id through this process's registry.
+
+        With multiple tokenizer workers each HTTP worker process holds its own
+        registry, so the resolution depends on which worker the router picks:
+        a worker that never served the original load would mint a fresh id and
+        die on the backend duplicate check. Fail loudly instead.
+        """
+        if obj.upsert and self.server_args.tokenizer_worker_num > 1:
+            raise ValueError(
+                "LoRA upsert is not supported with tokenizer_worker_num > 1: "
+                "each HTTP worker resolves lora_name against its own registry, "
+                "making upsert nondeterministic across workers."
+            )
+
     async def load_lora_adapter_from_tensors(
         self: TokenizerManager,
         obj: LoadLoRAAdapterFromTensorsReqInput,
@@ -653,6 +673,15 @@ async def load_lora_adapter_from_tensors(
             assert (
                 self.server_args.dp_size == 1 or self.server_args.enable_dp_attention
             ), "dp_size must be 1 or dp attention must be enabled for dynamic lora loading"
+            if obj.upsert:
+                # In-place refresh is only wired up on the from_distributed
+                # route (the disaggregated RL weight-sync path). Reject
+                # explicitly instead of dying later on the duplicate check
+                # with a fresh uuid.
+                raise ValueError(
+                    "upsert is not supported on the from_tensors route; use "
+                    "/load_lora_adapter_from_distributed to refresh an adapter in place."
+                )
             logger.info(
                 "Start load Lora adapter from tensors. Lora name=%s",
                 obj.lora_name,
@@ -707,6 +736,87 @@ async def load_lora_adapter_from_tensors(
                 error_message=str(e),
             )
 
+    async def load_lora_adapter_from_distributed(
+        self: TokenizerManager,
+        obj: LoadLoRAAdapterFromDistributedReqInput,
+        _: Optional[fastapi.Request] = None,
+    ) -> LoadLoRAAdapterFromDistributedReqOutput:
+        self.auto_create_handle_loop()
+
+        try:
+            if not self.server_args.enable_lora:
+                raise ValueError(
+                    "LoRA is not enabled. Please set `--enable-lora` to enable LoRA."
+                )
+
+            assert (
+                self.server_args.dp_size == 1
+            ), "dp_size must be 1 for dynamic lora loading"
+            logger.info(
+                "Start load Lora adapter from distributed. Lora name=%s, group=%s",
+                obj.lora_name,
+                obj.group_name,
+            )
+
+            async with self.lora_update_lock:
+                self._validate_lora_upsert_supported(obj)
+                # With upsert, a same-name adapter keeps its lora_id so the
+                # backend refreshes it in place instead of failing the
+                # duplicate check; otherwise this resolves to a fresh ref.
+                new_adapter, reused = await self.lora_registry.register_or_reuse(
+                    LoRARef(
+                        lora_name=obj.lora_name,
+                        lora_path="__distributed__",
+                        pinned=obj.pinned,
+                    ),
+                    upsert=obj.upsert,
+                )
+                obj.lora_id = new_adapter.lora_id
+                result = (await self.update_lora_adapter_communicator(obj))[0]
+
+                if result.success:
+                    if reused:
+                        await self.lora_registry.refresh(new_adapter)
+                    else:
+                        await self.lora_registry.register(new_adapter)
+                    self.lora_ref_cache[obj.lora_name] = new_adapter
+                if self.server_args.max_loaded_loras is not None:
+                    while (
+                        self.lora_registry.num_registered_loras
+                        > self.server_args.max_loaded_loras
+                    ):
+                        lru_lora_name = await self.lora_registry.lru_lora_name(
+                            exclude_pinned=True
+                        )
+                        if lru_lora_name is None:
+                            raise ValueError(
+                                "Didn't find any LoRA adapters when trying to evict LRU LoRA adapter. "
+                                f"LoRA registry is: {self.lora_registry._registry}"
+                            )
+
+                        logger.info(
+                            f"Unloading least recently used LoRA adapter '{lru_lora_name}' "
+                            f"(current number of adapters: {self.lora_registry.num_registered_loras}, "
+                            f"max allowed: {self.server_args.max_loaded_loras})"
+                        )
+
+                        unload_result = await self._unload_lora_adapter_locked(
+                            UnloadLoRAAdapterReqInput(lora_name=lru_lora_name)
+                        )
+                        if not unload_result.success:
+                            raise ValueError(
+                                f"Error while unloading LRU LoRA adapter '{lru_lora_name}': "
+                                f"{unload_result.error_message}"
+                            )
+                        del result.loaded_adapters[lru_lora_name]
+
+                return result
+        except ValueError as e:
+            return LoadLoRAAdapterFromDistributedReqOutput(
+                success=False,
+                error_message=str(e),
+            )
+
     async def unload_lora_adapter(
         self: TokenizerManager,
         obj: UnloadLoRAAdapterReqInput,
diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py
index 525d1f973d9f..cb63131f0e36 100644
--- a/python/sglang/srt/managers/tp_worker.py
+++ b/python/sglang/srt/managers/tp_worker.py
@@ -28,6 +28,7 @@
     GetWeightsByNameReqInput,
     InitWeightsSendGroupForRemoteInstanceReqInput,
     InitWeightsUpdateGroupReqInput,
+    LoadLoRAAdapterFromDistributedReqInput,
     LoadLoRAAdapterFromTensorsReqInput,
     LoadLoRAAdapterReqInput,
     SendWeightsToRemoteInstanceReqInput,
@@ -261,6 +262,22 @@ def load_lora_adapter_from_tensors(
             tensors,
             recv_req.config_dict,
             recv_req.added_tokens_config,
+            upsert=recv_req.upsert,
+        )
+        return result
+
+    def load_lora_adapter_from_distributed(
+        self, recv_req: LoadLoRAAdapterFromDistributedReqInput
+    ):
+        result = self.model_runner.load_lora_adapter_from_distributed(
+            recv_req.to_ref(),
+            recv_req.names,
+            recv_req.dtypes,
+            recv_req.shapes,
+            recv_req.config_dict,
+            recv_req.group_name,
+            recv_req.added_tokens_config,
+            upsert=recv_req.upsert,
         )
         return result
 
diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py
index 85ae1ca8ebda..2e32ee835622 100644
--- a/python/sglang/srt/model_executor/model_runner.py
+++ b/python/sglang/srt/model_executor/model_runner.py
@@ -83,6 +83,7 @@
 from sglang.srt.layers.utils.cp_utils import is_mla_prefill_cp_enabled
 from sglang.srt.lora.lora_manager import LoRAManager, init_lora_cuda_graph_moe_buffers
 from sglang.srt.lora.lora_registry import LoRARef
+from sglang.srt.managers.io_struct import LoRAUpdateOutput
 from sglang.srt.managers.schedule_batch import sanity_check_mm_pad_shift_value
 from sglang.srt.mem_cache import kv_cache_dtype
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
@@ -1048,11 +1049,77 @@ def load_lora_adapter(self, lora_ref: LoRARef):
         return self.lora_manager.load_lora_adapter(lora_ref)
 
     def load_lora_adapter_from_tensors(
-        self, lora_ref: LoRARef, tensors, config_dict, added_tokens_config=None
+        self,
+        lora_ref: LoRARef,
+        tensors,
+        config_dict,
+        added_tokens_config=None,
+        upsert: bool = False,
     ):
-        return self.lora_manager.load_lora_adapter_from_tensors(
-            lora_ref, tensors, config_dict, added_tokens_config
-        )
+        logger.info(f"LoRA adapter loading from tensors starts: {lora_ref}.")
+        result = self.lora_manager.load_lora_adapter_from_tensors(
+            lora_ref,
+            tensors,
+            config_dict,
+            added_tokens_config,
+            upsert=upsert,
+        )
+        logger.info(f"LoRA adapter loading from tensors completes: {lora_ref}.")
+        return result
+
+    def load_lora_adapter_from_distributed(
+        self,
+        lora_ref: LoRARef,
+        names,
+        dtypes,
+        shapes,
+        config_dict,
+        group_name,
+        added_tokens_config=None,
+        upsert: bool = False,
+    ):
+        """Load a new lora adapter whose weights are broadcast over the
+        `_model_update_group` process group (no CUDA IPC).
+        """
+        assert group_name in self._model_update_group, (
+            f"Group {group_name} not in {list(self._model_update_group.keys())}. "
+            "Please call `init_weights_update_group` first."
+        )
+
+        logger.info(f"LoRA adapter loading from distributed starts: {lora_ref}.")
+        try:
+            tensors = {}
+            handles = []
+            for name, dtype, shape in zip(names, dtypes, shapes):
+                target_dtype = (
+                    dtype if isinstance(dtype, torch.dtype) else getattr(torch, dtype)
+                )
+                weight = torch.empty(shape, dtype=target_dtype, device=self.device)
+                handles.append(
+                    torch.distributed.broadcast(
+                        weight,
+                        src=0,
+                        group=self._model_update_group[group_name],
+                        async_op=True,
+                    )
+                )
+                tensors[name] = weight
+            for handle in handles:
+                handle.wait()
+        except Exception as e:
+            error_msg = f"Failed to receive LoRA adapter weights from distributed: {e}."
+            logger.error(error_msg)
+            return LoRAUpdateOutput(success=False, error_message=error_msg)
+
+        result = self.lora_manager.load_lora_adapter_from_tensors(
+            lora_ref,
+            tensors,
+            config_dict,
+            added_tokens_config,
+            upsert=upsert,
+        )
+        logger.info(f"LoRA adapter loading from distributed completes: {lora_ref}.")
+        return result
 
     def unload_lora_adapter(self, lora_ref: LoRARef):
         """Unload a lora adapter that was previously loaded during initialization or dynamic loading."""
diff --git a/test/registered/unit/lora/test_lora_upsert.py b/test/registered/unit/lora/test_lora_upsert.py
new file mode 100644
index 000000000000..54478ae4128b
--- /dev/null
+++ b/test/registered/unit/lora/test_lora_upsert.py
@@ -0,0 +1,537 @@
+"""Unit tests for the LoRA upsert (in-place refresh) path.
+
+With upsert=True, loading an adapter that is already registered refreshes
+its weights in place, reusing the existing lora_id and memory-pool slot
+instead of failing with a duplicate error. Covers:
+
+  * LoRARegistry.get_lora_id / register_or_reuse / refresh
+  * LoRAManager.load_lora_adapter_from_tensors upsert semantics
+  * failed-upsert rollback (no half-updated live adapter)
+  * num_pinned_loras consistency across pinned flips
+  * LoRAManager.validate_new_adapter duplicate-name / starvation checks
+  * TokenizerControlMixin from_distributed id reuse; the from_tensors route
+    rejects upsert explicitly (only from_distributed supports in-place refresh)
+"""
+
+import asyncio
+import unittest
+from unittest.mock import AsyncMock, MagicMock, Mock
+
+import torch
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
+
+maybe_stub_sgl_kernel()
+
+from sglang.srt.lora.lora_manager import LoRAManager
+from sglang.srt.lora.lora_registry import LoRARef, LoRARegistry
+from sglang.srt.managers.io_struct import (
+    LoadLoRAAdapterFromDistributedReqInput,
+    LoadLoRAAdapterFromTensorsReqInput,
+)
+from sglang.srt.managers.tokenizer_manager import TokenizerManager
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+CONFIG_DICT = {"target_modules": ["q_proj"], "r": 8, "lora_alpha": 16}
+
+
+class TestLoRARegistryGetLoraId(CustomTestCase):
+    def test_returns_id_for_registered_adapter(self):
+        registry = LoRARegistry()
+        ref = LoRARef(lora_name="a", lora_path="/x")
+        asyncio.run(registry.register(ref))
+
+        self.assertEqual(asyncio.run(registry.get_lora_id("a")), ref.lora_id)
+
+    def test_returns_none_for_unregistered_adapter(self):
+        registry = LoRARegistry()
+        self.assertIsNone(asyncio.run(registry.get_lora_id("missing")))
+
+
+def _make_manager() -> LoRAManager:
+    """Create a LoRAManager via __new__ with only the fields the load path reads."""
+    manager = LoRAManager.__new__(LoRAManager)
+    manager.configs = {}
+    manager.loras = {}
+    manager.lora_refs = {}
+    manager.num_pinned_loras = 0
+    manager.max_loras_per_batch = 4
+    manager.base_hf_config = MagicMock(vocab_size=32000)
+    manager.lora_modules = []
+    manager.embed_tokens_module = None
+    manager.lm_head_module = None
+    manager.device = torch.device("cpu")
+    manager.memory_pool = MagicMock()
+    manager.memory_pool.can_support.return_value = True
+    manager.memory_pool.uid_to_buffer_id = {}
+    # Weight loading needs a real base model / backend; just build a stub.
+    manager._create_lora_adapter_from_tensors = Mock(
+        side_effect=lambda ref, config, tensors: MagicMock()
+    )
+    return manager
+
+
+class TestLoRAManagerUpsert(CustomTestCase):
+    def test_fresh_load_registers_adapter(self):
+        manager = _make_manager()
+        ref = LoRARef(lora_name="a", lora_path="__tensor__", pinned=True)
+
+        result = manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT)
+
+        self.assertTrue(result.success)
+        self.assertIn(ref.lora_id, manager.loras)
+        self.assertIs(manager.lora_refs[ref.lora_id], ref)
+        self.assertEqual(manager.num_pinned_loras, 1)
+
+    def test_duplicate_load_without_upsert_asserts(self):
+        manager = _make_manager()
+        ref = LoRARef(lora_name="a", lora_path="__tensor__")
+        manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT)
+
+        with self.assertRaises(AssertionError):
+            manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT)
+
+    def test_upsert_refreshes_loaded_adapter_in_place(self):
+        manager = _make_manager()
+        ref = LoRARef(lora_name="a", lora_path="__tensor__", pinned=True)
+        manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT)
+        # Simulate the adapter occupying a memory-pool buffer slot.
+        manager.memory_pool.uid_to_buffer_id = {ref.lora_id: 3}
+
+        new_config = dict(CONFIG_DICT, lora_alpha=32)
+        result = manager.load_lora_adapter_from_tensors(
+            ref, {}, new_config, upsert=True
+        )
+
+        self.assertTrue(result.success)
+        self.assertEqual(manager.configs[ref.lora_id].lora_alpha, 32)
+        # Weights are re-copied into the existing buffer slot.
+        manager.memory_pool.load_lora_weight_to_buffer.assert_called_once()
+        call = manager.memory_pool.load_lora_weight_to_buffer.call_args
+        self.assertEqual(call.args[0], ref.lora_id)
+        self.assertEqual(call.args[1], 3)
+        # The pinned slot is not double-counted.
+        self.assertEqual(manager.num_pinned_loras, 1)
+
+    def test_upsert_skips_pool_copy_when_not_resident(self):
+        manager = _make_manager()
+        ref = LoRARef(lora_name="a", lora_path="__tensor__")
+        manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT)
+
+        result = manager.load_lora_adapter_from_tensors(
+            ref, {}, CONFIG_DICT, upsert=True
+        )
+
+        self.assertTrue(result.success)
+        manager.memory_pool.load_lora_weight_to_buffer.assert_not_called()
+
+    def test_upsert_falls_back_to_register_when_not_loaded(self):
+        manager = _make_manager()
+        ref = LoRARef(lora_name="a", lora_path="__tensor__", pinned=True)
+
+        result = manager.load_lora_adapter_from_tensors(
+            ref, {}, CONFIG_DICT, upsert=True
+        )
+
+        self.assertTrue(result.success)
+        self.assertIn(ref.lora_id, manager.loras)
+        self.assertEqual(manager.num_pinned_loras, 1)
+
+
+class TestLoRARegistryRegisterOrReuse(CustomTestCase):
+    def test_upsert_reuses_id_of_registered_adapter(self):
+        registry = LoRARegistry()
+        existing = LoRARef(lora_name="a", lora_path="/x", pinned=False)
+        asyncio.run(registry.register(existing))
+
+        candidate = LoRARef(lora_name="a", lora_path="__tensor__", pinned=True)
+        resolved, reused = asyncio.run(registry.register_or_reuse(candidate, True))
+
+        self.assertTrue(reused)
+        self.assertEqual(resolved.lora_id, existing.lora_id)
+        self.assertEqual(resolved.lora_path, "__tensor__")
+        self.assertTrue(resolved.pinned)
+
+    def test_upsert_without_registered_adapter_keeps_fresh_id(self):
+        registry = LoRARegistry()
+        candidate = LoRARef(lora_name="a", lora_path="__tensor__")
+
+        resolved, reused = asyncio.run(registry.register_or_reuse(candidate, True))
+
+        self.assertFalse(reused)
+        self.assertIs(resolved, candidate)
+
+    def test_non_upsert_never_reuses(self):
+        registry = LoRARegistry()
+        asyncio.run(registry.register(LoRARef(lora_name="a", lora_path="/x")))
+
+        candidate = LoRARef(lora_name="a", lora_path="/x")
+        resolved, reused = asyncio.run(registry.register_or_reuse(candidate, False))
+
+        self.assertFalse(reused)
+        self.assertIs(resolved, candidate)
+
+    def test_refresh_replaces_ref_in_place(self):
+        registry = LoRARegistry()
+        existing = LoRARef(lora_name="a", lora_path="/x", pinned=False)
+        asyncio.run(registry.register(existing))
+
+        refreshed = LoRARef(
+            lora_id=existing.lora_id,
+            lora_name="a",
+            lora_path="__tensor__",
+            pinned=True,
+        )
+        asyncio.run(registry.refresh(refreshed))
+
+        self.assertEqual(registry.get_all_adapters()["a"].pinned, True)
+        self.assertEqual(asyncio.run(registry.get_lora_id("a")), existing.lora_id)
+
+    def test_refresh_rejects_id_mismatch(self):
+        registry = LoRARegistry()
+        asyncio.run(registry.register(LoRARef(lora_name="a", lora_path="/x")))
+
+        with self.assertRaises(AssertionError):
+            asyncio.run(
+                registry.refresh(LoRARef(lora_name="a", lora_path="__tensor__"))
+            )
+
+
+class TestUpsertRollback(CustomTestCase):
+    """A failed load/upsert must not leave a live adapter half-updated."""
+
+    def test_failed_fresh_load_leaves_no_state(self):
+        manager = _make_manager()
+        manager._create_lora_adapter_from_tensors = Mock(
+            side_effect=ValueError("bad tensors")
+        )
+        ref = LoRARef(lora_name="a", lora_path="__tensor__")
+
+        result = manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT)
+
+        self.assertFalse(result.success)
+        self.assertIn("bad tensors", result.error_message)
+        self.assertEqual(manager.configs, {})
+        self.assertEqual(manager.loras, {})
+        self.assertEqual(manager.lora_refs, {})
+
+    def test_failed_upsert_staging_keeps_old_adapter_serving(self):
+        manager = _make_manager()
+        ref = LoRARef(lora_name="a", lora_path="__tensor__")
+        manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT)
+        old_config = manager.configs[ref.lora_id]
+        old_lora = manager.loras[ref.lora_id]
+
+        manager._create_lora_adapter_from_tensors = Mock(
+            side_effect=ValueError("rank mismatch")
+        )
+        result = manager.load_lora_adapter_from_tensors(
+            ref, {}, CONFIG_DICT, upsert=True
+        )
+
+        self.assertFalse(result.success)
+        self.assertIs(manager.configs[ref.lora_id], old_config)
+        self.assertIs(manager.loras[ref.lora_id], old_lora)
+        manager.memory_pool.load_lora_weight_to_buffer.assert_not_called()
+
+    def test_failed_buffer_rewrite_restores_old_weights(self):
+        manager = _make_manager()
+        ref = LoRARef(lora_name="a", lora_path="__tensor__")
+        manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT)
+        old_config = manager.configs[ref.lora_id]
+        old_lora = manager.loras[ref.lora_id]
+        manager.memory_pool.uid_to_buffer_id = {ref.lora_id: 3}
+        manager.memory_pool.load_lora_weight_to_buffer.side_effect = [
+            RuntimeError("copy failed at layer k"),
+            None,  # the restore pass
+        ]
+
+        result = manager.load_lora_adapter_from_tensors(
+            ref, {}, CONFIG_DICT, upsert=True
+        )
+
+        self.assertFalse(result.success)
+        self.assertIn("copy failed", result.error_message)
+        # CPU-side state rolled back...
+        self.assertIs(manager.configs[ref.lora_id], old_config)
+        self.assertIs(manager.loras[ref.lora_id], old_lora)
+        # ...and the served buffer was rewritten back from the old adapter.
+        calls = manager.memory_pool.load_lora_weight_to_buffer.call_args_list
+        self.assertEqual(len(calls), 2)
+        self.assertIs(calls[1].args[2], old_lora)
+
+
+class TestUpsertPinnedAccounting(CustomTestCase):
+    def test_pinned_flip_updates_counter_both_ways(self):
+        manager = _make_manager()
+        unpinned = LoRARef(lora_name="a", lora_path="__tensor__", pinned=False)
+        manager.load_lora_adapter_from_tensors(unpinned, {}, CONFIG_DICT)
+        self.assertEqual(manager.num_pinned_loras, 0)
+
+        pinned = LoRARef(
+            lora_id=unpinned.lora_id,
+            lora_name="a",
+            lora_path="__tensor__",
+            pinned=True,
+        )
+        manager.load_lora_adapter_from_tensors(pinned, {}, CONFIG_DICT, upsert=True)
+        self.assertEqual(manager.num_pinned_loras, 1)
+
+        manager.load_lora_adapter_from_tensors(unpinned, {}, CONFIG_DICT, upsert=True)
+        self.assertEqual(manager.num_pinned_loras, 0)
+
+    def test_unload_after_pinned_flip_keeps_counter_consistent(self):
+        manager = _make_manager()
+        unpinned = LoRARef(lora_name="a", lora_path="__tensor__", pinned=False)
+        manager.load_lora_adapter_from_tensors(unpinned, {}, CONFIG_DICT)
+        pinned = LoRARef(
+            lora_id=unpinned.lora_id,
+            lora_name="a",
+            lora_path="__tensor__",
+            pinned=True,
+        )
+        manager.load_lora_adapter_from_tensors(pinned, {}, CONFIG_DICT, upsert=True)
+
+        result = manager.unload_lora_adapter(pinned)
+
+        self.assertTrue(result.success)
+        self.assertEqual(manager.num_pinned_loras, 0)
+
+    def test_pinned_refresh_allowed_at_pin_limit(self):
+        """Refreshing a pinned adapter adds no pinned slot; rejecting it would
+        freeze RL serving on the step-1 weights."""
+        manager = _make_manager()
+        manager.max_loras_per_batch = 2
+        ref = LoRARef(lora_name="a", lora_path="__tensor__", pinned=True)
+        manager.load_lora_adapter_from_tensors(ref, {}, CONFIG_DICT)
+        self.assertEqual(manager.num_pinned_loras, 1)
+
+        result = manager.load_lora_adapter_from_tensors(
+            ref, {}, CONFIG_DICT, upsert=True
+        )
+
+        self.assertTrue(result.success)
+        self.assertEqual(manager.num_pinned_loras, 1)
+
+    def test_fresh_pinned_load_still_rejected_at_pin_limit(self):
+        manager = _make_manager()
+        manager.max_loras_per_batch = 2
+        manager.load_lora_adapter_from_tensors(
+            LoRARef(lora_name="a", lora_path="__tensor__", pinned=True),
+            {},
+            CONFIG_DICT,
+        )
+
+        result = manager.load_lora_adapter_from_tensors(
+            LoRARef(lora_name="b", lora_path="__tensor__", pinned=True),
+            {},
+            CONFIG_DICT,
+        )
+
+        self.assertFalse(result.success)
+        self.assertIn("not allowed to pin all slots", result.error_message)
+
+
+class TestValidateNewAdapterDuplicates(CustomTestCase):
+    def test_duplicate_name_rejected_without_update(self):
+        manager = _make_manager()
+        existing = LoRARef(lora_name="a", lora_path="/x")
+        manager.lora_refs[existing.lora_id] = existing
+        config = MagicMock(lora_added_tokens_size=0, use_dora=False)
+
+        with self.assertRaisesRegex(ValueError, "already loaded"):
+            manager.validate_new_adapter(config, LoRARef(lora_name="a", lora_path="/y"))
+
+    def test_duplicate_name_allowed_for_update(self):
+        manager = _make_manager()
+        existing = LoRARef(lora_name="a", lora_path="/x")
+        manager.lora_refs[existing.lora_id] = existing
+        config = MagicMock(lora_added_tokens_size=0, use_dora=False)
+
+        manager.validate_new_adapter(config, existing, is_update=True)
+
+
+def _make_tokenizer_manager(tokenizer_worker_num: int = 1) -> TokenizerManager:
+    tm = TokenizerManager.__new__(TokenizerManager)
+    tm.server_args = MagicMock()
+    tm.server_args.enable_lora = True
+    tm.server_args.dp_size = 1
+    tm.server_args.max_loaded_loras = None
+    tm.server_args.tokenizer_worker_num = tokenizer_worker_num
+    tm.auto_create_handle_loop = Mock()
+    tm.lora_update_lock = asyncio.Lock()
+    tm.lora_registry = LoRARegistry()
+    tm.lora_ref_cache = {}
+    tm.update_lora_adapter_communicator = AsyncMock(
+        return_value=[MagicMock(success=True)]
+    )
+    return tm
+
+
+def _make_distributed_req(upsert: bool) -> LoadLoRAAdapterFromDistributedReqInput:
+    return LoadLoRAAdapterFromDistributedReqInput(
+        lora_name="a",
+        config_dict=CONFIG_DICT,
+        names=[],
+        dtypes=[],
+        shapes=[],
+        upsert=upsert,
+    )
+
+
+def _make_tensors_req(
+    upsert: bool, pinned: bool = False
+) -> LoadLoRAAdapterFromTensorsReqInput:
+    return LoadLoRAAdapterFromTensorsReqInput(
+        lora_name="a",
+        config_dict=CONFIG_DICT,
+        serialized_named_tensors=[],
+        pinned=pinned,
+        upsert=upsert,
+    )
+
+
+class TestLoadFromDistributedUpsert(CustomTestCase):
+    def test_upsert_reuses_existing_lora_id(self):
+        tm = _make_tokenizer_manager()
+        existing = LoRARef(lora_name="a", lora_path="__distributed__")
+        asyncio.run(tm.lora_registry.register(existing))
+
+        obj = _make_distributed_req(upsert=True)
+        result = asyncio.run(tm.load_lora_adapter_from_distributed(obj))
+
+        self.assertTrue(result.success)
+        self.assertEqual(obj.lora_id, existing.lora_id)
+        tm.update_lora_adapter_communicator.assert_awaited_once_with(obj)
+        # Not re-registered: still exactly one adapter with the original id.
+        self.assertEqual(tm.lora_registry.num_registered_loras, 1)
+        self.assertEqual(
+            asyncio.run(tm.lora_registry.get_lora_id("a")), existing.lora_id
+        )
+        self.assertEqual(tm.lora_ref_cache["a"].lora_id, existing.lora_id)
+
+    def test_upsert_registers_when_missing(self):
+        tm = _make_tokenizer_manager()
+
+        obj = _make_distributed_req(upsert=True)
+        result = asyncio.run(tm.load_lora_adapter_from_distributed(obj))
+
+        self.assertTrue(result.success)
+        self.assertIsNotNone(obj.lora_id)
+        self.assertEqual(asyncio.run(tm.lora_registry.get_lora_id("a")), obj.lora_id)
+
+    def test_non_upsert_duplicate_fails(self):
+        tm = _make_tokenizer_manager()
+        asyncio.run(
+            tm.lora_registry.register(
+                LoRARef(lora_name="a", lora_path="__distributed__")
+            )
+        )
+
+        obj = _make_distributed_req(upsert=False)
+        result = asyncio.run(tm.load_lora_adapter_from_distributed(obj))
+
+        self.assertFalse(result.success)
+        self.assertIn("already exists", result.error_message)
+
+    def test_upsert_refreshes_registered_ref(self):
+        # The registry ref (not just lora_ref_cache) must adopt the new
+        # metadata: LRU eviction reads ``pinned`` from the registry.
+        tm = _make_tokenizer_manager()
+        existing = LoRARef(lora_name="a", lora_path="__distributed__", pinned=False)
+        asyncio.run(tm.lora_registry.register(existing))
+
+        obj = _make_distributed_req(upsert=True)
+        obj.pinned = True
+        result = asyncio.run(tm.load_lora_adapter_from_distributed(obj))
+
+        self.assertTrue(result.success)
+        registered = tm.lora_registry.get_all_adapters()["a"]
+        self.assertEqual(registered.lora_id, existing.lora_id)
+        self.assertTrue(registered.pinned)
+
+    def test_failed_backend_load_keeps_registry_untouched(self):
+        tm = _make_tokenizer_manager()
+        existing = LoRARef(lora_name="a", lora_path="__distributed__", pinned=False)
+        asyncio.run(tm.lora_registry.register(existing))
+        tm.update_lora_adapter_communicator = AsyncMock(
+            return_value=[MagicMock(success=False, error_message="boom")]
+        )
+
+        obj = _make_distributed_req(upsert=True)
+        obj.pinned = True
+        result = asyncio.run(tm.load_lora_adapter_from_distributed(obj))
+
+        self.assertFalse(result.success)
+        self.assertIs(tm.lora_registry.get_all_adapters()["a"], existing)
+        self.assertNotIn("a", tm.lora_ref_cache)
+
+
+class TestLoadFromTensorsUpsertUnsupported(CustomTestCase):
+    """Only the from_distributed route supports in-place refresh; the
+    from_tensors route must reject upsert explicitly instead of minting a
+    fresh uuid and dying later on the backend duplicate check."""
+
+    def test_upsert_rejected_explicitly(self):
+        tm = _make_tokenizer_manager()
+        asyncio.run(
+            tm.lora_registry.register(LoRARef(lora_name="a", lora_path="__tensor__"))
+        )
+
+        obj = _make_tensors_req(upsert=True)
+        result = asyncio.run(tm.load_lora_adapter_from_tensors(obj))
+
+        self.assertFalse(result.success)
+        self.assertIn("not supported on the from_tensors route", result.error_message)
+        tm.update_lora_adapter_communicator.assert_not_awaited()
+
+    def test_non_upsert_load_still_works(self):
+        tm = _make_tokenizer_manager()
+
+        obj = _make_tensors_req(upsert=False)
+        result = asyncio.run(tm.load_lora_adapter_from_tensors(obj))
+
+        self.assertTrue(result.success)
+        self.assertEqual(asyncio.run(tm.lora_registry.get_lora_id("a")), obj.lora_id)
+
+    def test_non_upsert_duplicate_fails(self):
+        tm = _make_tokenizer_manager()
+        asyncio.run(
+            tm.lora_registry.register(LoRARef(lora_name="a", lora_path="__tensor__"))
+        )
+
+        obj = _make_tensors_req(upsert=False)
+        result = asyncio.run(tm.load_lora_adapter_from_tensors(obj))
+
+        self.assertFalse(result.success)
+        self.assertIn("already exists", result.error_message)
+
+
+class TestUpsertMultiTokenizerWorkerGuard(CustomTestCase):
+    """Upsert resolves names against a per-process registry; with >1 tokenizer
+    workers that resolution is nondeterministic, so it must fail loudly."""
+
+    def test_distributed_upsert_rejected_with_multiple_workers(self):
+        tm = _make_tokenizer_manager(tokenizer_worker_num=2)
+
+        result = asyncio.run(
+            tm.load_lora_adapter_from_distributed(_make_distributed_req(True))
+        )
+
+        self.assertFalse(result.success)
+        self.assertIn("tokenizer_worker_num", result.error_message)
+
+    def test_non_upsert_load_unaffected_by_multiple_workers(self):
+        tm = _make_tokenizer_manager(tokenizer_worker_num=2)
+
+        result = asyncio.run(
+            tm.load_lora_adapter_from_tensors(_make_tensors_req(False))
+        )
+
+        self.assertTrue(result.success)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

From 9e3f38aac3fc13f19db699e18f63a03725017e7e Mon Sep 17 00:00:00 2001
From: JD-ETH <jaedon.guo@gmail.com>
Date: Fri, 24 Jul 2026 21:39:52 -0700
Subject: [PATCH 19/43] [14/27] [sglang-miles] P2P weight update support and
 fixes (#21278, #22663)

Includes making Cohere2MoeConfig a dataclass so parameter mapping can
introspect config fields for P2P weight update.

Rebased onto v0.5.16, which extracted the remote-instance transfer engine out
of ModelRunner into RemoteInstanceWeightTransporter: the per-rank
RankParallelismConfig is now built in `init_engine()` and published from
`maybe_register_and_publish_weight_info()`, instead of the ModelRunner methods
this commit originally added. The hoisted deepseek expert_params_mapping also
picks up v0.5.16's broader `is_wint4afp8_or_wint4a16_config` predicate rather
than the old `quant_config.get_name() == "w4afp8"` check.

Co-authored-by: JensenFire <elevenji@outlook.com>
---
 python/sglang/srt/configs/cohere2_moe.py      |  10 +-
 .../sglang/srt/distributed/parallel_state.py  | 180 ++++++++++-
 .../engine_info_bootstrap_server.py           |  41 ++-
 python/sglang/srt/entrypoints/http_server.py  |  26 ++
 .../remote_instance_weight_transporter.py     |  50 +++
 .../srt/model_loader/parameter_mapper.py      | 261 ++++++++++++++++
 .../deepseek_common/deepseek_weight_loader.py |  66 ++--
 python/sglang/srt/models/deepseek_v2.py       |  47 +++
 python/sglang/srt/models/glm4.py              |  20 +-
 python/sglang/srt/models/glm4_moe.py          |  70 ++---
 python/sglang/srt/models/glm4_moe_lite.py     |  43 +++
 python/sglang/srt/models/glm4_moe_nextn.py    |  17 +
 python/sglang/srt/models/llama.py             |  28 +-
 python/sglang/srt/models/qwen2.py             |  20 +-
 python/sglang/srt/models/qwen3.py             |  32 +-
 python/sglang/srt/models/qwen3_moe.py         |  34 +-
 .../test_parallelism_context_integration.py   | 275 +++++++++++++++++
 test/srt/models/test_params_mapping.py        | 292 ++++++++++++++++++
 18 files changed, 1362 insertions(+), 150 deletions(-)
 create mode 100644 python/sglang/srt/model_loader/parameter_mapper.py
 create mode 100644 test/registered/disaggregation/test_parallelism_context_integration.py
 create mode 100644 test/srt/models/test_params_mapping.py

diff --git a/python/sglang/srt/configs/cohere2_moe.py b/python/sglang/srt/configs/cohere2_moe.py
index cd470bd69f49..8603c704644f 100644
--- a/python/sglang/srt/configs/cohere2_moe.py
+++ b/python/sglang/srt/configs/cohere2_moe.py
@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 """Cohere2Moe text config used by the Cohere Command-A Plus checkpoints."""
 
+from dataclasses import dataclass
+
 from transformers.configuration_utils import PreTrainedConfig
 from transformers.models.auto.configuration_auto import CONFIG_MAPPING
 
@@ -13,6 +15,7 @@ def strict(cls):  # type: ignore[misc]
 
 
 @strict
+@dataclass
 class Cohere2MoeConfig(PreTrainedConfig):
     model_type = "cohere2_moe"
     keys_to_ignore_at_inference = ["past_key_values"]
@@ -52,6 +55,9 @@ class Cohere2MoeConfig(PreTrainedConfig):
     rms_norm_eps: float | None = None
     sliding_window_pattern: int = 4
 
+    def validate_rope(self):
+        return super().validate_rope()
+
     def __post_init__(self, **kwargs):
         if self.num_key_value_heads is None:
             self.num_key_value_heads = self.num_attention_heads
@@ -82,7 +88,9 @@ def __post_init__(self, **kwargs):
             ]
             self.layer_types = prefix_layers + rest_layers
 
-        super().__post_init__(**kwargs)
+        post_init = getattr(super(), "__post_init__", None)
+        if post_init is not None:
+            post_init(**kwargs)
 
 
 try:
diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py
index 58a5a6469e51..03d0ee89554c 100644
--- a/python/sglang/srt/distributed/parallel_state.py
+++ b/python/sglang/srt/distributed/parallel_state.py
@@ -32,11 +32,11 @@
 import weakref
 from collections import namedtuple
 from contextlib import contextmanager, nullcontext
-from dataclasses import dataclass
+from dataclasses import asdict, dataclass
 from datetime import timedelta
 from multiprocessing import shared_memory
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import torch
 import torch.distributed
@@ -2858,3 +2858,179 @@ def monkey_patch_vllm_parallel_state(reverse: bool = False):
         setattr(vllm_parallel_state, "get_pp_group", get_pp_group)
         setattr(vllm_parallel_state, "get_tp_group", get_tp_group)
         setattr(vllm_parallel_state, "get_world_group", get_world_group)
+
+
+@dataclass
+class RankParallelismConfig:
+    """
+    Complete parallelism configuration for a single inference rank.
+
+    This configuration captures all the parallelism settings needed to recreate
+    a model shard outside of sglang. It supports:
+    - TP/PP/EP for model parallelism
+    - MoE-TP/Attn-TP/Attn-DP for MoE and DP attention.
+    """
+
+    tp_size: int = 1
+    tp_rank: int = 0
+    pp_size: int = 1
+    pp_rank: int = 0
+    ep_size: int = 1
+    ep_rank: int = 0
+    moe_tp_size: int = 1
+    moe_tp_rank: int = 0
+    attn_tp_size: int = 1
+    attn_tp_rank: int = 0
+    attn_dp_size: int = 1
+    attn_dp_rank: int = 0
+    attn_cp_size: int = 1
+    attn_cp_rank: int = 0
+    moe_dp_size: int = 1
+    moe_dp_rank: int = 0
+
+    world_size: int = 1
+    global_rank: int = 0
+    local_rank: int = 0
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "RankParallelismConfig":
+        """Create from dictionary, filtering unknown fields."""
+        import dataclasses
+
+        valid_fields = {f.name for f in dataclasses.fields(cls)}
+        filtered_data = {k: v for k, v in data.items() if k in valid_fields}
+        return cls(**filtered_data)
+
+    @classmethod
+    def from_parallel_state(cls, local_rank: int = 0) -> "RankParallelismConfig":
+        """Extract current parallelism settings from the global parallel state."""
+        tp_size = get_tensor_model_parallel_world_size()
+        tp_rank = get_tensor_model_parallel_rank()
+
+        # Import dp_attention lazily to avoid circular imports
+        from sglang.srt.layers.dp_attention import (
+            get_attention_cp_rank,
+            get_attention_cp_size,
+            get_attention_dp_rank,
+            get_attention_dp_size,
+            get_attention_tp_rank,
+            get_attention_tp_size,
+        )
+
+        return cls(
+            tp_size=tp_size,
+            tp_rank=tp_rank,
+            pp_size=get_pipeline_model_parallel_world_size(),
+            pp_rank=get_pipeline_model_parallel_rank(),
+            ep_size=get_moe_expert_parallel_world_size(),
+            ep_rank=get_moe_expert_parallel_rank(),
+            moe_tp_size=get_moe_tensor_parallel_world_size(),
+            moe_tp_rank=get_moe_tensor_parallel_rank(),
+            attn_tp_size=get_attention_tp_size(),
+            attn_tp_rank=get_attention_tp_rank(),
+            attn_dp_size=get_attention_dp_size(),
+            attn_dp_rank=get_attention_dp_rank(),
+            attn_cp_size=get_attention_cp_size(),
+            attn_cp_rank=get_attention_cp_rank(),
+            moe_dp_size=get_moe_data_parallel_world_size(),
+            moe_dp_rank=get_moe_data_parallel_rank(),
+            world_size=(
+                torch.distributed.get_world_size()
+                if torch.distributed.is_initialized()
+                else 1
+            ),
+            global_rank=(
+                torch.distributed.get_rank()
+                if torch.distributed.is_initialized()
+                else 0
+            ),
+            local_rank=local_rank,
+        )
+
+
+# Globals on parallel_state module to save/restore
+_PS_GLOBALS = ("_TP", "_PP", "_MOE_EP", "_MOE_TP", "_ATTN_TP", "_ATTN_CP", "_MOE_DP")
+# Globals on dp_attention module to save/restore
+_DA_GLOBALS = ("_ATTN_DP_RANK", "_ATTN_DP_SIZE", "_ENABLE_DP_ATTENTION_FLAG")
+
+
+class ParallelismContext:
+    """
+    Context manager for creating model replicas with specific parallelism settings.
+
+    Temporarily sets global variables to allow creating model shards outside of a
+    real distributed environment.
+    Usage:
+        with ParallelismContext(RankParallelismConfig.from_dict(parallelism_info)):
+            model = get_model(...)
+    """
+
+    def __init__(self, parallelism_config: RankParallelismConfig):
+        self.config = parallelism_config
+        self._original_globals: Dict[str, Any] = {}
+
+    def _create_mock_group(self, world_size: int, rank_in_group: int):
+        """Create a mock group coordinator with all necessary properties."""
+        mock_group = MagicMock()
+        mock_group.world_size = world_size
+        mock_group.rank_in_group = rank_in_group
+        mock_group.rank = rank_in_group
+        mock_group.local_rank = rank_in_group
+        mock_group.ranks = list(range(world_size))
+        mock_group.first_rank = 0
+        mock_group.last_rank = world_size - 1
+        mock_group.is_first_rank = rank_in_group == 0
+        mock_group.is_last_rank = rank_in_group == world_size - 1
+        mock_group.next_rank = mock_group.ranks[(rank_in_group + 1) % world_size]
+        mock_group.prev_rank = mock_group.ranks[(rank_in_group - 1) % world_size]
+        return mock_group
+
+    def __enter__(self):
+        conf = self.config
+
+        from sglang.srt.distributed import parallel_state
+        from sglang.srt.layers import dp_attention
+
+        # Save original globals
+        for name in _PS_GLOBALS:
+            self._original_globals[name] = getattr(parallel_state, name, None)
+        for name in _DA_GLOBALS:
+            self._original_globals[name] = getattr(dp_attention, name, None)
+
+        # Build and set mock group objects on parallel_state
+        _ps_new_values = {
+            "_TP": self._create_mock_group(conf.tp_size, conf.tp_rank),
+            "_PP": self._create_mock_group(conf.pp_size, conf.pp_rank),
+            "_MOE_EP": self._create_mock_group(conf.ep_size, conf.ep_rank),
+            "_MOE_TP": self._create_mock_group(conf.moe_tp_size, conf.moe_tp_rank),
+            "_ATTN_TP": self._create_mock_group(conf.attn_tp_size, conf.attn_tp_rank),
+            "_ATTN_CP": self._create_mock_group(conf.attn_cp_size, conf.attn_cp_rank),
+            "_MOE_DP": self._create_mock_group(conf.moe_dp_size, conf.moe_dp_rank),
+        }
+        for name, value in _ps_new_values.items():
+            setattr(parallel_state, name, value)
+
+        # Set dp_attention scalar globals
+        dp_attention._ATTN_DP_RANK = conf.attn_dp_rank
+        dp_attention._ATTN_DP_SIZE = conf.attn_dp_size
+        dp_attention._ENABLE_DP_ATTENTION_FLAG = conf.attn_dp_size > 1
+
+        logger.info(f"[ParallelismContext] Activated: {conf}")
+        return self
+
+    def __exit__(self, *args):
+        from sglang.srt.distributed import parallel_state
+        from sglang.srt.layers import dp_attention
+
+        # Restore original globals
+        for name in _PS_GLOBALS:
+            setattr(parallel_state, name, self._original_globals.get(name))
+        for name in _DA_GLOBALS:
+            setattr(dp_attention, name, self._original_globals.get(name))
+
+        logger.info("[ParallelismContext] Deactivated")
+        return False
diff --git a/python/sglang/srt/entrypoints/engine_info_bootstrap_server.py b/python/sglang/srt/entrypoints/engine_info_bootstrap_server.py
index 77de7fc7d030..88e48075a644 100644
--- a/python/sglang/srt/entrypoints/engine_info_bootstrap_server.py
+++ b/python/sglang/srt/entrypoints/engine_info_bootstrap_server.py
@@ -31,7 +31,8 @@ class EngineInfoBootstrapServer:
     accesses the collected info directly in-process; external consumers can
     query via HTTP GET.
 
-    Currently supports transfer engine memory registration info.
+    Currently supports transfer engine memory registration info and
+    per-rank parallelism configuration.
     """
 
     def __init__(self, host: str, port: int):
@@ -40,6 +41,8 @@ def __init__(self, host: str, port: int):
 
         # Storage: {tp_rank: (session_id, weights_info_dict)}
         self.transfer_engine_info: Dict[int, Tuple] = {}
+        # Storage: {tp_rank: parallelism_config_dict}
+        self.parallelism_config: Dict[int, dict] = {}
         self.lock = threading.Lock()
 
         app = FastAPI()
@@ -89,6 +92,38 @@ def get_transfer_engine_info(rank: int):
 
         config = uvicorn.Config(app, host=host, port=port, log_level="warning")
         self._server = uvicorn.Server(config)
+
+        @app.put("/register_parallelism_config")
+        def register_parallelism_config(data: dict):
+            try:
+                tp_rank = data["tp_rank"]
+                config = data["parallelism_config"]
+
+                with self.lock:
+                    self.parallelism_config[tp_rank] = config
+
+                logger.info(f"Registered parallelism config for tp_rank={tp_rank}")
+                return PlainTextResponse("OK")
+            except Exception as e:
+                logger.error(f"Failed to register parallelism config: {e}")
+                raise HTTPException(status_code=400, detail=str(e))
+
+        @app.get("/get_parallelism_config")
+        def get_parallelism_config(rank: int):
+            if rank < 0:
+                raise HTTPException(status_code=400, detail="Invalid rank parameter")
+
+            with self.lock:
+                config = self.parallelism_config.get(rank)
+
+            if config is None:
+                raise HTTPException(
+                    status_code=404,
+                    detail=f"No parallelism config for rank {rank}",
+                )
+
+            return config
+
         self._thread = threading.Thread(
             target=self._server.run,
             daemon=True,
@@ -103,3 +138,7 @@ def close(self):
     def get_transfer_engine_info(self, rank: int) -> Optional[Tuple]:
         """Direct in-process access for co-located HTTP server (no HTTP round-trip)."""
         return self.transfer_engine_info.get(rank)
+
+    def get_parallelism_config_info(self, rank: int) -> Optional[dict]:
+        """Direct in-process access for parallelism config (no HTTP round-trip)."""
+        return self.parallelism_config.get(rank)
diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py
index 06a652673e29..f5b8bf6277f1 100644
--- a/python/sglang/srt/entrypoints/http_server.py
+++ b/python/sglang/srt/entrypoints/http_server.py
@@ -1273,6 +1273,32 @@ async def remote_instance_transfer_engine_info(rank: int = None):
     )
 
 
+@app.get("/parallelism_config")
+@auth_level(AuthLevel.ADMIN_OPTIONAL)
+async def parallelism_config(rank: int = None):
+    """Get per-rank parallelism config from the bootstrap server."""
+    if rank is None or rank < 0:
+        return ORJSONResponse(
+            {"error": {"message": "Missing or invalid rank parameter"}},
+            status_code=HTTPStatus.BAD_REQUEST,
+        )
+
+    server_args = _global_state.tokenizer_manager.server_args
+    try:
+
+        resp = requests.get(
+            f"{server_args.engine_info_bootstrap_url}/get_parallelism_config",
+            params={"rank": rank},
+            timeout=5,
+        )
+        if resp.status_code == 200:
+            return resp.json()
+    except Exception:
+        pass
+
+    return Response(status_code=HTTPStatus.BAD_REQUEST)
+
+
 @app.post("/init_weights_update_group")
 @auth_level(AuthLevel.ADMIN_OPTIONAL)
 async def init_weights_update_group(
diff --git a/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py b/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py
index bf56487ad4a3..476e8070c6d9 100644
--- a/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py
+++ b/python/sglang/srt/model_executor/model_runner_components/remote_instance_weight_transporter.py
@@ -6,6 +6,7 @@
 
 import torch
 
+from sglang.srt.distributed.parallel_state import RankParallelismConfig
 from sglang.srt.environ import envs
 from sglang.srt.model_loader.remote_instance_weight_loader_utils import (
     RemoteInstanceWeightLoaderBackend,
@@ -26,6 +27,7 @@ class RemoteInstanceWeightTransporter:
     engine: Optional[Any] = None
     session_id: str = ""
     weight_info: Optional[dict[str, tuple[int, int, int]]] = None
+    parallelism_config: Optional[RankParallelismConfig] = None
     _nixl_manager: Optional[Any] = None
 
     @property
@@ -51,6 +53,9 @@ def init_engine(self):
         self.session_id = NetworkAddress(
             local_ip, self.engine.get_rpc_port()
         ).to_host_port_str()
+        self.parallelism_config = RankParallelismConfig.from_parallel_state(
+            self.tp_rank
+        )
 
     def maybe_register_and_publish_weight_info(self) -> None:
         if (
@@ -67,6 +72,51 @@ def maybe_register_and_publish_weight_info(self) -> None:
             self.weight_info = register_memory_region(self.model, self.engine)
             self._register_to_engine_info_bootstrap()
 
+        # The P2P weight-update client needs each rank's parallelism layout to
+        # map training-side parameters onto this rank's shards.
+        if (
+            self.server_args.remote_instance_weight_loader_use_transfer_engine()
+            and self.parallelism_config is not None
+        ):
+            self._register_parallelism_config_to_bootstrap()
+
+    def _bootstrap_url(self) -> str:
+        if self.server_args.dist_init_addr:
+            bootstrap_host = (
+                NetworkAddress.parse(self.server_args.dist_init_addr).resolved().host
+            )
+        else:
+            bootstrap_host = "127.0.0.1"
+        bootstrap_port = self.server_args.engine_info_bootstrap_port
+        return NetworkAddress(bootstrap_host, bootstrap_port).to_url()
+
+    def _register_parallelism_config_to_bootstrap(self) -> None:
+        """Register this rank's parallelism config with the EngineInfoBootstrapServer."""
+        import requests as http_requests
+
+        bootstrap_url = self._bootstrap_url()
+        url = f"{bootstrap_url}/register_parallelism_config"
+        payload = {
+            "tp_rank": self.tp_rank,
+            "parallelism_config": self.parallelism_config.to_dict(),
+        }
+        try:
+            resp = http_requests.put(url, json=payload, timeout=5)
+            if resp.status_code == 200:
+                logger.info(
+                    f"Registered parallelism config for tp_rank={self.tp_rank} "
+                    f"with bootstrap server at {bootstrap_url}"
+                )
+            else:
+                logger.error(
+                    f"Failed to register parallelism config for tp_rank={self.tp_rank}: "
+                    f"{resp.status_code}, {resp.text}"
+                )
+        except Exception as e:
+            logger.error(
+                f"Failed to register parallelism config for tp_rank={self.tp_rank}: {e}"
+            )
+
     def _register_to_engine_info_bootstrap(self: RemoteInstanceWeightTransporter):
         """Register transfer engine info with the EngineInfoBootstrapServer via HTTP PUT.
 
diff --git a/python/sglang/srt/model_loader/parameter_mapper.py b/python/sglang/srt/model_loader/parameter_mapper.py
new file mode 100644
index 000000000000..56056d570843
--- /dev/null
+++ b/python/sglang/srt/model_loader/parameter_mapper.py
@@ -0,0 +1,261 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Parameter mapping from HuggingFace checkpoint names to SGLang model parameters.
+
+This module provides utilities for translating weight names between HuggingFace
+checkpoint format and SGLang's internal parameter naming, handling:
+
+1. Stacked Parameter Fusion
+   - gate_proj + up_proj → gate_up_proj (num_shards=2)
+   - q_proj + k_proj + v_proj → qkv_proj (num_shards=3)
+   - q_a_proj + kv_a_proj_with_mqa → fused_qkv_a_proj_with_mqa (DeepSeek MLA)
+
+2. Expert Parameter Sharding (MoE models)
+   - experts.{id}.gate_proj + experts.{id}.up_proj → experts.w13_weight (num_shards=2)
+   - experts.{id}.down_proj → experts.w2_weight (num_shards=1)
+   - Handles expert parallelism: num_local_experts = n_routed // ep_size + shared
+
+3. Scale Remapping (Quantized models)
+   - k_proj.k_scale → attn.k_scale
+   - v_proj.v_scale → attn.v_scale
+   - Quark-specific: output_scale → per-component scales
+
+Supported Models:
+    Dense: Llama, Qwen2, Qwen3, GLM4
+    MoE:   DeepSeekV2/V3/R1, Qwen3-MoE, GLM4-MoE, GLM4-MoE-Lite (GLM-4.7)
+
+Example:
+    >>> mapper = ParameterMapper.from_model(model)
+    >>> result = mapper.map("model.layers.0.mlp.gate_proj.weight")
+    >>> result.sglang_name  # "model.layers.0.mlp.gate_up_proj.weight"
+    >>> result.shard_id     # 0
+    >>> result.num_shards   # 2
+"""
+
+from dataclasses import dataclass
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+StackedParamsEntry = Tuple[str, str, Union[int, str]]
+ExpertParamsEntry = Tuple[str, str, int, Union[int, str]]
+
+
+@dataclass
+class MappingResult:
+    """Result of mapping a HuggingFace checkpoint weight name to SGLang parameter."""
+
+    sglang_name: str
+    shard_id: Optional[Union[int, str]]
+    num_shards: int
+    expert_id: Optional[int]
+    num_local_experts: Optional[int]
+
+
+# Standard FP8 scale remapping patterns
+_SCALE_REMAP_PATTERNS: List[Tuple[str, str, str]] = [
+    (".k_scale", ".self_attn.k_proj.k_scale", ".self_attn.attn.k_scale"),
+    (".v_scale", ".self_attn.v_proj.v_scale", ".self_attn.attn.v_scale"),
+    (".k_scale", ".k_scale", ".attn.k_scale"),
+    (".v_scale", ".v_scale", ".attn.v_scale"),
+]
+
+# Quark quantization scale remapping
+_QUARK_SCALE_REMAP: Dict[str, str] = {
+    ".q_proj.output_scale": ".attn.q_scale",
+    ".k_proj.output_scale": ".attn.k_scale",
+    ".v_proj.output_scale": ".attn.v_scale",
+    "self_attn.prob_output_scale": ".attn.prob_scale",
+}
+
+
+class ParameterMapper:
+    """Maps HuggingFace checkpoint weight names to SGLang model parameters.
+
+    This class pre-computes lookup tables at initialization for efficient
+    repeated mapping. It handles:
+    - Stacked/fused parameter mapping (gate_up_proj, qkv_proj, etc.)
+    - Expert parameter mapping with shard information
+    - Scale remapping for quantized models
+    - Model-specific weight name mutations
+    """
+
+    def __init__(
+        self,
+        stacked_params_mapping: List[StackedParamsEntry],
+        expert_params_mapping: List[ExpertParamsEntry],
+        num_local_experts: int = 0,
+        mutate_weight_preload: Optional[Callable[[str], str]] = None,
+        custom_scale_remap: Optional[Callable[[str], str]] = None,
+    ):
+        """Initialize the parameter mapper with model-specific configuration.
+
+        Args:
+            stacked_params_mapping: List of (sglang_name, hf_name, shard_id) tuples.
+                Example: [("gate_up_proj", "gate_proj", 0), ("gate_up_proj", "up_proj", 1)]
+            expert_params_mapping: List of (sglang_name, hf_name, expert_id, shard_id) tuples.
+                Example: [("w13_weight", "experts.0.gate_proj.weight", 0, 0), ...]
+            num_local_experts: Number of experts in the current model rank.
+                For EP=1: num_local_experts = n_routed_experts + num_fused_shared_experts
+                For EP>1: num_local_experts = n_routed_experts // ep_size + num_fused_shared_experts
+            mutate_weight_preload: Optional function to transform weight names before mapping.
+                Used for shared expert fusion in DeepSeek (shared_experts → experts.{n_routed}).
+            custom_scale_remap: Optional function for model-specific scale remapping.
+                Used for DeepSeek k_proj/v_proj → attn_mqa scale mapping.
+        """
+        self.num_local_experts = num_local_experts
+        self._mutate_weight_preload = mutate_weight_preload
+        self._custom_scale_remap = custom_scale_remap
+
+        self._stacked_lookup, self._stacked_num_shards = self._build_stacked_lookup(
+            stacked_params_mapping
+        )
+        self._expert_lookup, self._expert_num_shards = self._build_expert_lookup(
+            expert_params_mapping
+        )
+
+    @staticmethod
+    def _build_stacked_lookup(
+        mapping: List[StackedParamsEntry],
+    ) -> Tuple[Dict[str, Tuple[str, Union[int, str]]], Dict[str, int]]:
+        """Build lookup table and num_shards from stacked params mapping."""
+        lookup: Dict[str, Tuple[str, Union[int, str]]] = {}
+        shard_counts: Dict[str, int] = {}
+
+        for sglang_name, hf_name, shard_id in mapping:
+            lookup[hf_name] = (sglang_name, shard_id)
+            shard_counts[sglang_name] = shard_counts.get(sglang_name, 0) + 1
+
+        return lookup, shard_counts
+
+    @staticmethod
+    def _build_expert_lookup(
+        mapping: List[ExpertParamsEntry],
+    ) -> Tuple[Dict[str, Tuple[str, int, Union[int, str]]], Dict[str, int]]:
+        """Build lookup table and num_shards from expert params mapping."""
+        lookup: Dict[str, Tuple[str, int, Union[int, str]]] = {}
+        shard_counts: Dict[str, int] = {}
+
+        for sglang_name, hf_name, expert_id, shard_id in mapping:
+            lookup[hf_name] = (sglang_name, expert_id, shard_id)
+
+        for sglang_name, _, _, shard_id in mapping:
+            key = sglang_name
+            if key not in shard_counts:
+                unique_shards = set(
+                    s_id for s_name, _, _, s_id in mapping if s_name == sglang_name
+                )
+                shard_counts[key] = len(unique_shards)
+
+        return lookup, shard_counts
+
+    def _apply_scale_remap(self, name: str) -> str:
+        """Apply standard and Quark scale remapping patterns."""
+        for suffix, pattern, replacement in _SCALE_REMAP_PATTERNS:
+            if name.endswith(suffix) and pattern in name:
+                return name.replace(pattern, replacement)
+
+        for quark_suffix, replacement in _QUARK_SCALE_REMAP.items():
+            if name.endswith(quark_suffix):
+                return name.replace(quark_suffix, replacement)
+
+        return name
+
+    def map(self, hf_weight_name: str) -> MappingResult:
+        """Map a HuggingFace checkpoint weight name to SGLang parameter info.
+
+        Args:
+            hf_weight_name: The weight name from HuggingFace checkpoint.
+
+        Returns:
+            MappingResult with mapped name and sharding information.
+        """
+        name = hf_weight_name
+
+        if self._mutate_weight_preload is not None:
+            name = self._mutate_weight_preload(name)
+
+        if "scale" in name:
+            if self._custom_scale_remap is not None:
+                remapped = self._custom_scale_remap(name)
+                if remapped != name:
+                    name = remapped
+                else:
+                    name = self._apply_scale_remap(name)
+            else:
+                name = self._apply_scale_remap(name)
+
+        for hf_pattern, (
+            sglang_name,
+            expert_id,
+            shard_id,
+        ) in self._expert_lookup.items():
+            if hf_pattern in name:
+                mapped_name = name.replace(hf_pattern, sglang_name)
+                return MappingResult(
+                    sglang_name=mapped_name,
+                    shard_id=shard_id,
+                    num_shards=self._expert_num_shards.get(sglang_name, 1),
+                    expert_id=expert_id,
+                    num_local_experts=self.num_local_experts,
+                )
+
+        for hf_pattern, (sglang_name, shard_id) in self._stacked_lookup.items():
+            if hf_pattern in name:
+                mapped_name = name.replace(hf_pattern, sglang_name)
+                return MappingResult(
+                    sglang_name=mapped_name,
+                    shard_id=shard_id,
+                    num_shards=self._stacked_num_shards.get(sglang_name, 1),
+                    expert_id=None,
+                    num_local_experts=None,
+                )
+
+        return MappingResult(
+            sglang_name=name,
+            shard_id=None,
+            num_shards=1,
+            expert_id=None,
+            num_local_experts=None,
+        )
+
+    @classmethod
+    def from_model(cls, model) -> "ParameterMapper":
+        """Create a ParameterMapper from a model instance; currently supports
+        DeepseekV2ForCausalLM, Glm4ForCausalLM, Glm4MoeForCausalLM,
+        Glm4MoeLiteForCausalLM, LlamaForCausalLM, Qwen2ForCausalLM,
+        Qwen3ForCausalLM, Qwen3MoeForCausalLM."""
+        stacked_mapping = list(getattr(model, "stacked_params_mapping", []) or [])
+        expert_mapping = list(getattr(model, "expert_params_mapping", []) or [])
+
+        num_local_experts = 0
+        if hasattr(model, "num_local_experts"):
+            num_local_experts = model.num_local_experts
+        elif expert_mapping:
+            expert_ids = set(entry[2] for entry in expert_mapping)
+            num_local_experts = len(expert_ids)
+
+        mutate_fn = None
+        if hasattr(model, "mutate_weight_preload"):
+            mutate_fn = model.mutate_weight_preload
+
+        scale_fn = None
+        if hasattr(model, "custom_scale_remap"):
+            scale_fn = model.custom_scale_remap
+
+        return cls(
+            stacked_params_mapping=stacked_mapping,
+            expert_params_mapping=expert_mapping,
+            num_local_experts=num_local_experts,
+            mutate_weight_preload=mutate_fn,
+            custom_scale_remap=scale_fn,
+        )
diff --git a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py
index 165119e2d8af..3ffbbf9b6b7c 100644
--- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py
+++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py
@@ -25,7 +25,6 @@
 from sglang.srt.distributed.parallel_state import GroupCoordinator
 from sglang.srt.environ import envs
 from sglang.srt.layers import deep_gemm_wrapper
-from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.quantization.fp8_utils import (
     block_quant_dequant,
@@ -58,7 +57,6 @@
     _use_aiter_gfx95,
     awq_dequantize_func,
     enable_nextn_moe_bf16_cast_to_fp8,
-    is_wint4afp8_or_wint4a16_config,
 )
 from sglang.srt.utils import bind_or_assign, get_bool_env_var, log_info_on_rank0
 
@@ -148,6 +146,18 @@ class DeepseekV2WeightLoaderMixin:
     quant_config: Optional[QuantizationConfig]
     pp_group: GroupCoordinator
     num_fused_shared_experts: int
+    # Weight mapping relationships determined at model initialization time.
+    fuse_qkv_a_proj: bool
+    stacked_params_mapping: List[Tuple[str, str, int]]
+    expert_params_mapping: List[Tuple[str, str, int, int]]
+
+    def mutate_weight_preload(self, name: str) -> str:
+        """Override in subclass for model-specific weight name mutations."""
+        return name
+
+    def custom_scale_remap(self, name: str) -> str:
+        """Override in subclass for model-specific scale remapping."""
+        return name
 
     def do_load_weights(
         self,
@@ -166,33 +176,7 @@ def do_load_weights(
             weights, NVFP4_CKPT_FP8_ATTN_QUANT_MODULES, nextn_conf
         )
 
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("gate_up_proj", "gate_proj", 0),
-            ("gate_up_proj", "up_proj", 1),
-        ]
-
-        # Params for weights, fp8 weight scales, fp8 activation scales
-        # (param_name, weight_name, expert_id, shard_id)
-        expert_params_mapping = FusedMoE.make_expert_params_mapping(
-            ckpt_gate_proj_name="gate_proj",
-            ckpt_down_proj_name="down_proj",
-            ckpt_up_proj_name="up_proj",
-            num_experts=self.config.n_routed_experts + self.num_fused_shared_experts,
-        )
-        # Params for special naming rules in mixed-precision models, for example:
-        # model.layers.xx.mlp.experts.xx.w1.input_scale. For details,
-        # see https://huggingface.co/Barrrrry/DeepSeek-R1-W4AFP8/blob/main.
-        if is_wint4afp8_or_wint4a16_config(self.quant_config):
-            expert_params_mapping += FusedMoE.make_expert_input_scale_params_mapping(
-                num_experts=self.config.n_routed_experts
-            )
-
-        # Fuse q_a_proj and kv_a_proj_with_mqa along output dimension when q_lora_rank is not None
-        fuse_qkv_a_proj = hasattr(self.config, "q_lora_rank") and (
-            self.config.q_lora_rank is not None
-        )
-        cached_a_proj = {} if fuse_qkv_a_proj else None
+        cached_a_proj = {} if self.fuse_qkv_a_proj else None
 
         pending_indexer_wk: Dict[str, Dict[str, torch.Tensor]] = {}
 
@@ -217,11 +201,7 @@ def do_load_weights(
                     )
                 ):
                     continue
-                if self.num_fused_shared_experts > 0 and "mlp.shared_experts" in name:
-                    name = name.replace(
-                        "mlp.shared_experts",
-                        f"mlp.experts.{self.config.n_routed_experts}",
-                    )
+                name = self.mutate_weight_preload(name)
 
                 weight_names.append(name)
 
@@ -270,7 +250,7 @@ def do_load_weights(
                 ):
                     continue
 
-                for param_name, weight_name, shard_id in stacked_params_mapping:
+                for param_name, weight_name, shard_id in self.stacked_params_mapping:
                     # Skip non-stacked layers and experts (experts handled below).
                     if weight_name not in name:
                         continue
@@ -284,6 +264,10 @@ def do_load_weights(
                     # for mlp.experts[0].gate_gate_up_proj, which breaks load.
                     if ("mlp.experts." in name) and name not in params_dict:
                         continue
+                    # q_a_proj / kv_a_proj_with_mqa must bypass stacked loading
+                    # and use the cache+concat fused A-proj path below.
+                    if param_name == "fused_qkv_a_proj_with_mqa":
+                        continue
                     name = name.replace(weight_name, param_name)
                     # Skip loading extra bias for GPTQ models.
                     if name.endswith(".bias") and name not in params_dict:
@@ -299,7 +283,7 @@ def do_load_weights(
                     )
                     break
                 else:
-                    for mapping in expert_params_mapping:
+                    for mapping in self.expert_params_mapping:
                         param_name, weight_name, expert_id, shard_id = mapping
                         if weight_name not in name:
                             continue
@@ -336,7 +320,7 @@ def do_load_weights(
                         # Skip loading norm if not last rank in pipeline parallelism
                         if ".norm." in name and not self.pp_group.is_last_rank:
                             continue
-                        if fuse_qkv_a_proj and (
+                        if self.fuse_qkv_a_proj and (
                             "q_a_proj" in name or "kv_a_proj_with_mqa" in name
                         ):
                             cached_a_proj[name] = _clone_if_runai_streamed_tensor(
@@ -406,13 +390,7 @@ def do_load_weights(
                             if (
                                 "k_scale" in name or "v_scale" in name
                             ) and name not in params_dict:
-                                # modelopt attn kv scale is named differently
-                                for scale in ["k_scale", "v_scale"]:
-                                    if scale in name:
-                                        name = name.replace(
-                                            f"{scale[0]}_proj", "attn_mqa"
-                                        )
-                                        break
+                                name = self.custom_scale_remap(name)
                             if name not in params_dict:
                                 # modelopt ckpt contains not needed weights for MTP module:
                                 # model.decoder.self_attn.attn_mqa.v_scale and
diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py
index a406c08231fb..7acaf038edb2 100644
--- a/python/sglang/srt/models/deepseek_v2.py
+++ b/python/sglang/srt/models/deepseek_v2.py
@@ -2732,6 +2732,37 @@ def __init__(
             self.lm_head = PPMissingLayer()
         self.logits_processor = LogitsProcessor(config)
 
+        # Stacked params mapping for unified weight loading API
+        self.stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+        # Add A-proj fusion mapping when q_lora_rank is enabled
+        # q_a_proj + kv_a_proj_with_mqa -> fused_qkv_a_proj_with_mqa
+        if self.fuse_qkv_a_proj:
+            self.stacked_params_mapping.extend(
+                [
+                    ("fused_qkv_a_proj_with_mqa", "q_a_proj", 0),
+                    ("fused_qkv_a_proj_with_mqa", "kv_a_proj_with_mqa", 1),
+                ]
+            )
+        self.expert_params_mapping = FusedMoE.make_expert_params_mapping(
+            ckpt_gate_proj_name="gate_proj",
+            ckpt_down_proj_name="down_proj",
+            ckpt_up_proj_name="up_proj",
+            num_experts=self.config.n_routed_experts + self.num_fused_shared_experts,
+        )
+        # Params for special naming rules in mixed-precision models, for example:
+        # model.layers.xx.mlp.experts.xx.w1.input_scale. For details,
+        # see https://huggingface.co/Barrrrry/DeepSeek-R1-W4AFP8/blob/main.
+        if is_wint4afp8_or_wint4a16_config(self.quant_config):
+            self.expert_params_mapping += (
+                FusedMoE.make_expert_input_scale_params_mapping(
+                    num_experts=self.config.n_routed_experts
+                )
+            )
+
         self._routed_experts_weights_of_layer = LazyValue(
             lambda: {
                 layer_id: layer.mlp.get_moe_weights()
@@ -2754,6 +2785,22 @@ def __init__(
         q_lora_rank = config.q_lora_rank if hasattr(config, "q_lora_rank") else None
         get_attn_tp_context().init_context(q_lora_rank, is_deepseek_dsa(config))
 
+    def mutate_weight_preload(self, name: str) -> str:
+        """DeepSeek V2: shared expert fusion."""
+        if self.num_fused_shared_experts > 0 and "mlp.shared_experts" in name:
+            return name.replace(
+                "mlp.shared_experts",
+                f"mlp.experts.{self.config.n_routed_experts}",
+            )
+        return name
+
+    def custom_scale_remap(self, name: str) -> str:
+        """DeepSeek V2: k_proj -> attn_mqa when k_scale in name, v_proj -> attn_mqa when v_scale in name."""
+        for scale in ["k_scale", "v_scale"]:
+            if scale in name:
+                return name.replace(f"{scale[0]}_proj", "attn_mqa")
+        return name
+
     @property
     def routed_experts_weights_of_layer(self):
         return self._routed_experts_weights_of_layer.value
diff --git a/python/sglang/srt/models/glm4.py b/python/sglang/srt/models/glm4.py
index 00a5057d0e3e..91f8fceaea8d 100644
--- a/python/sglang/srt/models/glm4.py
+++ b/python/sglang/srt/models/glm4.py
@@ -458,6 +458,17 @@ def __init__(
 
         self.logits_processor = LogitsProcessor(config)
         self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
+
+        # Stacked params mapping for unified weight loading API
+        self.stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
+
         # For EAGLE3 support
         self.capture_aux_hidden_states = False
 
@@ -552,14 +563,7 @@ def end_layer(self):
         return self.model.end_layer
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            (".qkv_proj", ".q_proj", "q"),
-            (".qkv_proj", ".k_proj", "k"),
-            (".qkv_proj", ".v_proj", "v"),
-            (".gate_up_proj", ".up_proj", 1),
-            (".gate_up_proj", ".gate_proj", 0),
-        ]
+        stacked_params_mapping = self.stacked_params_mapping
 
         params_dict = dict(self.named_parameters())
         for name, loaded_weight in weights:
diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py
index 758aeaa524db..cfe3a96d8ee1 100644
--- a/python/sglang/srt/models/glm4_moe.py
+++ b/python/sglang/srt/models/glm4_moe.py
@@ -15,7 +15,6 @@
 """Inference-only GLM-4.5, GLM-4.6 and GLM-4.7 model compatible with HuggingFace weights"""
 
 import logging
-import re
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 import torch
@@ -1174,9 +1173,37 @@ def __init__(
         )
         self.logits_processor = LogitsProcessor(config)
 
+        # Stacked params mapping for unified weight loading API
+        self.stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+        self.expert_params_mapping = FusedMoE.make_expert_params_mapping(
+            ckpt_gate_proj_name="gate_proj",
+            ckpt_down_proj_name="down_proj",
+            ckpt_up_proj_name="up_proj",
+            num_experts=self.config.n_routed_experts + self.num_fused_shared_experts,
+        )
+
         # For EAGLE3 support
         self.capture_aux_hidden_states = False
 
+    def mutate_weight_preload(self, name: str) -> str:
+        """GLM4-MoE: shared expert fusion."""
+        if self.num_fused_shared_experts > 0 and "mlp.shared_experts" in name:
+            return name.replace(
+                "mlp.shared_experts",
+                f"mlp.experts.{self.config.n_routed_experts}",
+            )
+        return name
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.model.embed_tokens
+
     def determine_num_fused_shared_experts(self):
         if get_server_args().disable_shared_experts_fusion:
             return
@@ -1269,43 +1296,8 @@ def load_weights(
             else:
                 raise ValueError("num_nextn_predict_layers is not in the config")
 
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-            ("gate_up_proj", "gate_proj", 0),
-            ("gate_up_proj", "up_proj", 1),
-        ]
-
-        if self.num_fused_shared_experts > 0:
-            assert self.num_fused_shared_experts == 1
-
-            def iter_weights_with_fused_shared_experts(
-                weights: Iterable[Tuple[str, torch.Tensor]],
-            ) -> Iterable[Tuple[str, torch.Tensor]]:
-
-                pattern = re.compile(
-                    r"^model\.layers\.(\d+)\.mlp\.shared_experts\.(.+)$"
-                )
-                for name, weight in weights:
-                    match = pattern.match(name)
-                    if match:
-                        layer_id = int(match.group(1))
-                        suffix = match.group(2)
-                        name = f"model.layers.{layer_id}.mlp.experts.{self.config.n_routed_experts}.{suffix}"
-                    yield name, weight
-
-            weights = iter_weights_with_fused_shared_experts(weights)
-
-        # Params for weights, fp8 weight scales, fp8 activation scales
-        # (param_name, weight_name, expert_id, shard_id)
-        expert_params_mapping = FusedMoE.make_expert_params_mapping(
-            ckpt_gate_proj_name="gate_proj",
-            ckpt_down_proj_name="down_proj",
-            ckpt_up_proj_name="up_proj",
-            num_experts=self.config.n_routed_experts + self.num_fused_shared_experts,
-        )
+        stacked_params_mapping = self.stacked_params_mapping
+        expert_params_mapping = self.expert_params_mapping
 
         if is_nextn:
             nextn_layer_prefix = f"model.layers.{nextn_layer_id}"
@@ -1326,6 +1318,8 @@ def iter_weights_with_fused_shared_experts(
         for name, loaded_weight in weights:
             weight_names.append(name)
 
+            name = self.mutate_weight_preload(name)
+
             if not is_nextn:
                 if hasattr(self.config, "num_nextn_predict_layers"):
                     num_nextn_layers = self.config.num_nextn_predict_layers
diff --git a/python/sglang/srt/models/glm4_moe_lite.py b/python/sglang/srt/models/glm4_moe_lite.py
index 0ed6b272be85..4d4e934f408a 100644
--- a/python/sglang/srt/models/glm4_moe_lite.py
+++ b/python/sglang/srt/models/glm4_moe_lite.py
@@ -920,6 +920,33 @@ def __init__(
         )
         self.capture_aux_hidden_states = False
 
+        # Weight loading mappings for ParameterMapper compatibility
+        self.stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+        # Add A-proj fusion mapping when q_lora_rank is enabled (MLA)
+        self.fuse_qkv_a_proj = hasattr(config, "q_lora_rank") and (
+            config.q_lora_rank is not None
+        )
+        if self.fuse_qkv_a_proj:
+            self.stacked_params_mapping.extend(
+                [
+                    ("fused_qkv_a_proj_with_mqa", "q_a_proj", 0),
+                    ("fused_qkv_a_proj_with_mqa", "kv_a_proj_with_mqa", 1),
+                ]
+            )
+        self.expert_params_mapping = FusedMoE.make_expert_params_mapping(
+            ckpt_gate_proj_name="gate_proj",
+            ckpt_down_proj_name="down_proj",
+            ckpt_up_proj_name="up_proj",
+            num_experts=config.n_routed_experts + self.num_fused_shared_experts,
+        )
+
     @property
     def routed_experts_weights_of_layer(self):
         return self._routed_experts_weights_of_layer.value
@@ -1041,6 +1068,22 @@ def set_dflash_layers_to_capture(self, layer_ids: List[int]):
         self.capture_aux_hidden_states = True
         self.model.layers_to_capture = [val + 1 for val in layer_ids]
 
+    def mutate_weight_preload(self, name: str) -> str:
+        """GLM4-MoE-Lite: shared expert fusion."""
+        if self.num_fused_shared_experts > 0 and "mlp.shared_experts" in name:
+            return name.replace(
+                "mlp.shared_experts",
+                f"mlp.experts.{self.config.n_routed_experts}",
+            )
+        return name
+
+    def custom_scale_remap(self, name: str) -> str:
+        """GLM4-MoE-Lite: k_proj/v_proj -> attn_mqa for MLA kv scale."""
+        for s in ["k_scale", "v_scale"]:
+            if s in name:
+                return name.replace(f"{s[0]}_proj", "attn_mqa")
+        return name
+
     def load_weights(
         self,
         weights: Iterable[Tuple[str, torch.Tensor]],
diff --git a/python/sglang/srt/models/glm4_moe_nextn.py b/python/sglang/srt/models/glm4_moe_nextn.py
index 3126fd026846..24e206b3cf9a 100644
--- a/python/sglang/srt/models/glm4_moe_nextn.py
+++ b/python/sglang/srt/models/glm4_moe_nextn.py
@@ -25,6 +25,7 @@
 from sglang.srt.layers.dp_attention import is_dp_attention_enabled
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.logits_processor import LogitsProcessor
+from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
@@ -145,6 +146,22 @@ def __init__(
             0 if get_server_args().disable_shared_experts_fusion else 1
         )
 
+        # Weight loading mappings (must match parent for load_weights to work)
+        self.stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+        self.expert_params_mapping = FusedMoE.make_expert_params_mapping(
+            ckpt_gate_proj_name="gate_proj",
+            ckpt_down_proj_name="down_proj",
+            ckpt_up_proj_name="up_proj",
+            num_experts=self.config.n_routed_experts + self.num_fused_shared_experts,
+        )
+
     @torch.no_grad()
     def forward(
         self,
diff --git a/python/sglang/srt/models/llama.py b/python/sglang/srt/models/llama.py
index a7eec70af5ef..1e4772b5c19a 100644
--- a/python/sglang/srt/models/llama.py
+++ b/python/sglang/srt/models/llama.py
@@ -510,9 +510,21 @@ def __init__(
             (".gate_up_proj", ".gate_proj", 0),
             (".gate_up_proj", ".up_proj", 1),
         ]
+        # Llama-specific scale remapping patterns (suffix, pattern, replacement)
+        self._llama_scale_remap_patterns = [
+            (".activation_scale", ".activation_scale", ".input_scale"),
+            (".weight_scale_inv", ".weight_scale_inv", ".weight_scale"),
+        ]
 
         self.capture_aux_hidden_states = False
 
+    def custom_scale_remap(self, name: str) -> str:
+        """Llama: activation_scale->input_scale, weight_scale_inv->weight_scale."""
+        for suffix, pattern, replacement in self._llama_scale_remap_patterns:
+            if name.endswith(suffix) and pattern in name:
+                return name.replace(pattern, replacement)
+        return name
+
     def _init_model(
         self,
         config: LlamaConfig,
@@ -623,23 +635,11 @@ def get_num_params(self):
         return len(params_dict)
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            (".qkv_proj", ".q_proj", "q"),
-            (".qkv_proj", ".k_proj", "k"),
-            (".qkv_proj", ".v_proj", "v"),
-            (".gate_up_proj", ".gate_proj", 0),
-            (".gate_up_proj", ".up_proj", 1),
-        ]
 
         params_dict = dict(self.named_parameters())
 
         for name, loaded_weight in weights:
-            if name.endswith(".activation_scale"):
-                name = name.replace(".activation_scale", ".input_scale")
-            if name.endswith(".weight_scale_inv"):
-                name = name.replace(".weight_scale_inv", ".weight_scale")
-
+            name = self.custom_scale_remap(name)
             layer_id = get_layer_id(name)
             if (
                 layer_id is not None
@@ -666,7 +666,7 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
                 if name is None:
                     continue
 
-            for param_name, weight_name, shard_id in stacked_params_mapping:
+            for param_name, weight_name, shard_id in self.stacked_params_mapping:
                 if weight_name not in name:
                     continue
                 name = name.replace(weight_name, param_name)
diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py
index c79fa483106b..49cdd00ac236 100644
--- a/python/sglang/srt/models/qwen2.py
+++ b/python/sglang/srt/models/qwen2.py
@@ -469,6 +469,17 @@ def __init__(
 
         self.logits_processor = LogitsProcessor(config)
         self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
+
+        # Stacked params mapping for unified weight loading API
+        self.stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+
         # For EAGLE3 support
         self.capture_aux_hidden_states = False
 
@@ -563,14 +574,7 @@ def end_layer(self):
         return self.model.end_layer
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-            ("gate_up_proj", "gate_proj", 0),
-            ("gate_up_proj", "up_proj", 1),
-        ]
+        stacked_params_mapping = self.stacked_params_mapping
 
         params_dict = dict(self.named_parameters())
         for name, loaded_weight in weights:
diff --git a/python/sglang/srt/models/qwen3.py b/python/sglang/srt/models/qwen3.py
index c9db1b8c39d9..4961a41bded1 100644
--- a/python/sglang/srt/models/qwen3.py
+++ b/python/sglang/srt/models/qwen3.py
@@ -39,6 +39,7 @@
     should_force_bfloat16_dense_tensor_math,
 )
 from sglang.srt.utils import add_prefix, get_bool_env_var, is_cuda, is_hip, is_npu
+from sglang.srt.utils.hf_transformers_utils import get_rope_config
 
 Qwen3Config = None
 
@@ -329,16 +330,7 @@ def __init__(
     ) -> None:
         super().__init__()
         self.hidden_size = config.hidden_size
-        if (
-            hasattr(config, "rope_parameters")
-            and config.rope_parameters
-            and "rope_theta" in config.rope_parameters
-        ):
-            rope_theta = config.rope_parameters["rope_theta"]
-            rope_scaling = config.rope_parameters
-        else:
-            rope_theta = getattr(config, "rope_theta", 1000000)
-            rope_scaling = getattr(config, "rope_scaling", None)
+        rope_theta, rope_scaling = get_rope_config(config)
         max_position_embeddings = getattr(config, "max_position_embeddings", 32768)
         head_dim = getattr(config, "head_dim", None)
         self.self_attn = Qwen3Attention(
@@ -491,6 +483,16 @@ def __init__(
             config, quant_config=quant_config, prefix=add_prefix("model", prefix)
         )
 
+        # Stacked params mapping for unified weight loading API
+        self.stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+
         # handle the lm head on different pp ranks
         if self.pp_group.is_last_rank:
             if self.pp_group.world_size == 1 and config.tie_word_embeddings:
@@ -602,15 +604,7 @@ def end_layer(self):
         return self.model.end_layer
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-            ("gate_up_proj", "gate_proj", 0),
-            ("gate_up_proj", "up_proj", 1),
-        ]
-
+        stacked_params_mapping = self.stacked_params_mapping
         params_dict = dict(self.named_parameters())
         for name, loaded_weight in weights:
             if not name.startswith("model.") and (
diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py
index c277c11ffa6e..1d4cb28e9ee3 100644
--- a/python/sglang/srt/models/qwen3_moe.py
+++ b/python/sglang/srt/models/qwen3_moe.py
@@ -982,6 +982,23 @@ def __init__(
             use_attn_tp_group=get_server_args().enable_dp_lm_head,
         )
         self.logits_processor = LogitsProcessor(config)
+
+        # Stacked params mapping for unified weight loading API
+        self.stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+        self.expert_params_mapping = FusedMoE.make_expert_params_mapping(
+            ckpt_gate_proj_name="gate_proj",
+            ckpt_down_proj_name="down_proj",
+            ckpt_up_proj_name="up_proj",
+            num_experts=self.config.num_experts,
+        )
+
         self.capture_aux_hidden_states = False
 
         self.attn_cp_size = get_parallel().attn_cp_size
@@ -1121,21 +1138,8 @@ def set_dflash_layers_to_capture(self, layer_ids: List[int]):
     def load_weights(
         self, weights: Iterable[Tuple[str, torch.Tensor]], is_mtp: bool = False
     ):
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-            ("gate_up_proj", "gate_proj", 0),
-            ("gate_up_proj", "up_proj", 1),
-        ]
-
-        expert_params_mapping = FusedMoE.make_expert_params_mapping(
-            ckpt_gate_proj_name="gate_proj",
-            ckpt_down_proj_name="down_proj",
-            ckpt_up_proj_name="up_proj",
-            num_experts=self.config.num_experts,
-        )
+        stacked_params_mapping = self.stacked_params_mapping
+        expert_params_mapping = self.expert_params_mapping
 
         # Pre-define `params_dict` to avoid repeated expensive traversal of model parameters.
         params_dict = dict(self.named_parameters())
diff --git a/test/registered/disaggregation/test_parallelism_context_integration.py b/test/registered/disaggregation/test_parallelism_context_integration.py
new file mode 100644
index 000000000000..e4dde076fc78
--- /dev/null
+++ b/test/registered/disaggregation/test_parallelism_context_integration.py
@@ -0,0 +1,275 @@
+"""
+Integration tests for ParallelismContext with real sglang servers.
+
+Tests that ParallelismContext can instantiate models with correct tensor parallel
+sharding by comparing parameter names and sizes against a running sglang server.
+
+Run with:
+    pytest test/registered/distributed/test_parallelism_context_integration.py -v
+
+Full test suite (non-CI):
+    - TP=2 small model (Qwem2.5-1.5B-Instruct)
+    - EP=2 small MOE model (DeepSeek-Coder-V2-Lite-Instruct)
+    - MLA model with hybrid dp attention (DeepSeek-Coder-V2-Lite-Instruct)
+
+CI test (reduced):
+    - TP=2 small model only
+"""
+
+import dataclasses
+import gc
+from typing import Dict, List, Tuple
+
+import pytest
+import requests
+import torch
+
+from sglang.srt.distributed.parallel_state import RankParallelismConfig
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.test_utils import (
+    DEFAULT_MLA_MODEL_NAME_FOR_TEST,
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    is_in_ci,
+    popen_launch_server,
+)
+from sglang.utils import terminate_process
+
+register_cuda_ci(est_time=145, stage="extra-a", runner_config="2-gpu-large")
+register_amd_ci(est_time=72, suite="stage-b-test-2-gpu-large-amd")
+
+
+def get_transfer_engine_info(url: str, rank: int) -> Dict:
+    """Get transfer engine info (parameter names and sizes) for a rank."""
+    response = requests.get(
+        f"{url}/remote_instance_transfer_engine_info",
+        params={"rank": rank},
+    )
+    response.raise_for_status()
+    return response.json()
+
+
+def get_parallelism_config(url: str, rank: int) -> Dict:
+    """Get parallelism config for a rank."""
+    response = requests.get(f"{url}/parallelism_config", params={"rank": rank})
+    response.raise_for_status()
+    return response.json()
+
+
+def get_server_info(url: str) -> Dict:
+    """Get server info."""
+    response = requests.get(f"{url}/server_info")
+    response.raise_for_status()
+    return response.json()
+
+
+def verify_model_params_match_for_rank(
+    url: str,
+    rank: int,
+    server_info: Dict,
+    test_gpu_id: int,
+):
+    """Verify model parameters match for a specific rank by recreating a model shard."""
+    transfer_info = get_transfer_engine_info(url, rank)
+    server_weights_info = transfer_info["remote_instance_transfer_engine_info"][1]
+
+    # Get parallelism config from running server
+    parallelism_config_data = get_parallelism_config(url, rank)
+    parallelism_config = RankParallelismConfig.from_dict(parallelism_config_data)
+    # Get server args from server info
+    from sglang.srt.server_args import ServerArgs
+
+    valid_fields = {f.name for f in dataclasses.fields(ServerArgs)}
+    filtered_info = {k: v for k, v in server_info.items() if k in valid_fields}
+    filtered_info.pop("model_config", None)
+    server_args = ServerArgs(**filtered_info)
+
+    from sglang.srt import server_args as server_args_module
+    from sglang.srt.distributed.parallel_state import ParallelismContext
+
+    original_global_server_args = server_args_module._global_server_args
+
+    try:
+        # In a Mock ParallelismContext, instantiate the model for this rank.
+        # Use a separate GPU (test_gpu_id) to avoid memory conflicts with the running server.
+        server_args_module._global_server_args = server_args
+        with ParallelismContext(parallelism_config):
+            from sglang.srt.configs.device_config import DeviceConfig
+            from sglang.srt.configs.load_config import LoadConfig
+            from sglang.srt.configs.model_config import ModelConfig
+            from sglang.srt.model_loader import get_model
+
+            model_config = ModelConfig.from_server_args(server_args)
+            load_config = LoadConfig(load_format="dummy")
+            device_config = DeviceConfig(device="cuda", gpu_id=test_gpu_id)
+
+            torch.cuda.set_device(test_gpu_id)
+            model = get_model(
+                model_config=model_config,
+                load_config=load_config,
+                device_config=device_config,
+            )
+            model_params = {}
+            for name, param in model.named_parameters():
+                model_params[name] = param.numel() * param.element_size()
+
+            # Verify all server parameters exist in model with same size
+            mismatches = []
+            missing = []
+            for param_name, (ptr, numel, elem_size) in server_weights_info.items():
+                expected_size = numel * elem_size
+                if param_name not in model_params:
+                    missing.append(param_name)
+                elif model_params[param_name] != expected_size:
+                    mismatches.append(
+                        f"{param_name}: model={model_params[param_name]}, server={expected_size}"
+                    )
+
+            assert not missing, f"Rank {rank}: Missing parameters: {missing}"
+            assert not mismatches, f"Rank {rank}: Size mismatches: {mismatches}"
+            del model
+            torch.cuda.empty_cache()
+
+    finally:
+        server_args_module._global_server_args = original_global_server_args
+
+
+TEST_CONFIGS: List[Tuple[str, str, int, List[str], int]] = [
+    # Basic TP=2 test (CI only)
+    (
+        "tp2_small",
+        DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN,
+        2,
+        [],
+        2,
+    ),
+    # EP=2: MoE experts split across 2 groups, moe_tp=1 per group
+    (
+        "mla_ep2",
+        DEFAULT_MLA_MODEL_NAME_FOR_TEST,
+        2,
+        ["--ep-size", "2"],
+        2,
+    ),
+    (
+        "mla_dp2_tp4",
+        DEFAULT_MLA_MODEL_NAME_FOR_TEST,
+        4,
+        ["--enable-dp-attention", "--dp", "2"],
+        4,
+    ),
+    (
+        "mla_dp2_ep2_tp4",
+        DEFAULT_MLA_MODEL_NAME_FOR_TEST,
+        4,
+        ["--enable-dp-attention", "--dp", "2", "--ep-size", "2"],
+        4,
+    ),
+    (
+        "mla_dp2_ep4_tp4",
+        DEFAULT_MLA_MODEL_NAME_FOR_TEST,
+        4,
+        ["--enable-dp-attention", "--dp", "2", "--ep-size", "4"],
+        4,
+    ),
+    (
+        "mla_dp4_ep2_tp4",
+        DEFAULT_MLA_MODEL_NAME_FOR_TEST,
+        4,
+        ["--enable-dp-attention", "--dp", "4", "--ep-size", "2"],
+        4,
+    ),
+]
+
+
+def get_test_configs():
+    if is_in_ci():
+        return [TEST_CONFIGS[0]]
+    else:
+        return TEST_CONFIGS
+
+
+def _get_test_params():
+    """Generate pytest parameters based on test configs."""
+    configs = get_test_configs()
+    params = []
+    ids = []
+    for (
+        test_id,
+        model_name,
+        tp_size,
+        extra_args,
+        min_gpus,
+    ) in configs:
+        params.append(
+            pytest.param(
+                (model_name, tp_size, extra_args, min_gpus),
+                id=test_id,
+            )
+        )
+    return params
+
+
+class TestParallelismContextIntegration:
+    """
+    Test that ParallelismContext can instantiate models with the same
+    parameter names and sizes as the sglang server engine.
+    """
+
+    @pytest.mark.parametrize("config", _get_test_params())
+    def test_model_instantiation_matches_server(self, config):
+        """
+        Test that a model instantiated with ParallelismContext has the same
+        parameter names and sizes as the model in the sglang server.
+
+        This test:
+        1. Starts a server with specified parallelism config
+        2. Gets transfer_engine_info for all ranks (contains param names and sizes)
+        3. Gets parallelism_config and server_info
+        4. Uses ParallelismContext to instantiate a model for each rank
+        5. Compares the parameter names and sizes
+        """
+        model_name, tp_size, extra_args, min_gpus = config
+        url = DEFAULT_URL_FOR_TEST
+
+        # Need min_gpus for server + 1 extra GPU for test model instantiation
+        required_gpus = min_gpus + 1
+        if torch.cuda.device_count() < required_gpus:
+            pytest.skip(
+                f"Need at least {required_gpus} GPUs (server={min_gpus} + test=1), have {torch.cuda.device_count()}"
+            )
+        test_gpu_id = min_gpus  # e.g., if server uses 0-1, test uses 2
+
+        # Build server args
+        other_args = [
+            "--tp-size",
+            str(tp_size),
+            "--remote-instance-weight-loader-start-seed-via-transfer-engine",
+            "--trust-remote-code",
+        ]
+        other_args.extend(extra_args)
+
+        process = None
+        try:
+            process = popen_launch_server(
+                model_name,
+                url,
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                other_args=other_args,
+            )
+            server_info = get_server_info(url)
+
+            for rank in range(tp_size):
+                verify_model_params_match_for_rank(url, rank, server_info, test_gpu_id)
+
+        finally:
+            if process is not None:
+                terminate_process(process)
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
diff --git a/test/srt/models/test_params_mapping.py b/test/srt/models/test_params_mapping.py
new file mode 100644
index 000000000000..e064266cab6a
--- /dev/null
+++ b/test/srt/models/test_params_mapping.py
@@ -0,0 +1,292 @@
+"""Unit tests for ParameterMapper."""
+
+from types import SimpleNamespace
+
+import pytest
+
+from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
+from sglang.srt.model_loader.parameter_mapper import ParameterMapper
+
+_DEEPSEEK_N_ROUTED = 4
+_DEEPSEEK_N_LOCAL = _DEEPSEEK_N_ROUTED + 1  # +1 fused shared expert
+_QWEN3MOE_N = 4
+_GLM4LITE_N_ROUTED = 4
+_GLM4LITE_N_LOCAL = _GLM4LITE_N_ROUTED + 1  # +1 fused shared expert
+
+
+def _make_model(**kwargs):
+    """Create a stub model object for ParameterMapper.from_model()."""
+    return SimpleNamespace(**kwargs)
+
+
+def _deepseek_mutate(name):
+    if "mlp.shared_experts" in name:
+        return name.replace("mlp.shared_experts", f"mlp.experts.{_DEEPSEEK_N_ROUTED}")
+    return name
+
+
+def _deepseek_scale_remap(name):
+    for s in ["k_scale", "v_scale"]:
+        if s in name:
+            return name.replace(f"{s[0]}_proj", "attn_mqa")
+    return name
+
+
+def _glm4lite_mutate(name):
+    if "mlp.shared_experts" in name:
+        return name.replace("mlp.shared_experts", f"mlp.experts.{_GLM4LITE_N_ROUTED}")
+    return name
+
+
+_LLAMA_SCALE_PATTERNS = [
+    (".activation_scale", ".activation_scale", ".input_scale"),
+    (".weight_scale_inv", ".weight_scale_inv", ".weight_scale"),
+]
+
+
+def _llama_scale_remap(name):
+    for suffix, pattern, replacement in _LLAMA_SCALE_PATTERNS:
+        if name.endswith(suffix) and pattern in name:
+            return name.replace(pattern, replacement)
+    return name
+
+
+@pytest.fixture
+def qwen_mapper():
+    """Qwen2/Qwen3 (dense): QKV fusion, gate/up fusion, no experts."""
+    return ParameterMapper.from_model(
+        _make_model(
+            stacked_params_mapping=[
+                ("qkv_proj", "q_proj", "q"),
+                ("qkv_proj", "k_proj", "k"),
+                ("qkv_proj", "v_proj", "v"),
+                ("gate_up_proj", "gate_proj", 0),
+                ("gate_up_proj", "up_proj", 1),
+            ],
+        )
+    )
+
+
+@pytest.fixture
+def llama_mapper():
+    """Llama/GLM4 (dense): dot-prefixed stacked params, custom scale remap."""
+    return ParameterMapper.from_model(
+        _make_model(
+            stacked_params_mapping=[
+                (".qkv_proj", ".q_proj", "q"),
+                (".qkv_proj", ".k_proj", "k"),
+                (".qkv_proj", ".v_proj", "v"),
+                (".gate_up_proj", ".gate_proj", 0),
+                (".gate_up_proj", ".up_proj", 1),
+            ],
+            custom_scale_remap=_llama_scale_remap,
+        )
+    )
+
+
+@pytest.fixture
+def qwen3moe_mapper():
+    """Qwen3-MoE: QKV fusion + experts, no shared expert fusion."""
+    return ParameterMapper.from_model(
+        _make_model(
+            stacked_params_mapping=[
+                ("qkv_proj", "q_proj", "q"),
+                ("qkv_proj", "k_proj", "k"),
+                ("qkv_proj", "v_proj", "v"),
+                ("gate_up_proj", "gate_proj", 0),
+                ("gate_up_proj", "up_proj", 1),
+            ],
+            expert_params_mapping=FusedMoE.make_expert_params_mapping(
+                ckpt_gate_proj_name="gate_proj",
+                ckpt_down_proj_name="down_proj",
+                ckpt_up_proj_name="up_proj",
+                num_experts=_QWEN3MOE_N,
+            ),
+        )
+    )
+
+
+@pytest.fixture
+def deepseek_mapper():
+    """DeepSeek V2/V3: MLA A-proj fusion, shared expert fusion, custom scale remap."""
+    return ParameterMapper.from_model(
+        _make_model(
+            stacked_params_mapping=[
+                ("gate_up_proj", "gate_proj", 0),
+                ("gate_up_proj", "up_proj", 1),
+                ("fused_qkv_a_proj_with_mqa", "q_a_proj", 0),
+                ("fused_qkv_a_proj_with_mqa", "kv_a_proj_with_mqa", 1),
+            ],
+            expert_params_mapping=FusedMoE.make_expert_params_mapping(
+                ckpt_gate_proj_name="gate_proj",
+                ckpt_down_proj_name="down_proj",
+                ckpt_up_proj_name="up_proj",
+                num_experts=_DEEPSEEK_N_LOCAL,
+            ),
+            mutate_weight_preload=_deepseek_mutate,
+            custom_scale_remap=_deepseek_scale_remap,
+        )
+    )
+
+
+@pytest.fixture
+def glm4lite_mapper():
+    """GLM4-MoE-Lite (GLM-4.7): QKV fusion, MLA A-proj fusion, shared expert fusion, custom scale remap."""
+    return ParameterMapper.from_model(
+        _make_model(
+            stacked_params_mapping=[
+                ("qkv_proj", "q_proj", "q"),
+                ("qkv_proj", "k_proj", "k"),
+                ("qkv_proj", "v_proj", "v"),
+                ("gate_up_proj", "gate_proj", 0),
+                ("gate_up_proj", "up_proj", 1),
+                ("fused_qkv_a_proj_with_mqa", "q_a_proj", 0),
+                ("fused_qkv_a_proj_with_mqa", "kv_a_proj_with_mqa", 1),
+            ],
+            expert_params_mapping=FusedMoE.make_expert_params_mapping(
+                ckpt_gate_proj_name="gate_proj",
+                ckpt_down_proj_name="down_proj",
+                ckpt_up_proj_name="up_proj",
+                num_experts=_GLM4LITE_N_LOCAL,
+            ),
+            mutate_weight_preload=_glm4lite_mutate,
+            custom_scale_remap=_deepseek_scale_remap,
+        )
+    )
+
+
+# ── Helpers ──────────────────────────────────────────────────────────────────
+
+
+def to_expect(name, shard=None, n=1, expert=None, n_exp=None):
+    """Shorthand for expected MappingResult fields."""
+    return (name, shard, n, expert, n_exp)
+
+
+def _assert(mapper, ckpt, expected):
+    r = mapper.map(ckpt)
+    name, shard, n, expert, n_exp = expected
+    assert (
+        r.sglang_name,
+        r.shard_id,
+        r.num_shards,
+        r.expert_id,
+        r.num_local_experts,
+    ) == (name, shard, n, expert, n_exp), f"map({ckpt!r}) = {r}"
+
+
+# ── Tests ────────────────────────────────────────────────────────────────────
+
+# fmt: off
+_QWEN_CASES = [
+    # QKV fusion (Qwen2, Qwen3, GLM4-MoE)
+    ("layers.0.attn.q_proj.weight",     to_expect("layers.0.attn.qkv_proj.weight", "q", 3)),
+    ("layers.0.attn.k_proj.weight",     to_expect("layers.0.attn.qkv_proj.weight", "k", 3)),
+    ("layers.0.attn.v_proj.weight",     to_expect("layers.0.attn.qkv_proj.weight", "v", 3)),
+    # Gate/Up fusion
+    ("layers.0.mlp.gate_proj.weight",   to_expect("layers.0.mlp.gate_up_proj.weight", 0, 2)),
+    ("layers.0.mlp.up_proj.weight",     to_expect("layers.0.mlp.gate_up_proj.weight", 1, 2)),
+    # Pass-through
+    ("layers.0.mlp.down_proj.weight",   to_expect("layers.0.mlp.down_proj.weight")),
+    ("embed_tokens.weight",             to_expect("embed_tokens.weight")),
+    # Standard scale remap (no custom_scale_remap)
+    ("model.layers.0.self_attn.k_scale", to_expect("model.layers.0.self_attn.attn.k_scale")),
+    ("model.layers.0.self_attn.v_scale", to_expect("model.layers.0.self_attn.attn.v_scale")),
+]
+
+_LLAMA_CASES = [
+    # Dot-prefixed QKV fusion (Llama, GLM4)
+    ("model.layers.0.self_attn.q_proj.weight", to_expect("model.layers.0.self_attn.qkv_proj.weight", "q", 3)),
+    ("model.layers.0.self_attn.k_proj.weight", to_expect("model.layers.0.self_attn.qkv_proj.weight", "k", 3)),
+    # Dot-prefixed gate/up
+    ("model.layers.0.mlp.gate_proj.weight",    to_expect("model.layers.0.mlp.gate_up_proj.weight", 0, 2)),
+    # Llama-specific scale remap + stacked (scales follow their weights)
+    ("model.layers.0.mlp.gate_proj.activation_scale", to_expect("model.layers.0.mlp.gate_up_proj.input_scale", 0, 2)),
+    ("model.layers.0.mlp.gate_proj.weight_scale_inv", to_expect("model.layers.0.mlp.gate_up_proj.weight_scale", 0, 2)),
+    # Pass-through
+    ("model.layers.0.mlp.down_proj.weight",    to_expect("model.layers.0.mlp.down_proj.weight")),
+]
+
+_QWEN3MOE_CASES = [
+    # QKV fusion
+    ("model.layers.0.self_attn.q_proj.weight",          to_expect("model.layers.0.self_attn.qkv_proj.weight", "q", 3)),
+    # Expert mapping (no shared expert fusion)
+    ("model.layers.0.mlp.experts.0.gate_proj.weight",   to_expect("model.layers.0.mlp.experts.w13_weight", "w1", 2, 0, _QWEN3MOE_N)),
+    ("model.layers.0.mlp.experts.3.down_proj.weight",   to_expect("model.layers.0.mlp.experts.w2_weight", "w2", 1, 3, _QWEN3MOE_N)),
+    # shared_experts falls through to stacked mapping (no mutate_weight_preload)
+    ("model.layers.0.mlp.shared_experts.gate_proj.weight", to_expect("model.layers.0.mlp.shared_experts.gate_up_proj.weight", 0, 2)),
+]
+
+_DEEPSEEK_CASES = [
+    # MLA A-proj fusion
+    ("model.layers.0.self_attn.q_a_proj.weight",            to_expect("model.layers.0.self_attn.fused_qkv_a_proj_with_mqa.weight", 0, 2)),
+    ("model.layers.0.self_attn.kv_a_proj_with_mqa.weight",  to_expect("model.layers.0.self_attn.fused_qkv_a_proj_with_mqa.weight", 1, 2)),
+    # Shared expert fusion via mutate_weight_preload
+    ("model.layers.0.mlp.shared_experts.gate_proj.weight",   to_expect("model.layers.0.mlp.experts.w13_weight", "w1", 2, _DEEPSEEK_N_ROUTED, _DEEPSEEK_N_LOCAL)),
+    ("model.layers.0.mlp.shared_experts.down_proj.weight",   to_expect("model.layers.0.mlp.experts.w2_weight", "w2", 1, _DEEPSEEK_N_ROUTED, _DEEPSEEK_N_LOCAL)),
+    # Custom scale remap (k_proj/v_proj -> attn_mqa, NOT double-remapped)
+    ("model.layers.0.self_attn.k_proj.k_scale",             to_expect("model.layers.0.self_attn.attn_mqa.k_scale")),
+    ("model.layers.0.self_attn.v_proj.v_scale",             to_expect("model.layers.0.self_attn.attn_mqa.v_scale")),
+    # kv_b_proj pass-through (decomposed in post_load_weights)
+    ("model.layers.0.self_attn.kv_b_proj.weight",           to_expect("model.layers.0.self_attn.kv_b_proj.weight")),
+]
+
+_GLM4LITE_CASES = [
+    # QKV fusion (GLM-4.7 uses standard QKV unlike DeepSeek which uses MLA-only)
+    ("model.layers.0.self_attn.q_proj.weight",              to_expect("model.layers.0.self_attn.qkv_proj.weight", "q", 3)),
+    ("model.layers.0.self_attn.k_proj.weight",              to_expect("model.layers.0.self_attn.qkv_proj.weight", "k", 3)),
+    ("model.layers.0.self_attn.v_proj.weight",              to_expect("model.layers.0.self_attn.qkv_proj.weight", "v", 3)),
+    # MLA A-proj fusion (GLM-4.7 also uses MLA with q_lora_rank)
+    ("model.layers.0.self_attn.q_a_proj.weight",            to_expect("model.layers.0.self_attn.fused_qkv_a_proj_with_mqa.weight", 0, 2)),
+    ("model.layers.0.self_attn.kv_a_proj_with_mqa.weight",  to_expect("model.layers.0.self_attn.fused_qkv_a_proj_with_mqa.weight", 1, 2)),
+    # Gate/Up fusion (non-expert layers)
+    ("model.layers.0.mlp.gate_proj.weight",                 to_expect("model.layers.0.mlp.gate_up_proj.weight", 0, 2)),
+    ("model.layers.0.mlp.up_proj.weight",                   to_expect("model.layers.0.mlp.gate_up_proj.weight", 1, 2)),
+    # Expert mapping
+    ("model.layers.0.mlp.experts.0.gate_proj.weight",       to_expect("model.layers.0.mlp.experts.w13_weight", "w1", 2, 0, _GLM4LITE_N_LOCAL)),
+    ("model.layers.0.mlp.experts.0.up_proj.weight",         to_expect("model.layers.0.mlp.experts.w13_weight", "w3", 2, 0, _GLM4LITE_N_LOCAL)),
+    ("model.layers.0.mlp.experts.3.down_proj.weight",       to_expect("model.layers.0.mlp.experts.w2_weight", "w2", 1, 3, _GLM4LITE_N_LOCAL)),
+    # Shared expert fusion via mutate_weight_preload
+    ("model.layers.0.mlp.shared_experts.gate_proj.weight",  to_expect("model.layers.0.mlp.experts.w13_weight", "w1", 2, _GLM4LITE_N_ROUTED, _GLM4LITE_N_LOCAL)),
+    ("model.layers.0.mlp.shared_experts.down_proj.weight",  to_expect("model.layers.0.mlp.experts.w2_weight", "w2", 1, _GLM4LITE_N_ROUTED, _GLM4LITE_N_LOCAL)),
+    # Custom scale remap (same as DeepSeek: k_proj/v_proj -> attn_mqa)
+    ("model.layers.0.self_attn.k_proj.k_scale",             to_expect("model.layers.0.self_attn.attn_mqa.k_scale")),
+    ("model.layers.0.self_attn.v_proj.v_scale",             to_expect("model.layers.0.self_attn.attn_mqa.v_scale")),
+    # Pass-through
+    ("model.layers.0.mlp.down_proj.weight",                 to_expect("model.layers.0.mlp.down_proj.weight")),
+    ("model.layers.0.self_attn.kv_b_proj.weight",           to_expect("model.layers.0.self_attn.kv_b_proj.weight")),
+]
+# fmt: on
+
+
+@pytest.mark.parametrize("ckpt,expected", _QWEN_CASES, ids=[c[0] for c in _QWEN_CASES])
+def test_qwen(qwen_mapper, ckpt, expected):
+    _assert(qwen_mapper, ckpt, expected)
+
+
+@pytest.mark.parametrize(
+    "ckpt,expected", _LLAMA_CASES, ids=[c[0] for c in _LLAMA_CASES]
+)
+def test_llama(llama_mapper, ckpt, expected):
+    _assert(llama_mapper, ckpt, expected)
+
+
+@pytest.mark.parametrize(
+    "ckpt,expected", _QWEN3MOE_CASES, ids=[c[0] for c in _QWEN3MOE_CASES]
+)
+def test_qwen3moe(qwen3moe_mapper, ckpt, expected):
+    _assert(qwen3moe_mapper, ckpt, expected)
+
+
+@pytest.mark.parametrize(
+    "ckpt,expected", _DEEPSEEK_CASES, ids=[c[0] for c in _DEEPSEEK_CASES]
+)
+def test_deepseek(deepseek_mapper, ckpt, expected):
+    _assert(deepseek_mapper, ckpt, expected)
+
+
+@pytest.mark.parametrize(
+    "ckpt,expected", _GLM4LITE_CASES, ids=[c[0] for c in _GLM4LITE_CASES]
+)
+def test_glm4lite(glm4lite_mapper, ckpt, expected):
+    _assert(glm4lite_mapper, ckpt, expected)

From 37f5a22f8f67d4a5dfbd9e0956863c05b55b4667 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Fri, 24 Jul 2026 21:40:20 -0700
Subject: [PATCH 20/43] [15/27] [sglang-miles] rollout indexer replay: raw
 seq-relative topk capture + meta_info threading

---
 .../srt/layers/attention/dsa/dsa_indexer.py   | 169 +++++++++++++-----
 .../layers/attention/dsa/dsa_topk_backend.py  |  42 +++--
 .../srt/layers/attention/dsa_backend.py       |   2 +
 .../srt/managers/detokenizer_manager.py       |   1 +
 python/sglang/srt/managers/io_struct.py       |   7 +
 .../srt/managers/multi_tokenizer_mixin.py     |   2 +
 .../scheduler_components/output_streamer.py   |  10 ++
 .../sglang/srt/managers/tokenizer_manager.py  |   3 +
 8 files changed, 177 insertions(+), 59 deletions(-)

diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py
index 196639f18146..ba1f1792739f 100644
--- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py
@@ -39,6 +39,7 @@
 )
 from sglang.srt.runtime_context import get_parallel, get_server_args
 from sglang.srt.state_capturer.indexer_topk import (
+    get_global_indexer_capturer,
     maybe_capture_indexer_topk,
 )
 from sglang.srt.utils import (
@@ -1033,7 +1034,14 @@ def _get_topk_paged(
 
         # NOTE(dark): logits should be cleaned in topk_transform
         self._mask_init_and_local_tokens(logits, seqlens_32)
-        topk_result = metadata.topk_transform(logits, self.index_topk)
+        capture = get_global_indexer_capturer() is not None
+        if capture:
+            topk_result, raw_result = metadata.topk_transform(
+                logits, self.index_topk, return_raw_indices=True
+            )
+        else:
+            topk_result = metadata.topk_transform(logits, self.index_topk)
+            raw_result = None
         # Restore possible padding exist in the hidden states.
         if not _is_hip and q_offset < q_fp8.shape[0]:
             pad_len = q_fp8.shape[0] - q_offset
@@ -1044,7 +1052,9 @@ def _get_topk_paged(
                 device=topk_result.device,
             )
             topk_result = torch.cat([topk_result, padding], dim=0)
-        return topk_result
+            if capture:
+                raw_result = torch.cat([raw_result, padding], dim=0)
+        return topk_result, raw_result
 
     def _get_mqa_logits_budget_bytes(self, device_index: int) -> int:
         free_mem_fraction = self._mqa_logits_free_mem_fraction()
@@ -1154,8 +1164,10 @@ def _get_topk_ragged(
             topk_result = torch.full(
                 (token_nums, self.index_topk), -1, device=device, dtype=torch.int32
             )
+        capture = get_global_indexer_capturer() is not None
+        raw_result = torch.full_like(topk_result, -1) if capture else None
         if batch_size == 0:
-            return topk_result
+            return topk_result, raw_result
 
         ks, ke = metadata.get_indexer_kvcache_range()
 
@@ -1222,9 +1234,17 @@ def _get_topk_ragged(
             assert logits.shape[1] == k_offset
 
             self._mask_init_and_local_tokens(logits, seq_lens_expanded, ks)
-            raw_topk_result = metadata.topk_transform(logits, self.index_topk, ks=ks)
-            topk_result[:q_offset] = raw_topk_result
-            return topk_result
+            if capture:
+                transformed, raw = metadata.topk_transform(
+                    logits, self.index_topk, ks=ks, return_raw_indices=True
+                )
+                topk_result[:q_offset] = transformed
+                raw_result[:q_offset] = raw
+            else:
+                topk_result[:q_offset] = metadata.topk_transform(
+                    logits, self.index_topk, ks=ks
+                )
+            return topk_result, raw_result
 
         bytes_per_row = k_offset * self._MQA_LOGITS_BYTES_PER_ELEM
         max_rows = max(1, int(logits_budget_bytes // max(bytes_per_row, 1)))
@@ -1290,19 +1310,32 @@ def _get_topk_ragged(
                 cu_seqlens_q_chunk = cu_seqlens_q_full[start:end]
                 batch_idx_chunk = token_to_batch_idx[start:end]
 
-            raw_topk_chunk = metadata.topk_transform(
-                logits_chunk,
-                self.index_topk,
-                ks=ks[start:end],
-                cu_seqlens_q=cu_seqlens_q_chunk,
-                ke_offset=lengths_chunk,
-                batch_idx_list=batch_idx_chunk,
-                topk_indices_offset_override=topk_offset_chunk,
-            )
-            topk_result[start:end] = raw_topk_chunk
+            if capture:
+                transformed, raw = metadata.topk_transform(
+                    logits_chunk,
+                    self.index_topk,
+                    ks=ks[start:end],
+                    cu_seqlens_q=cu_seqlens_q_chunk,
+                    ke_offset=lengths_chunk,
+                    batch_idx_list=batch_idx_chunk,
+                    topk_indices_offset_override=topk_offset_chunk,
+                    return_raw_indices=True,
+                )
+                topk_result[start:end] = transformed
+                raw_result[start:end] = raw
+            else:
+                topk_result[start:end] = metadata.topk_transform(
+                    logits_chunk,
+                    self.index_topk,
+                    ks=ks[start:end],
+                    cu_seqlens_q=cu_seqlens_q_chunk,
+                    ke_offset=lengths_chunk,
+                    batch_idx_list=batch_idx_chunk,
+                    topk_indices_offset_override=topk_offset_chunk,
+                )
             start = end
 
-        return topk_result
+        return topk_result, raw_result
 
     def _forward_cuda_k_only(
         self,
@@ -1368,7 +1401,7 @@ def _forward_cuda_k_only(
 
         # MHA doesn't need topk_indices
         if not return_indices:
-            return None
+            return None, None
 
         # MLA: use dummy logits with topk kernel's fast path to generate indices
         # When length <= 2048, naive_topk_cuda directly generates [0,1,...,length-1,-1,...]
@@ -1379,13 +1412,19 @@ def _forward_cuda_k_only(
             dtype=torch.float32,
             device=x_meta.device,
         )
-        raw_topk_result = metadata.topk_transform(dummy_logits, self.index_topk)
+        if get_global_indexer_capturer() is not None:
+            raw_topk_result, raw_result = metadata.topk_transform(
+                dummy_logits, self.index_topk, return_raw_indices=True
+            )
+        else:
+            raw_topk_result = metadata.topk_transform(dummy_logits, self.index_topk)
+            raw_result = None
         if topk_result is not None:
             # PCG/BCG: fill the valid prefix of the padded static buffer and
             # leave padded rows at the -1 sentinel.
             topk_result[: raw_topk_result.shape[0]] = raw_topk_result
-            return None
-        return raw_topk_result
+            return None, raw_result
+        return raw_topk_result, raw_result
 
     def _get_topk_ragged_with_cp(
         self,
@@ -1409,6 +1448,8 @@ def _get_topk_ragged_with_cp(
         assert page_size == 64, "only support page size 64"
         assert len(weights.shape) == 3
         weights = weights.squeeze(-1)
+        capture = get_global_indexer_capturer() is not None
+        raw_result = None
         k_fp8_list = []
         k_scale_list = []
         ks_list = []
@@ -1482,14 +1523,25 @@ def _get_topk_ragged_with_cp(
                     ke,
                     clean_logits=False,
                 )
-            topk_result = metadata.topk_transform(
-                logits,
-                self.index_topk,
-                ks=ks,
-                cu_seqlens_q=actual_seq_q,
-                ke_offset=ke_offset,
-                batch_idx_list=batch_idx_list,
-            )
+            if capture:
+                topk_result, raw_result = metadata.topk_transform(
+                    logits,
+                    self.index_topk,
+                    ks=ks,
+                    cu_seqlens_q=actual_seq_q,
+                    ke_offset=ke_offset,
+                    batch_idx_list=batch_idx_list,
+                    return_raw_indices=True,
+                )
+            else:
+                topk_result = metadata.topk_transform(
+                    logits,
+                    self.index_topk,
+                    ks=ks,
+                    cu_seqlens_q=actual_seq_q,
+                    ke_offset=ke_offset,
+                    batch_idx_list=batch_idx_list,
+                )
         else:
             kv_len = (
                 forward_batch.seq_lens_cpu[0].item()
@@ -1532,15 +1584,25 @@ def _get_topk_ragged_with_cp(
             actual_seq_q = torch.tensor([actual_seq_q], dtype=torch.int32).to(
                 device="cuda", non_blocking=True
             )
-            topk_result = metadata.topk_transform(
-                logits,
-                self.index_topk,
-                ks=ks,
-                cu_seqlens_q=actual_seq_q,
-                ke_offset=ke_offset,
-            )
+            if capture:
+                topk_result, raw_result = metadata.topk_transform(
+                    logits,
+                    self.index_topk,
+                    ks=ks,
+                    cu_seqlens_q=actual_seq_q,
+                    ke_offset=ke_offset,
+                    return_raw_indices=True,
+                )
+            else:
+                topk_result = metadata.topk_transform(
+                    logits,
+                    self.index_topk,
+                    ks=ks,
+                    cu_seqlens_q=actual_seq_q,
+                    ke_offset=ke_offset,
+                )
 
-        return topk_result
+        return topk_result, raw_result
 
     def forward_indexer(
         self,
@@ -1724,6 +1786,15 @@ def forward_xpu(
             x, q_lora, positions, forward_batch, layer_id, return_indices
         )
 
+    def _capture_and_return(self, layer_id, topk_result, raw_result):
+        # Capture the model's natural (sequence-relative) topk for rollout R3
+        # replay; the attention kernel still receives the transformed
+        # (paged/ragged kv-cache) indices in topk_result.
+        maybe_capture_indexer_topk(
+            layer_id, raw_result if raw_result is not None else topk_result
+        )
+        return topk_result
+
     def forward_cuda(
         self,
         x: torch.Tensor,
@@ -1777,7 +1848,7 @@ def forward_cuda(
 
         # Optimization: fast path when skipping topk computation
         if skip_logits_computation and (not self.dsa_enable_prefill_cp):
-            topk_result = self._forward_cuda_k_only(
+            topk_result, raw_result = self._forward_cuda_k_only(
                 x,
                 positions,
                 forward_batch,
@@ -1788,7 +1859,7 @@ def forward_cuda(
                 return_indices,
             )
             topk_result = _broadcast_indexer_topk_from_rank0(topk_result)
-            return maybe_capture_indexer_topk(layer_id, topk_result)
+            return self._capture_and_return(layer_id, topk_result, raw_result)
 
         # When weights_proj is LoRA-wrapped, use an eager module call so the
         # wrapper owns base+delta and no LoRA kernel runs under torch.compile.
@@ -1968,6 +2039,7 @@ def forward_cuda(
             else:
                 weights = self._get_logits_head_gate(x_for_gate, q_scale)
 
+        raw_result = None
         if _is_cuda or _is_hip:
             # In piecewise/breakable CUDA graph, any access to seq_lens_cpu
             # creates a Dynamo shape guard. These graph modes never have empty
@@ -1986,14 +2058,14 @@ def forward_cuda(
                         device=x_meta.device,
                     )
                     topk_result = _broadcast_indexer_topk_from_rank0(topk_result)
-                    return maybe_capture_indexer_topk(layer_id, topk_result)
+                    return self._capture_and_return(layer_id, topk_result, None)
 
             if (
                 forward_batch.forward_mode.is_decode_or_idle()
                 or forward_batch.forward_mode.is_target_verify()
                 or forward_batch.forward_mode.is_draft_extend_v2()
             ):
-                topk_result = self._get_topk_paged(
+                topk_result, raw_result = self._get_topk_paged(
                     forward_batch, layer_id, q_fp8, weights, metadata
                 )
             else:
@@ -2020,7 +2092,7 @@ def forward_cuda(
                     weights_prev, weights_next = torch.split(
                         weights, (weights.shape[0] + 1) // 2, dim=0
                     )
-                    topk_result_prev = self._get_topk_ragged_with_cp(
+                    topk_result_prev, raw_prev = self._get_topk_ragged_with_cp(
                         forward_batch,
                         layer_id,
                         q_fp8_prev,
@@ -2030,7 +2102,7 @@ def forward_cuda(
                         actual_seq_q_prev,
                     )
 
-                    topk_result_next = self._get_topk_ragged_with_cp(
+                    topk_result_next, raw_next = self._get_topk_ragged_with_cp(
                         forward_batch,
                         layer_id,
                         q_fp8_next,
@@ -2040,8 +2112,13 @@ def forward_cuda(
                         actual_seq_q_next,
                     )
                     topk_result = torch.cat([topk_result_prev, topk_result_next], dim=0)
+                    raw_result = (
+                        torch.cat([raw_prev, raw_next], dim=0)
+                        if raw_prev is not None
+                        else None
+                    )
                     topk_result = _broadcast_indexer_topk_from_rank0(topk_result)
-                    return maybe_capture_indexer_topk(layer_id, topk_result)
+                    return self._capture_and_return(layer_id, topk_result, raw_result)
                 else:
                     # In-graph (PCG/BCG) non-CP prefill is handled earlier by the
                     # graph DSA split-op dispatch, so only the eager path reaches
@@ -2050,7 +2127,7 @@ def forward_cuda(
                         "Internal error: in-graph DSA prefill must go through the "
                         "graph DSA split-op dispatch"
                     )
-                    topk_result = self._get_topk_ragged(
+                    topk_result, raw_result = self._get_topk_ragged(
                         enable_dual_stream,
                         forward_batch,
                         layer_id,
@@ -2067,7 +2144,7 @@ def forward_cuda(
                 layer_id=layer_id,
             )
         topk_result = _broadcast_indexer_topk_from_rank0(topk_result)
-        return maybe_capture_indexer_topk(layer_id, topk_result)
+        return self._capture_and_return(layer_id, topk_result, raw_result)
 
     def forward_npu(
         self,
diff --git a/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py b/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py
index f9db4b88e6a0..525dd66d484d 100644
--- a/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py
+++ b/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py
@@ -84,9 +84,21 @@ def topk_transform(
         row_starts: Optional[torch.Tensor] = None,
         batch_idx_list: Optional[List[int]] = None,
         force_unfused_topk: bool = False,
+        return_raw_indices: bool = False,
     ) -> torch.Tensor:
         if not envs.SGLANG_DSA_FUSE_TOPK.get() or force_unfused_topk:
-            return self.topk_func(logits, lengths, topk, row_starts=row_starts)
+            result = self.topk_func(logits, lengths, topk, row_starts=row_starts)
+            return (result, result) if return_raw_indices else result
+
+        # Sequence-relative selection, before the kv-cache coordinate remap done
+        # by the fused transform below. Rollout R3 replay compares against the
+        # model's natural (sequence-relative) topk, so capture must use these raw
+        # indices rather than the transformed (paged/ragged) result.
+        raw_indices = (
+            self.topk_func(logits, lengths, topk, row_starts=row_starts)
+            if return_raw_indices
+            else None
+        )
 
         # Decode-shaped PAGED top-k (plain decode AND spec verify / draft-extend,
         # whose expanded rows match the same shape) routes to the DeepSeek-V4 top-k
@@ -108,7 +120,8 @@ def topk_transform(
             == logits.shape[0]
             == attn_metadata.real_page_table.shape[0]
         ):
-            return _topk_transform_v2_paged(logits, lengths, topk, attn_metadata)
+            result = _topk_transform_v2_paged(logits, lengths, topk, attn_metadata)
+            return (result, raw_indices) if return_raw_indices else result
 
         # The legacy transforms below read attn_metadata.page_table_1 (page_size=1),
         # which is always present here: the fold only drops it for the decode case
@@ -127,7 +140,7 @@ def topk_transform(
                     if batch_idx_list is not None
                     else attn_metadata.page_table_1
                 )
-                return fast_topk_transform_fused(
+                result = fast_topk_transform_fused(
                     score=logits,
                     lengths=lengths,
                     page_table_size_1=page_table_size_1,
@@ -135,22 +148,22 @@ def topk_transform(
                     topk=topk,
                     row_starts=row_starts,
                 )
-            if topk_transform_method == TopkTransformMethod.RAGGED:
+            elif topk_transform_method == TopkTransformMethod.RAGGED:
                 if topk_indices_offset is None:
                     raise RuntimeError(
                         "RAGGED topk_transform requires topk_indices_offset; "
                         "expected extend-without-speculative metadata."
                     )
-                return fast_topk_transform_ragged_fused(
+                result = fast_topk_transform_ragged_fused(
                     score=logits,
                     lengths=lengths,
                     topk_indices_offset=topk_indices_offset,
                     topk=topk,
                     row_starts=row_starts,
                 )
-            raise RuntimeError(f"Unsupported {topk_transform_method = }.")
-
-        if self.is_flashinfer():
+            else:
+                raise RuntimeError(f"Unsupported {topk_transform_method = }.")
+        elif self.is_flashinfer():
             import flashinfer
 
             if topk_transform_method == TopkTransformMethod.PAGED:
@@ -162,7 +175,7 @@ def topk_transform(
                     device=logits.device,
                     num_rows=logits.shape[0],
                 )
-                return flashinfer.top_k_page_table_transform(
+                result = flashinfer.top_k_page_table_transform(
                     logits.contiguous(),
                     attn_metadata.page_table_1.contiguous(),
                     lengths.contiguous(),
@@ -173,13 +186,13 @@ def topk_transform(
                     dsa_graph_safe=True,
                     row_starts=local_row_starts,
                 )
-            if topk_transform_method == TopkTransformMethod.RAGGED:
+            elif topk_transform_method == TopkTransformMethod.RAGGED:
                 if topk_indices_offset is None:
                     raise RuntimeError(
                         "RAGGED topk_transform requires topk_indices_offset; "
                         "expected extend-without-speculative metadata."
                     )
-                return flashinfer.top_k_ragged_transform(
+                result = flashinfer.top_k_ragged_transform(
                     logits.contiguous(),
                     topk_indices_offset.contiguous(),
                     lengths.contiguous(),
@@ -189,9 +202,12 @@ def topk_transform(
                     dsa_graph_safe=True,
                     row_starts=row_starts,
                 )
-            raise RuntimeError(f"Unsupported {topk_transform_method = }.")
+            else:
+                raise RuntimeError(f"Unsupported {topk_transform_method = }.")
+        else:
+            raise RuntimeError(f"Unsupported {self = } for SGLANG_DSA_FUSE_TOPK.")
 
-        raise RuntimeError(f"Unsupported {self = } for SGLANG_DSA_FUSE_TOPK.")
+        return (result, raw_indices) if return_raw_indices else result
 
 
 def _topk_unfused(
diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py
index 2606b608157f..f8bfa0eadb3f 100644
--- a/python/sglang/srt/layers/attention/dsa_backend.py
+++ b/python/sglang/srt/layers/attention/dsa_backend.py
@@ -294,6 +294,7 @@ def topk_transform(
         ke_offset: Optional[torch.Tensor] = None,
         batch_idx_list: Optional[List[int]] = None,
         topk_indices_offset_override: Optional[torch.Tensor] = None,
+        return_raw_indices: bool = False,
     ) -> torch.Tensor:
         if topk_indices_offset_override is not None:
             cu_topk_indices_offset = topk_indices_offset_override
@@ -323,6 +324,7 @@ def topk_transform(
             row_starts=ks,
             batch_idx_list=batch_idx_list,
             force_unfused_topk=self.force_unfused_topk,
+            return_raw_indices=return_raw_indices,
         )
 
 
diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py
index 7b3ca52e5640..4aab66b43359 100644
--- a/python/sglang/srt/managers/detokenizer_manager.py
+++ b/python/sglang/srt/managers/detokenizer_manager.py
@@ -474,6 +474,7 @@ def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOutput):
             output_hidden_states=recv_obj.output_hidden_states,
             routed_experts=routed_experts,
             indexer_topk=indexer_topk,
+            indexer_topk_num_layers=recv_obj.indexer_topk_num_layers,
             customized_info=recv_obj.customized_info,
             placeholder_tokens_idx=None,
             placeholder_tokens_val=None,
diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py
index 1fefc17c44d6..ca742e642741 100644
--- a/python/sglang/srt/managers/io_struct.py
+++ b/python/sglang/srt/managers/io_struct.py
@@ -1281,6 +1281,10 @@ class BatchTokenIDOutput(BaseBatchReq, kw_only=True):
     # Pickled Optional[List[SchedulerReqTimeStats]]
     time_stats: Optional[PickleWrapper] = None
 
+    # Number of indexer layers, set when indexer_topk is non-empty. The
+    # rollout-indexer-replay consumer needs the layer count to reshape the
+    # flattened per-token topk back into (token, layer, topk).
+    indexer_topk_num_layers: Optional[int] = None
     # Multimodal prompt token counts (image/audio/video). None when not applicable.
     image_tokens: Optional[List[int]] = None
     audio_tokens: Optional[List[int]] = None
@@ -1363,6 +1367,9 @@ class BatchStrOutput(BaseBatchReq, kw_only=True):
     # Pickled Optional[List[SchedulerReqTimeStats]]
     time_stats: Optional[PickleWrapper] = None
 
+    # Number of indexer layers, set when indexer_topk is non-empty (see
+    # BatchTokenIDOutput.indexer_topk_num_layers).
+    indexer_topk_num_layers: Optional[int] = None
     # Multimodal prompt token counts (image/audio/video). None when not applicable.
     image_tokens: Optional[List[int]] = None
     audio_tokens: Optional[List[int]] = None
diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py
index 8a1ba183999a..d4e3ccb5eca6 100644
--- a/python/sglang/srt/managers/multi_tokenizer_mixin.py
+++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py
@@ -247,6 +247,7 @@ def _handle_output_by_index(output, i):
             indexer_topk=_extract_field_by_index(
                 output, "indexer_topk", i, check_length=False
             ),
+            indexer_topk_num_layers=getattr(output, "indexer_topk_num_layers", None),
             retraction_counts=_extract_field_by_index(output, "retraction_counts", i),
             placeholder_tokens_idx=None,
             placeholder_tokens_val=None,
@@ -355,6 +356,7 @@ def _handle_output_by_index(output, i):
             indexer_topk=_extract_field_by_index(
                 output, "indexer_topk", i, check_length=False
             ),
+            indexer_topk_num_layers=getattr(output, "indexer_topk_num_layers", None),
             customized_info=_extract_field_by_index(
                 output, "customized_info", i, check_length=False
             ),
diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py
index e03d172e8643..28f4e2ba86ab 100644
--- a/python/sglang/srt/managers/scheduler_components/output_streamer.py
+++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py
@@ -561,6 +561,15 @@ def to_payload(
         if not (self.rids or is_idle_batch):
             return None
         dp_ranks = [dp_rank] * len(self.rids) if self.rids else None
+        indexer_topk_num_layers = None
+        if self.return_indexer_topk:
+            from sglang.srt.state_capturer.indexer_topk import (
+                get_global_indexer_capturer,
+            )
+
+            _idx_cap = get_global_indexer_capturer()
+            if _idx_cap is not None:
+                indexer_topk_num_layers = _idx_cap.num_indexer_layers
         return BatchTokenIDOutput(
             rids=self.rids,
             http_worker_ipcs=self.http_worker_ipcs,
@@ -605,6 +614,7 @@ def to_payload(
             output_hidden_states=self.output_hidden_states,
             routed_experts=self.routed_experts,
             indexer_topk=self.indexer_topk,
+            indexer_topk_num_layers=indexer_topk_num_layers,
             customized_info=(
                 wrap_as_pickle(self.customized_info) if self.customized_info else None
             ),
diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py
index 60b35ea021dd..f8a40fe4901e 100644
--- a/python/sglang/srt/managers/tokenizer_manager.py
+++ b/python/sglang/srt/managers/tokenizer_manager.py
@@ -2021,6 +2021,9 @@ async def _handle_batch_output(
                     if isinstance(val, torch.Tensor):
                         val = pybase64.b64encode(val.numpy().tobytes()).decode("utf-8")
                     meta_info["indexer_topk"] = val
+                    n = getattr(recv_obj, "indexer_topk_num_layers", None)
+                    if n is not None:
+                        meta_info["indexer_topk_num_layers"] = n
             if getattr(recv_obj, "dp_ranks", None):
                 meta_info["dp_rank"] = recv_obj.dp_ranks[i]
 

From bb17e99b24a6610f83838fc0bcf0fd8ca3738f40 Mon Sep 17 00:00:00 2001
From: Jiajun Li <guapisolo@gmail.com>
Date: Fri, 24 Jul 2026 23:01:05 -0700
Subject: [PATCH 21/43] [16/27] [sglang-miles] RL weight-update sessions +
 distributed update for spec draft worker(s) (#27749, #28575, #18565, #22663,
 #28001, #29675, #27750)

Squash of the spec-draft distributed-update work and the weight-processing
session work: the latter rewrites iter_draft_runners()/get_model_runners() the
former introduces, so they cannot be applied independently.

  - Draft runners never join the update group, so the target receives the
    broadcast once and the weights are loaded into every selected runner.
  - Begin/EndWeightUpdate session chain: engine / http_server /
    tokenizer_control_mixin (pause-aware locking) -> scheduler ->
    SchedulerWeightUpdaterManager, with {target,draft,all} runner selectors.
  - loader.py: public post_load_weights plus restore_weight/postprocess_weight
    via _apply_quant_method_hook (skips LoRA wrappers).
  - weight checker: skip_tensor_list + role-prefixed overall_checksum over
    get_model_runners(selector), replacing _get_draft_model_runner.

Reimplemented onto v0.5.16, which had already extracted weight updating into
ModelRunner.weight_updater (WeightUpdater): the receive/load split lands there
as receive_weights_from_distributed() + load_weights() rather than on
ModelRunner, and the scheduler fan-out drives runner.weight_updater.*. The
worker-level update_weights_from_{distributed,tensor} entry points on
TpModelWorker / EAGLEWorkerV2 / NGRAMWorker are dropped: the scheduler now owns
the fan-out, and leaving them would be a second path that updates one runner
only. iter_runners() replaces the ad-hoc draft-runner discovery on every spec v2
worker (incl. DFlashWorkerV2). The unit test moves to test/registered/rl/, since
v0.5.16 retired test/srt/ from CI collection.

Co-authored-by: Yueming Yuan <yym022502@gmail.com>
Co-authored-by: JD-ETH <jaedon.guo@gmail.com>
Co-authored-by: maocheng23 <35615230+maocheng23@users.noreply.github.com>
---
 python/sglang/srt/entrypoints/engine.py       |  19 ++
 python/sglang/srt/entrypoints/http_server.py  |  32 +++
 .../compressed_tensors/compressed_tensors.py  |   4 +
 python/sglang/srt/managers/io_struct.py       |  34 ++-
 python/sglang/srt/managers/scheduler.py       |  10 +
 .../scheduler_components/weight_updater.py    | 198 +++++++++++----
 .../srt/managers/tokenizer_control_mixin.py   |  40 ++++
 python/sglang/srt/managers/tp_worker.py       |  32 +--
 .../sglang/srt/model_executor/model_runner.py |  31 ++-
 .../model_runner_components/weight_updater.py | 121 +++++-----
 python/sglang/srt/model_loader/loader.py      |  42 +++-
 .../srt/speculative/dflash_worker_v2.py       |   8 +-
 .../sglang/srt/speculative/eagle_worker_v2.py |  33 +--
 .../multi_layer_eagle_worker_v2.py            |   5 +-
 python/sglang/srt/speculative/ngram_worker.py |  19 +-
 python/sglang/srt/utils/weight_checker.py     |  79 ++++--
 ...t_distributed_weight_update_spec_worker.py | 225 ++++++++++++++++++
 17 files changed, 721 insertions(+), 211 deletions(-)
 create mode 100644 test/registered/rl/test_distributed_weight_update_spec_worker.py

diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py
index e85f3dc7a7eb..25a2a6213032 100644
--- a/python/sglang/srt/entrypoints/engine.py
+++ b/python/sglang/srt/entrypoints/engine.py
@@ -59,9 +59,11 @@
 )
 from sglang.srt.managers.detokenizer_manager import run_detokenizer_process
 from sglang.srt.managers.io_struct import (
+    BeginWeightUpdateReqInput,
     CloseSessionReqInput,
     DestroyWeightsUpdateGroupReqInput,
     EmbeddingReqInput,
+    EndWeightUpdateReqInput,
     GenerateReqInput,
     GetWeightsByNameReqInput,
     InitWeightsUpdateGroupReqInput,
@@ -1044,6 +1046,23 @@ def destroy_weights_update_group(
             self.tokenizer_manager.destroy_weights_update_group(obj, None)
         )
 
+    def begin_weight_update(self, selector: str = "all"):
+        """Open a weight-update session: unpack in-place-quantized weights on the
+        selected runners so update_weights_from_{distributed,tensor} can load into
+        them. Must be closed with end_weight_update()."""
+        obj = BeginWeightUpdateReqInput(selector=selector)
+        return self.loop.run_until_complete(
+            self.tokenizer_manager.begin_weight_update(obj, None)
+        )
+
+    def end_weight_update(self):
+        """Close the session opened by begin_weight_update() and finalize quantized
+        weights into kernel layout."""
+        obj = EndWeightUpdateReqInput()
+        return self.loop.run_until_complete(
+            self.tokenizer_manager.end_weight_update(obj, None)
+        )
+
     def update_weights_from_distributed(
         self,
         names: list[str],
diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py
index f5b8bf6277f1..af29ae5f7e1a 100644
--- a/python/sglang/srt/entrypoints/http_server.py
+++ b/python/sglang/srt/entrypoints/http_server.py
@@ -113,6 +113,7 @@
 from sglang.srt.managers.io_struct import (
     AbortReq,
     AttachHiCacheStorageReqInput,
+    BeginWeightUpdateReqInput,
     CheckWeightsReqInput,
     CloseSessionReqInput,
     ConfigureLoggingReq,
@@ -120,6 +121,7 @@
     DestroyWeightsUpdateGroupReqInput,
     DumperControlReqInput,
     EmbeddingReqInput,
+    EndWeightUpdateReqInput,
     GenerateReqInput,
     GetWeightsByNameReqInput,
     InitWeightsSendGroupForRemoteInstanceReqInput,
@@ -1353,6 +1355,36 @@ async def update_weights_from_tensor(
     )
 
 
+@app.post("/begin_weight_update")
+@auth_level(AuthLevel.ADMIN_OPTIONAL)
+async def begin_weight_update(
+    obj: Annotated[BeginWeightUpdateReqInput, Body()], request: Request
+):
+    """Open a weight-update session so in-place-quantized weights become loadable."""
+    success, message = await _global_state.tokenizer_manager.begin_weight_update(
+        obj, request
+    )
+    return ORJSONResponse(
+        {"success": success, "message": message},
+        status_code=HTTPStatus.OK if success else HTTPStatus.BAD_REQUEST,
+    )
+
+
+@app.post("/end_weight_update")
+@auth_level(AuthLevel.ADMIN_OPTIONAL)
+async def end_weight_update(
+    obj: Annotated[EndWeightUpdateReqInput, Body()], request: Request
+):
+    """Close the weight-update session and finalize quantized weights."""
+    success, message = await _global_state.tokenizer_manager.end_weight_update(
+        obj, request
+    )
+    return ORJSONResponse(
+        {"success": success, "message": message},
+        status_code=HTTPStatus.OK if success else HTTPStatus.BAD_REQUEST,
+    )
+
+
 @app.post("/update_weights_from_distributed")
 @auth_level(AuthLevel.ADMIN_OPTIONAL)
 async def update_weights_from_distributed(
diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py
index f9d3e2ea0203..051db4d66b1b 100644
--- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py
+++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py
@@ -1021,6 +1021,10 @@ def __init__(self, quantization_config: CompressedTensorsConfig):
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         layer.scheme.process_weights_after_loading(layer)
 
+    def restore_weights_before_loading(self, layer: torch.nn.Module) -> None:
+        if hasattr(layer.scheme, "restore_weights_before_loading"):
+            layer.scheme.restore_weights_before_loading(layer)
+
     def create_weights(
         self,
         layer: torch.nn.Module,
diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py
index ca742e642741..f5f5b15dbec2 100644
--- a/python/sglang/srt/managers/io_struct.py
+++ b/python/sglang/srt/managers/io_struct.py
@@ -1583,6 +1583,9 @@ class UpdateWeightsFromDistributedReqInput(BaseReq, kw_only=True):
     weight_version: Optional[str] = None
     # Optional format specification for loading
     load_format: Optional[str] = None
+    # Which model runners to update: "target" (target model only), "draft" (draft
+    # worker(s) only), or "all" (default).
+    selector: Literal["target", "draft", "all"] = "all"
     # Whether to call torch.cuda.empty_cache() during flush
     torch_empty_cache: bool = False
 
@@ -1609,8 +1612,9 @@ class UpdateWeightsFromTensorReqInput(BaseReq, kw_only=True):
     abort_all_requests: bool = False
     # Optional: Update weight version along with weights
     weight_version: Optional[str] = None
-    # Optional: Determine whether to disable updating the draft model
-    disable_draft_model: Optional[bool] = None
+    # Which model runners to update: "target" (target model only), "draft" (draft
+    # worker(s) only), or "all" (default).
+    selector: Literal["target", "draft", "all"] = "all"
     # Whether to call torch.cuda.empty_cache() during flush
     torch_empty_cache: bool = False
 
@@ -1758,9 +1762,35 @@ class ResumeMemoryOccupationReqOutput(BaseReq, kw_only=True):
     pass
 
 
+class BeginWeightUpdateReqInput(BaseReq, kw_only=True):
+    """Open a weight-update session: unpack in-place-quantized weights on the
+    selected runners so fresh weights can be loaded into them."""
+
+    selector: Literal["target", "draft", "all"] = "all"
+
+
+class BeginWeightUpdateReqOutput(BaseReq, kw_only=True):
+    success: bool
+    message: str
+
+
+class EndWeightUpdateReqInput(BaseReq, kw_only=True):
+    """Close the weight-update session opened by BeginWeightUpdateReqInput."""
+
+
+class EndWeightUpdateReqOutput(BaseReq, kw_only=True):
+    success: bool
+    message: str
+
+
 class CheckWeightsReqInput(BaseReq, kw_only=True):
     action: str = "checksum"
     allow_quant_error: bool = False
+    # Substrings of tensor names to exclude from reset/compare/checksum.
+    skip_tensor_list: Optional[List[str]] = None
+    # Which model runners to update: "target" (target model only), "draft" (draft
+    # worker(s) only), or "all" (default).
+    selector: Literal["target", "draft", "all"] = "all"
 
 
 # Wire versions of the pydantic ParallelismInfo/ChecksumInfo in
diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
index 327a7ee414a4..f397c02cd126 100644
--- a/python/sglang/srt/managers/scheduler.py
+++ b/python/sglang/srt/managers/scheduler.py
@@ -89,6 +89,7 @@
     AttachHiCacheStorageReqOutput,
     BatchTokenizedEmbeddingReqInput,
     BatchTokenizedGenerateReqInput,
+    BeginWeightUpdateReqInput,
     CheckWeightsReqInput,
     ClearHiCacheReqInput,
     ClearHiCacheReqOutput,
@@ -100,6 +101,7 @@
     DetachHiCacheStorageReqOutput,
     DumperControlReqInput,
     DumperControlReqOutput,
+    EndWeightUpdateReqInput,
     ExpertDistributionReq,
     ExpertDistributionReqOutput,
     ExpertDistributionReqType,
@@ -1367,6 +1369,14 @@ def init_request_dispatcher(self):
                     SendWeightsToRemoteInstanceReqInput,
                     self.send_weights_to_remote_instance,
                 ),
+                (
+                    BeginWeightUpdateReqInput,
+                    self.weight_updater.begin_weight_update,
+                ),
+                (
+                    EndWeightUpdateReqInput,
+                    self.weight_updater.end_weight_update,
+                ),
                 (
                     UpdateWeightsFromDistributedReqInput,
                     self.weight_updater.update_weights_from_distributed,
diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py
index 7013ce04fe8a..62fc51554027 100644
--- a/python/sglang/srt/managers/scheduler_components/weight_updater.py
+++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py
@@ -1,12 +1,11 @@
 from __future__ import annotations
 
-import hashlib
 import logging
 import time
 import traceback
 from contextlib import contextmanager
 from dataclasses import dataclass, field
-from typing import Any, Callable, Dict, Iterator, Optional, Tuple
+from typing import Any, Callable, Dict, Iterator, List, Optional, Set, Tuple
 
 import msgspec
 import torch
@@ -19,11 +18,15 @@
 )
 from sglang.srt.disaggregation.utils import DisaggregationMode
 from sglang.srt.managers.io_struct import (
+    BeginWeightUpdateReqInput,
+    BeginWeightUpdateReqOutput,
     ChecksumInfo,
     CheckWeightsReqInput,
     CheckWeightsReqOutput,
     DestroyWeightsUpdateGroupReqInput,
     DestroyWeightsUpdateGroupReqOutput,
+    EndWeightUpdateReqInput,
+    EndWeightUpdateReqOutput,
     GetWeightsByNameReqInput,
     GetWeightsByNameReqOutput,
     InitWeightsUpdateGroupReqInput,
@@ -41,35 +44,40 @@
     UpdateWeightsFromTensorReqInput,
     UpdateWeightsFromTensorReqOutput,
 )
+from sglang.srt.utils import MultiprocessingSerializer
+from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions
+from sglang.srt.utils.weight_checker import overall_checksum
 
 logger = logging.getLogger(__name__)
 
 
-def _get_draft_model_runner(draft_worker):
-    # DFlash / FrozenKVMTP workers expose draft_model_runner directly
-    runner = getattr(draft_worker, "draft_model_runner", None)
-    if runner is not None:
-        return runner
-    # EAGLEWorkerV2: _draft_worker.draft_runner
-    inner = getattr(draft_worker, "_draft_worker", None)
-    if inner is not None:
-        runner = getattr(inner, "draft_runner", None)
-        if runner is not None:
-            return runner
-    return None
-
-
-def _merge_checksum_payloads(target: Dict, draft: Dict) -> Dict:
-    merged_checksums = dict(target["checksums"])
-    for name, chk in draft["checksums"].items():
-        merged_checksums[f"draft.{name}"] = chk
-    h = hashlib.sha256()
-    for name in sorted(merged_checksums):
-        h.update(name.encode())
-        h.update(merged_checksums[name].encode())
-    target["checksums"] = merged_checksums
-    target["per_gpu_checksum"] = h.hexdigest()
-    return target
+def _merge_checksum_payloads(role_payloads: List[Tuple[str, Dict]]) -> Dict:
+    merged: Dict[str, str] = {}
+    parallelism_infos = []
+    for role, p in role_payloads:
+        for name, chk in p["checksums"].items():
+            # Only non-target roles are prefixed, so target keys stay stable.
+            key = name if role == "" else f"{role}.{name}"
+            if key in merged:
+                raise ValueError(f"checksum key collision: {key}")
+            merged[key] = chk
+        parallelism_infos.append({"role": role or "target", **p["parallelism_info"]})
+    return {
+        "checksums": merged,
+        "per_gpu_checksum": overall_checksum(merged),
+        "parallelism_info": parallelism_infos,
+    }
+
+
+def _parse_runner_selector(selector: str) -> Set[str]:
+    """Map a {target, draft, all} weight-op selector to the set of roles it covers."""
+    if selector == "all":
+        return {"target", "draft"}
+    if selector in ("target", "draft"):
+        return {selector}
+    raise ValueError(
+        f"invalid selector {selector!r}; expected 'target', 'draft', or 'all'"
+    )
 
 
 @dataclass(kw_only=True, slots=True)
@@ -84,6 +92,11 @@ class SchedulerWeightUpdaterManager:
     metrics_collector: Optional[Any] = None
     offload_tags: set = field(default_factory=set)
     stashed_model_static_state: Any = None
+    _weight_update_in_progress: bool = False
+    _weight_update_loaded: bool = False
+    # Runner selector for the open session, recorded at begin_weight_update and
+    # reused by end_weight_update so the same set is restored and finalized.
+    _weight_update_selector: str = "all"
 
     @contextmanager
     def _observe_weight_load(self, source: str) -> Iterator[None]:
@@ -135,29 +148,87 @@ def destroy_weights_update_group(
         success, message = self.tp_worker.destroy_weights_update_group(recv_req)
         return DestroyWeightsUpdateGroupReqOutput(success=success, message=message)
 
+    def iter_weight_update_workers(
+        self, selector: str = "all"
+    ) -> List[Tuple[str, Any]]:
+        """Resolve a {target, draft, all} selector to (role, worker) pairs, target
+        first. This is the worker-level inclusion decision; each worker then
+        contributes its own runners via iter_runners()."""
+        parsed = _parse_runner_selector(selector)
+        workers: List[Tuple[str, Any]] = []
+        if "target" in parsed:
+            workers.append(("target", self.tp_worker))
+        if "draft" in parsed and self.draft_worker is not None:
+            workers.append(("draft", self.draft_worker))
+        return workers
+
+    def get_model_runners(self, selector: str = "all") -> List[Tuple[str, Any]]:
+        """Resolve a {target, draft, all} selector to (role, ModelRunner) pairs,
+        target first. Role is "" for the target runner; draft roles come from the
+        draft worker's iter_runners()."""
+        runners: List[Tuple[str, Any]] = []
+        for _, worker in self.iter_weight_update_workers(selector):
+            runners += worker.iter_runners()
+        return runners
+
     def update_weights_from_distributed(
         self,
         recv_req: UpdateWeightsFromDistributedReqInput,
     ) -> Tuple[bool, str]:
-        """Update the online model parameter."""
+        """Update the online model parameter, fanning out to the selected runners."""
+        assert (
+            self._weight_update_in_progress
+        ), "update_weights_from_distributed requires an open begin_weight_update session"
         with self._observe_weight_load("distributed"):
-            success, message = self.tp_worker.update_weights_from_distributed(recv_req)
+            # Only the target (main) model joined this process's update group, so it
+            # receives the broadcast once; the received weights are then loaded into
+            # each selected runner locally. Draft runners never join the group.
+            try:
+                weights = self.tp_worker.model_runner.weight_updater.receive_weights_from_distributed(
+                    recv_req.names,
+                    recv_req.dtypes,
+                    recv_req.shapes,
+                    recv_req.group_name,
+                    recv_req.load_format,
+                )
+                for _, runner in self.get_model_runners(recv_req.selector):
+                    runner.weight_updater.load_weights(weights)
+                success, message = True, "Succeeded to update parameter online."
+            except Exception as e:
+                success = False
+                message = (
+                    f"Failed to update parameter online: {e}. The full weights of the "
+                    "ModelRunner are partially updated. Please discard the whole weights."
+                )
+                logger.error(message)
             if success:
+                self._weight_update_loaded = True
                 self.flush_cache_after_weight_update(recv_req)
-            else:
-                logger.error(message)
             return UpdateWeightsFromDistributedReqOutput(
                 success=success, message=message
             )
 
     def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
-        """Update the online model parameter from tensors."""
+        """Update the online model parameter from tensors, fanning out to the
+        selected runners."""
+        assert (
+            self._weight_update_in_progress
+        ), "update_weights_from_tensor requires an open begin_weight_update session"
         with self._observe_weight_load("tensor"):
-            if recv_req.disable_draft_model:
-                worker = self.tp_worker
-            else:
-                worker = self.draft_worker or self.tp_worker
-            success, message = worker.update_weights_from_tensor(recv_req)
+            monkey_patch_torch_reductions()
+            named_tensors = MultiprocessingSerializer.deserialize(
+                recv_req.serialized_named_tensors[self.tp_worker.ps.tp_rank]
+            )
+            success, message = True, "Success"
+            for _, runner in self.get_model_runners(recv_req.selector):
+                success, message = runner.weight_updater.update_weights_from_tensor(
+                    named_tensors=named_tensors,
+                    load_format=recv_req.load_format,
+                )
+                if not success:
+                    break
+            if success:
+                self._weight_update_loaded = True
             if success:
                 self.flush_cache_after_weight_update(recv_req)
             else:
@@ -183,6 +254,36 @@ def get_weights_by_name(self, recv_req: GetWeightsByNameReqInput):
         parameter = self.tp_worker.get_weights_by_name(recv_req)
         return GetWeightsByNameReqOutput(parameter=parameter)
 
+    def begin_weight_update(self, recv_req: BeginWeightUpdateReqInput):
+        """Begin a weight-update session: restore in-place-packed weights to a
+        loadable state on the selected runners (target and/or draft), so the draft
+        model is prepared identically to the target. The selector is recorded and
+        reused by end_weight_update so the same set is finalized."""
+        assert (
+            not self._weight_update_in_progress
+        ), "begin_weight_update called while a weight-update session is already open"
+        self._weight_update_selector = recv_req.selector
+        for _, runner in self.get_model_runners(recv_req.selector):
+            runner.begin_weight_update()
+        self._weight_update_in_progress = True
+        self._weight_update_loaded = False
+        torch.distributed.barrier(group=self.tp_cpu_group)
+        return BeginWeightUpdateReqOutput(success=True, message="Success")
+
+    def end_weight_update(self, recv_req: EndWeightUpdateReqInput):
+        """End the weight-update session on the runners begin_weight_update opened
+        (its recorded selector): quant finalize on each, plus model.post_load_weights
+        only when load_weights was bypassed this session (e.g. P2P/RDMA)."""
+        assert (
+            self._weight_update_in_progress
+        ), "end_weight_update called without begin_weight_update"
+        run_post_load = not self._weight_update_loaded
+        for _, runner in self.get_model_runners(self._weight_update_selector):
+            runner.end_weight_update(run_post_load=run_post_load)
+        self._weight_update_in_progress = False
+        torch.distributed.barrier(group=self.tp_cpu_group)
+        return EndWeightUpdateReqOutput(success=True, message="Success")
+
     def release_memory_occupation(self, recv_req: ReleaseMemoryOccupationReqInput):
         assert (
             self.is_fully_idle()
@@ -270,19 +371,16 @@ def resume_memory_occupation(self, recv_req: ResumeMemoryOccupationReqInput):
 
     def check_weights(self, recv_req: CheckWeightsReqInput):
         try:
-            payload = self.tp_worker.model_runner.check_weights(
-                action=recv_req.action, allow_quant_error=recv_req.allow_quant_error
-            )
-
-            if self.draft_worker is not None:
-                draft_runner = _get_draft_model_runner(self.draft_worker)
-                if draft_runner is not None:
-                    draft_payload = draft_runner.check_weights(
-                        action=recv_req.action,
-                        allow_quant_error=recv_req.allow_quant_error,
-                    )
-                    if payload is not None and draft_payload is not None:
-                        payload = _merge_checksum_payloads(payload, draft_payload)
+            role_payloads = []
+            for role, runner in self.get_model_runners(recv_req.selector):
+                p = runner.check_weights(
+                    action=recv_req.action,
+                    allow_quant_error=recv_req.allow_quant_error,
+                    skip_tensor_list=recv_req.skip_tensor_list,
+                )
+                if p is not None:
+                    role_payloads.append((role, p))
+            payload = _merge_checksum_payloads(role_payloads) if role_payloads else None
 
             tp_size = torch.distributed.get_world_size(group=self.tp_cpu_group)
             if tp_size > 1 and payload is not None:
diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py
index f5e3d5807e23..3848d136ae14 100644
--- a/python/sglang/srt/managers/tokenizer_control_mixin.py
+++ b/python/sglang/srt/managers/tokenizer_control_mixin.py
@@ -15,6 +15,8 @@
     AddExternalCorpusReqOutput,
     AttachHiCacheStorageReqInput,
     AttachHiCacheStorageReqOutput,
+    BeginWeightUpdateReqInput,
+    BeginWeightUpdateReqOutput,
     ChecksumInfo,
     CheckWeightsReqInput,
     CheckWeightsReqOutput,
@@ -27,6 +29,8 @@
     DetachHiCacheStorageReqOutput,
     DumperControlReqInput,
     DumperControlReqOutput,
+    EndWeightUpdateReqInput,
+    EndWeightUpdateReqOutput,
     ExpertDistributionReq,
     ExpertDistributionReqOutput,
     ExpertDistributionReqType,
@@ -119,6 +123,8 @@
     ("get_internal_state", GetInternalStateReqOutput),
     ("set_internal_state", SetInternalStateReqOutput),
     ("expert_distribution", ExpertDistributionReqOutput),
+    ("begin_weight_update", BeginWeightUpdateReqOutput),
+    ("end_weight_update", EndWeightUpdateReqOutput),
     ("update_lora_adapter", LoRAUpdateOutput),
     ("dumper_control", DumperControlReqOutput),
     ("scale_elastic_ep", ScaleElasticEPReqOutput),
@@ -417,6 +423,40 @@ async def destroy_weights_update_group(
         results = await self.destroy_weights_update_group_communicator(obj)
         return FanOutCommunicator.merge_results(results)
 
+    async def _weight_update_session_call(
+        self: TokenizerManager, communicator, obj
+    ) -> Tuple[bool, str]:
+        """Run one weight-update session RPC under the same pause-aware locking as
+        update_weights_from_distributed: while the engine is paused the writer lock
+        is already held by whoever paused it, so taking it again would deadlock."""
+        self.auto_create_handle_loop()
+        async with self.is_pause_cond:
+            is_paused = self.is_pause
+            if is_paused:
+                results = await communicator(obj)
+        if not is_paused:
+            async with self.model_update_lock.writer_lock:
+                results = await communicator(obj)
+        return FanOutCommunicator.merge_results(results)
+
+    async def begin_weight_update(
+        self: TokenizerManager,
+        obj: BeginWeightUpdateReqInput,
+        request: Optional[fastapi.Request] = None,
+    ) -> Tuple[bool, str]:
+        return await self._weight_update_session_call(
+            self.begin_weight_update_communicator, obj
+        )
+
+    async def end_weight_update(
+        self: TokenizerManager,
+        obj: EndWeightUpdateReqInput,
+        request: Optional[fastapi.Request] = None,
+    ) -> Tuple[bool, str]:
+        return await self._weight_update_session_call(
+            self.end_weight_update_communicator, obj
+        )
+
     async def update_weights_from_distributed(
         self: TokenizerManager,
         obj: UpdateWeightsFromDistributedReqInput,
diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py
index cb63131f0e36..e5e51b010efd 100644
--- a/python/sglang/srt/managers/tp_worker.py
+++ b/python/sglang/srt/managers/tp_worker.py
@@ -34,9 +34,7 @@
     SendWeightsToRemoteInstanceReqInput,
     UnloadLoRAAdapterReqInput,
     UpdateWeightFromDiskReqInput,
-    UpdateWeightsFromDistributedReqInput,
     UpdateWeightsFromIPCReqInput,
-    UpdateWeightsFromTensorReqInput,
 )
 from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.scheduler import GenerationBatchResult
@@ -160,31 +158,6 @@ def send_weights_to_remote_instance(
         )
         return success, message
 
-    def update_weights_from_distributed(
-        self, recv_req: UpdateWeightsFromDistributedReqInput
-    ):
-        success, message = (
-            self.model_runner.weight_updater.update_weights_from_distributed(
-                recv_req.names,
-                recv_req.dtypes,
-                recv_req.shapes,
-                recv_req.group_name,
-                recv_req.load_format,
-            )
-        )
-        return success, message
-
-    def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
-
-        monkey_patch_torch_reductions()
-        success, message = self.model_runner.weight_updater.update_weights_from_tensor(
-            named_tensors=MultiprocessingSerializer.deserialize(
-                recv_req.serialized_named_tensors[self.ps.tp_rank]
-            ),
-            load_format=recv_req.load_format,
-        )
-        return success, message
-
     def update_weights_from_ipc(self, recv_req: UpdateWeightsFromIPCReqInput):
         """Update weights from IPC for checkpoint-engine integration."""
         success, message = self.model_runner.weight_updater.update_weights_from_ipc(
@@ -489,6 +462,11 @@ def _init_dllm_algorithm(self):
     def model_runner(self) -> ModelRunner:
         return self._model_runner
 
+    def iter_runners(self) -> List[Tuple[str, ModelRunner]]:
+        """(role, runner) pairs this worker owns for weight ops. The target worker
+        owns one runner and uses the empty role so its checksum keys stay unprefixed."""
+        return [("", self._model_runner)]
+
     def register_hicache_layer_transfer_counter(self, counter: LayerDoneCounter):
         self.hicache_layer_transfer_counter = counter
 
diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py
index 2e32ee835622..475681e0a762 100644
--- a/python/sglang/srt/model_executor/model_runner.py
+++ b/python/sglang/srt/model_executor/model_runner.py
@@ -20,7 +20,7 @@
 import logging
 import time
 from dataclasses import dataclass
-from typing import Optional, Union
+from typing import List, Optional, Union
 
 import torch
 import torch.distributed as dist
@@ -160,6 +160,11 @@
     EagerRunner,
     get_batch_sizes_to_capture,
 )
+from sglang.srt.model_loader.loader import (
+    post_load_weights,
+    postprocess_weight,
+    restore_weight,
+)
 from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import get_server_args
 from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
@@ -1638,9 +1643,29 @@ def compute_logprobs_only(
             forward_batch.token_ids_logprobs,
         )
 
-    def check_weights(self, action: str, allow_quant_error: bool = False):
+    def begin_weight_update(self) -> None:
+        """Begin a weight-update session: restore in-place-packed weights to a
+        loadable state (no-op for schemes that don't repack)."""
+        restore_weight(self.model, torch.device(self.device))
+
+    def end_weight_update(self, run_post_load: bool) -> None:
+        """End the weight-update session: optionally run model.post_load_weights
+        (when load_weights was bypassed this session, e.g. P2P/RDMA), then finalize
+        quantized weights into kernel layout."""
+        if run_post_load:
+            post_load_weights(self.model)
+        postprocess_weight(self.model, torch.device(self.device))
+
+    def check_weights(
+        self,
+        action: str,
+        allow_quant_error: bool = False,
+        skip_tensor_list: Optional[List[str]] = None,
+    ):
         return self._weight_checker.handle(
-            action=action, allow_quant_error=allow_quant_error
+            action=action,
+            allow_quant_error=allow_quant_error,
+            skip_tensor_list=skip_tensor_list,
         )
 
     def _expand_eplb_metadata_for_scale(
diff --git a/python/sglang/srt/model_executor/model_runner_components/weight_updater.py b/python/sglang/srt/model_executor/model_runner_components/weight_updater.py
index 8ea10ef04cfd..abd11a99f813 100644
--- a/python/sglang/srt/model_executor/model_runner_components/weight_updater.py
+++ b/python/sglang/srt/model_executor/model_runner_components/weight_updater.py
@@ -180,7 +180,11 @@ def model_load_weights(model, iter):
         logger.info("Update weights end.")
         return True, "Succeeded to update model weights."
 
-    def update_weights_from_distributed(
+    def load_weights(self: WeightUpdater, weights) -> None:
+        """Load an in-memory list of (name, tensor) weights into this runner's model."""
+        self.get_model().load_weights(weights)
+
+    def receive_weights_from_distributed(
         self: WeightUpdater,
         names,
         dtypes,
@@ -188,82 +192,79 @@ def update_weights_from_distributed(
         group_name,
         load_format: Optional[str] = None,
     ):
-        """
-        Update specific parameter in the model weights online
-        through `_model_update_group` process group.
+        """Receive one weight broadcast from the training engine over this runner's
+        `_model_update_group` and return the named tensors WITHOUT loading them.
 
-        Args:
-            name: the name of the parameter to be updated.
-            dtype: the data type of the parameter to be updated.
-            shape: the shape of the parameter to be updated.
+        Only the runner that joined the group (the target / main model) can receive;
+        the caller loads the result into each runner it wants updated. Speculative
+        draft runners never join the group, so they are fed from here.
         """
-
         assert group_name in self._model_update_group, (
             f"Group {group_name} not in {list(self._model_update_group.keys())}. "
             "Please call `init_weights_update_group` first."
         )
 
         if load_format == "flattened_bucket":
-            return self._update_bucketed_weights_from_distributed(
+            return self._receive_bucketed_weights_from_distributed(
                 names, dtypes, shapes, group_name
             )
-        try:
-            weights = []
-            handles = []
-            for name, dtype, shape in zip(names, dtypes, shapes):
-                target_dtype = (
-                    dtype if isinstance(dtype, torch.dtype) else getattr(torch, dtype)
-                )
-                weight = torch.empty(shape, dtype=target_dtype, device=self.device)
-                handles.append(
-                    torch.distributed.broadcast(
-                        weight,
-                        src=0,
-                        group=self._model_update_group[group_name],
-                        async_op=True,
-                    )
-                )
-                weights.append((name, weight))
-            for handle in handles:
-                handle.wait()
-
-            self.get_model().load_weights(weights)
-            return True, "Succeeded to update parameter online."
 
-        except Exception as e:
-            error_msg = (
-                f"Failed to update parameter online: {e}. "
-                f"The full weights of the ModelRunner are partially updated. "
-                f"Please discard the whole weights."
+        weights = []
+        handles = []
+        for name, dtype, shape in zip(names, dtypes, shapes):
+            target_dtype = (
+                dtype if isinstance(dtype, torch.dtype) else getattr(torch, dtype)
             )
-            logger.error(error_msg)
-            return False, error_msg
+            weight = torch.empty(shape, dtype=target_dtype, device=self.device)
+            handles.append(
+                torch.distributed.broadcast(
+                    weight,
+                    src=0,
+                    group=self._model_update_group[group_name],
+                    async_op=True,
+                )
+            )
+            weights.append((name, weight))
+        for handle in handles:
+            handle.wait()
+        return weights
 
-    def _update_bucketed_weights_from_distributed(
+    def _receive_bucketed_weights_from_distributed(
         self: WeightUpdater, names, dtypes, shapes, group_name
     ):
+        named_tensors = []
+        for name, dtype, shape in zip(names, dtypes, shapes):
+            target_dtype = (
+                dtype if isinstance(dtype, torch.dtype) else getattr(torch, dtype)
+            )
+            named_tensors.append(
+                (name, torch.empty(shape, dtype=target_dtype, device=self.device))
+            )
+        bucket = FlattenedTensorBucket(named_tensors=named_tensors)
+        flattened_tensor = bucket.get_flattened_tensor()
+        torch.distributed.broadcast(
+            flattened_tensor,
+            src=0,
+            group=self._model_update_group[group_name],
+        )
+        return bucket.reconstruct_tensors()
+
+    def update_weights_from_distributed(
+        self: WeightUpdater,
+        names,
+        dtypes,
+        shapes,
+        group_name,
+        load_format: Optional[str] = None,
+    ):
+        """Receive and load one weight broadcast into this runner's model."""
         try:
-            named_tensors = []
-            for name, dtype, shape in zip(names, dtypes, shapes):
-                target_dtype = (
-                    dtype if isinstance(dtype, torch.dtype) else getattr(torch, dtype)
-                )
-                named_tensors.append(
-                    (
-                        name,
-                        torch.empty(shape, dtype=target_dtype, device=self.device),
-                    )
-                )
-            bucket = FlattenedTensorBucket(named_tensors=named_tensors)
-            flattened_tensor = bucket.get_flattened_tensor()
-            torch.distributed.broadcast(
-                flattened_tensor,
-                src=0,
-                group=self._model_update_group[group_name],
+            weights = self.receive_weights_from_distributed(
+                names, dtypes, shapes, group_name, load_format
             )
-            reconstructed_tensors = bucket.reconstruct_tensors()
-            self.get_model().load_weights(reconstructed_tensors)
-            return True, f"Succeeded to update parameter online."
+            self.load_weights(weights)
+            return True, "Succeeded to update parameter online."
+
         except Exception as e:
             error_msg = (
                 f"Failed to update parameter online: {e}. "
diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py
index 13c0c07d43e0..b1339733ab4c 100644
--- a/python/sglang/srt/model_loader/loader.py
+++ b/python/sglang/srt/model_loader/loader.py
@@ -321,7 +321,7 @@ def _initialize_model(
     return model_class(**kwargs)
 
 
-def _post_load_weights(model: nn.Module) -> None:
+def post_load_weights(model: nn.Module) -> None:
     # Loaders that bypass `model.load_weights()` (dummy / sharded state / remote instance /
     # remote fs) must trigger the model's post-load fixup explicitly; `model.load_weights()`
     # would normally do it internally. NextN subclasses override the method to fill in
@@ -330,6 +330,36 @@ def _post_load_weights(model: nn.Module) -> None:
         model.post_load_weights()
 
 
+def _apply_quant_method_hook(model: nn.Module, target_device, hook_name: str) -> None:
+    """Run a quant_method hook (restore/process) on every quantized module.
+
+    LoRA wrappers forward `quant_method` for forward-path dispatch but don't own
+    the packed params; skip them so the inner base layer (yielded separately by
+    `named_modules`) handles it.
+    """
+    from sglang.srt.lora.layers import BaseLayerWithLoRA
+
+    for _, module in model.named_modules():
+        if isinstance(module, BaseLayerWithLoRA):
+            continue
+        quant_method = getattr(module, "quant_method", None)
+        if quant_method is not None and hasattr(quant_method, hook_name):
+            with device_loading_context(module, target_device):
+                getattr(quant_method, hook_name)(module)
+
+
+def restore_weight(model: nn.Module, target_device) -> None:
+    """Undo in-place quant packing so fresh weights can be loaded
+    (no-op for schemes that don't repack, e.g. plain fp8)."""
+    _apply_quant_method_hook(model, target_device, "restore_weights_before_loading")
+
+
+def postprocess_weight(model: nn.Module, target_device) -> None:
+    """Finalize quantized weights into kernel layout (Marlin repack, UE8M0 requant,
+    transpose, ...)."""
+    _apply_quant_method_hook(model, target_device, "process_weights_after_loading")
+
+
 class BaseModelLoader(ABC):
     """Base class for model loaders."""
 
@@ -1436,7 +1466,7 @@ def load_model(
             # random values to the weights.
             initialize_dummy_weights(model)
 
-            _post_load_weights(model)
+            post_load_weights(model)
 
             for _, module in model.named_modules():
                 quant_method = getattr(module, "quant_method", None)
@@ -1588,7 +1618,7 @@ def load_model(
             if state_dict:
                 raise ValueError(f"Missing keys {tuple(state_dict)} in loaded state!")
 
-            _post_load_weights(model)
+            post_load_weights(model)
 
         return model.eval()
 
@@ -2357,7 +2387,7 @@ def load_model_from_remote_instance_by_nccl(
                 )
             current_platform.synchronize()
 
-            _post_load_weights(model)
+            post_load_weights(model)
         end_get_weights_tic = time.time()
         logger.debug(
             f"finish getting all weights from remote instance, time used: {(end_get_weights_tic - start_get_weights_tic):.4f}s"
@@ -2420,7 +2450,7 @@ def load_model_from_remote_instance_by_transfer_engine(
             logger.error(f"batch transfer failed, error: {ret}")
             return False
 
-        _post_load_weights(model)
+        post_load_weights(model)
 
         return True
 
@@ -2510,7 +2540,7 @@ def _load_model_from_remote_kv(
         if state_dict:
             raise ValueError(f"Missing keys {tuple(state_dict)} in loaded state!")
 
-        _post_load_weights(model)
+        post_load_weights(model)
 
     def _load_model_from_remote_fs(
         self, model, client, model_config: ModelConfig, device_config: DeviceConfig
diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py
index a1095828f93a..d35ade470318 100644
--- a/python/sglang/srt/speculative/dflash_worker_v2.py
+++ b/python/sglang/srt/speculative/dflash_worker_v2.py
@@ -1,7 +1,7 @@
 import logging
 import math
 from dataclasses import replace
-from typing import List, Optional
+from typing import TYPE_CHECKING, List, Optional, Tuple
 
 import torch
 
@@ -53,6 +53,9 @@
 _is_npu = is_npu()
 
 
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
 logger = logging.getLogger(__name__)
 
 _FusedKVMaterializeHelper = None
@@ -154,6 +157,9 @@ class DFlashWorkerV2(BaseSpecWorker):
     scheduler runs it synchronously when overlap is disabled.
     """
 
+    def iter_runners(self) -> List[Tuple[str, "ModelRunner"]]:
+        return [("draft", self.draft_model_runner)]
+
     def __init__(
         self,
         server_args: ServerArgs,
diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py
index f14399aa0fd2..c1523267a7a3 100644
--- a/python/sglang/srt/speculative/eagle_worker_v2.py
+++ b/python/sglang/srt/speculative/eagle_worker_v2.py
@@ -2,7 +2,7 @@
 import logging
 import time
 from dataclasses import replace
-from typing import List, Optional
+from typing import TYPE_CHECKING, List, Optional, Tuple
 
 import torch
 
@@ -28,7 +28,6 @@
     speculative_moe_a2a_backend_context,
     speculative_moe_backend_context,
 )
-from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput
 from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.scheduler import GenerationBatchResult
 from sglang.srt.managers.tp_worker import TpModelWorker
@@ -92,7 +91,6 @@
     maybe_detect_oob,
 )
 from sglang.srt.utils.common import (
-    MultiprocessingSerializer,
     empty_context,
     fast_topk,
     get_available_gpu_memory,
@@ -104,7 +102,6 @@
     is_xpu,
     log_info_on_rank0,
 )
-from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions
 
 _is_cpu = is_cpu()
 _is_npu = is_npu()
@@ -114,6 +111,9 @@
 _is_xpu = is_xpu()
 
 
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
 logger = logging.getLogger(__name__)
 
 
@@ -1023,6 +1023,9 @@ def _draft_extend_for_decode(
 
 
 class EAGLEWorkerV2(BaseSpecWorker):
+    def iter_runners(self) -> List[Tuple[str, "ModelRunner"]]:
+        return [("draft", self.draft_runner)]
+
     def __init__(
         self,
         server_args: ServerArgs,
@@ -1514,25 +1517,3 @@ def verify(self, batch: ScheduleBatch):
             metadata_ready_pre_pad=False,
             finalize_tree_path=True,
         )
-
-    def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
-        monkey_patch_torch_reductions()
-        named_tensors = MultiprocessingSerializer.deserialize(
-            recv_req.serialized_named_tensors[self.ps.tp_rank]
-        )
-        success, message = (
-            self.draft_worker.draft_runner.weight_updater.update_weights_from_tensor(
-                named_tensors=named_tensors,
-                load_format=recv_req.load_format,
-            )
-        )
-        if not success:
-            return success, message
-
-        success, message = (
-            self.target_worker.model_runner.weight_updater.update_weights_from_tensor(
-                named_tensors=named_tensors,
-                load_format=recv_req.load_format,
-            )
-        )
-        return success, message
diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
index 4f3720776161..b6c5b29cc8a1 100644
--- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
+++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
@@ -16,7 +16,7 @@
 
 import logging
 from dataclasses import replace
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING, List, Tuple
 
 import torch
 
@@ -875,6 +875,9 @@ def _draft_extend_for_decode(
 
 
 class MultiLayerEagleWorkerV2(BaseSpecWorker):
+    def iter_runners(self) -> List[Tuple[str, ModelRunner]]:
+        return [(f"draft_step_{i}", r) for i, r in enumerate(self.draft_runner_list)]
+
     def __init__(
         self,
         server_args: ServerArgs,
diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py
index 01ca29f55f82..a08ff9b20e21 100644
--- a/python/sglang/srt/speculative/ngram_worker.py
+++ b/python/sglang/srt/speculative/ngram_worker.py
@@ -1,5 +1,5 @@
 import logging
-from typing import List, Optional
+from typing import TYPE_CHECKING, List, Optional, Tuple
 
 import numpy as np
 import torch
@@ -32,6 +32,9 @@
 
 _is_cpu = is_cpu()
 
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
 logger = logging.getLogger(__name__)
 
 
@@ -39,6 +42,11 @@
 
 
 class NGRAMWorker(BaseSpecWorker):
+    def iter_runners(self) -> List[Tuple[str, "ModelRunner"]]:
+        # NGRAM shares the target's model_runner -- no independent draft weights
+        # (the n-gram corpus is a CPU lookup structure built from token streams).
+        return []
+
     def alloc_memory_pool(self, **kwargs):
         # The target memory pool does not exist yet when __init__ runs.
         self.req_to_token_pool, self.token_to_kv_pool_allocator = (
@@ -115,15 +123,6 @@ def clear_cache_pool(self):
         self.ngram_corpus.reset()
         self._prev_decode_rids = set()
 
-    def update_weights_from_tensor(self, recv_req):
-        # NGRAM has no draft weights of its own — the n-gram corpus is a CPU
-        # lookup structure built from request token streams — and its
-        # `model_runner` is shared with the target worker. The scheduler
-        # mixin dispatches via `self.draft_worker or self.tp_worker`, so
-        # without this method any caller of `update_weights_from_tensor`
-        # under `--speculative-algorithm NGRAM` raises AttributeError.
-        return self.target_worker.update_weights_from_tensor(recv_req)
-
     def add_external_corpus(self, corpus_id: str, token_chunks: list[list[int]]) -> int:
         return self.ngram_corpus.load_external_corpus_named(corpus_id, token_chunks)
 
diff --git a/python/sglang/srt/utils/weight_checker.py b/python/sglang/srt/utils/weight_checker.py
index 8d2c088386eb..f67e61ee4e2f 100644
--- a/python/sglang/srt/utils/weight_checker.py
+++ b/python/sglang/srt/utils/weight_checker.py
@@ -1,7 +1,7 @@
 import hashlib
 import logging
 import time
-from typing import Any, Callable, Dict, Iterable, NamedTuple, Optional, Set
+from typing import Any, Callable, Dict, Iterable, List, NamedTuple, Optional, Set
 
 import torch
 import torch.distributed as dist
@@ -63,24 +63,50 @@ def _is_non_persistent_buffer_name(name: str) -> bool:
     return any(pat in name for pat in _NON_PERSISTENT_BUFFER_PATTERNS)
 
 
+def _is_skip_weight_check(name, param, skip_tensor_list=None) -> bool:
+    # One skip group shared by reset / compare / checksum; the _skip_weight_check
+    # flag is set on kv-cache k/v_scale and process_weights_after_loading placeholders.
+    return (
+        _is_non_persistent_buffer_name(name)
+        or getattr(param, "_skip_weight_check", False)
+        or any(pat in name for pat in (skip_tensor_list or ()))
+    )
+
+
+def overall_checksum(checksums: Dict[str, str]) -> str:
+    h = hashlib.sha256()
+    for name in sorted(checksums):
+        h.update(name.encode())
+        h.update(checksums[name].encode())
+    return h.hexdigest()
+
+
 class WeightChecker:
     def __init__(self, *, get_model: Callable[[], Any], ps: Any):
         self._get_model = get_model
         self._ps = ps
         self._snapshot_tensors = None
 
-    def handle(self, action: str, allow_quant_error: bool = False) -> Optional[Dict]:
+    def handle(
+        self,
+        action: str,
+        allow_quant_error: bool = False,
+        skip_tensor_list: Optional[List[str]] = None,
+    ) -> Optional[Dict]:
         logger.info(
-            f"[WeightChecker] handle action={action} allow_quant_error={allow_quant_error}"
+            f"[WeightChecker] handle action={action} "
+            f"allow_quant_error={allow_quant_error} skip_tensor_list={skip_tensor_list}"
         )
         if action == "snapshot":
             return self._snapshot()
         elif action == "reset_tensors":
-            return self._reset_tensors()
+            return self._reset_tensors(skip_tensor_list)
         elif action == "compare":
-            return self._compare(allow_quant_error=allow_quant_error)
+            return self._compare(
+                allow_quant_error=allow_quant_error, skip_tensor_list=skip_tensor_list
+            )
         elif action == "checksum":
-            return self._compute_checksum()
+            return self._compute_checksum(skip_tensor_list)
         else:
             raise Exception(f"Unsupported {action=}")
 
@@ -93,21 +119,32 @@ def _snapshot(self):
             named_tensors
         ), f"should not have duplicated tensor name"
 
-    def _reset_tensors(self):
+    def _skip_compare_names(
+        self, skip_tensor_list: Optional[List[str]] = None
+    ) -> Set[str]:
+        return {
+            name
+            for name, param in self._model_state()
+            if _is_skip_weight_check(name, param, skip_tensor_list)
+        }
+
+    def _reset_tensors(self, skip_tensor_list: Optional[List[str]] = None):
         for name, param in self._model_state():
-            if _is_non_persistent_buffer_name(name):
+            # Skip exactly what compare/checksum skip, so reset only poisons
+            # tensors compare will verify.
+            if _is_skip_weight_check(name, param, skip_tensor_list):
                 continue
             param.copy_(_random_like(param))
 
-    def _compare(self, allow_quant_error: bool = False):
+    def _compare(
+        self,
+        allow_quant_error: bool = False,
+        skip_tensor_list: Optional[List[str]] = None,
+    ):
         assert self._snapshot_tensors is not None
 
         quantized_set = _build_quantized_set(self._get_model())
-        skip_compare_names = {
-            name
-            for name, param in self._model_state()
-            if getattr(param, "_skip_weight_check", False)
-        }
+        skip_compare_names = self._skip_compare_names(skip_tensor_list)
         _check_tensors(
             expect_tensors=_build_check_entries(
                 self._snapshot_tensors, skip_compare_names, quantized_set
@@ -118,16 +155,12 @@ def _compare(self, allow_quant_error: bool = False):
             allow_quant_error=allow_quant_error,
         )
 
-    def _compute_checksum(self) -> Dict:
+    def _compute_checksum(self, skip_tensor_list: Optional[List[str]] = None) -> Dict:
         torch.cuda.synchronize()
         start = time.perf_counter()
 
         quantized_set = _build_quantized_set(self._get_model())
-        skip_compare_names = {
-            name
-            for name, param in self._model_state()
-            if getattr(param, "_skip_weight_check", False)
-        }
+        skip_compare_names = self._skip_compare_names(skip_tensor_list)
 
         # Hash the dequantized weight so two (qweight, scale) pairs with the same
         # bf16 hash equal.
@@ -138,11 +171,7 @@ def _compute_checksum(self) -> Dict:
             if should_compare:
                 checksums[name] = _hash_tensor(comparable.dequantize().data)
 
-        h = hashlib.sha256()
-        for name in sorted(checksums):
-            h.update(name.encode())
-            h.update(checksums[name].encode())
-        overall = h.hexdigest()
+        overall = overall_checksum(checksums)
 
         torch.cuda.synchronize()
         elapsed = time.perf_counter() - start
diff --git a/test/registered/rl/test_distributed_weight_update_spec_worker.py b/test/registered/rl/test_distributed_weight_update_spec_worker.py
new file mode 100644
index 000000000000..43f534092831
--- /dev/null
+++ b/test/registered/rl/test_distributed_weight_update_spec_worker.py
@@ -0,0 +1,225 @@
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+import pytest
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+from sglang.srt.managers.io_struct import (
+    BeginWeightUpdateReqInput,
+    EndWeightUpdateReqInput,
+    UpdateWeightsFromDistributedReqInput,
+)
+from sglang.srt.managers.scheduler_components.weight_updater import (
+    SchedulerWeightUpdaterManager,
+)
+
+
+def _distributed_req(selector="all"):
+    return UpdateWeightsFromDistributedReqInput(
+        names=["model.layers.0.weight"],
+        dtypes=["float32"],
+        shapes=[[1]],
+        group_name="weight_update_group",
+        flush_cache=False,
+        selector=selector,
+    )
+
+
+def _manager(tp_worker, draft_worker):
+    # metrics_collector defaults to None, so _observe_weight_load is a no-op; the
+    # reqs below set flush_cache=False so flush_cache is never called either.
+    manager = SchedulerWeightUpdaterManager(
+        tp_worker=tp_worker,
+        draft_worker=draft_worker,
+        tp_cpu_group=object(),
+        memory_saver_adapter=Mock(),
+        flush_cache=Mock(return_value=True),
+        is_fully_idle=Mock(return_value=True),
+    )
+    # update_weights_from_* assert an open begin_weight_update session.
+    manager._weight_update_in_progress = True
+    return manager
+
+
+def test_scheduler_distributed_update_receives_once_on_target_loads_into_each():
+    # Default selector ("all"): only the target (main model) owns the update group,
+    # so it receives the broadcast once; that single weights object is then loaded
+    # into every selected runner — receive once on the target, load into each.
+    weights = object()
+    target_runner = Mock()
+    target_runner.weight_updater.receive_weights_from_distributed.return_value = weights
+    draft_runner = Mock()
+    manager = _manager(
+        tp_worker=SimpleNamespace(
+            model_runner=target_runner,
+            iter_runners=lambda: [("", target_runner)],
+        ),
+        draft_worker=SimpleNamespace(iter_runners=lambda: [("draft", draft_runner)]),
+    )
+
+    output = manager.update_weights_from_distributed(_distributed_req())
+
+    assert output.success is True
+    target_runner.weight_updater.receive_weights_from_distributed.assert_called_once_with(
+        ["model.layers.0.weight"],
+        ["float32"],
+        [[1]],
+        "weight_update_group",
+        None,
+    )
+    # The single received weights object is loaded into every selected runner.
+    target_runner.weight_updater.load_weights.assert_called_once_with(weights)
+    draft_runner.weight_updater.load_weights.assert_called_once_with(weights)
+
+
+def test_scheduler_distributed_update_target_only_selector_skips_draft():
+    # selector="target": the target still receives once, but the draft worker is
+    # never enumerated and no draft runner is loaded.
+    weights = object()
+    target_runner = Mock()
+    target_runner.weight_updater.receive_weights_from_distributed.return_value = weights
+    draft_worker = Mock()
+    manager = _manager(
+        tp_worker=SimpleNamespace(
+            model_runner=target_runner,
+            iter_runners=lambda: [("", target_runner)],
+        ),
+        draft_worker=draft_worker,
+    )
+
+    output = manager.update_weights_from_distributed(
+        _distributed_req(selector="target")
+    )
+
+    assert output.success is True
+    target_runner.weight_updater.receive_weights_from_distributed.assert_called_once()
+    target_runner.weight_updater.load_weights.assert_called_once_with(weights)
+    draft_worker.iter_runners.assert_not_called()
+
+
+def _session_manager(target_runner, draft_runner):
+    return _manager(
+        tp_worker=SimpleNamespace(iter_runners=lambda: [("", target_runner)]),
+        draft_worker=SimpleNamespace(iter_runners=lambda: [("draft", draft_runner)]),
+    )
+
+
+def test_begin_weight_update_restores_target_and_draft():
+    # The session begins on every runner (target + draft): the draft model is
+    # restored to a loadable state identically to the target.
+    target_runner = Mock()
+    draft_runner = Mock()
+    manager = _session_manager(target_runner, draft_runner)
+    manager._weight_update_in_progress = False
+
+    with patch("torch.distributed.barrier"):
+        output = manager.begin_weight_update(BeginWeightUpdateReqInput())
+
+    assert output.success is True
+    target_runner.begin_weight_update.assert_called_once_with()
+    draft_runner.begin_weight_update.assert_called_once_with()
+    assert manager._weight_update_in_progress is True
+    assert manager._weight_update_loaded is False
+
+
+def test_end_weight_update_runs_post_load_on_both_when_load_was_bypassed():
+    # No load_weights happened this session (e.g. P2P/RDMA), so end runs
+    # post_load_weights then quant finalize on BOTH target and draft.
+    target_runner = Mock()
+    draft_runner = Mock()
+    manager = _session_manager(target_runner, draft_runner)
+    manager._weight_update_loaded = False
+
+    with patch("torch.distributed.barrier"):
+        output = manager.end_weight_update(EndWeightUpdateReqInput())
+
+    assert output.success is True
+    target_runner.end_weight_update.assert_called_once_with(run_post_load=True)
+    draft_runner.end_weight_update.assert_called_once_with(run_post_load=True)
+    assert manager._weight_update_in_progress is False
+
+
+def test_end_weight_update_skips_post_load_on_both_when_weights_loaded():
+    # A distributed/tensor load happened this session, so post_load is skipped on
+    # both runners; only quant finalize runs.
+    target_runner = Mock()
+    draft_runner = Mock()
+    manager = _session_manager(target_runner, draft_runner)
+    manager._weight_update_loaded = True
+
+    with patch("torch.distributed.barrier"):
+        manager.end_weight_update(EndWeightUpdateReqInput())
+
+    target_runner.end_weight_update.assert_called_once_with(run_post_load=False)
+    draft_runner.end_weight_update.assert_called_once_with(run_post_load=False)
+
+
+def test_model_runner_begin_end_wire_to_loader_hooks():
+    # ModelRunner.begin/end delegate to the loader: begin restores; end runs
+    # post_load only when requested, always finalizes quant layout.
+    import sglang.srt.model_executor.model_runner as mr
+
+    runner = SimpleNamespace(model=object(), device="cpu")
+
+    with patch.object(mr, "restore_weight") as restore:
+        mr.ModelRunner.begin_weight_update(runner)
+    restore.assert_called_once()
+
+    with patch.object(mr, "post_load_weights") as post_load, patch.object(
+        mr, "postprocess_weight"
+    ) as postprocess:
+        mr.ModelRunner.end_weight_update(runner, run_post_load=True)
+    post_load.assert_called_once()
+    postprocess.assert_called_once()
+
+    with patch.object(mr, "post_load_weights") as post_load, patch.object(
+        mr, "postprocess_weight"
+    ) as postprocess:
+        mr.ModelRunner.end_weight_update(runner, run_post_load=False)
+    post_load.assert_not_called()
+    postprocess.assert_called_once()
+
+
+def test_begin_weight_update_selector_restores_only_selected_and_is_recorded():
+    # begin(selector="draft") opens the session on the draft only; the target is
+    # untouched, and the selector is recorded for end to reuse.
+    target_runner = Mock()
+    draft_runner = Mock()
+    manager = _session_manager(target_runner, draft_runner)
+    manager._weight_update_in_progress = False
+
+    with patch("torch.distributed.barrier"):
+        manager.begin_weight_update(BeginWeightUpdateReqInput(selector="draft"))
+
+    target_runner.begin_weight_update.assert_not_called()
+    draft_runner.begin_weight_update.assert_called_once_with()
+    assert manager._weight_update_selector == "draft"
+
+
+def test_end_weight_update_reuses_session_selector_from_begin():
+    # end has no selector of its own; it finalizes exactly the set begin opened.
+    target_runner = Mock()
+    draft_runner = Mock()
+    manager = _session_manager(target_runner, draft_runner)
+    manager._weight_update_in_progress = False
+
+    with patch("torch.distributed.barrier"):
+        manager.begin_weight_update(BeginWeightUpdateReqInput(selector="draft"))
+        manager.end_weight_update(EndWeightUpdateReqInput())
+
+    target_runner.end_weight_update.assert_not_called()
+    draft_runner.end_weight_update.assert_called_once()
+
+
+def test_begin_weight_update_rejects_reentry():
+    # A second begin while a session is open would leave the first session's
+    # restored runners unfinalized — reject it loudly.
+    manager = _session_manager(Mock(), Mock())
+    manager._weight_update_in_progress = True
+
+    with patch("torch.distributed.barrier"):
+        with pytest.raises(AssertionError, match="already open"):
+            manager.begin_weight_update(BeginWeightUpdateReqInput())

From 94442dcbc679cf386195732d3ca4783f23e35819 Mon Sep 17 00:00:00 2001
From: Yueming Yuan <yym022502@gmail.com>
Date: Fri, 24 Jul 2026 23:01:14 -0700
Subject: [PATCH 22/43] [17/27] [sglang-miles] exclude shared skip-topk layer
 indexer weights from RL weight check (#29339)

---
 python/sglang/srt/models/deepseek_v2.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py
index 7acaf038edb2..80671b39b783 100644
--- a/python/sglang/srt/models/deepseek_v2.py
+++ b/python/sglang/srt/models/deepseek_v2.py
@@ -1686,6 +1686,10 @@ def __init__(
                     self.skip_topk = dsa_layer_skips_topk(config, layer_id)
                     self.next_skip_topk = dsa_layer_skips_topk(config, layer_id + 1)
 
+            if self.skip_topk:
+                for p in self.indexer.parameters():
+                    p._skip_weight_check = True
+
         self.kv_b_proj = ColumnParallelLinear(
             self.kv_lora_rank,
             self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),

From 30bce16ad0ee482f7306a0ce157f3acd9430dd8c Mon Sep 17 00:00:00 2001
From: Yueming Yuan <yym022502@gmail.com>
Date: Fri, 24 Jul 2026 23:04:02 -0700
Subject: [PATCH 23/43] [18/27] [sglang-miles] Fix stale _attn_sink_local cache
 after RL weight updates (#30421)

The per-rank attn_sink slice is built once and cached; an RL weight update
rewrites attn_sink in place, so decode keeps serving the pre-update sink.
Refresh the cache from post_load_weights instead.

Extended past the original PR for v0.5.16, which added deepseek_v4_dspark.py
with the same lazily-cached slice: refresh_attn_sink_cache moves from MQALayer
up to MqaAttentionBase (with an overridable _attn_sink_pad_width, since DSpark
pads q to _PAD_NUM_HEADS rather than padded_num_heads), and
DeepseekV4ForCausalLMDSpark gains the matching post_load_weights hook. Without
this, DSpark + RL would keep the same stale-sink bug this commit fixes for the
non-DSpark path.
---
 python/sglang/srt/models/deepseek_v4.py       | 46 ++++++++++++-------
 .../sglang/srt/models/deepseek_v4_dspark.py   | 19 ++++++--
 python/sglang/srt/models/deepseek_v4_nextn.py |  1 +
 3 files changed, 45 insertions(+), 21 deletions(-)

diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py
index c123d642cbaa..bccd5e19ae57 100644
--- a/python/sglang/srt/models/deepseek_v4.py
+++ b/python/sglang/srt/models/deepseek_v4.py
@@ -453,9 +453,15 @@ def __init__(
         self.fuse_wqa_wkv = fuse
 
         self.attn_sink = nn.Parameter(torch.empty(self.n_heads, dtype=torch.float32))
+        # FlashMLA's fp8 sparse decode kernel only specializes h_q for {64, 128}.
+        # Pad the per-rank heads to 64 (not the full n_heads) when they fit, to
+        # dispatch the cheaper decode::head64 variant; attn_sink is sliced to
+        # this rank and padded to match.
+        self.padded_num_heads = 64 if self.n_local_heads <= 64 else self.n_heads
         self._attn_sink_local: Optional[torch.Tensor] = (
             self.attn_sink if self.attn_tp_size == 1 else None
         )
+
         if fuse:
             self.wqkv_a = ReplicatedLinear(
                 self.hidden_size,
@@ -545,6 +551,26 @@ def __init__(
         self.register_buffer("freqs_cis", freqs_cis, persistent=False)
         self.freqs_cis: torch.Tensor
 
+    def _attn_sink_pad_width(self) -> int:
+        """Width of the per-rank attn_sink buffer; subclasses that pad q differently
+        override this so the sink matches their q layout."""
+        return self.padded_num_heads
+
+    def refresh_attn_sink_cache(self):
+        if self._attn_sink_local is self.attn_sink:
+            return
+        if self._attn_sink_local is None:
+            self._attn_sink_local = self.attn_sink.new_zeros(
+                self._attn_sink_pad_width()
+            )
+        self._attn_sink_local[: self.n_local_heads].copy_(
+            self.attn_sink[
+                self.attn_tp_rank
+                * self.n_local_heads : (self.attn_tp_rank + 1)
+                * self.n_local_heads
+            ]
+        )
+
 
 class MQALayer(MqaAttentionBase):
     def __init__(
@@ -1108,30 +1134,17 @@ def forward(
 
         tp_slice, q_padded, q_out = slice(None), None, None
         if self.tp_size > 1:
-            # FlashMLA's fp8 sparse decode kernel only specializes h_q for {64, 128}.
-            # Pad the per-rank heads to 64 (not the full n_heads) when they fit, to
-            # dispatch the cheaper decode::head64 variant; attn_sink is sliced to
-            # this rank and padded to match.
-            padded_num_heads = 64 if self.n_local_heads <= 64 else self.n_heads
             # Only [0:n_local_heads] is written below. Uninitialized padded TP
             # heads inject NaN into attention on gfx942 (fnuz), so zero-init
             # there; other archs tolerate new_empty and skip the per-forward
             # memset.
             if _is_gfx942_supported:
-                q_padded = x.new_zeros(x.shape[0], padded_num_heads, self.head_dim)
+                q_padded = x.new_zeros(x.shape[0], self.padded_num_heads, self.head_dim)
             else:
-                q_padded = x.new_empty(x.shape[0], padded_num_heads, self.head_dim)
+                q_padded = x.new_empty(x.shape[0], self.padded_num_heads, self.head_dim)
             tp_slice = slice(0, self.n_local_heads)
             q_out = q_padded[:, tp_slice, :]
-            if self._attn_sink_local is None:
-                # Build once on the first forward (post weight load); a per-call
-                # rebuild would replay a fill+copy per layer in the decode graph.
-                rank = self.tp_rank
-                sink = self.attn_sink.new_zeros(padded_num_heads)
-                sink[: self.n_local_heads] = self.attn_sink[
-                    rank * self.n_local_heads : (rank + 1) * self.n_local_heads
-                ]
-                self._attn_sink_local = sink
+            assert self._attn_sink_local is not None
 
         if enable_multi_stream:
             # Multi-stream path always fuses cache write into the K kernel,
@@ -2592,6 +2605,7 @@ def post_load_weights(self, is_nextn=False, weight_names=None):
             ):
                 self_attn.indexer.compressor.apply_ape_hotfix()
             layer.refresh_mhc_norm_weight_cache()
+            self_attn.refresh_attn_sink_cache()
 
     @staticmethod
     def remap_weight_name_to_dpsk_hf_format(
diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py
index 98a411afd18d..cabcc7d8e4cc 100644
--- a/python/sglang/srt/models/deepseek_v4_dspark.py
+++ b/python/sglang/srt/models/deepseek_v4_dspark.py
@@ -121,15 +121,16 @@ def kv_proj_only(self, x: torch.Tensor) -> torch.Tensor:
         kv, _ = self.wkv(x)
         return kv
 
+    def _attn_sink_pad_width(self) -> int:
+        # DSpark pads q to _PAD_NUM_HEADS only when the rank has fewer heads than
+        # that, so the sink follows the same width rather than padded_num_heads.
+        return max(self.n_local_heads, _PAD_NUM_HEADS)
+
     def _local_attn_sink(self) -> torch.Tensor:
         if self.attn_tp_size == 1:
             return self.attn_sink
         if self._attn_sink_local is None:
-            rank = self.attn_tp_rank
-            num_heads = self.n_local_heads
-            sink = self.attn_sink.new_zeros(max(num_heads, _PAD_NUM_HEADS))
-            sink[:num_heads] = self.attn_sink[rank * num_heads : (rank + 1) * num_heads]
-            self._attn_sink_local = sink
+            self.refresh_attn_sink_cache()
         return self._attn_sink_local
 
     def _store_block_kv(
@@ -838,6 +839,14 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> None:
             params_dict=params_dict, loaded_params=loaded_params
         )
 
+    def post_load_weights(self):
+        # The per-rank attn_sink slice is cached on first use; an RL weight update
+        # rewrites attn_sink in place, so re-slice it here or decode keeps serving
+        # the pre-update sink.
+        for module in self.modules():
+            if isinstance(module, DSparkAttention):
+                module.refresh_attn_sink_cache()
+
     def _assert_confidence_head_loaded(
         self, *, params_dict: dict, loaded_params: set
     ) -> None:
diff --git a/python/sglang/srt/models/deepseek_v4_nextn.py b/python/sglang/srt/models/deepseek_v4_nextn.py
index 1dd326c6e717..2c5120c10c4c 100644
--- a/python/sglang/srt/models/deepseek_v4_nextn.py
+++ b/python/sglang/srt/models/deepseek_v4_nextn.py
@@ -278,6 +278,7 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
 
     def post_load_weights(self, is_nextn=False, weight_names=None):
         super().post_load_weights(is_nextn=True, weight_names=weight_names)
+        self.model.decoder.self_attn.refresh_attn_sink_cache()
 
 
 EntryClass = [DeepseekV4ForCausalLMNextN]

From 627661641b8e614321a3e9300283b23f3c36ac59 Mon Sep 17 00:00:00 2001
From: "Ethan (Yusheng) Su" <yushengsu@radixark.ai>
Date: Fri, 24 Jul 2026 23:04:02 -0700
Subject: [PATCH 24/43] [19/27] [lora] Support GDN in_proj_ba adapters for
 Qwen3.5 (#30499)

---
 python/sglang/srt/lora/lora.py      | 60 +++++++++++++++++++++++++++++
 python/sglang/srt/lora/utils.py     |  4 ++
 python/sglang/srt/models/qwen3_5.py |  4 ++
 3 files changed, 68 insertions(+)

diff --git a/python/sglang/srt/lora/lora.py b/python/sglang/srt/lora/lora.py
index c154e1c5581b..99848d822e5f 100644
--- a/python/sglang/srt/lora/lora.py
+++ b/python/sglang/srt/lora/lora.py
@@ -210,6 +210,8 @@ def _normalize_weights(self):
             self._normalize_in_proj(layer.weights)
             # Stack in_proj_q + in_proj_k + in_proj_v + in_proj_z → in_proj_qkvz for GDN layers
             self._normalize_in_proj_qkvz(layer.weights)
+            # Stack in_proj_b + in_proj_a → in_proj_ba for GDN layers
+            self._normalize_in_proj_ba(layer.weights)
             weight_names = list(layer.weights.keys())
             self.normalize_gate_up_proj(weight_names, layer.weights)
             weight_names = list(layer.weights.keys())
@@ -450,6 +452,25 @@ def _normalize_in_proj_qkvz(self, weights: Dict[str, torch.Tensor]):
                 weights.pop(k_name)
                 weights.pop(v_name)
                 weights.pop(z_name)
+            elif "in_proj_qkv." in weight_name:
+                # 2-way split (Megatron-Bridge adapter export): in_proj_qkv
+                # covers the q|k|v rows and in_proj_z the z rows, with one
+                # shared lora_A. The stacked buffer expects one A block per
+                # slice, so repeat the qkv A block 3x (q, k, v) before
+                # appending the z block; B rows concatenate directly.
+                z_name = weight_name.replace("in_proj_qkv.", "in_proj_z.")
+                if z_name not in weights:
+                    continue
+                qkvz_name = weight_name.replace("in_proj_qkv.", "in_proj_qkvz.")
+                cat_dim = weights[weight_name].dim() - 2
+                qkv_w = weights[weight_name]
+                if "lora_A" in weight_name:
+                    repeat_dims = [1] * qkv_w.dim()
+                    repeat_dims[cat_dim] = 3
+                    qkv_w = qkv_w.repeat(*repeat_dims)
+                weights[qkvz_name] = torch.cat((qkv_w, weights[z_name]), cat_dim)
+                weights.pop(weight_name)
+                weights.pop(z_name)
             elif "in_proj_qkvz" in weight_name and "lora_A" in weight_name:
                 # Already-merged adapter: replicate the shared A across the 4
                 # stacked slots the buffer expects (q, k, v, z).
@@ -459,6 +480,45 @@ def _normalize_in_proj_qkvz(self, weights: Dict[str, torch.Tensor]):
                 weights[weight_name] = weights[weight_name].repeat(*repeat_dims)
             # else (in_proj_qkvz lora_B, or unrelated): no-op.
 
+    def _normalize_in_proj_ba(self, weights: Dict[str, torch.Tensor]):
+        """Normalize in_proj_ba weights for GDN (GatedDeltaNet) layers like
+        Qwen3.5.
+
+        Two adapter formats are handled:
+
+        1. Split: ``in_proj_b + in_proj_a`` (HF checkpoint naming, also the
+           Megatron-Bridge adapter export) are present as separate weights →
+           concatenate them into ``in_proj_ba`` (B rows b|a; A blocks b, a).
+
+        2. Already-merged: the adapter has a single ``in_proj_ba`` weight
+           (PEFT trained against SGLang's fused Linear). The stacked buffer
+           expects two per-slice ``A`` blocks, so repeat ``lora_A`` 2x along
+           the rank dim. ``lora_B`` is already full-output-dim and matches
+           the buffer directly.
+        """
+        for weight_name in list(weights.keys()):
+            # NB: match with the trailing dot so the merged "in_proj_ba."
+            # names don't take the split branch.
+            if "in_proj_b." in weight_name:
+                a_name = weight_name.replace("in_proj_b.", "in_proj_a.")
+                if a_name not in weights:
+                    continue
+                ba_name = weight_name.replace("in_proj_b.", "in_proj_ba.")
+                cat_dim = weights[weight_name].dim() - 2
+                weights[ba_name] = torch.cat(
+                    (weights[weight_name], weights[a_name]), cat_dim
+                )
+                weights.pop(weight_name)
+                weights.pop(a_name)
+            elif "in_proj_ba" in weight_name and "lora_A" in weight_name:
+                # Already-merged adapter: replicate the shared A across the 2
+                # stacked slots the buffer expects (b, a).
+                ndim = weights[weight_name].dim()
+                repeat_dims = [1] * ndim
+                repeat_dims[ndim - 2] = 2
+                weights[weight_name] = weights[weight_name].repeat(*repeat_dims)
+            # else (in_proj_ba lora_B, or unrelated): no-op.
+
     def normalize_gate_up_proj(
         self, weight_names: List[str], weights: Dict[str, torch.Tensor]
     ):
diff --git a/python/sglang/srt/lora/utils.py b/python/sglang/srt/lora/utils.py
index 16459ea78739..4950daa616a6 100644
--- a/python/sglang/srt/lora/utils.py
+++ b/python/sglang/srt/lora/utils.py
@@ -253,6 +253,8 @@ def get_normalized_target_modules(
         "v_proj": "qkv_proj",
         "gate_proj": "gate_up_proj",
         "up_proj": "gate_up_proj",
+        "in_proj_b": "in_proj_ba",
+        "in_proj_a": "in_proj_ba",
         "out_proj": "out_proj",
         "embed_tokens": "embed_tokens",
         "vocab_emb": "embed_tokens",
@@ -291,6 +293,7 @@ def get_stacked_multiply(
     stacked_rank = {
         "qkv_proj": 3,
         "in_proj_qkvz": 4,  # GDN packed input projection
+        "in_proj_ba": 2,  # GDN packed b/a input projection
         "gate_up_proj": 2,
         "gate_up_proj_moe": 2,
         "gate_up_proj_shared_moe": 2,
@@ -352,6 +355,7 @@ def get_target_module_name(full_module_name: str, target_modules: Set[str]) -> s
         "out_proj",
         "in_proj",
         "in_proj_qkvz",
+        "in_proj_ba",
         "up_proj",
         "gate_up_proj",
         "down_proj",
diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py
index fe26e906e1ae..0e12141dc864 100644
--- a/python/sglang/srt/models/qwen3_5.py
+++ b/python/sglang/srt/models/qwen3_5.py
@@ -1138,6 +1138,7 @@ class Qwen3_5ForCausalLM(nn.Module):
         "o_proj",
         "out_proj",
         "in_proj_qkvz",
+        "in_proj_ba",
         "gate_up_proj",
         "down_proj",
         "lm_head",
@@ -1163,6 +1164,9 @@ def get_hidden_dim(self, module_name: str, layer_idx: int):
             key_dim = config.linear_key_head_dim * config.linear_num_key_heads
             value_dim = config.linear_value_head_dim * config.linear_num_value_heads
             return config.hidden_size, key_dim * 2 + value_dim * 2
+        elif module_name == "in_proj_ba":
+            # b + a projections: one scalar per linear-attention value head each
+            return config.hidden_size, config.linear_num_value_heads * 2
         elif module_name == "gate_up_proj":
             # MoE: shared expert uses shared_expert_intermediate_size
             # Dense: regular MLP uses intermediate_size

From e0bd0e68e16c3db1005658aa279238f239b31568 Mon Sep 17 00:00:00 2001
From: Nan Jiang <59716405+nanjiangwill@users.noreply.github.com>
Date: Fri, 24 Jul 2026 23:04:31 -0700
Subject: [PATCH 25/43] [20/27] [RL] Add /pull_weights: engine-side pull of
 published weights into a host-local checkpoint (#30366, #28524)

Folds in the two fixes #28524 made to this endpoint -- Annotated[..., Body()]
on the handler and keyword construction of PullWeightsReqOutput, which is a
kw_only msgspec Struct and would raise on the positional form. The rest of
#28524 (dumper / dump-comparator) is already upstream in v0.5.16 and is dropped.
---
 python/pyproject.toml                         |   3 +-
 python/sglang/srt/entrypoints/http_server.py  |  14 +
 python/sglang/srt/managers/io_struct.py       |  16 +
 python/sglang/srt/managers/scheduler.py       |   5 +
 .../scheduler_components/weight_updater.py    |  40 ++
 .../srt/managers/tokenizer_control_mixin.py   |  12 +
 python/sglang/srt/server_args.py              |   4 +
 .../srt/weight_sync/local_checkpoint.py       | 350 ++++++++++++++++++
 8 files changed, 443 insertions(+), 1 deletion(-)
 create mode 100644 python/sglang/srt/weight_sync/local_checkpoint.py

diff --git a/python/pyproject.toml b/python/pyproject.toml
index f0146cebae60..b934e83351f8 100755
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -87,7 +87,8 @@ dependencies = [
   "uvloop",
   "watchfiles",
   "xgrammar==0.2.1",
-  "zstandard",
+  "xxhash",  # /pull_weights delta checksum
+  "zstandard",  # /pull_weights delta codec
 ]
 
 [[tool.uv.index]]
diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py
index af29ae5f7e1a..4c2ff20458c5 100644
--- a/python/sglang/srt/entrypoints/http_server.py
+++ b/python/sglang/srt/entrypoints/http_server.py
@@ -133,6 +133,7 @@
     ParseFunctionCallReq,
     PauseGenerationReqInput,
     ProfileReq,
+    PullWeightsReqInput,
     ReleaseMemoryOccupationReqInput,
     ResumeMemoryOccupationReqInput,
     SendWeightsToRemoteInstanceReqInput,
@@ -1200,6 +1201,19 @@ async def update_weights_from_disk(
         )
 
 
+@app.post("/pull_weights")
+@auth_level(AuthLevel.ADMIN_OPTIONAL)
+async def pull_weights(obj: Annotated[PullWeightsReqInput, Body()], request: Request):
+    """Have every host of this deployment pull published weight deltas into its
+    local checkpoint (materialized from the model path on first use)."""
+    success, message = await _global_state.tokenizer_manager.pull_weights(obj, request)
+
+    content = {"success": success, "message": message}
+    return ORJSONResponse(
+        content, status_code=HTTPStatus.OK if success else HTTPStatus.BAD_REQUEST
+    )
+
+
 @app.post("/init_weights_send_group_for_remote_instance")
 @auth_level(AuthLevel.ADMIN_OPTIONAL)
 async def init_weights_send_group_for_remote_instance(
diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py
index f5f5b15dbec2..507d573be770 100644
--- a/python/sglang/srt/managers/io_struct.py
+++ b/python/sglang/srt/managers/io_struct.py
@@ -1569,6 +1569,22 @@ class UpdateWeightFromDiskReqOutput(BaseReq, kw_only=True):
     num_paused_requests: int = 0
 
 
+class PullWeightsReqInput(BaseReq, kw_only=True):
+    # Host-local checkpoint dir the pulled weights land in; seeded from the
+    # server's model path when the published stream has no full version.
+    local_checkpoint_dir: str
+    # Shared dir the publisher writes weight_v{N:06d}/ version dirs under; each
+    # version is a full HF checkpoint or a delta against the previous version.
+    source_dir: str
+    # The version to bring the local checkpoint up to.
+    target_version: int
+
+
+class PullWeightsReqOutput(BaseReq, kw_only=True):
+    success: bool
+    message: str
+
+
 class UpdateWeightsFromDistributedReqInput(BaseReq, kw_only=True):
     names: List[str]
     dtypes: List[str]
diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
index f397c02cd126..92fbe2c2769f 100644
--- a/python/sglang/srt/managers/scheduler.py
+++ b/python/sglang/srt/managers/scheduler.py
@@ -125,6 +125,7 @@
     OpenSessionReqInput,
     PauseGenerationReqInput,
     ProfileReq,
+    PullWeightsReqInput,
     ReleaseMemoryOccupationReqInput,
     RemoveExternalCorpusReqInput,
     RemoveExternalCorpusReqOutput,
@@ -1405,6 +1406,10 @@ def init_request_dispatcher(self):
                     CheckWeightsReqInput,
                     self.weight_updater.check_weights,
                 ),
+                (
+                    PullWeightsReqInput,
+                    self.weight_updater.pull_weights,
+                ),
                 (SlowDownReqInput, self.slow_down),
                 (
                     ProfileReq,
diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py
index 62fc51554027..c68699b827ef 100644
--- a/python/sglang/srt/managers/scheduler_components/weight_updater.py
+++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py
@@ -31,6 +31,8 @@
     GetWeightsByNameReqOutput,
     InitWeightsUpdateGroupReqInput,
     InitWeightsUpdateGroupReqOutput,
+    PullWeightsReqInput,
+    PullWeightsReqOutput,
     ReleaseMemoryOccupationReqInput,
     ReleaseMemoryOccupationReqOutput,
     ResumeMemoryOccupationReqInput,
@@ -135,6 +137,44 @@ def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput):
                 success=success, message=message, num_paused_requests=0
             )
 
+    def pull_weights(self, recv_req: PullWeightsReqInput):
+        """Sync this host's local checkpoint up to recv_req.target_version.
+
+        Every rank runs the pull; a per-host file lock collapses co-located
+        ranks to one pull. Success is gathered across the TP group (all nodes),
+        so the reply only reports success once every host holds a verified
+        checkpoint.
+        """
+        from sglang.srt.weight_sync import local_checkpoint
+
+        server_args = self.tp_worker.model_runner.server_args
+        try:
+            local_checkpoint.pull(
+                local_checkpoint_dir=recv_req.local_checkpoint_dir,
+                base_dir=server_args.model_path,
+                source_dir=recv_req.source_dir,
+                target_version=recv_req.target_version,
+                pre_read_hook=server_args.custom_pull_weights_pre_read_hook,
+            )
+            success, message = True, "Success."
+        except Exception:
+            success, message = False, traceback.format_exc()
+            logger.error(message)
+
+        tp_size = (
+            torch.distributed.get_world_size(group=self.tp_cpu_group)
+            if torch.distributed.is_initialized()
+            else 1
+        )
+        if tp_size > 1:
+            results = [None] * tp_size
+            torch.distributed.all_gather_object(
+                results, (success, message), group=self.tp_cpu_group
+            )
+            success = all(ok for ok, _ in results)
+            message = "; ".join(msg for ok, msg in results if not ok) or message
+        return PullWeightsReqOutput(success=success, message=message)
+
     def init_weights_update_group(self, recv_req: InitWeightsUpdateGroupReqInput):
         """Initialize the online model parameter update group."""
         success, message = self.tp_worker.init_weights_update_group(recv_req)
diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py
index 3848d136ae14..1d3ff8d6ddd2 100644
--- a/python/sglang/srt/managers/tokenizer_control_mixin.py
+++ b/python/sglang/srt/managers/tokenizer_control_mixin.py
@@ -57,6 +57,8 @@
     ProfileReq,
     ProfileReqOutput,
     ProfileReqType,
+    PullWeightsReqInput,
+    PullWeightsReqOutput,
     ReleaseMemoryOccupationReqInput,
     ReleaseMemoryOccupationReqOutput,
     RemoveExternalCorpusReqInput,
@@ -111,6 +113,7 @@
     ("release_memory_occupation", ReleaseMemoryOccupationReqOutput),
     ("resume_memory_occupation", ResumeMemoryOccupationReqOutput),
     ("check_weights", CheckWeightsReqOutput),
+    ("pull_weights", PullWeightsReqOutput),
     ("slow_down", SlowDownReqOutput),
     ("flush_cache", FlushCacheReqOutput),
     ("add_external_corpus", AddExternalCorpusReqOutput),
@@ -916,6 +919,15 @@ async def resume_memory_occupation(
         self.auto_create_handle_loop()
         await self.resume_memory_occupation_communicator(obj)
 
+    async def pull_weights(
+        self: TokenizerManager,
+        obj: PullWeightsReqInput,
+        request: Optional[fastapi.Request] = None,
+    ) -> Tuple[bool, str]:
+        self.auto_create_handle_loop()
+        results = await self.pull_weights_communicator(obj)
+        return FanOutCommunicator.merge_results(results)
+
     async def check_weights(
         self: TokenizerManager,
         obj: CheckWeightsReqInput,
diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index c44a6f04b61a..04ba960b9ad2 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -2626,6 +2626,10 @@ class ServerArgs:
             nargs="*",
         ),
     ] = None
+    custom_pull_weights_pre_read_hook: A[
+        Optional[str],
+        "Import path of a hook(source_dir, target_version) that /pull_weights calls before reading the published weights. POSIX shared filesystems need no hook; object-store-backed mounts often lack cross-host read-after-write consistency, so another host's writes only become visible after an explicit refresh.",
+    ] = None
     weight_loader_disable_mmap: A[
         bool,
         "Disable mmap while loading weight using safetensors.",
diff --git a/python/sglang/srt/weight_sync/local_checkpoint.py b/python/sglang/srt/weight_sync/local_checkpoint.py
new file mode 100644
index 000000000000..8f59802f5590
--- /dev/null
+++ b/python/sglang/srt/weight_sync/local_checkpoint.py
@@ -0,0 +1,350 @@
+"""Host-local pull of published weights (the /pull_weights endpoint).
+
+A trainer publishes each weight sync as a version directory ``weight_v{N:06d}/``
+under a shared ``source_dir``. Each version is a canonical HF checkpoint
+directory of one of two kinds, distinguished by its index metadata:
+
+- **full**: an ordinary checkpoint. Pulling it copies it into the host-local
+  ``local_checkpoint_dir``, replacing whatever is there — no history needed.
+- **delta** (index metadata carries ``delta_encoding``): safetensors files
+  holding zstd-compressed per-tensor diffs against version N-1, plus per-tensor
+  checksums of the new state. Pulling it patches the local checkpoint in place.
+
+Version 0 is the engine's own base checkpoint (``model_path``). Every host of a
+(possibly multi-node) deployment runs the same pull; the engine then reloads the
+local checkpoint through the ordinary ``update_weights_from_disk`` path.
+
+``pull()`` is safe to call concurrently from every scheduler rank on a host: a
+per-host file lock serializes the work and an applied-version marker makes the
+extra calls no-ops.
+"""
+
+from __future__ import annotations
+
+import fcntl
+import glob
+import importlib
+import json
+import logging
+import mmap
+import os
+import shutil
+import struct
+import threading
+import zlib
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager
+from typing import Optional
+
+import numpy as np
+import zstandard
+
+logger = logging.getLogger(__name__)
+
+# The delta-apply phases (decompress, XOR/scatter, checksum) are memory-bandwidth
+# bound and release the GIL, so a thread pool over tensors recovers the
+# bandwidth one thread leaves idle.
+NUM_WORKERS = min(32, (os.cpu_count() or 8))
+
+# Per-checkpoint dir holding the applied-version marker and the pull lock.
+SYNC_DIR = ".weight_sync"
+
+
+def pull(
+    local_checkpoint_dir: str,
+    base_dir: str,
+    source_dir: str,
+    target_version: int,
+    pre_read_hook: Optional[str] = None,
+) -> None:
+    """Bring the host-local checkpoint up to ``target_version``.
+
+    Seeds from the newest full checkpoint at or below the target — the engine's
+    own base (``base_dir``) for a pure-delta stream, a published full version
+    otherwise — then applies the remaining deltas in order. A local checkpoint
+    already past the seed point just continues its delta chain. Raises on any
+    per-tensor checksum mismatch (fail loud, never serve bad weights).
+    """
+    # Object-store-backed shared filesystems lack cross-host read-after-write
+    # consistency: the publisher's files only appear here after an explicit
+    # refresh, which the deployment supplies as this hook. POSIX shared
+    # filesystems (NFS, Lustre, ...) need none.
+    if target_version > 0 and pre_read_hook:
+        _load_hook(pre_read_hook)(source_dir, target_version)
+    with _pull_lock(local_checkpoint_dir):
+        applied = _read_applied_version(local_checkpoint_dir)  # None on a fresh host
+        # Scan back from the target for the newest full version. Stop at the
+        # local state — below it a reset can never be needed (or, on a fresh
+        # host, at 0 = the engine's base).
+        floor = applied if applied is not None else 0
+        start = target_version
+        while start > floor and _is_delta(_version_dir(source_dir, start)):
+            start -= 1
+        if applied is None or start > applied:
+            seed_dir = base_dir if start == 0 else _version_dir(source_dir, start)
+            _reset_checkpoint(seed_dir, local_checkpoint_dir, start)
+        else:
+            start = applied
+        for version in range(start + 1, target_version + 1):
+            _apply_delta(local_checkpoint_dir, _version_dir(source_dir, version))
+
+
+def _load_hook(path: str):
+    module_path, _, name = path.rpartition(".")
+    return getattr(importlib.import_module(module_path), name)
+
+
+def _version_dir(source_dir: str, version: int) -> str:
+    return os.path.join(source_dir, f"weight_v{version:06d}")
+
+
+def _is_delta(version_dir: str) -> bool:
+    """A version is a delta iff its index metadata declares an encoding; an
+    ordinary HF checkpoint (with or without an index) is a full version."""
+    if not os.path.isdir(version_dir):
+        raise FileNotFoundError(f"published weight version missing: {version_dir}")
+    try:
+        with open(os.path.join(version_dir, "model.safetensors.index.json")) as f:
+            return "delta_encoding" in json.load(f).get("metadata", {})
+    except FileNotFoundError:
+        return False
+
+
+class _Adler32:
+    """adler32 behind the incremental .update / .hexdigest interface the hash objects expose."""
+
+    def __init__(self):
+        self._value = 1
+
+    def update(self, data) -> None:
+        self._value = zlib.adler32(data, self._value)
+
+    def hexdigest(self) -> str:
+        return f"{self._value:08x}"
+
+
+def _new_hasher(algorithm: str):
+    if algorithm == "xxh3-128":
+        import xxhash
+
+        return xxhash.xxh3_128()
+    if algorithm == "blake3":
+        import blake3
+
+        return blake3.blake3()
+    if algorithm == "adler32":
+        return _Adler32()
+    raise KeyError(f"unknown checksum algorithm {algorithm!r}")
+
+
+def _checksum(algorithm: str, buf) -> str:
+    hasher = _new_hasher(algorithm)
+    hasher.update(buf)
+    return hasher.hexdigest()
+
+
+@contextmanager
+def _pull_lock(local_checkpoint_dir: str):
+    sync = os.path.join(local_checkpoint_dir, SYNC_DIR)
+    os.makedirs(sync, exist_ok=True)
+    with open(os.path.join(sync, "lock"), "w") as f:
+        fcntl.flock(f, fcntl.LOCK_EX)
+        try:
+            yield
+        finally:
+            fcntl.flock(f, fcntl.LOCK_UN)
+
+
+def _read_applied_version(local_checkpoint_dir: str) -> Optional[int]:
+    try:
+        with open(os.path.join(local_checkpoint_dir, SYNC_DIR, "state.json")) as f:
+            return int(json.load(f)["version"])
+    except FileNotFoundError:
+        return None
+
+
+def _write_applied_version(local_checkpoint_dir: str, version: int) -> None:
+    path = os.path.join(local_checkpoint_dir, SYNC_DIR, "state.json")
+    tmp = path + ".tmp"
+    with open(tmp, "w") as f:
+        json.dump({"version": f"{version:06d}"}, f)
+        f.flush()
+        os.fsync(f.fileno())
+    os.replace(tmp, path)
+
+
+def _drop_page_cache(path: str) -> None:
+    """Evict a file from the page cache (POSIX_FADV_DONTNEED)."""
+    if not hasattr(os, "posix_fadvise"):  # POSIX-only (absent on macOS/Windows)
+        return
+    try:
+        fd = os.open(path, os.O_RDONLY)
+        try:
+            os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)
+        finally:
+            os.close(fd)
+    except OSError:
+        pass
+
+
+def _reset_checkpoint(src_dir: str, local_checkpoint_dir: str, version: int) -> None:
+    """Make local_checkpoint_dir an exact copy of the full checkpoint in src_dir
+    (files the new checkpoint doesn't have — e.g. differently-sharded old ones —
+    are pruned). Later deltas chain on top of this state."""
+    logger.info(
+        "Pulling full checkpoint v%d %s -> %s", version, src_dir, local_checkpoint_dir
+    )
+    os.makedirs(local_checkpoint_dir, exist_ok=True)
+    src_files = [entry for entry in os.scandir(src_dir) if entry.is_file()]
+    for entry in src_files:
+        shutil.copy2(entry.path, os.path.join(local_checkpoint_dir, entry.name))
+        # don't let the source evict the local copy we keep resident
+        _drop_page_cache(entry.path)
+    names = {entry.name for entry in src_files}
+    for entry in os.scandir(local_checkpoint_dir):
+        if entry.is_file() and entry.name not in names:
+            os.remove(entry.path)
+    # a truncated copy (e.g. an object-store mount surfacing metadata before
+    # bytes) must fail loud, not serve bad weights
+    for entry in src_files:
+        copied = os.path.getsize(os.path.join(local_checkpoint_dir, entry.name))
+        if copied != entry.stat().st_size:
+            raise RuntimeError(
+                f"size mismatch copying {entry.name}: src {entry.stat().st_size} != local {copied}"
+            )
+    _write_applied_version(local_checkpoint_dir, version)
+
+
+def _tensor_locations(ckpt_dir: str) -> dict:
+    """Map each tensor name to (file, byte offset, nbytes) by reading every safetensors header."""
+    locations = {}
+    for path in glob.glob(os.path.join(ckpt_dir, "*.safetensors")):
+        with open(path, "rb") as f:
+            (header_len,) = struct.unpack("<Q", f.read(8))
+            header = json.loads(f.read(header_len))
+        for name, info in header.items():
+            if name == "__metadata__":
+                continue
+            begin, end = info["data_offsets"]
+            locations[name] = (path, 8 + header_len + begin, end - begin)
+    return locations
+
+
+def _apply_delta(local_checkpoint_dir: str, version_dir: str) -> None:
+    """Apply one version's delta in place: decompress + apply + checksum each tensor across a thread
+    pool (each writes a distinct mmap region, so the writes don't conflict). Any mismatch raises.
+    """
+    with open(os.path.join(version_dir, "model.safetensors.index.json")) as f:
+        meta = json.load(f)["metadata"]
+    applied = _read_applied_version(local_checkpoint_dir)
+    if applied == int(meta["version"]):
+        return
+    if applied != int(meta["base_version"]):
+        raise RuntimeError(
+            f"out-of-order delta: local at {applied}, delta builds on {meta['base_version']}"
+        )
+    if meta["compression_format"] != "zstd":
+        raise NotImplementedError(
+            f"compression {meta['compression_format']!r} not supported"
+        )
+    encoding = meta["delta_encoding"]
+    algorithm = meta["checksum_format"]
+    locations = _tensor_locations(local_checkpoint_dir)
+    open_mmaps = {}
+    mismatches = []
+    lock = threading.Lock()
+    file_bytes = []  # keep alive: items hold zero-copy views into these
+    items = []  # (name, compressed_view, path, offset, nbytes, want_checksum)
+    try:
+        for delta_file in sorted(glob.glob(os.path.join(version_dir, "*.safetensors"))):
+            with open(delta_file, "rb") as f:
+                blob = f.read()
+            file_bytes.append(blob)
+            (header_len,) = struct.unpack("<Q", blob[:8])
+            header = json.loads(blob[8 : 8 + header_len])
+            want_checksums = header.get("__metadata__", {})
+            view = memoryview(blob)
+            for name, info in header.items():
+                if name == "__metadata__":
+                    continue
+                begin, end = info["data_offsets"]
+                path, offset, nbytes = locations[name]
+                if path not in open_mmaps:
+                    fh = open(path, "r+b")
+                    open_mmaps[path] = (fh, mmap.mmap(fh.fileno(), 0))
+                data_start = 8 + header_len
+                items.append(
+                    (
+                        name,
+                        view[data_start + begin : data_start + end],
+                        path,
+                        offset,
+                        nbytes,
+                        want_checksums.get(name),
+                    )
+                )
+
+        # prefetch into page cache (evicted during the rollout) so the apply
+        # doesn't fault from cold storage
+        for _, mm in open_mmaps.values():
+            try:
+                mm.madvise(mmap.MADV_WILLNEED)
+            except (OSError, AttributeError, ValueError):
+                pass
+
+        def apply_xor(item) -> None:
+            name, compressed, path, offset, nbytes, want = item
+            region = np.ndarray(
+                (nbytes,), dtype=np.uint8, buffer=open_mmaps[path][1], offset=offset
+            )
+            hasher = _new_hasher(algorithm)
+            reader = zstandard.ZstdDecompressor().stream_reader(compressed)
+            pos = 0
+            # 2 MB chunks stay L2-resident across decompress -> XOR -> checksum
+            while pos < nbytes:
+                block = reader.read(min(2 << 20, nbytes - pos))
+                if not block:
+                    break
+                chunk = np.frombuffer(block, dtype=np.uint8)
+                region[pos : pos + chunk.size] ^= chunk
+                hasher.update(region[pos : pos + chunk.size])
+                pos += chunk.size
+            if hasher.hexdigest() != want:
+                with lock:
+                    mismatches.append(name)
+
+        def apply_overwrite(item) -> None:
+            name, compressed, path, offset, nbytes, want = item
+            delta = np.frombuffer(
+                zstandard.ZstdDecompressor().decompress(compressed), dtype=np.uint8
+            )
+            region = np.ndarray(
+                (nbytes,), dtype=np.uint8, buffer=open_mmaps[path][1], offset=offset
+            )
+            count = int.from_bytes(delta[:4], "little")
+            positions = np.frombuffer(delta[4 : 4 + 4 * count], dtype="<u4")
+            region[positions] = delta[4 + 4 * count :]
+            if _checksum(algorithm, region) != want:
+                with lock:
+                    mismatches.append(name)
+
+        if encoding == "xor":
+            apply_tensor = apply_xor
+        elif encoding == "overwrite":
+            apply_tensor = apply_overwrite
+        else:
+            raise NotImplementedError(f"delta encoding {encoding!r} not supported")
+        with ThreadPoolExecutor(max_workers=NUM_WORKERS) as pool:
+            list(pool.map(apply_tensor, items))
+        # no msync: the engine reads these pages via the shared cache; durability
+        # isn't needed (a host that loses the cache rebuilds from base)
+    finally:
+        for fh, mm in open_mmaps.values():
+            mm.close()
+            fh.close()
+    if mismatches:
+        raise RuntimeError(
+            f"checksum mismatch for {len(mismatches)} tensors after applying {version_dir}: "
+            f"{sorted(mismatches)[:20]}"
+        )
+    _write_applied_version(local_checkpoint_dir, int(meta["version"]))

From 4b642641102ea9414290d66199a4f4d182aa6962 Mon Sep 17 00:00:00 2001
From: Mathew Han <49226490+mathewjhan@users.noreply.github.com>
Date: Fri, 24 Jul 2026 23:04:31 -0700
Subject: [PATCH 26/43] [21/27] feat(sglang-miles): Support aborting requests
 by rid prefix - multi-lora needs (#30912)

---
 python/sglang/srt/entrypoints/http_server.py  |   2 +-
 python/sglang/srt/managers/io_struct.py       |   7 +
 .../sglang/srt/managers/tokenizer_manager.py  |  68 ++-
 .../managers/test_abort_request_prefix.py     | 440 ++++++++++++++++++
 4 files changed, 508 insertions(+), 9 deletions(-)
 create mode 100644 test/registered/unit/managers/test_abort_request_prefix.py

diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py
index 4c2ff20458c5..5b7de28d519c 100644
--- a/python/sglang/srt/entrypoints/http_server.py
+++ b/python/sglang/srt/entrypoints/http_server.py
@@ -1632,7 +1632,7 @@ async def abort_request(obj: Annotated[AbortReq, Body()], request: Request):
     """Abort a request."""
     try:
         _global_state.tokenizer_manager.abort_request(
-            rid=obj.rid, abort_all=obj.abort_all
+            rid=obj.rid, abort_all=obj.abort_all, prefix=obj.prefix
         )
         return Response(status_code=200)
     except Exception as e:
diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py
index 507d573be770..de60b2068882 100644
--- a/python/sglang/srt/managers/io_struct.py
+++ b/python/sglang/srt/managers/io_struct.py
@@ -1851,6 +1851,13 @@ class AbortReq(BaseReq, kw_only=True):
     # The finished reason data (from BaseFinishReason.to_json())
     finished_reason: Optional[FinishReasonDict] = None
     abort_message: Optional[str] = None
+    # Treat ``rid`` as a namespace prefix. Note this only relaxes the
+    # tokenizer-side gate (which otherwise requires an exact live rid, with a
+    # single tokenizer worker): once the request reaches the scheduler,
+    # matching is prefix-based (``rid.startswith``) regardless of this flag,
+    # because batch requests derive child rids as ``f"{rid}_{i}"`` and an
+    # abort for the parent rid must cover them.
+    prefix: bool = False
 
     def __post_init__(self):
         # FIXME: This is a hack to keep the same with the old code
diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py
index f8a40fe4901e..0b435574b208 100644
--- a/python/sglang/srt/managers/tokenizer_manager.py
+++ b/python/sglang/srt/managers/tokenizer_manager.py
@@ -182,6 +182,12 @@ class ReqState:
     last_completion_tokens: int = 1
     ttft_observed: bool = False
 
+    # An abort matched this rid while the request was still tokenizer-held
+    # (parked at the pause gate / model-update lock / tokenization). The
+    # scheduler never saw the rid, so the dispatch path must resolve the
+    # request as aborted instead of sending it.
+    abort_before_dispatch: bool = False
+
     # For streaming output
     last_output_offset: int = 0
 
@@ -1364,10 +1370,35 @@ def _should_use_batch_tokenization(self, batch_size, requests) -> bool:
             )
         )
 
+    def _abort_instead_of_dispatch(
+        self,
+        tokenized_obj: Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput],
+    ) -> bool:
+        """Resolve a request whose rid an abort matched while it was still
+        tokenizer-held (see abort_request): the scheduler never saw the rid, so
+        it must be finished here instead of dispatched. Returns True if the
+        request was aborted."""
+        state = self.rid_to_state.get(tokenized_obj.rid)
+        if (
+            state is None
+            or not state.abort_before_dispatch
+            or is_health_check_generate_req(tokenized_obj)
+        ):
+            return False
+        self._handle_abort_req(
+            AbortReq(
+                rid=tokenized_obj.rid,
+                abort_message="Aborted by AbortReq before dispatch to scheduler",
+            )
+        )
+        return True
+
     def _send_one_request(
         self,
         tokenized_obj: Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput],
     ):
+        if self._abort_instead_of_dispatch(tokenized_obj):
+            return
         tokenized_obj.time_stats.set_api_server_dispatch_time()
         tokenized_obj = wrap_shm_features(tokenized_obj)
         time_stats = tokenized_obj.time_stats
@@ -1383,6 +1414,13 @@ def _send_batch_request(
         ],
     ):
         """Send a batch of tokenized requests as a single batched request to the scheduler."""
+        tokenized_objs = [
+            tokenized_obj
+            for tokenized_obj in tokenized_objs
+            if not self._abort_instead_of_dispatch(tokenized_obj)
+        ]
+        if not tokenized_objs:
+            return
         set_time_batch(tokenized_objs, "set_api_server_dispatch_time")
         time_stats = [tokenized_obj.time_stats for tokenized_obj in tokenized_objs]
         for tokenized_obj in tokenized_objs:
@@ -1708,18 +1746,32 @@ async def _handle_batch_request(
                     except StopAsyncIteration:
                         pass
 
-    def abort_request(self, rid: str = "", abort_all: bool = False):
+    def abort_request(
+        self, rid: str = "", abort_all: bool = False, prefix: bool = False
+    ):
         # Empty rid would startswith-match every request on the scheduler.
         if not abort_all and not rid:
             logger.warning("Ignore abort_request with empty rid and abort_all=False")
             return
-        if (
-            not abort_all
-            and self.server_args.tokenizer_worker_num == 1
-            and rid not in self.rid_to_state
-        ):
-            return
-        req = AbortReq(rid=rid, abort_all=abort_all)
+        if not abort_all and self.server_args.tokenizer_worker_num == 1:
+            if prefix:
+                if not any(r.startswith(rid) for r in self.rid_to_state):
+                    return
+            elif rid not in self.rid_to_state:
+                return
+        # A matching request can still be tokenizer-held: rid_to_state is
+        # populated in _init_req_state, several awaits (pause gate,
+        # model_update_lock, tokenization) before the scheduler learns the rid
+        # in _send_one_request. The scheduler-side abort cannot match those, so
+        # flag them here; the dispatch path resolves flagged requests as
+        # aborted instead of sending them. Already-dispatched requests are
+        # unaffected (the flag is only read at dispatch).
+        for tracked_rid, state in self.rid_to_state.items():
+            if abort_all or (
+                tracked_rid.startswith(rid) if prefix else tracked_rid == rid
+            ):
+                state.abort_before_dispatch = True
+        req = AbortReq(rid=rid, abort_all=abort_all, prefix=prefix)
         self._dispatch_to_scheduler(req)
         if self.enable_metrics:
             # TODO: also use custom_labels from the request
diff --git a/test/registered/unit/managers/test_abort_request_prefix.py b/test/registered/unit/managers/test_abort_request_prefix.py
new file mode 100644
index 000000000000..5734193e845e
--- /dev/null
+++ b/test/registered/unit/managers/test_abort_request_prefix.py
@@ -0,0 +1,440 @@
+"""Unit tests for abort-by-rid-prefix.
+
+Tokenizer side: with prefix=True, the rid is treated as a prefix — the
+early-return gate matches any tracked rid starting with it, matching
+tokenizer-held (not yet dispatched) requests are resolved locally, and the
+AbortReq is forwarded with prefix=True.
+
+Scheduler side: matching is prefix-based (``rid.startswith``) regardless of
+the flag, because batch requests derive child rids as ``f"{rid}_{i}"``.
+"""
+
+import asyncio
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock, Mock, patch
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
+
+maybe_stub_sgl_kernel()
+
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.managers.io_struct import AbortReq
+from sglang.srt.managers.schedule_batch import FINISH_ABORT
+from sglang.srt.managers.scheduler import Scheduler
+from sglang.srt.managers.tokenizer_manager import ReqState, TokenizerManager
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+def _make_tokenizer_manager(rids=(), tokenizer_worker_num=1) -> TokenizerManager:
+    """Create a TokenizerManager with mocked dependencies, bypassing __init__."""
+    tm = TokenizerManager.__new__(TokenizerManager)
+    tm.server_args = MagicMock()
+    tm.server_args.tokenizer_worker_num = tokenizer_worker_num
+    tm.enable_metrics = False
+    tm.rid_to_state = {rid: Mock() for rid in rids}
+    tm.send_to_scheduler = MagicMock()
+    tm.tokenizer_ipc_name = None
+    # The IPC boundary: sock_send's wire format varies (pickle/msgpack), so
+    # tests observe dispatched objects here instead of on the zmq socket.
+    tm._dispatch_to_scheduler = MagicMock()
+    return tm
+
+
+def _sent_req(tm) -> AbortReq:
+    tm._dispatch_to_scheduler.assert_called_once()
+    return tm._dispatch_to_scheduler.call_args.args[0]
+
+
+class TestAbortRequestPrefix(CustomTestCase):
+    def test_prefix_match_sends_abort(self):
+        tm = _make_tokenizer_manager(rids=["job-1-seq-0", "job-1-seq-1", "other"])
+        tm.abort_request(rid="job-1", prefix=True)
+
+        req = _sent_req(tm)
+        self.assertEqual(req.rid, "job-1")
+        self.assertTrue(req.prefix)
+        self.assertFalse(req.abort_all)
+
+    def test_prefix_without_match_is_ignored(self):
+        tm = _make_tokenizer_manager(rids=["other-1", "other-2"])
+        tm.abort_request(rid="job-1", prefix=True)
+
+        tm._dispatch_to_scheduler.assert_not_called()
+
+    def test_prefix_requires_full_prefix_not_substring(self):
+        tm = _make_tokenizer_manager(rids=["seq-job-1"])
+        tm.abort_request(rid="job-1", prefix=True)
+
+        tm._dispatch_to_scheduler.assert_not_called()
+
+    def test_exact_match_still_works_without_prefix(self):
+        tm = _make_tokenizer_manager(rids=["job-1"])
+        tm.abort_request(rid="job-1")
+
+        req = _sent_req(tm)
+        self.assertEqual(req.rid, "job-1")
+        self.assertFalse(req.prefix)
+
+    def test_exact_mode_does_not_prefix_match(self):
+        # rid is only a prefix of a tracked request; exact mode must ignore it.
+        tm = _make_tokenizer_manager(rids=["job-1-seq-0"])
+        tm.abort_request(rid="job-1")
+
+        tm._dispatch_to_scheduler.assert_not_called()
+
+    def test_empty_rid_is_ignored(self):
+        # An empty rid would prefix-match every request on the scheduler.
+        tm = _make_tokenizer_manager(rids=["job-1"])
+        tm.abort_request(rid="", prefix=True)
+
+        tm._dispatch_to_scheduler.assert_not_called()
+
+    def test_multi_tokenizer_worker_skips_local_check(self):
+        # With >1 tokenizer workers, rid_to_state is not authoritative; the
+        # abort must be forwarded even if this worker tracks no matching rid.
+        tm = _make_tokenizer_manager(rids=[], tokenizer_worker_num=2)
+        tm.abort_request(rid="job-1", prefix=True)
+
+        req = _sent_req(tm)
+        self.assertEqual(req.rid, "job-1")
+        self.assertTrue(req.prefix)
+
+
+def _make_state(rid: str) -> ReqState:
+    obj = SimpleNamespace(rid=rid, stream=False, return_logprob=False)
+    return ReqState([], False, asyncio.Event(), obj, MagicMock())
+
+
+def _make_tokenizer_manager_with_states(rids) -> TokenizerManager:
+    tm = _make_tokenizer_manager()
+    tm.rid_to_state = {rid: _make_state(rid) for rid in rids}
+    return tm
+
+
+class TestAbortTokenizerHeldRequests(CustomTestCase):
+    """An abort must also cover requests admitted to the tokenizer (rid_to_state
+    populated in _init_req_state) but not yet dispatched to the scheduler —
+    e.g. parked at the pause gate during a weight update. The scheduler cannot
+    match those rids, so abort_request flags them and the dispatch path
+    resolves them as aborted instead of sending them."""
+
+    def test_prefix_abort_flags_matching_states(self):
+        tm = _make_tokenizer_manager_with_states(
+            ["job-1-seq-0", "job-1-seq-1", "other"]
+        )
+        tm.abort_request(rid="job-1", prefix=True)
+
+        self.assertTrue(tm.rid_to_state["job-1-seq-0"].abort_before_dispatch)
+        self.assertTrue(tm.rid_to_state["job-1-seq-1"].abort_before_dispatch)
+        self.assertFalse(tm.rid_to_state["other"].abort_before_dispatch)
+        # The AbortReq is still forwarded for already-dispatched requests.
+        self.assertEqual(_sent_req(tm).rid, "job-1")
+
+    def test_exact_abort_flags_only_exact_state(self):
+        tm = _make_tokenizer_manager_with_states(["job-1", "job-1-seq-0"])
+        tm.abort_request(rid="job-1")
+
+        self.assertTrue(tm.rid_to_state["job-1"].abort_before_dispatch)
+        self.assertFalse(tm.rid_to_state["job-1-seq-0"].abort_before_dispatch)
+
+    def test_abort_all_flags_every_state(self):
+        tm = _make_tokenizer_manager_with_states(["a", "b"])
+        tm.abort_request(abort_all=True)
+
+        self.assertTrue(tm.rid_to_state["a"].abort_before_dispatch)
+        self.assertTrue(tm.rid_to_state["b"].abort_before_dispatch)
+
+    def test_dispatch_resolves_flagged_request_locally(self):
+        tm = _make_tokenizer_manager_with_states(["job-1-seq-0"])
+        tm.server_args.weight_version = "v0"
+        state = tm.rid_to_state["job-1-seq-0"]
+        tm.abort_request(rid="job-1", prefix=True)
+        tm._dispatch_to_scheduler.reset_mock()
+
+        tm._send_one_request(SimpleNamespace(rid="job-1-seq-0"))
+
+        # Never dispatched; resolved as aborted so _wait_one_response returns.
+        tm._dispatch_to_scheduler.assert_not_called()
+        self.assertTrue(state.finished)
+        self.assertTrue(state.event.is_set())
+        self.assertNotIn("job-1-seq-0", tm.rid_to_state)
+        finish_reason = state.out_list[-1]["meta_info"]["finish_reason"]
+        self.assertEqual(finish_reason["type"], "abort")
+
+    def test_dispatch_sends_unflagged_request(self):
+        tm = _make_tokenizer_manager_with_states(["job-1", "other"])
+        tm.abort_request(rid="job-1")
+        tm._dispatch_to_scheduler.reset_mock()
+
+        tokenized_obj = MagicMock()
+        tokenized_obj.rid = "other"
+        with unittest.mock.patch(
+            "sglang.srt.managers.tokenizer_manager.wrap_shm_features",
+            side_effect=lambda obj: obj,
+        ):
+            tm._send_one_request(tokenized_obj)
+
+        tm._dispatch_to_scheduler.assert_called_once_with(tokenized_obj)
+        self.assertIn("other", tm.rid_to_state)
+
+    def test_batch_dispatch_filters_flagged_requests(self):
+        tm = _make_tokenizer_manager_with_states(["job-1-seq-0", "job-1-seq-1"])
+        tm.server_args.weight_version = "v0"
+        tm.abort_request(rid="job-1", prefix=True)
+        tm._dispatch_to_scheduler.reset_mock()
+
+        tm._send_batch_request(
+            [SimpleNamespace(rid="job-1-seq-0"), SimpleNamespace(rid="job-1-seq-1")]
+        )
+
+        tm._dispatch_to_scheduler.assert_not_called()
+        self.assertEqual(tm.rid_to_state, {})
+
+
+class FakeReq:
+    def __init__(self, rid: str):
+        self.rid = rid
+        self.mamba_pool_idx = None
+        self.to_finish = None
+
+    def finished(self) -> bool:
+        return False
+
+
+def _make_scheduler(waiting_rids=(), running_rids=(), chunked_rid=None):
+    sched = SimpleNamespace()
+    sched.chunked_req = FakeReq(chunked_rid) if chunked_rid is not None else None
+    sched.waiting_queue = [FakeReq(rid) for rid in waiting_rids]
+    sched.enable_hicache_storage = False
+    sched.disaggregation_mode = DisaggregationMode.NULL
+    sched.grammar_manager = MagicMock()
+    sched.running_batch = SimpleNamespace(reqs=[FakeReq(rid) for rid in running_rids])
+    sched.cur_batch = None
+    sched.ipc_channels = MagicMock()
+    return sched
+
+
+class TestSchedulerAbortMatching(CustomTestCase):
+    """Scheduler-side matching semantics for AbortReq (see io_struct.AbortReq:
+    always ``rid.startswith``, so batch children ``f"{rid}_{i}"`` are covered)."""
+
+    def test_prefix_abort_isolates_namespaces(self):
+        sched = _make_scheduler(
+            waiting_rids=["A::1", "A::2", "B::1"],
+            running_rids=["A::3", "B::2"],
+        )
+        Scheduler.abort_request(sched, AbortReq(rid="A::", prefix=True))
+
+        self.assertEqual([req.rid for req in sched.waiting_queue], ["B::1"])
+        running = {req.rid: req for req in sched.running_batch.reqs}
+        self.assertIsInstance(running["A::3"].to_finish, FINISH_ABORT)
+        self.assertIsNone(running["B::2"].to_finish)
+        # Every waiting-queue abort echoes back to the tokenizer for cleanup.
+        aborted_rids = {
+            call.args[0].rid
+            for call in sched.ipc_channels.send_to_tokenizer.send_output.call_args_list
+        }
+        self.assertEqual(aborted_rids, {"A::1", "A::2"})
+
+    def test_matching_is_prefix_based_even_without_prefix_flag(self):
+        # Pre-existing scheduler semantics: batch requests derive child rids as
+        # f"{rid}_{i}", so an exact-mode abort for the parent must cover them.
+        sched = _make_scheduler(waiting_rids=["job-1_0", "job-1_1", "job-2_0"])
+        Scheduler.abort_request(sched, AbortReq(rid="job-1", prefix=False))
+
+        self.assertEqual([req.rid for req in sched.waiting_queue], ["job-2_0"])
+
+    def test_chunked_request_is_prefix_matched(self):
+        sched = _make_scheduler(chunked_rid="A::9")
+        Scheduler.abort_request(sched, AbortReq(rid="A::", prefix=True))
+
+        self.assertIs(sched._pending_chunked_abort_req, sched.chunked_req)
+
+    def test_abort_all_covers_everything(self):
+        sched = _make_scheduler(waiting_rids=["A::1"], running_rids=["B::1"])
+        Scheduler.abort_request(sched, AbortReq(abort_all=True))
+
+        self.assertEqual(sched.waiting_queue, [])
+        self.assertIsInstance(sched.running_batch.reqs[0].to_finish, FINISH_ABORT)
+
+
+def _make_disagg_req(rid: str, pending_bootstrap: bool = False) -> FakeReq:
+    req = FakeReq(rid)
+    req.pending_bootstrap = pending_bootstrap
+    req.disagg_kv_sender = Mock()
+    return req
+
+
+def _make_prefill_scheduler(waiting_rids=(), bootstrap_rids=(), inflight_rids=()):
+    sched = _make_scheduler()
+    sched.disaggregation_mode = DisaggregationMode.PREFILL
+    sched.waiting_queue = [
+        _make_disagg_req(rid, pending_bootstrap=True) for rid in waiting_rids
+    ]
+    sched.req_to_metadata_buffer_idx_allocator = MagicMock()
+    sched.disagg_prefill_bootstrap_queue = SimpleNamespace(
+        queue=[_make_disagg_req(rid) for rid in bootstrap_rids]
+    )
+    sched.disagg_prefill_inflight_queue = [
+        _make_disagg_req(rid) for rid in inflight_rids
+    ]
+    return sched
+
+
+def _make_decode_req(rid: str) -> SimpleNamespace:
+    return SimpleNamespace(req=FakeReq(rid), kv_receiver=Mock())
+
+
+def _make_retracted_req(rid: str) -> SimpleNamespace:
+    return SimpleNamespace(rid=rid, kv_cache_cpu=object())
+
+
+def _make_decode_scheduler(
+    waiting_rids=(), prealloc_rids=(), transfer_rids=(), retracted_rids=()
+):
+    sched = _make_scheduler(waiting_rids=waiting_rids)
+    sched.disaggregation_mode = DisaggregationMode.DECODE
+    sched.tree_cache = MagicMock()
+    sched.disagg_decode_prealloc_queue = SimpleNamespace(
+        queue=[_make_decode_req(rid) for rid in prealloc_rids],
+        retracted_queue=[_make_retracted_req(rid) for rid in retracted_rids],
+    )
+    sched.disagg_decode_transfer_queue = SimpleNamespace(
+        queue=[_make_decode_req(rid) for rid in transfer_rids]
+    )
+    return sched
+
+
+def _echoed_rids(sched) -> set:
+    return {
+        call.args[0].rid
+        for call in sched.ipc_channels.send_to_tokenizer.send_output.call_args_list
+    }
+
+
+class TestSchedulerDisaggPrefillAbort(CustomTestCase):
+    """PREFILL-side disaggregation abort matching: the bootstrap and in-flight
+    queues hold requests the waiting queue no longer tracks, and the waiting
+    queue itself must release the metadata buffer slot and abort a
+    still-bootstrapping KV sender."""
+
+    def test_bootstrap_and_inflight_queues_prefix_matched(self):
+        sched = _make_prefill_scheduler(
+            bootstrap_rids=["A::1", "B::1"], inflight_rids=["A::2", "B::2"]
+        )
+        Scheduler.abort_request(sched, AbortReq(rid="A::", prefix=True))
+
+        bootstrap = {r.rid: r for r in sched.disagg_prefill_bootstrap_queue.queue}
+        bootstrap["A::1"].disagg_kv_sender.abort.assert_called_once()
+        bootstrap["B::1"].disagg_kv_sender.abort.assert_not_called()
+        inflight = {r.rid: r for r in sched.disagg_prefill_inflight_queue}
+        inflight["A::2"].disagg_kv_sender.abort.assert_called_once()
+        inflight["B::2"].disagg_kv_sender.abort.assert_not_called()
+
+    def test_waiting_queue_releases_metadata_buffer(self):
+        sched = _make_prefill_scheduler(waiting_rids=["A::1", "B::1"])
+        with patch(
+            "sglang.srt.managers.scheduler.maybe_release_metadata_buffer"
+        ) as release:
+            Scheduler.abort_request(sched, AbortReq(rid="A::", prefix=True))
+
+        self.assertEqual([r.rid for r in sched.waiting_queue], ["B::1"])
+        release.assert_called_once()
+        released_req, allocator = release.call_args.args
+        self.assertEqual(released_req.rid, "A::1")
+        self.assertIs(allocator, sched.req_to_metadata_buffer_idx_allocator)
+        self.assertEqual(_echoed_rids(sched), {"A::1"})
+
+    def test_waiting_queue_pending_bootstrap_gates_sender_abort(self):
+        pending = _make_disagg_req("A::1", pending_bootstrap=True)
+        bootstrapped = _make_disagg_req("A::2", pending_bootstrap=False)
+        sched = _make_prefill_scheduler()
+        sched.waiting_queue = [pending, bootstrapped]
+        with patch("sglang.srt.managers.scheduler.maybe_release_metadata_buffer"):
+            Scheduler.abort_request(sched, AbortReq(rid="A::", prefix=True))
+
+        pending.disagg_kv_sender.abort.assert_called_once()
+        bootstrapped.disagg_kv_sender.abort.assert_not_called()
+
+    def test_abort_all_covers_prefill_queues(self):
+        sched = _make_prefill_scheduler(bootstrap_rids=["A::1"], inflight_rids=["B::1"])
+        Scheduler.abort_request(sched, AbortReq(abort_all=True))
+
+        sched.disagg_prefill_bootstrap_queue.queue[
+            0
+        ].disagg_kv_sender.abort.assert_called_once()
+        sched.disagg_prefill_inflight_queue[
+            0
+        ].disagg_kv_sender.abort.assert_called_once()
+
+
+class TestSchedulerDisaggDecodeAbort(CustomTestCase):
+    """DECODE-side disaggregation abort matching: prealloc/transfer queues
+    abort their KV receivers, the retracted queue frees CPU KV cache and
+    echoes the abort back to the tokenizer, and waiting-queue requests
+    release their preallocated KV cache."""
+
+    def test_prealloc_and_transfer_queues_prefix_matched(self):
+        sched = _make_decode_scheduler(
+            prealloc_rids=["A::1", "B::1"], transfer_rids=["A::2", "B::2"]
+        )
+        Scheduler.abort_request(sched, AbortReq(rid="A::", prefix=True))
+
+        prealloc = {d.req.rid: d for d in sched.disagg_decode_prealloc_queue.queue}
+        prealloc["A::1"].kv_receiver.abort.assert_called_once()
+        prealloc["B::1"].kv_receiver.abort.assert_not_called()
+        transfer = {d.req.rid: d for d in sched.disagg_decode_transfer_queue.queue}
+        transfer["A::2"].kv_receiver.abort.assert_called_once()
+        transfer["B::2"].kv_receiver.abort.assert_not_called()
+
+    def test_retracted_queue_frees_cpu_cache_and_echoes(self):
+        sched = _make_decode_scheduler(retracted_rids=["A::1", "B::1"])
+        aborted = sched.disagg_decode_prealloc_queue.retracted_queue[0]
+        Scheduler.abort_request(sched, AbortReq(rid="A::", prefix=True))
+
+        self.assertEqual(
+            [d.rid for d in sched.disagg_decode_prealloc_queue.retracted_queue],
+            ["B::1"],
+        )
+        self.assertFalse(hasattr(aborted, "kv_cache_cpu"))
+        self.assertEqual(_echoed_rids(sched), {"A::1"})
+
+    def test_waiting_queue_releases_kv_cache(self):
+        sched = _make_decode_scheduler(waiting_rids=["A::1", "B::1"])
+        with patch("sglang.srt.managers.scheduler.release_kv_cache") as release:
+            Scheduler.abort_request(sched, AbortReq(rid="A::", prefix=True))
+
+        self.assertEqual([r.rid for r in sched.waiting_queue], ["B::1"])
+        release.assert_called_once()
+        self.assertEqual(release.call_args.args[0].rid, "A::1")
+
+    def test_exact_mode_is_still_prefix_matched_in_disagg_queues(self):
+        # Same load-bearing semantics as the non-disagg queues: batch children
+        # derive rids as f"{rid}_{i}", so exact-mode must cover them here too.
+        sched = _make_decode_scheduler(prealloc_rids=["job-1_0"])
+        Scheduler.abort_request(sched, AbortReq(rid="job-1", prefix=False))
+
+        sched.disagg_decode_prealloc_queue.queue[
+            0
+        ].kv_receiver.abort.assert_called_once()
+
+    def test_abort_all_covers_decode_queues(self):
+        sched = _make_decode_scheduler(
+            prealloc_rids=["A::1"], transfer_rids=["B::1"], retracted_rids=["C::1"]
+        )
+        Scheduler.abort_request(sched, AbortReq(abort_all=True))
+
+        sched.disagg_decode_prealloc_queue.queue[
+            0
+        ].kv_receiver.abort.assert_called_once()
+        sched.disagg_decode_transfer_queue.queue[
+            0
+        ].kv_receiver.abort.assert_called_once()
+        self.assertEqual(sched.disagg_decode_prealloc_queue.retracted_queue, [])
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

From d0c484a477d81c4d2379f2bcaee35e8de9256efd Mon Sep 17 00:00:00 2001
From: Zhichen Zeng <zczeng@uw.edu>
Date: Fri, 24 Jul 2026 23:04:47 -0700
Subject: [PATCH 27/43] [22/27] [sglang-miles] Fix flush_cache() no-op after
 pause_generation in retract (#31962)

---
 python/sglang/srt/managers/scheduler.py       |   7 +-
 .../scheduler_components/weight_updater.py    |   6 +-
 ...est_scheduler_flush_cache_after_retract.py | 109 ++++++++++++++++++
 3 files changed, 116 insertions(+), 6 deletions(-)
 create mode 100644 test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py

diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
index 92fbe2c2769f..92acb3a31827 100644
--- a/python/sglang/srt/managers/scheduler.py
+++ b/python/sglang/srt/managers/scheduler.py
@@ -3706,7 +3706,7 @@ def on_idle(self):
         # sleep until next event
         self.maybe_sleep_on_idle()
 
-    def is_fully_idle(self, for_health_check=False) -> bool:
+    def is_fully_idle(self, for_health_check=False, ignore_waiting=False) -> bool:
         # Health check piggybacks on running requests in process_output.
         # Only running_batch + waiting_queue guarantee active GPU processing;
         # disagg queues (bootstrap/prealloc/transfer) may have items without
@@ -3723,7 +3723,8 @@ def is_fully_idle(self, for_health_check=False) -> bool:
         )
 
         # Waiting queues: waiting + bootstrapping + preallocation + kv transfer (decode)
-        idle &= len(self.waiting_queue) == 0
+        if not ignore_waiting:
+            idle &= len(self.waiting_queue) == 0
 
         if not for_health_check:
             # Grammar queue and prefill inflight queue may not produce batch
@@ -3868,7 +3869,7 @@ def detach_hicache_storage_wrapped(
 
     def flush_cache(self, empty_cache: bool = True):
         """Flush memory pools (e.g., KV cache, Mamba cache) and optionally empty device allocator cache."""
-        if self.is_fully_idle():
+        if self.is_fully_idle(ignore_waiting=self._engine_paused):
             self.cur_batch_for_debug = None
             self.last_batch = None
             self.tree_cache.reset()
diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py
index c68699b827ef..150aa18d3b34 100644
--- a/python/sglang/srt/managers/scheduler_components/weight_updater.py
+++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py
@@ -325,8 +325,9 @@ def end_weight_update(self, recv_req: EndWeightUpdateReqInput):
         return EndWeightUpdateReqOutput(success=True, message="Success")
 
     def release_memory_occupation(self, recv_req: ReleaseMemoryOccupationReqInput):
-        assert (
-            self.is_fully_idle()
+        scheduler = self.scheduler
+        assert self.is_fully_idle(
+            ignore_waiting=scheduler is not None and scheduler._engine_paused
         ), "release_memory_occupation should be called only when server is idle."
 
         tags = recv_req.tags
@@ -338,7 +339,6 @@ def release_memory_occupation(self, recv_req: ReleaseMemoryOccupationReqInput):
             self.offload_tags.add(tag)
 
         if GPU_MEMORY_TYPE_KV_CACHE in tags:
-            scheduler = self.scheduler
             if scheduler is not None:
                 if scheduler.disaggregation_mode == DisaggregationMode.DECODE:
                     for queue_name in (
diff --git a/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py b/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py
new file mode 100644
index 000000000000..ff3f771eb9da
--- /dev/null
+++ b/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py
@@ -0,0 +1,109 @@
+import unittest
+from collections import deque
+from unittest.mock import MagicMock
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import maybe_stub_sgl_kernel
+
+maybe_stub_sgl_kernel()
+
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.managers.scheduler import Scheduler
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+register_cpu_ci(est_time=6, suite="base-c-test-cpu")
+
+
+def _make_req(is_retracted: bool) -> MagicMock:
+    req = MagicMock()
+    req.is_retracted = is_retracted
+    return req
+
+
+class TestSchedulerFlushCacheAfterRetract(unittest.TestCase):
+    """Regression coverage for flush_cache() no-oping while the engine is paused.
+
+    With generation paused and running_batch drained, nothing in waiting_queue
+    holds KV/pool state: neither requests pause_generation(mode="retract") just
+    moved there, nor requests that were never scheduled at all. The latter are
+    always present under high concurrency (in-flight targets exceed engine
+    parallelism), where gating the flush on them starves weight updates until
+    the caller's retry budget times out (reproduced in RL training at 512
+    in-flight requests). See RETRACT_MODE_FLUSH_CACHE_FIX.md in miles."""
+
+    def _new_scheduler(self) -> Scheduler:
+        scheduler = Scheduler.__new__(Scheduler)
+        scheduler._engine_paused = True
+        scheduler.enable_overlap = False
+        scheduler.last_batch = None
+        scheduler.cur_batch = None
+        scheduler.chunked_req = None
+        scheduler.disaggregation_mode = DisaggregationMode.NULL
+        scheduler.running_batch = MagicMock()
+        scheduler.running_batch.is_empty.return_value = True
+        scheduler.result_queue = deque()
+        scheduler.waiting_queue = []
+        scheduler.dllm_manager = MagicMock()
+        scheduler.dllm_manager.any_staging_reqs.return_value = False
+        scheduler._pp_microbatches_drained = MagicMock(return_value=True)
+        scheduler.grammar_manager = MagicMock()
+        scheduler.grammar_manager.grammar_queue = []
+        scheduler.enable_hierarchical_cache = False
+        scheduler.enable_hisparse = False
+        scheduler.tree_cache = MagicMock()
+        scheduler.req_to_token_pool = MagicMock()
+        scheduler.token_to_kv_pool_allocator = MagicMock()
+        scheduler.draft_worker = None
+        scheduler.metrics_reporter = MagicMock()
+        return scheduler
+
+    def test_flush_cache_succeeds_when_only_retracted_reqs_are_waiting(self):
+        scheduler = self._new_scheduler()
+        scheduler.waiting_queue = [_make_req(is_retracted=True)]
+
+        self.assertTrue(scheduler.flush_cache())
+        scheduler.tree_cache.reset.assert_called_once()
+        scheduler.req_to_token_pool.clear.assert_called_once()
+        scheduler.token_to_kv_pool_allocator.clear.assert_called_once()
+
+    def test_flush_cache_succeeds_with_never_scheduled_reqs_while_paused(self):
+        """A request that was queued but never scheduled holds no KV/pool state
+        either; a paused-engine flush must not starve on it."""
+        scheduler = self._new_scheduler()
+        scheduler.waiting_queue = [_make_req(is_retracted=False)]
+
+        self.assertTrue(scheduler.flush_cache())
+        scheduler.tree_cache.reset.assert_called_once()
+
+    def test_flush_cache_succeeds_on_mixed_queue_while_paused(self):
+        scheduler = self._new_scheduler()
+        scheduler.waiting_queue = [
+            _make_req(is_retracted=True),
+            _make_req(is_retracted=False),
+        ]
+
+        self.assertTrue(scheduler.flush_cache())
+        scheduler.tree_cache.reset.assert_called_once()
+
+    def test_is_fully_idle_default_ignores_nothing(self):
+        """Regression guard: every other is_fully_idle() caller passes no args,
+        so a non-empty waiting_queue must still report not-idle by default."""
+        scheduler = self._new_scheduler()
+        scheduler.waiting_queue = [_make_req(is_retracted=True)]
+
+        self.assertFalse(scheduler.is_fully_idle())
+        self.assertTrue(scheduler.is_fully_idle(ignore_waiting=True))
+
+    def test_flush_cache_does_not_ignore_waiting_reqs_while_unpaused(self):
+        """ignore_waiting is gated on _engine_paused inside flush_cache — a
+        flush issued without pausing first must not be silently permissive."""
+        scheduler = self._new_scheduler()
+        scheduler._engine_paused = False
+        scheduler.waiting_queue = [_make_req(is_retracted=True)]
+
+        self.assertFalse(scheduler.flush_cache())
+        scheduler.tree_cache.reset.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()

From dd2c725374221112dbd80e7ad87dee3bf0d1466b Mon Sep 17 00:00:00 2001
From: Yueming Yuan <yym022502@gmail.com>
Date: Sat, 25 Jul 2026 03:03:08 -0700
Subject: [PATCH 28/43] [23/27] [sglang-miles] check_weights: make the wire
 ChecksumInfo match the merged per-role payload

`_merge_checksum_payloads` folds each runner's payload into one, tagging every
entry with its role, so `parallelism_info` is a LIST of per-role entries. That
shape was never validated on v0.5.15 -- v0.5.16 added the msgspec wire structs
plus `msgspec.convert(p, ChecksumInfo)` in `check_weights`, and the wire struct
still declared a single `ParallelismInfo` with no `role`, so every
check_weights('checksum') call would have raised on conversion.

This is the shape miles already consumes: `checksum_utils._gpu_rank` iterates
`parallelism_info` and reads `role_info["rank"]` per role, and its tests cover
the target+draft case sharing one GPU rank.
---
 python/sglang/srt/managers/io_struct.py       |  8 ++-
 .../managers/test_msgpack_ipc_roundtrip.py    | 72 +++++++++++++------
 2 files changed, 56 insertions(+), 24 deletions(-)

diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py
index de60b2068882..19823719a90c 100644
--- a/python/sglang/srt/managers/io_struct.py
+++ b/python/sglang/srt/managers/io_struct.py
@@ -1813,6 +1813,9 @@ class CheckWeightsReqInput(BaseReq, kw_only=True):
 # sglang.srt.utils.weight_checker. Not array_like: the payload is read by field
 # name and re-serialized to JSON, so it must stay a {field: value} map.
 class ParallelismInfo(msgspec.Struct, kw_only=True):
+    # Which runner this describes: "target", or a draft role such as "draft" /
+    # "draft_step_0". One entry per runner the checksum covers.
+    role: str
     tp_rank: int
     tp_size: int
     dp_rank: int
@@ -1826,7 +1829,10 @@ class ParallelismInfo(msgspec.Struct, kw_only=True):
 class ChecksumInfo(msgspec.Struct, kw_only=True):
     checksums: Dict[str, str]
     per_gpu_checksum: str
-    parallelism_info: ParallelismInfo
+    # One entry per role the checksum covers: the target model plus, under
+    # speculative decoding, each draft runner. All roles on a rank share the GPU
+    # rank; consumers key off it to merge the shards.
+    parallelism_info: List[ParallelismInfo]
 
 
 class CheckWeightsReqOutput(BaseReq, kw_only=True):
diff --git a/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py b/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py
index e3d3e95d067b..d8c2a4eb524a 100644
--- a/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py
+++ b/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py
@@ -63,9 +63,17 @@ def _contains_dataclass(obj) -> bool:
     return False
 
 
-def _parallelism_info() -> ParallelismInfo:
+def _parallelism_info(role: str = "target") -> ParallelismInfo:
     return ParallelismInfo(
-        tp_rank=0, tp_size=2, dp_rank=0, dp_size=1, pp_rank=0, pp_size=1, rank=0, size=2
+        role=role,
+        tp_rank=0,
+        tp_size=2,
+        dp_rank=0,
+        dp_size=1,
+        pp_rank=0,
+        pp_size=1,
+        rank=0,
+        size=2,
     )
 
 
@@ -73,7 +81,8 @@ def _checksum_info(tag: str) -> ChecksumInfo:
     return ChecksumInfo(
         checksums={f"model.layers.{tag}": "deadbeef"},
         per_gpu_checksum="cafef00d",
-        parallelism_info=_parallelism_info(),
+        # One entry per role: the target model plus each speculative draft runner.
+        parallelism_info=[_parallelism_info("target"), _parallelism_info("draft")],
     )
 
 
@@ -192,29 +201,46 @@ def test_check_weights_multi_rank_payload(self):
         self.assertEqual(len(decoded.payload), 2)
         as_dict = msgspec_to_builtins(decoded.payload[0])
         self.assertEqual(as_dict["per_gpu_checksum"], "cafef00d")
-        self.assertIn("tp_rank", as_dict["parallelism_info"])
+        self.assertEqual(
+            [pi["role"] for pi in as_dict["parallelism_info"]], ["target", "draft"]
+        )
+        self.assertIn("tp_rank", as_dict["parallelism_info"][0])
 
     def test_check_weights_producer_conversion(self):
-        # Mirrors weight_updater.check_weights: WeightChecker returns
-        # ChecksumInfo.model_dump() (a dict), converted to the msgspec struct via
-        # msgspec.convert, and the result round-trips as the payload.
-        pydantic_checksum = PydanticChecksumInfo(
-            checksums={"model.layers.0": "deadbeef"},
-            per_gpu_checksum="cafef00d",
-            parallelism_info=PydanticParallelismInfo(
-                tp_rank=0,
-                tp_size=2,
-                dp_rank=0,
-                dp_size=1,
-                pp_rank=0,
-                pp_size=1,
-                rank=0,
-                size=2,
-            ),
+        # Mirrors weight_updater.check_weights end to end: each runner's
+        # WeightChecker returns a single-role ChecksumInfo.model_dump(),
+        # _merge_checksum_payloads folds the roles into one per-role payload, and
+        # that is what msgspec.convert has to accept.
+        from sglang.srt.managers.scheduler_components.weight_updater import (
+            _merge_checksum_payloads,
+        )
+
+        def _dump(role_tag):
+            return PydanticChecksumInfo(
+                checksums={f"model.layers.0{role_tag}": "deadbeef"},
+                per_gpu_checksum="cafef00d",
+                parallelism_info=PydanticParallelismInfo(
+                    tp_rank=0,
+                    tp_size=2,
+                    dp_rank=0,
+                    dp_size=1,
+                    pp_rank=0,
+                    pp_size=1,
+                    rank=0,
+                    size=2,
+                ),
+            ).model_dump()
+
+        merged = _merge_checksum_payloads([("", _dump("")), ("draft", _dump(""))])
+        converted = msgspec.convert(merged, ChecksumInfo)
+        # Draft keys are role-prefixed so they never collide with the target's.
+        self.assertEqual(
+            sorted(converted.checksums), ["draft.model.layers.0", "model.layers.0"]
+        )
+        self.assertEqual(
+            [pi.role for pi in converted.parallelism_info], ["target", "draft"]
         )
-        converted = msgspec.convert(pydantic_checksum.model_dump(), ChecksumInfo)
-        self.assertEqual(converted.per_gpu_checksum, "cafef00d")
-        self.assertEqual(converted.parallelism_info.tp_rank, 0)
+        self.assertEqual(converted.parallelism_info[0].tp_rank, 0)
         output = CheckWeightsReqOutput(success=True, message="ok", payload=[converted])
         self.assertEqual(_round_trip(output), output)
 

From 76875851dc9c34123f861c483fe835d2aabf242a Mon Sep 17 00:00:00 2001
From: JD-ETH <jaedon.guo@gmail.com>
Date: Sat, 25 Jul 2026 04:28:14 -0700
Subject: [PATCH 29/43] [24/27] [sglang-miles] Port RankParallelismConfig /
 ParallelismContext to v0.5.16's parallel state

Two follow-ups to [14/24]'s P2P weight-update support, both caused by v0.5.16
retiring the module-level dp-attention accessors. Found by CI: every GPU stage
that initialises the remote-instance transfer engine died at startup with

    remote_instance_weight_transporter.py:56 init_engine
      -> RankParallelismConfig.from_parallel_state
    ImportError: cannot import name 'get_attention_cp_rank'
                 from 'sglang.srt.layers.dp_attention'

1. `from_parallel_state` imported get_attention_{tp,dp,cp}_{rank,size} from
   dp_attention. v0.5.16 removed all six; the attention-side ranks and sizes now
   live on the runtime ParallelState, so read them from `get_parallel()`.

2. `ParallelismContext` set `dp_attention._ENABLE_DP_ATTENTION_FLAG`, which
   v0.5.16 replaced with `get_flags().dp.enabled`. That assignment would not have
   raised -- it would have created a stale attribute nobody reads, leaving
   `is_dp_attention_enabled()` reporting the ambient value instead of the
   context's and silently building a replica with the wrong sharding. The flag is
   now saved, set and restored on the runtime flags instead, and dropped from
   _DA_GLOBALS.
---
 .../sglang/srt/distributed/parallel_state.py  | 41 +++++++++++--------
 1 file changed, 23 insertions(+), 18 deletions(-)

diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py
index 03d0ee89554c..eeccece89475 100644
--- a/python/sglang/srt/distributed/parallel_state.py
+++ b/python/sglang/srt/distributed/parallel_state.py
@@ -2911,15 +2911,12 @@ def from_parallel_state(cls, local_rank: int = 0) -> "RankParallelismConfig":
         tp_size = get_tensor_model_parallel_world_size()
         tp_rank = get_tensor_model_parallel_rank()
 
-        # Import dp_attention lazily to avoid circular imports
-        from sglang.srt.layers.dp_attention import (
-            get_attention_cp_rank,
-            get_attention_cp_size,
-            get_attention_dp_rank,
-            get_attention_dp_size,
-            get_attention_tp_rank,
-            get_attention_tp_size,
-        )
+        # v0.5.16 retired the get_attention_{tp,dp,cp}_* helpers from dp_attention;
+        # the attention-side ranks/sizes now live on the runtime ParallelState.
+        # Imported lazily to avoid a circular import.
+        from sglang.srt.runtime_context import get_parallel
+
+        ps = get_parallel()
 
         return cls(
             tp_size=tp_size,
@@ -2930,12 +2927,12 @@ def from_parallel_state(cls, local_rank: int = 0) -> "RankParallelismConfig":
             ep_rank=get_moe_expert_parallel_rank(),
             moe_tp_size=get_moe_tensor_parallel_world_size(),
             moe_tp_rank=get_moe_tensor_parallel_rank(),
-            attn_tp_size=get_attention_tp_size(),
-            attn_tp_rank=get_attention_tp_rank(),
-            attn_dp_size=get_attention_dp_size(),
-            attn_dp_rank=get_attention_dp_rank(),
-            attn_cp_size=get_attention_cp_size(),
-            attn_cp_rank=get_attention_cp_rank(),
+            attn_tp_size=ps.attn_tp_size,
+            attn_tp_rank=ps.attn_tp_rank,
+            attn_dp_size=ps.attn_dp_size,
+            attn_dp_rank=ps.attn_dp_rank,
+            attn_cp_size=ps.attn_cp_size,
+            attn_cp_rank=ps.attn_cp_rank,
             moe_dp_size=get_moe_data_parallel_world_size(),
             moe_dp_rank=get_moe_data_parallel_rank(),
             world_size=(
@@ -2954,8 +2951,10 @@ def from_parallel_state(cls, local_rank: int = 0) -> "RankParallelismConfig":
 
 # Globals on parallel_state module to save/restore
 _PS_GLOBALS = ("_TP", "_PP", "_MOE_EP", "_MOE_TP", "_ATTN_TP", "_ATTN_CP", "_MOE_DP")
-# Globals on dp_attention module to save/restore
-_DA_GLOBALS = ("_ATTN_DP_RANK", "_ATTN_DP_SIZE", "_ENABLE_DP_ATTENTION_FLAG")
+# Globals on dp_attention module to save/restore. v0.5.16 moved the dp-attention
+# enable flag out of this module onto the runtime flags (get_flags().dp.enabled),
+# so it is saved/restored separately below rather than as a module attribute.
+_DA_GLOBALS = ("_ATTN_DP_RANK", "_ATTN_DP_SIZE")
 
 
 class ParallelismContext:
@@ -2994,6 +2993,7 @@ def __enter__(self):
 
         from sglang.srt.distributed import parallel_state
         from sglang.srt.layers import dp_attention
+        from sglang.srt.runtime_context import get_flags
 
         # Save original globals
         for name in _PS_GLOBALS:
@@ -3017,7 +3017,10 @@ def __enter__(self):
         # Set dp_attention scalar globals
         dp_attention._ATTN_DP_RANK = conf.attn_dp_rank
         dp_attention._ATTN_DP_SIZE = conf.attn_dp_size
-        dp_attention._ENABLE_DP_ATTENTION_FLAG = conf.attn_dp_size > 1
+        # is_dp_attention_enabled() reads this, not a dp_attention module global.
+        dp_flags = get_flags().dp
+        self._original_dp_enabled = dp_flags.enabled
+        dp_flags.enabled = conf.attn_dp_size > 1
 
         logger.info(f"[ParallelismContext] Activated: {conf}")
         return self
@@ -3025,12 +3028,14 @@ def __enter__(self):
     def __exit__(self, *args):
         from sglang.srt.distributed import parallel_state
         from sglang.srt.layers import dp_attention
+        from sglang.srt.runtime_context import get_flags
 
         # Restore original globals
         for name in _PS_GLOBALS:
             setattr(parallel_state, name, self._original_globals.get(name))
         for name in _DA_GLOBALS:
             setattr(dp_attention, name, self._original_globals.get(name))
+        get_flags().dp.enabled = self._original_dp_enabled
 
         logger.info("[ParallelismContext] Deactivated")
         return False

From 47325f4cf661497d611701221cd627ceb3419894 Mon Sep 17 00:00:00 2001
From: Yueming Yuan <yym022502@gmail.com>
Date: Sat, 25 Jul 2026 04:56:29 -0700
Subject: [PATCH 30/43] [25/27] [sglang-miles] Repoint three symbols v0.5.16
 renamed or moved

Three stale references left by the rebase. The cherry-picks applied cleanly
because the surrounding lines matched, so nothing surfaced until the code ran.

1. model_runner.forward: `decode_cuda_graph_runner.num_tokens_per_bs` ->
   `.captured_req_width`. v0.5.16 renamed the per-request capture width (and
   `ModelRunner.decode_num_tokens_per_bs` -> `decode_num_tokens_per_req`), so
   every decode with a captured graph raised AttributeError.
   From [2/25], R3 DeepEP/MTP support.

2. lora/backend/base_backend: `dp_attention.get_attention_tp_size()` ->
   `get_parallel().attn_tp_size`. Same removal as [24/25]; this one is a lazy
   import so it only fired on the MoE-LoRA DP-attention path.

3. lora/backend/base_backend: `layers.utils.cp_utils` ->
   `layers.cp.padding` for get_cp_padding_align_size. The module moved in
   v0.5.16; every other caller in the tree already uses the new path.

Found 1 and 2 in CI; found 3 by then resolving every sglang symbol imported by
lines this branch adds against the v0.5.16 tree, which is how it should have been
checked in the first place.
---
 python/sglang/srt/lora/backend/base_backend.py   | 11 ++++++-----
 python/sglang/srt/model_executor/model_runner.py | 13 ++++++-------
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/python/sglang/srt/lora/backend/base_backend.py b/python/sglang/srt/lora/backend/base_backend.py
index 45443b5da440..46042bfa1990 100644
--- a/python/sglang/srt/lora/backend/base_backend.py
+++ b/python/sglang/srt/lora/backend/base_backend.py
@@ -27,12 +27,13 @@ def get_gathered_moe_num_tokens(forward_batch: ForwardBatch, num_tokens: int) ->
     global_num_tokens = getattr(forward_batch, "global_num_tokens_cpu", None)
     if not global_num_tokens:
         return num_tokens
-    from sglang.srt.layers.dp_attention import get_attention_tp_size
+    # Local import: a module-level cp import here is circular (see forward_batch_info).
+    # v0.5.16 moved this helper layers.utils.cp_utils -> layers.cp.padding.
+    from sglang.srt.layers.cp.padding import get_cp_padding_align_size
+    from sglang.srt.runtime_context import get_parallel
 
-    # Local import: a module-level cp_utils import here is circular (see forward_batch_info).
-    from sglang.srt.layers.utils.cp_utils import get_cp_padding_align_size
-
-    attn_tp_size = get_attention_tp_size()
+    # v0.5.16 retired dp_attention.get_attention_tp_size(); it lives on ParallelState now.
+    attn_tp_size = get_parallel().attn_tp_size
     cp_align_size = get_cp_padding_align_size()
     upper = max(
         ceil_align(ceil_align(t, attn_tp_size), cp_align_size)
diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py
index 475681e0a762..9bd491cd5cc5 100644
--- a/python/sglang/srt/model_executor/model_runner.py
+++ b/python/sglang/srt/model_executor/model_runner.py
@@ -1369,14 +1369,13 @@ def forward(
         output.expert_distribution_metrics = recorder_outputs.get("metrics")
 
         no_copy_to_cpu = not self.server_args.disable_overlap_schedule
-        # In speculative decoding, num_tokens_per_bs > 1, so pass the actual
-        # number of tokens per DP rank in CUDA graph, not the batch size.
+        # In speculative decoding more than one token is captured per request, so
+        # pass the actual number of tokens per DP rank in CUDA graph, not the batch
+        # size. v0.5.16 renamed this width num_tokens_per_bs -> captured_req_width.
         cuda_graph_num_tokens = None
-        if getattr(self.decode_cuda_graph_runner, "bs", None):
-            cuda_graph_num_tokens = (
-                self.decode_cuda_graph_runner.bs
-                * self.decode_cuda_graph_runner.num_tokens_per_bs
-            )
+        runner = self.decode_cuda_graph_runner
+        if getattr(runner, "bs", None):
+            cuda_graph_num_tokens = runner.bs * runner.captured_req_width
 
         if (
             not self.is_draft_worker

From 119689c08f0ba5e32c85db1b17dd9b532f440183 Mon Sep 17 00:00:00 2001
From: Yueming Yuan <yym022502@gmail.com>
Date: Sat, 25 Jul 2026 05:58:02 -0700
Subject: [PATCH 31/43] [26/27] [sglang-miles] Repoint two more attributes
 v0.5.16 relocated

Same class as [25/26]: attributes that still exist but now live somewhere else,
so the cherry-picks applied cleanly and only failed when the line ran.

1. `TpModelWorker.tp_rank` -> `self.ps.tp_rank` in
   `load_lora_adapter_from_tensors`. v0.5.16 keeps the rank on the ParallelState
   the worker holds, and defines no `tp_rank` attribute at all, so every LoRA
   load-from-tensors raised AttributeError. From [11/26] (MoE-LoRA).
   The two `[LORA-CHECK]` log lines in the same method have the same defect --
   they are v0.5.16's own code, latent because they only run on a checksum
   mismatch -- and are fixed here too rather than left as a trap.

2. `ModelRunner._model_update_group` -> `self.weight_updater._model_update_group`
   in `load_lora_adapter_from_distributed`. v0.5.16 moved the update-group
   registry onto the WeightUpdater component; the assert and the broadcast both
   read it. From [13/26]. Not yet reached by CI -- found by the scan below.

Found 2 by extending the pre-push checks: for every `self.X` this branch adds,
flag the ones never assigned or declared anywhere in the same file. That is a
tight signal for a relocated attribute (22 hits over 94 changed files, and the
rest are genuine mixin/base-class attributes).
---
 python/sglang/srt/managers/tp_worker.py          | 6 +++---
 python/sglang/srt/model_executor/model_runner.py | 8 +++++---
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py
index e5e51b010efd..0e90d5a4791d 100644
--- a/python/sglang/srt/managers/tp_worker.py
+++ b/python/sglang/srt/managers/tp_worker.py
@@ -189,7 +189,7 @@ def load_lora_adapter_from_tensors(
         # consumer means the CUDA-IPC ref counter on the producer's
         # bucket drops cleanly each cycle.
         monkey_patch_torch_reductions()
-        serialized = recv_req.serialized_named_tensors[self.tp_rank]
+        serialized = recv_req.serialized_named_tensors[self.ps.tp_rank]
         if recv_req.load_format == "flattened_bucket":
             flattened_data = MultiprocessingSerializer.deserialize(serialized)
             bucket = FlattenedTensorBucket(
@@ -223,12 +223,12 @@ def load_lora_adapter_from_tensors(
             extra = [n for n in tensors if n not in exp]
             if mismatch or missing or extra:
                 raise RuntimeError(
-                    f"[LORA-CHECK] rank{self.tp_rank} adapter sync MISMATCH of {len(exp)} expected: "
+                    f"[LORA-CHECK] rank{self.ps.tp_rank} adapter sync MISMATCH of {len(exp)} expected: "
                     f"{len(mismatch)} value-diff {mismatch[:5]}, {len(missing)} missing {missing[:5]}, "
                     f"{len(extra)} extra {extra[:5]}"
                 )
             logger.info(
-                f"[LORA-CHECK] rank{self.tp_rank} adapter sync OK: {len(exp)}/{len(exp)} tensors match (sha256)"
+                f"[LORA-CHECK] rank{self.ps.tp_rank} adapter sync OK: {len(exp)}/{len(exp)} tensors match (sha256)"
             )
         result = self.model_runner.load_lora_adapter_from_tensors(
             recv_req.to_ref(),
diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py
index 9bd491cd5cc5..6b32cb645fca 100644
--- a/python/sglang/srt/model_executor/model_runner.py
+++ b/python/sglang/srt/model_executor/model_runner.py
@@ -1086,8 +1086,10 @@ def load_lora_adapter_from_distributed(
         """Load a new lora adapter whose weights are broadcast over the
         `_model_update_group` process group (no CUDA IPC).
         """
-        assert group_name in self._model_update_group, (
-            f"Group {group_name} not in {list(self._model_update_group.keys())}. "
+        # v0.5.16 moved the update-group registry onto the WeightUpdater component.
+        update_groups = self.weight_updater._model_update_group
+        assert group_name in update_groups, (
+            f"Group {group_name} not in {list(update_groups.keys())}. "
             "Please call `init_weights_update_group` first."
         )
 
@@ -1104,7 +1106,7 @@ def load_lora_adapter_from_distributed(
                     torch.distributed.broadcast(
                         weight,
                         src=0,
-                        group=self._model_update_group[group_name],
+                        group=update_groups[group_name],
                         async_op=True,
                     )
                 )

From a3839f78922c15088186b4bd772013046c4a0104 Mon Sep 17 00:00:00 2001
From: Jiajun Li <guapisolo@gmail.com>
Date: Sat, 25 Jul 2026 06:32:10 -0700
Subject: [PATCH 32/43] [27/27] [sglang-miles] Fix iter_runners reaching the
 draft runner on the wrong object

`EAGLEWorkerV2.iter_runners` returned `self.draft_runner` and
`MultiLayerEagleWorkerV2.iter_runners` returned `self.draft_runner_list`, but
both attributes live on the inner EagleDraftWorker, not on the *V2 wrapper --
the wrappers reach it through the `BaseSpecWorker.draft_worker` property, which
is what the original [16/27] code did before I transcribed it. Every
check_weights / weight-update fan-out under EAGLE spec decoding failed with

    AttributeError: 'EAGLEWorkerV2' object has no attribute 'draft_runner'.
                    Did you mean: 'draft_worker'?

DFlashWorkerV2 is unaffected: it does own `draft_model_runner` directly.

My own transcription error while porting [16/27] onto v0.5.16, not a v0.5.16
change -- the attribute sits in the same class in v0.5.15.
---
 python/sglang/srt/speculative/eagle_worker_v2.py             | 2 +-
 python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py | 5 ++++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py
index c1523267a7a3..198b8e8c78c7 100644
--- a/python/sglang/srt/speculative/eagle_worker_v2.py
+++ b/python/sglang/srt/speculative/eagle_worker_v2.py
@@ -1024,7 +1024,7 @@ def _draft_extend_for_decode(
 
 class EAGLEWorkerV2(BaseSpecWorker):
     def iter_runners(self) -> List[Tuple[str, "ModelRunner"]]:
-        return [("draft", self.draft_runner)]
+        return [("draft", self.draft_worker.draft_runner)]
 
     def __init__(
         self,
diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
index b6c5b29cc8a1..cf490bd6de55 100644
--- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
+++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
@@ -876,7 +876,10 @@ def _draft_extend_for_decode(
 
 class MultiLayerEagleWorkerV2(BaseSpecWorker):
     def iter_runners(self) -> List[Tuple[str, ModelRunner]]:
-        return [(f"draft_step_{i}", r) for i, r in enumerate(self.draft_runner_list)]
+        return [
+            (f"draft_step_{i}", r)
+            for i, r in enumerate(self.draft_worker.draft_runner_list)
+        ]
 
     def __init__(
         self,

From 5138eb28845d137b860cddbff258de0754387165 Mon Sep 17 00:00:00 2001
From: Jiajun Li <jiajun.li@radixark.ai>
Date: Wed, 29 Jul 2026 13:22:44 -0700
Subject: [PATCH 33/43] fix: honor weight-check skips for quantized entries
 (#32809)

---
 python/sglang/srt/utils/weight_checker.py         |  6 +++++-
 test/registered/rl/test_weight_checker_e2e.py     |  7 ++++++-
 test/registered/unit/utils/test_weight_checker.py | 12 ++++++++++++
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/python/sglang/srt/utils/weight_checker.py b/python/sglang/srt/utils/weight_checker.py
index f67e61ee4e2f..6fbf56e03ef4 100644
--- a/python/sglang/srt/utils/weight_checker.py
+++ b/python/sglang/srt/utils/weight_checker.py
@@ -319,7 +319,11 @@ def _build_check_entries(
             continue  # compared via its weight's comparable
         if name in quantized_set:
             qw = quantized_set[name]
-            yield CheckEntry(name, True, qw.comparable_cls(tensor, raw[qw.scale_name]))
+            yield CheckEntry(
+                name,
+                name not in skip_compare_names,
+                qw.comparable_cls(tensor, raw[qw.scale_name]),
+            )
         else:
             should_compare = name not in skip_compare_names and (
                 not _is_non_persistent_buffer_name(name)
diff --git a/test/registered/rl/test_weight_checker_e2e.py b/test/registered/rl/test_weight_checker_e2e.py
index bff4ad102304..6c90f667a7e6 100644
--- a/test/registered/rl/test_weight_checker_e2e.py
+++ b/test/registered/rl/test_weight_checker_e2e.py
@@ -160,7 +160,12 @@ def test_e_checksum_returns_ranks_with_hashes(self):
         self.assertIn("checksums", first)
         self.assertIn("parallelism_info", first)
 
-        info = first["parallelism_info"]
+        parallelism_info = first["parallelism_info"]
+        self.assertIsInstance(parallelism_info, list)
+        self.assertGreaterEqual(len(parallelism_info), 1)
+
+        info = parallelism_info[0]
+        self.assertEqual(info["role"], "target")
         for key in (
             "tp_rank",
             "tp_size",
diff --git a/test/registered/unit/utils/test_weight_checker.py b/test/registered/unit/utils/test_weight_checker.py
index 79fe66aabf24..54eaebe0027f 100644
--- a/test/registered/unit/utils/test_weight_checker.py
+++ b/test/registered/unit/utils/test_weight_checker.py
@@ -267,6 +267,18 @@ def test_substring_match_not_endswith(self):
             [("weird.cos_sin_cache.foo.bar", False, RawComparable(t))],
         )
 
+    def test_skip_set_marks_quantized_entry_not_compared(self):
+        qweight, sf_fp32, _ = _build_fp8_quant_pair()
+        raw = {"x.weight": qweight, "x.weight_scale_inv": sf_fp32}
+        quantized_set = {
+            "x.weight": QuantizedWeight(Fp8BlockComparable, "x.weight_scale_inv")
+        }
+        ref = Fp8BlockComparable(qweight, sf_fp32)
+        _assert_entries_close(
+            _build_check_entries(raw, {"x.weight"}, quantized_set),
+            [("x.weight", False, ref)],
+        )
+
     # --- fp8 quant pair (real dequant on real fp8 tensors) ---
 
     def test_fp8_quant_pair_yields_lazy_pair(self):

From 49a5ef85eb6f808fe555b21c3a04f06b08ba6719 Mon Sep 17 00:00:00 2001
From: Xinyu Jiang <xinyuj2@andrew.cmu.edu>
Date: Mon, 27 Jul 2026 19:48:05 -0700
Subject: [PATCH 34/43] [Fix] Make RowParallelLinear k-size tuple-aware for FP8
 (#30742)

(cherry picked from commit f7ea06e1e519c569c1d6cb7e1c7b1128cbdbc39f)
---
 python/sglang/srt/layers/linear.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py
index b2385a95753e..f56ed7322c50 100644
--- a/python/sglang/srt/layers/linear.py
+++ b/python/sglang/srt/layers/linear.py
@@ -1611,7 +1611,18 @@ def forward(self, input_, skip_all_reduce=False, forward_batch=None):
                 get_tp_group(), disabled=not is_allocation_symmetric()
             )
         with symm_ctx:
-            if should_use_tp_invariant_row_linear(input_parallel.shape[-1]):
+            from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod
+
+            is_tuple = isinstance(input_parallel, tuple)
+            is_fp8 = isinstance(self.quant_method, Fp8LinearMethod)
+            k_size = (
+                input_parallel[0].shape[-1] if is_tuple else input_parallel.shape[-1]
+            )
+            if should_use_tp_invariant_row_linear(k_size):
+                if is_fp8:
+                    raise NotImplementedError(
+                        "FP8 tp-invariant row-linear not yet supported"
+                    )
                 output_parallel = torch.ops.tp_inv_ops.matmul_tp_inv(
                     input_parallel, self.weight.t(), bias_
                 )

From 3003d70f680d41c59d1b7acbf65cb47795dfd19e Mon Sep 17 00:00:00 2001
From: "Ethan (Yusheng) Su" <yushengsu@radixark.ai>
Date: Wed, 29 Jul 2026 13:57:25 -0700
Subject: [PATCH 35/43] [sglang-miles] Warn instead of silently dropping
 one-sided MoE expert LoRA targets (#32376)

Carried onto v0.5.16. The target-module gate is kept as-is: v0.5.16 added an
is_shared_fused_moe arm to the FusedMoEWithLoRA branch, and dropping the
redundant re-check there (as the original commit did) would let a shared-MoE
module through unchecked, which is beyond this fix's intent.
---
 python/sglang/srt/lora/lora_manager.py | 35 +++++++++++++++++++++++---
 1 file changed, 32 insertions(+), 3 deletions(-)

diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py
index f03db5093751..c888ecec49f6 100644
--- a/python/sglang/srt/lora/lora_manager.py
+++ b/python/sglang/srt/lora/lora_manager.py
@@ -978,6 +978,8 @@ def init_lora_modules(self):
 
         self.embed_tokens_module: Optional[BaseLayerWithLoRA] = None
         self.lm_head_module: Optional[BaseLayerWithLoRA] = None
+        # One MoE layer's worth of warning is enough; the condition is model-wide.
+        warned_one_sided_moe_targets = False
 
         # When tie_word_embeddings=True, lm_head is the same Python object as
         # embed_tokens. PyTorch's named_modules() deduplicates by object identity,
@@ -1069,9 +1071,36 @@ def init_lora_modules(self):
                 )
                 continue
 
-            if isinstance(module, (FusedMoE, InklingBatchDenseMLP)) and all(
-                x in self.target_modules for x in ["gate_up_proj", "down_proj"]
-            ):
+            if isinstance(module, (FusedMoE, InklingBatchDenseMLP)):
+                expert_projections = ("gate_up_proj", "down_proj")
+                missing = [
+                    x for x in expert_projections if x not in self.target_modules
+                ]
+                if len(missing) == 1 and not warned_one_sided_moe_targets:
+                    warned_one_sided_moe_targets = True
+                    # The MoE-LoRA hooks inject a delta after gate_up and after
+                    # down together, so wrapping for only one of them is not
+                    # implemented and the layer stays unwrapped. Adapters that only
+                    # train the dense/shared MLP still work, but one that carries
+                    # expert weights for the targeted projection has them loaded
+                    # into the pool and never applied. Warn once -- silently
+                    # dropping part of an adapter reads as a serving/training
+                    # mismatch with no symptom.
+                    present = next(x for x in expert_projections if x != missing[0])
+                    logger.warning(
+                        "LoRA target modules include %r but not %r, so MoE expert layers "
+                        "(e.g. %s) get no adapter at all: expert LoRA needs both projections. "
+                        "Any %r expert weights in a loaded adapter will be ignored. Add %r to "
+                        "the target modules to enable expert LoRA -- an adapter that does not "
+                        "train it keeps zero-filled buffers and contributes no delta.",
+                        present,
+                        missing[0],
+                        module_name,
+                        present,
+                        missing[0],
+                    )
+                if missing:
+                    continue
                 layer_id = get_layer_id(module_name)
                 if layer_id is None:
                     # FusedMoE submodules outside the decoder layer hierarchy

From 71db1014ec0a0fec28ad3850ce7ab0e75c5603dc Mon Sep 17 00:00:00 2001
From: Yuzhen Zhou <82826991+zyzshishui@users.noreply.github.com>
Date: Thu, 30 Jul 2026 16:46:14 -0700
Subject: [PATCH 36/43] [sglang-miles] Back up the CUDA graph pool across TMS
 pause/resume (#32718)

---
 .../model_executor/runner_backend/full_cuda_graph_backend.py  | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py
index e475ed9ca9fe..a39382cd42f7 100644
--- a/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py
+++ b/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py
@@ -103,6 +103,10 @@ def capture_one(
             graph_ctx = partial(
                 self._memory_saver_adapter.cuda_graph,
                 tag=GPU_MEMORY_TYPE_CUDA_GRAPH,
+                # The pool holds capture-time-initialized device state that replays
+                # read but never rewrite, so pausing must preserve the pool contents
+                # instead of discarding them.
+                enable_cpu_backup=True,
             )
         else:
             graph_ctx = self._device_module.graph

From 3fe50eddf881b4f33c60c1867af84eeb1626d91d Mon Sep 17 00:00:00 2001
From: Zhichen Zeng <zczeng@uw.edu>
Date: Thu, 30 Jul 2026 17:37:56 -0700
Subject: [PATCH 37/43] [Cherry-pick of #32861] Fix Inkling tool-call parsing
 recovery, content handling, and streaming (#32958)

Co-authored-by: ispobock <ispobaoke@gmail.com>
---
 .../srt/function_call/inkling_detector.py     | 370 ++++++++----------
 python/sglang/srt/parser/reasoning_parser.py  |   9 +-
 .../test_function_call_parser.py              | 234 +++++++++--
 .../unit/parser/test_reasoning_parser.py      |  11 +
 4 files changed, 381 insertions(+), 243 deletions(-)

diff --git a/python/sglang/srt/function_call/inkling_detector.py b/python/sglang/srt/function_call/inkling_detector.py
index 845aa47dc667..6340b690c587 100644
--- a/python/sglang/srt/function_call/inkling_detector.py
+++ b/python/sglang/srt/function_call/inkling_detector.py
@@ -1,11 +1,8 @@
 import json
 import logging
-import re
 from collections.abc import Mapping
 from typing import List, Optional
 
-from partial_json_parser.core.exceptions import MalformedJSON
-from partial_json_parser.core.options import Allow
 from xgrammar import StructuralTag
 
 from sglang.srt.entrypoints.openai.protocol import Tool
@@ -16,9 +13,9 @@
     ToolCallItem,
     _GetInfoFunc,
 )
-from sglang.srt.function_call.utils import _is_complete_json, _partial_json_loads
 from sglang.srt.parser.inkling_tokenizer import (
     CONTENT_INVOKE_TOOL_JSON,
+    CONTENT_INVOKE_TOOL_TEXT,
     END_MESSAGE,
     INKLING_CONTROL_TOKENS,
     INKLING_SPECIAL_TOKEN_IDS,
@@ -28,6 +25,11 @@
 logger = logging.getLogger(__name__)
 
 
+def _reject_nonfinite_number(value: str) -> float:
+    # Strict tool-payload parsing rejects NaN/Infinity; only recovery accepts them.
+    raise ValueError(f"{value} is not a valid JSON number")
+
+
 class InklingDetector(BaseFormatDetector):
     """
     Detector for Inkling structured tool calls.
@@ -40,176 +42,189 @@ def __init__(self):
         super().__init__()
         self.bot_token = CONTENT_INVOKE_TOOL_JSON
         self.eot_token = END_MESSAGE
-        self.tool_call_regex = re.compile(
-            re.escape(self.bot_token) + r"\s*(.*?)\s*" + re.escape(self.eot_token),
-            re.DOTALL,
-        )
-        self._current_header_name: str | None = None
+        # Streaming: index of the next call to emit; once a call fails to frame,
+        # the rest of the response streams through verbatim.
+        self._stream_call_index = 0
+        self._raw_passthrough = False
 
     def has_tool_call(self, text: str) -> bool:
-        return self.bot_token in text
+        return self.bot_token in text or CONTENT_INVOKE_TOOL_TEXT in text
 
     def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult:
-        if self.bot_token not in text:
+        if not self.has_tool_call(text):
             return StreamingParseResult(normal_text=self._clean_normal_text(text))
 
-        try:
-            calls: list[ToolCallItem] = []
-            for match in self.tool_call_regex.finditer(text):
-                try:
-                    payload = json.loads(match.group(1).strip())
-                except json.JSONDecodeError as exc:
-                    logger.warning("Invalid Inkling tool call JSON: %s", exc)
-                    continue
-                if not isinstance(payload, Mapping):
-                    logger.warning("Invalid Inkling tool call payload: %s", payload)
-                    continue
-                _, header_name = self._split_trailing_tool_header(text[: match.start()])
-                call = self._tool_call_item(
-                    payload, tools, len(calls), header_name=header_name
-                )
-                if call is not None:
-                    calls.append(call)
-
-            if not calls:
-                # Every candidate call was rejected (bad payload or a
-                # header/payload name mismatch). Match the framework contract
-                # every other detector follows: normal_text is only the content
-                # BEFORE the tool marker — the rejected tool-call region is
-                # dropped, never regurgitated as visible content.
-                prefix, _ = self._split_trailing_tool_header(
-                    text[: text.find(self.bot_token)]
-                )
-                return StreamingParseResult(normal_text=self._clean_normal_text(prefix))
-
-            normal_prefix, _ = self._split_trailing_tool_header(
-                text[: text.find(self.bot_token)]
-            )
-            normal_text = self._clean_normal_text(normal_prefix)
+        parsed = self._parse_canonical(text, tools)
+        if parsed is not None:
+            normal_text, calls = parsed
             return StreamingParseResult(normal_text=normal_text, calls=calls)
-        except Exception as exc:
-            logger.error("Error in Inkling detect_and_parse: %s", exc, exc_info=True)
-            prefix, _ = self._split_trailing_tool_header(
-                text[: text.find(self.bot_token)]
-            )
-            return StreamingParseResult(normal_text=self._clean_normal_text(prefix))
 
-    def parse_streaming_increment(
-        self, new_text: str, tools: List[Tool]
-    ) -> StreamingParseResult:
-        self._buffer += new_text
-        current_text = self._buffer
-
-        if self.bot_token not in current_text:
-            header_start = self._pending_tool_header_start(current_text)
-            if header_start is not None:
-                safe_text = current_text[:header_start]
-                self._buffer = current_text[header_start:]
-                return StreamingParseResult(
-                    normal_text=self._clean_normal_text(safe_text)
-                )
-            # Hold back a partial prefix of ANY token _clean_normal_text
-            # strips — emitting a split control token leaks its first half as
-            # visible text (the completed token would have been stripped).
-            partial_len = max(
-                self._ends_with_partial_token(current_text, token)
-                for token in INKLING_CONTROL_TOKENS
-            )
-            if partial_len:
-                safe_text = current_text[:-partial_len]
-                self._buffer = current_text[-partial_len:]
+        # Canonical framing failed. Recover a single call from the last tool
+        # marker; if that fails too, surface the whole visible payload as text.
+        recovered = self._recover_last_json_call(text, tools)
+        if recovered is not None:
+            return StreamingParseResult(normal_text="", calls=[recovered])
+        return StreamingParseResult(normal_text=self._clean_normal_text(text))
+
+    def _parse_canonical(
+        self, text: str, tools: List[Tool]
+    ) -> tuple[str, list[ToolCallItem]] | None:
+        """Extract every tool call under strict framing.
+
+        Returns (visible_text, calls) when EVERY marker frames a valid call, or
+        None if any marker is malformed/unterminated — one bad call fails the
+        whole batch, matching streaming.
+        """
+        calls: list[ToolCallItem] = []
+        normal_parts: list[str] = []
+        pos = 0
+        while pos < len(text):
+            marker_pos, marker_token, is_json = self._next_tool_marker(text, pos)
+            if marker_pos is None:
+                normal_parts.append(text[pos:])
+                break
+
+            prefix, _ = self._split_trailing_tool_header(text[pos:marker_pos])
+            normal_parts.append(prefix)
+
+            body_start = marker_pos + len(marker_token)
+            eot = text.find(self.eot_token, body_start)
+            if eot == -1:
+                return None  # unterminated -> recovery reads through EOS
+            body = text[body_start:eot]
+
+            if is_json:
+                call = self._canonical_json_call(body, tools, len(calls))
+                if call is None:
+                    return None
             else:
-                safe_text = current_text
-                self._buffer = ""
-            return StreamingParseResult(normal_text=self._clean_normal_text(safe_text))
-
-        bot_pos = current_text.find(self.bot_token)
-        if bot_pos > 0:
-            normal_text, self._current_header_name = self._split_trailing_tool_header(
-                current_text[:bot_pos]
-            )
-            self._buffer = current_text[bot_pos:]
-            normal_text = self._clean_normal_text(normal_text)
-            if normal_text:
-                return StreamingParseResult(normal_text=normal_text)
-            current_text = self._buffer
-
-        if not hasattr(self, "_tool_indices"):
-            self._tool_indices = self._get_tool_indices(tools)
-
-        start_idx = len(self.bot_token)
-        while start_idx < len(current_text) and current_text[start_idx].isspace():
-            start_idx += 1
-
-        flags = Allow.ALL if self.current_tool_name_sent else Allow.ALL & ~Allow.STR
+                call = self._text_tool_call(body, len(calls))
+            calls.append(call)
+            pos = eot + len(self.eot_token)
+
+        return self._clean_normal_text("".join(normal_parts)), calls
+
+    def _next_tool_marker(self, text: str, start: int) -> tuple[int | None, str, bool]:
+        """Earliest json/text tool marker at or after ``start`` (is_json flag)."""
+        json_pos = text.find(self.bot_token, start)
+        text_pos = text.find(CONTENT_INVOKE_TOOL_TEXT, start)
+        if json_pos == -1 and text_pos == -1:
+            return None, "", True
+        if text_pos == -1 or (json_pos != -1 and json_pos <= text_pos):
+            return json_pos, self.bot_token, True
+        return text_pos, CONTENT_INVOKE_TOOL_TEXT, False
+
+    def _canonical_json_call(
+        self, body: str, tools: List[Tool], call_index: int
+    ) -> ToolCallItem | None:
         try:
-            payload, end_idx = _partial_json_loads(current_text[start_idx:], flags)
-        except (MalformedJSON, json.JSONDecodeError):
-            return StreamingParseResult()
+            payload = json.loads(body.strip(), parse_constant=_reject_nonfinite_number)
+        except (json.JSONDecodeError, ValueError):
+            return None
         if not isinstance(payload, Mapping):
-            return StreamingParseResult()
-
-        calls: list[ToolCallItem] = []
-        name = payload.get("name")
-        if (
-            not self.current_tool_name_sent
-            and isinstance(name, str)
-            and (self._current_header_name is None or self._current_header_name == name)
-        ):
-            self._ensure_current_tool()
-            calls.append(
-                ToolCallItem(
-                    tool_index=self.current_tool_id,
-                    name=name,
-                    parameters="",
-                )
-            )
-            self.current_tool_name_sent = True
-            self.prev_tool_call_arr[self.current_tool_id] = {
-                "name": name,
-                "arguments": {},
-            }
+            return None
+        return self._tool_call_item(payload, tools, call_index)
 
-        json_text = current_text[start_idx : start_idx + end_idx]
-        if not _is_complete_json(json_text):
-            return StreamingParseResult(calls=calls)
+    def _recover_last_json_call(
+        self, text: str, tools: List[Tool]
+    ) -> ToolCallItem | None:
+        """Recover one call from the last json marker, reading through the next
+        end token or EOS. Requires a nonempty name; accepts NaN/Infinity."""
+        last = text.rfind(self.bot_token)
+        if last == -1:
+            return None
+        body_start = last + len(self.bot_token)
+        eot = text.find(self.eot_token, body_start)
+        candidate = text[body_start:eot] if eot != -1 else text[body_start:]
+        candidate = self._clean_normal_text(candidate).strip()
+        try:
+            payload = json.loads(candidate)
+        except (json.JSONDecodeError, ValueError):
+            return None
+        if not isinstance(payload, Mapping) or not payload.get("name"):
+            return None
+        return self._tool_call_item(payload, tools, 0)
 
-        call = self._tool_call_item(
-            payload,
-            tools,
-            self.current_tool_id,
-            header_name=self._current_header_name,
+    def _text_tool_call(self, body: str, call_index: int) -> ToolCallItem:
+        # Headerless raw-text invocation: no structured name/args on the wire.
+        return ToolCallItem(
+            tool_index=call_index,
+            name="",
+            parameters=json.dumps({"text": self._clean_normal_text(body)}),
         )
-        if call is None:
-            self._abandon_current_tool()
+
+    def parse_streaming_increment(
+        self, new_text: str, tools: List[Tool]
+    ) -> StreamingParseResult:
+        self._buffer += new_text
+        if self._raw_passthrough:
+            out = self._clean_normal_text(self._buffer)
             self._buffer = ""
-            return StreamingParseResult(calls=calls)
-
-        if self.current_tool_id == -1:
-            self._ensure_current_tool()
-
-        args = json.loads(call.parameters)
-        self.prev_tool_call_arr[self.current_tool_id] = {
-            "name": call.name,
-            "arguments": args,
-        }
-        sent = self.streamed_args_for_tool[self.current_tool_id]
-        remaining_args = call.parameters[len(sent) :]
-        if remaining_args:
-            calls.append(
-                ToolCallItem(
-                    tool_index=self.current_tool_id,
-                    name=None,
-                    parameters=remaining_args,
-                )
-            )
-            self.streamed_args_for_tool[self.current_tool_id] += remaining_args
+            return StreamingParseResult(normal_text=out)
 
-        self._buffer = self._remaining_after_call(current_text, start_idx + end_idx)
-        self.current_tool_id += 1
-        self.current_tool_name_sent = False
-        self._current_header_name = None
-        return StreamingParseResult(calls=calls)
+        normal_parts: list[str] = []
+        calls: list[ToolCallItem] = []
+        while self._buffer:
+            marker_pos, marker_token, is_json = self._next_tool_marker(self._buffer, 0)
+            if marker_pos is None:
+                safe, hold = self._split_safe_text(self._buffer)
+                normal_parts.append(self._clean_normal_text(safe))
+                self._buffer = hold
+                break
+
+            prefix, _ = self._split_trailing_tool_header(self._buffer[:marker_pos])
+            body_start = marker_pos + len(marker_token)
+            eot = self._buffer.find(self.eot_token, body_start)
+            if eot == -1:
+                # Region incomplete — emit only the text before it and hold the
+                # marker + partial body so name/args stay atomic until the close.
+                normal_parts.append(self._clean_normal_text(prefix))
+                self._buffer = self._buffer[marker_pos:]
+                break
+
+            body = self._buffer[body_start:eot]
+            if is_json:
+                call = self._canonical_json_call(body, tools, self._stream_call_index)
+                if call is None:
+                    return self._stream_recover_or_fallback(tools)
+            else:
+                call = self._text_tool_call(body, self._stream_call_index)
+            normal_parts.append(self._clean_normal_text(prefix))
+            calls.append(call)
+            self._stream_call_index += 1
+            self._buffer = self._buffer[eot + len(self.eot_token) :]
+
+        return StreamingParseResult(normal_text="".join(normal_parts), calls=calls)
+
+    def _stream_recover_or_fallback(self, tools: List[Tool]) -> StreamingParseResult:
+        # A call failed to frame: recover one call from the last marker, then
+        # pass the rest of the response through as text.
+        self._raw_passthrough = True
+        recovered = self._recover_last_json_call(self._buffer, tools)
+        buffered = self._buffer
+        self._buffer = ""
+        if recovered is not None:
+            recovered.tool_index = self._stream_call_index
+            self._stream_call_index += 1
+            return StreamingParseResult(calls=[recovered])
+        return StreamingParseResult(normal_text=self._clean_normal_text(buffered))
+
+    def _split_safe_text(self, text: str) -> tuple[str, str]:
+        """Split off text safe to emit now from a tail that may be a forming
+        tool header or a split control token."""
+        header_start = self._pending_tool_header_start(text)
+        if header_start is not None:
+            return text[:header_start], text[header_start:]
+        partial_len = max(
+            (
+                self._ends_with_partial_token(text, token)
+                for token in INKLING_CONTROL_TOKENS
+            ),
+            default=0,
+        )
+        if partial_len:
+            return text[:-partial_len], text[-partial_len:]
+        return text, ""
 
     def structure_info(self) -> _GetInfoFunc:
         def info(name: str) -> StructureInfo:
@@ -278,21 +293,12 @@ def _tool_call_item(
         payload: Mapping[str, object],
         tools: List[Tool],
         call_index: int,
-        *,
-        header_name: str | None = None,
     ) -> ToolCallItem | None:
         name = payload.get("name")
         args = payload.get("args")
         if not isinstance(name, str) or not isinstance(args, Mapping):
             logger.warning("Invalid Inkling tool call payload: %s", payload)
             return None
-        if header_name is not None and header_name != name:
-            logger.warning(
-                "Inkling tool header %r does not match payload name %r",
-                header_name,
-                name,
-            )
-            return None
 
         if not hasattr(self, "_tool_indices"):
             self._tool_indices = self._get_tool_indices(tools)
@@ -309,28 +315,6 @@ def _tool_call_item(
             parameters=json.dumps(args, ensure_ascii=False),
         )
 
-    def _ensure_current_tool(self) -> None:
-        if self.current_tool_id == -1:
-            self.current_tool_id = 0
-        while len(self.prev_tool_call_arr) <= self.current_tool_id:
-            self.prev_tool_call_arr.append({})
-        while len(self.streamed_args_for_tool) <= self.current_tool_id:
-            self.streamed_args_for_tool.append("")
-
-    def _abandon_current_tool(self) -> None:
-        """Discard the in-flight call after a rejected payload.
-
-        Resetting ``current_tool_id`` to -1 here would collide the NEXT valid
-        call with tool index 0 (``_ensure_current_tool`` maps -1 -> 0) and
-        slice its arguments against index 0's already-streamed args. Keep the
-        counter: an unannounced slot is simply reused; an announced slot is
-        abandoned by advancing past it.
-        """
-        if self.current_tool_name_sent:
-            self.current_tool_id += 1
-        self.current_tool_name_sent = False
-        self._current_header_name = None
-
     def _split_trailing_tool_header(self, text: str) -> tuple[str, str | None]:
         message_pos = self._pending_tool_header_start(text)
         if message_pos is None:
@@ -350,14 +334,6 @@ def _pending_tool_header_start(self, text: str) -> int | None:
             return None
         return message_pos
 
-    def _remaining_after_call(self, text: str, end_idx: int) -> str:
-        remaining = text[end_idx:]
-        if remaining.startswith(self.eot_token):
-            return remaining[len(self.eot_token) :]
-        if self.eot_token in remaining:
-            return remaining.split(self.eot_token, 1)[1]
-        return remaining
-
     def _clean_normal_text(self, text: str) -> str:
         for token in INKLING_CONTROL_TOKENS:
             text = text.replace(token, "")
diff --git a/python/sglang/srt/parser/reasoning_parser.py b/python/sglang/srt/parser/reasoning_parser.py
index 838f40005a22..a3b6541eba94 100644
--- a/python/sglang/srt/parser/reasoning_parser.py
+++ b/python/sglang/srt/parser/reasoning_parser.py
@@ -7,6 +7,7 @@
 from sglang.srt.parser.harmony_parser import HarmonyParser
 from sglang.srt.parser.inkling_tokenizer import (
     CONTENT_INVOKE_TOOL_JSON,
+    CONTENT_INVOKE_TOOL_TEXT,
     CONTENT_MODEL_END_SAMPLING,
     CONTENT_TEXT,
     CONTENT_THINKING,
@@ -854,12 +855,12 @@ def flush_reasoning() -> None:
                     # a real header can only follow an end token. Preserve it
                     # instead of rerouting the rest of the block into a header.
                     emit(token)
-            elif token == CONTENT_INVOKE_TOOL_JSON:
+            elif token in (CONTENT_INVOKE_TOOL_JSON, CONTENT_INVOKE_TOOL_TEXT):
+                # Preserve the tool-invocation framing (json and headerless raw
+                # text) in content so the tool-call detector receives it.
                 flush_reasoning()
                 if self._kind == "header":
-                    content.extend(
-                        (MESSAGE_MODEL, self._pending_header, CONTENT_INVOKE_TOOL_JSON)
-                    )
+                    content.extend((MESSAGE_MODEL, self._pending_header, token))
                     self._pending_header = ""
                 else:
                     content.append(token)
diff --git a/test/registered/unit/function_call/test_function_call_parser.py b/test/registered/unit/function_call/test_function_call_parser.py
index 7dfa25925727..7ac9b5f27bab 100644
--- a/test/registered/unit/function_call/test_function_call_parser.py
+++ b/test/registered/unit/function_call/test_function_call_parser.py
@@ -85,30 +85,32 @@ def test_streaming_header_is_buffered_until_the_tool_kind(self):
         self.assertEqual(name, "weather")
         self.assertEqual(json.loads(parameters), {"city": "SF"})
 
-    def test_mismatched_header_is_rejected(self):
+    def test_header_name_is_ignored_and_payload_name_wins(self):
+        """The message header is author metadata, not a name check: a header
+        that differs from the payload name still yields a call named by the
+        payload."""
         detector = InklingDetector()
         source = (
             "<|message_model|>other<|content_invoke_tool_json|>"
             '{"name":"weather","args":{}}<|end_message|>'
         )
         result = detector.detect_and_parse(source, self.tools)
-        self.assertEqual(result.calls, [])
+        self.assertEqual(len(result.calls), 1)
+        self.assertEqual(result.calls[0].name, "weather")
+        self.assertEqual(result.normal_text, "")
 
-    def test_rejected_call_does_not_leak_protocol_tokens(self):
-        """Bug regression: the no-surviving-calls path returned the RAW text,
-        so a rejected call (e.g. header/payload mismatch) leaked <|...|>
-        protocol tokens into user-visible content."""
+    def test_raw_fallback_strips_protocol_tokens(self):
+        """When a payload cannot be parsed or recovered, the visible text is
+        surfaced as content with the <|...|> special tokens stripped."""
         detector = InklingDetector()
         source = (
-            "<|message_model|>other<|content_invoke_tool_json|>"
-            '{"name":"weather","args":{}}<|end_message|>'
+            "<|message_model|>weather<|content_invoke_tool_json|>"
+            "{not json at all<|end_message|>"
         )
         result = detector.detect_and_parse(source, self.tools)
+        self.assertEqual(result.calls, [])
         self.assertNotIn("<|", result.normal_text)
-        # Framework parity: the rejected tool-call REGION is dropped entirely
-        # (normal_text = content before the marker), like every other detector
-        # — the JSON payload must not surface as visible content either.
-        self.assertEqual(result.normal_text, "")
+        self.assertIn("{not json at all", result.normal_text)
 
     def test_headerless_legacy_tool_call_still_parses(self):
         """Spec tolerance: a bare <|content_invoke_tool_json|> block with no
@@ -149,33 +151,94 @@ def test_streaming_two_sequential_tool_calls_get_distinct_indices(self):
         self.assertEqual(json.loads(args_by_index[0]), {"city": "SF"})
         self.assertEqual(json.loads(args_by_index[1]), {"city": "NY"})
 
-    def test_streaming_rejection_does_not_collide_tool_indices(self):
-        """Bug regression: a rejected mid-stream call reset current_tool_id to
-        -1, so the NEXT valid call re-announced as tool_index 0 — colliding
-        with the first call's index and slicing its arguments against index
-        0's already-streamed args."""
+    def test_streaming_two_complete_tool_calls_in_one_delta_both_emit(self):
+        """Bug regression: parse_streaming_increment parsed one call per delta
+        and re-buffered the rest, relying on the NEXT delta to drain it. Two
+        complete calls arriving in a single (e.g. final) delta left the second
+        stranded in the buffer with no stream-end flush, so only the first was
+        emitted."""
+        detector = InklingDetector()
+        source = (
+            "<|message_model|>weather<|content_invoke_tool_json|>"
+            '{"name":"weather","args":{"city":"SF"}}<|end_message|>'
+            "<|message_model|>weather<|content_invoke_tool_json|>"
+            '{"name":"weather","args":{"city":"NY"}}<|end_message|>'
+        )
+        args_by_index: dict = {}
+        for call in detector.parse_streaming_increment(source, self.tools).calls:
+            args_by_index[call.tool_index] = (
+                args_by_index.get(call.tool_index, "") + call.parameters
+            )
+        self.assertEqual(sorted(args_by_index), [0, 1])
+        self.assertEqual(json.loads(args_by_index[0]), {"city": "SF"})
+        self.assertEqual(json.loads(args_by_index[1]), {"city": "NY"})
+
+    def test_streaming_text_then_tool_call_in_one_delta_emits_both(self):
+        """Bug regression: a delta carrying visible text followed by a complete
+        tool call emitted only the text and stranded the call in the buffer
+        (the drain loop stopped after the leading-text run), so a final such
+        delta dropped the call. The drain must continue past leading text."""
+        detector = InklingDetector()
+        source = (
+            "Sure, let me check.<|message_model|>weather<|content_invoke_tool_json|>"
+            '{"name":"weather","args":{"city":"SF"}}<|end_message|>'
+        )
+        result = detector.parse_streaming_increment(source, self.tools)
+        self.assertIn("Sure, let me check.", result.normal_text)
+        names = [c.name for c in result.calls if c.name]
+        self.assertEqual(names, ["weather"])
+        args = "".join(c.parameters for c in result.calls)
+        self.assertEqual(json.loads(args), {"city": "SF"})
+
+    def test_streaming_differing_headers_all_stream(self):
+        """The header is author metadata, not a name gate: three calls with
+        differing headers all stream, indexed 0/1/2 by payload name."""
+        detector = InklingDetector()
+        source = (
+            "<|message_model|>weather<|content_invoke_tool_json|>"
+            '{"name":"weather","args":{"city":"SF"}}<|end_message|>'
+            "<|message_model|>other<|content_invoke_tool_json|>"
+            '{"name":"weather","args":{"city":"XX"}}<|end_message|>'
+            "<|message_model|>weather<|content_invoke_tool_json|>"
+            '{"name":"weather","args":{"city":"NY"}}<|end_message|>'
+        )
+        args_by_index: dict = {}
+        for call in detector.parse_streaming_increment(source, self.tools).calls:
+            args_by_index[call.tool_index] = (
+                args_by_index.get(call.tool_index, "") + call.parameters
+            )
+        self.assertEqual(sorted(args_by_index), [0, 1, 2])
+        self.assertEqual(json.loads(args_by_index[0]), {"city": "SF"})
+        self.assertEqual(json.loads(args_by_index[1]), {"city": "XX"})
+        self.assertEqual(json.loads(args_by_index[2]), {"city": "NY"})
+
+    def test_streaming_malformed_call_switches_to_raw_passthrough(self):
+        """A call that fails to frame switches the stream to raw passthrough:
+        earlier calls stay emitted (streaming cannot un-emit), and everything
+        after the failure is surfaced as content, never as further calls."""
         detector = InklingDetector()
         chunks = [
             "<|message_model|>weather<|content_invoke_tool_json|>",
             '{"name":"weather","args":{"city":"SF"}}<|end_message|>',
-            # header/payload mismatch -> rejected
-            "<|message_model|>other<|content_invoke_tool_json|>",
-            '{"name":"weather","args":{"city":"NY"}}<|end_message|>',
-            # valid again
+            # unrecoverable -> raw passthrough from here on
+            "<|message_model|>weather<|content_invoke_tool_json|>",
+            "{not json at all<|end_message|>",
+            # would-be call, now passthrough text
             "<|message_model|>weather<|content_invoke_tool_json|>",
             '{"name":"weather","args":{"city":"LA"}}<|end_message|>',
         ]
-        args_by_index: dict = {}
+        calls: list = []
+        normal_text = ""
         for chunk in chunks:
-            for call in detector.parse_streaming_increment(chunk, self.tools).calls:
-                args_by_index[call.tool_index] = (
-                    args_by_index.get(call.tool_index, "") + call.parameters
-                )
-        self.assertEqual(json.loads(args_by_index[0]), {"city": "SF"})
-        self.assertEqual(len(args_by_index), 2)
-        second_index = max(args_by_index)
-        self.assertGreater(second_index, 0)
-        self.assertEqual(json.loads(args_by_index[second_index]), {"city": "LA"})
+            result = detector.parse_streaming_increment(chunk, self.tools)
+            normal_text += result.normal_text
+            calls.extend(result.calls)
+        self.assertEqual(len(calls), 1)
+        self.assertEqual(calls[0].name, "weather")
+        self.assertEqual(json.loads(calls[0].parameters), {"city": "SF"})
+        self.assertNotIn("<|", normal_text)
+        self.assertIn("{not json at all", normal_text)
+        self.assertIn("LA", normal_text)
 
     def test_undeclared_tool_name_is_surfaced(self):
         """A call to a tool absent from the request's tool list surfaces as a
@@ -212,8 +275,9 @@ def test_undeclared_tool_name_surfaces_in_streaming(self):
         self.assertEqual(json.loads(parameters), {"query": "q"})
         self.assertNotIn("<|", normal_text)
 
-    def test_malformed_json_does_not_leak_protocol_tokens(self):
-        """Malformed JSON must drop the protocol region and its tool header."""
+    def test_malformed_json_surfaces_as_raw_fallback(self):
+        """Malformed JSON that also fails recovery surfaces the visible payload
+        as content (special tokens stripped), not a tool call."""
         detector = InklingDetector()
         source = (
             "<|message_model|>weather<|content_invoke_tool_json|>"
@@ -221,10 +285,12 @@ def test_malformed_json_does_not_leak_protocol_tokens(self):
         )
         result = detector.detect_and_parse(source, self.tools)
         self.assertEqual(result.calls, [])
-        self.assertEqual(result.normal_text, "")
+        self.assertNotIn("<|", result.normal_text)
+        self.assertIn("{not json at all", result.normal_text)
 
-    def test_parser_does_not_restore_malformed_tool_call_as_text(self):
-        """The parser wrapper must preserve the detector's sanitized fallback."""
+    def test_parser_preserves_raw_fallback_text(self):
+        """The parser wrapper preserves the detector's raw fallback, so the
+        visible prefix plus the failed payload reach the caller as content."""
         from sglang.srt.function_call.function_call_parser import FunctionCallParser
 
         source = (
@@ -235,8 +301,9 @@ def test_parser_does_not_restore_malformed_tool_call_as_text(self):
         normal_text, calls = FunctionCallParser(self.tools, "inkling").parse_non_stream(
             source
         )
-        self.assertEqual(normal_text, "Visible prefix.")
         self.assertEqual(calls, [])
+        self.assertTrue(normal_text.startswith("Visible prefix."))
+        self.assertIn("{not json at all", normal_text)
 
     def test_parser_preserves_text_without_tool_call_marker(self):
         from sglang.srt.function_call.function_call_parser import FunctionCallParser
@@ -248,7 +315,10 @@ def test_parser_preserves_text_without_tool_call_marker(self):
         self.assertEqual(normal_text, source)
         self.assertEqual(calls, [])
 
-    def test_malformed_call_does_not_discard_an_earlier_valid_call(self):
+    def test_one_malformed_call_fails_the_whole_batch(self):
+        """All-or-nothing: a single unrecoverable call fails canonical framing
+        for the whole response, so even an earlier valid call is discarded and
+        the visible text is surfaced as content."""
         source = (
             "<|message_model|>weather<|content_invoke_tool_json|>"
             '{"name":"weather","args":{"city":"SF"}}<|end_message|>'
@@ -256,10 +326,10 @@ def test_malformed_call_does_not_discard_an_earlier_valid_call(self):
             "{not json at all<|end_message|>"
         )
         result = InklingDetector().detect_and_parse(source, self.tools)
-        self.assertEqual(result.normal_text, "")
-        self.assertEqual(len(result.calls), 1)
-        self.assertEqual(result.calls[0].name, "weather")
-        self.assertEqual(json.loads(result.calls[0].parameters), {"city": "SF"})
+        self.assertEqual(result.calls, [])
+        self.assertNotIn("<|", result.normal_text)
+        self.assertIn('{"name":"weather","args":{"city":"SF"}}', result.normal_text)
+        self.assertIn("{not json at all", result.normal_text)
 
     def test_clean_normal_text_strips_the_full_control_alphabet(self):
         """Fall-through text is cleaned against the whole shared control-token
@@ -278,6 +348,86 @@ def test_structural_tag_uses_the_canonical_header(self):
         self.assertEqual(info.trigger, header)
         self.assertTrue(info.begin.startswith(header + '{"name":"weather"'))
 
+    def test_content_after_tool_call_is_preserved(self):
+        """A tool call followed by a text block returns both: the call plus the
+        trailing visible content, not just the prefix before the marker."""
+        source = (
+            "<|message_model|>weather<|content_invoke_tool_json|>"
+            '{"name":"weather","args":{"city":"SF"}}<|end_message|>'
+            "<|message_model|><|content_text|>Here you go.<|end_message|>"
+        )
+        result = InklingDetector().detect_and_parse(source, self.tools)
+        self.assertEqual(len(result.calls), 1)
+        self.assertEqual(result.calls[0].name, "weather")
+        self.assertEqual(result.normal_text, "Here you go.")
+
+    def test_empty_name_is_allowed_on_the_canonical_path(self):
+        source = "<|content_invoke_tool_json|>" '{"name":"","args":{}}<|end_message|>'
+        result = InklingDetector().detect_and_parse(source, self.tools)
+        self.assertEqual(len(result.calls), 1)
+        self.assertEqual(result.calls[0].name, "")
+
+    def test_recovery_uses_only_the_last_marker(self):
+        """Canonical framing fails on the garbage payload; recovery reads only
+        the payload after the LAST marker."""
+        source = (
+            "<|message_model|>weather<|content_invoke_tool_json|>garbage"
+            "<|content_invoke_tool_json|>"
+            '{"name":"weather","args":{"city":"SF"}}<|end_message|>'
+        )
+        result = InklingDetector().detect_and_parse(source, self.tools)
+        self.assertEqual(len(result.calls), 1)
+        self.assertEqual(result.calls[0].name, "weather")
+        self.assertEqual(json.loads(result.calls[0].parameters), {"city": "SF"})
+
+    def test_recovery_requires_a_nonempty_name(self):
+        """Recovery (unlike the canonical path) rejects an empty name, falling
+        through to raw text."""
+        source = (
+            "<|message_model|>weather<|content_invoke_tool_json|>bad"
+            '<|content_invoke_tool_json|>{"name":"","args":{}}<|end_message|>'
+        )
+        result = InklingDetector().detect_and_parse(source, self.tools)
+        self.assertEqual(result.calls, [])
+        self.assertNotIn("<|", result.normal_text)
+
+    def test_nonfinite_numbers_rejected_canonically_but_recovered(self):
+        """NaN/Infinity are not valid canonical JSON, so the strict pass fails;
+        recovery accepts them."""
+        source = (
+            "<|content_invoke_tool_json|>"
+            '{"name":"weather","args":{"v":NaN}}<|end_message|>'
+        )
+        result = InklingDetector().detect_and_parse(source, self.tools)
+        self.assertEqual(len(result.calls), 1)
+        self.assertEqual(result.calls[0].name, "weather")
+
+    def test_streaming_name_not_emitted_before_end_message(self):
+        """Atomicity: the tool name is withheld until the closing marker, so a
+        call that never completes never leaks an orphan name delta."""
+        detector = InklingDetector()
+        pre = detector.parse_streaming_increment(
+            '<|message_model|>weather<|content_invoke_tool_json|>{"name":"wea',
+            self.tools,
+        )
+        self.assertEqual(pre.calls, [])
+        post = detector.parse_streaming_increment(
+            'ther","args":{"city":"SF"}}<|end_message|>', self.tools
+        )
+        self.assertEqual(len(post.calls), 1)
+        self.assertEqual(post.calls[0].name, "weather")
+        self.assertEqual(json.loads(post.calls[0].parameters), {"city": "SF"})
+
+    def test_raw_text_tool_invocation_surfaces_as_a_call(self):
+        """A headerless <|content_invoke_tool_text|> block reaches the tool loop
+        as a call carrying the raw body, instead of being dropped."""
+        source = "<|content_invoke_tool_text|>search the web<|end_message|>"
+        result = InklingDetector().detect_and_parse(source, self.tools)
+        self.assertEqual(len(result.calls), 1)
+        self.assertEqual(
+            json.loads(result.calls[0].parameters), {"text": "search the web"}
+        )
+
 
 class TestPythonicDetector(unittest.TestCase):
     def setUp(self):
diff --git a/test/registered/unit/parser/test_reasoning_parser.py b/test/registered/unit/parser/test_reasoning_parser.py
index e5616b27f87e..a2e5b5ab8fcf 100644
--- a/test/registered/unit/parser/test_reasoning_parser.py
+++ b/test/registered/unit/parser/test_reasoning_parser.py
@@ -195,6 +195,17 @@ def test_tool_header_is_preserved_for_the_tool_parser(self):
             content += detector.parse_streaming_increment(char).normal_text
         self.assertEqual(content, source)
 
+    def test_raw_text_tool_framing_is_preserved_for_the_tool_parser(self):
+        """The headerless <|content_invoke_tool_text|> block must survive into
+        content so the tool-call detector can surface it, rather than being
+        swallowed as header data."""
+        detector = InklingDetector()
+        source = "<|message_model|><|content_invoke_tool_text|>search<|end_message|>"
+        result = detector.detect_and_parse(source)
+        self.assertIn("<|content_invoke_tool_text|>", result.normal_text)
+        self.assertIn("search", result.normal_text)
+        self.assertEqual(result.reasoning_text, "")
+
     def test_quoted_message_model_token_inside_content_is_preserved(self):
         """Bug regression: the header branch flipped to header state on ANY
         <|message_model|> occurrence, so a literal token the model wrote

From 0b45fbc90bdf4eb021f97496555580aeb14bb78f Mon Sep 17 00:00:00 2001
From: Shi Dong <hello@dongshi.me>
Date: Sun, 2 Aug 2026 17:35:02 -0700
Subject: [PATCH 38/43] [sglang-miles] Release the LoRA reference on every
 request-cleanup path

A request acquires a reference in the LoRA registry before dispatch and is
expected to hand it back when it finishes. Two cleanup paths dropped the
rid_to_state entry without releasing: the abort echo in `_handle_abort_req`,
and the pre-dispatch failure path `_discard_pending_req_states` that a
400 Bad Request unwinds through.

One leaked reference keeps the registry's per-adapter counter above zero, so
`/unload_lora_adapter` waits in `wait_for_zero()` forever and every later
adapter swap hangs. That is fatal for colocated RL training, where each
weight sync unloads and reloads the trained adapter.

Couple removal and release into one helper, `_release_req_state()`, and use it
at all four cleanup sites so the two steps cannot drift apart again.
---
 .../sglang/srt/managers/tokenizer_manager.py  |  32 +++--
 .../test_tokenizer_manager_rid_cleanup.py     | 120 ++++++++++++++++++
 2 files changed, 139 insertions(+), 13 deletions(-)

diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py
index 0b435574b208..74bb013b5877 100644
--- a/python/sglang/srt/managers/tokenizer_manager.py
+++ b/python/sglang/srt/managers/tokenizer_manager.py
@@ -1490,6 +1490,7 @@ async def _handle_abort_finish_reason(
             finish_reason.get("type") == "abort"
             and finish_reason.get("status_code") == HTTPStatus.BAD_REQUEST
         ):
+            self._release_req_state(state.obj.rid)
             if not is_stream:
                 raise ValueError(finish_reason["message"])
             return out
@@ -1502,12 +1503,7 @@ async def _handle_abort_finish_reason(
         ):
             # Delete the key to prevent resending abort request to the scheduler and
             # to ensure aborted request state is cleaned up.
-            if state.obj.rid in self.rid_to_state:
-                del self.rid_to_state[state.obj.rid]
-
-            # Mark ongoing LoRA request as finished.
-            if self.enable_lora and state.obj.lora_path:
-                await self.lora_registry.release(state.obj.lora_id)
+            self._release_req_state(state.obj.rid)
             if not is_stream:
                 raise fastapi.HTTPException(
                     status_code=finish_reason["status_code"],
@@ -2217,11 +2213,7 @@ async def _handle_batch_output(
                         )
                     )
 
-                del self.rid_to_state[rid]
-
-                # Mark ongoing LoRA request as finished.
-                if self.enable_lora and state.obj.lora_path:
-                    asyncio.create_task(self.lora_registry.release(state.obj.lora_id))
+                self._release_req_state(rid)
 
             if out_dict is not None:
                 state.out_list.append(out_dict)
@@ -2879,7 +2871,7 @@ def _handle_abort_req(self, recv_obj: AbortReq):
             "output_ids": output_ids,
             "meta_info": meta_info,
         }
-        del self.rid_to_state[recv_obj.rid]
+        self._release_req_state(recv_obj.rid)
 
         state.out_list.append(out)
         state.event.set()
@@ -3080,6 +3072,20 @@ def _init_req_state(
                 time_stats.init_trace_ctx(rid, bootstrap_room, external_trace_header)
             time_stats.set_created_time(created_time)
 
+    def _release_req_state(self, rid: str) -> None:
+        """Drop the request's rid_to_state entry and release its LoRA reference.
+
+        Removal and release belong together: popping first keeps the release at
+        exactly one per request when several cleanup paths race for the same
+        rid, and a reference that is never handed back leaves
+        ``/unload_lora_adapter`` waiting on the registry counter forever.
+        """
+        state = self.rid_to_state.pop(rid, None)
+        if state is None:
+            return
+        if self.enable_lora and state.obj.lora_path:
+            asyncio.create_task(self.lora_registry.release(state.obj.lora_id))
+
     def _discard_pending_req_states(self, obj):
         """Drop rid_to_state entries created by _init_req_state for *obj*.
 
@@ -3092,7 +3098,7 @@ def _discard_pending_req_states(self, obj):
         else:
             rids = obj.rid
         for rid in rids:
-            self.rid_to_state.pop(rid, None)
+            self._release_req_state(rid)
 
     def _should_dispatch_to_encoder(
         self, obj: Union[GenerateReqInput, EmbeddingReqInput]
diff --git a/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py b/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py
index b023a22bcdda..c5b228032cce 100644
--- a/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py
+++ b/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py
@@ -10,10 +10,12 @@
   - _handle_batch_output cleans up rid_to_state on finished requests
   - _init_req_state rejects duplicate rids
   - Resubmission succeeds after cleanup
+  - Every cleanup path hands the request's LoRA reference back to the registry
 """
 
 import asyncio
 import unittest
+from http import HTTPStatus
 from unittest.mock import AsyncMock, MagicMock, Mock
 
 import msgspec
@@ -521,5 +523,123 @@ async def drive():
             self.assertNotIn(r, tm.rid_to_state)
 
 
+class _RecordingLoRARegistry:
+    """Stands in for LoRARegistry, recording which adapter ids were released."""
+
+    def __init__(self):
+        self.released = []
+
+    async def release(self, lora_id):
+        self.released.append(lora_id)
+
+
+def _enable_lora(tm: TokenizerManager) -> TokenizerManager:
+    tm.enable_lora = True
+    tm.lora_registry = _RecordingLoRARegistry()
+    return tm
+
+
+def _make_lora_req_state(rid: str, lora_id: str = "lora-1") -> ReqState:
+    state = _make_req_state(rid)
+    state.obj.lora_path = "adapter"
+    state.obj.lora_id = lora_id
+    return state
+
+
+class TestLoraReferenceRelease(CustomTestCase):
+    """Every path that drops a request must release its LoRA reference.
+
+    The registry counts in-flight requests per adapter and
+    /unload_lora_adapter waits for that count to reach zero, so a single
+    reference that is never handed back wedges every later adapter swap.
+    """
+
+    def test_abort_releases_reference(self):
+        tm = _enable_lora(_make_tokenizer_manager())
+        rid = "lora_abort_rid"
+        tm.rid_to_state[rid] = _make_lora_req_state(rid)
+
+        async def drive():
+            tm._handle_abort_req(_make_abort_req(rid))
+            await asyncio.sleep(0)
+
+        asyncio.run(drive())
+
+        self.assertNotIn(rid, tm.rid_to_state)
+        self.assertEqual(tm.lora_registry.released, ["lora-1"])
+
+    def test_bad_request_abort_releases_reference(self):
+        tm = _enable_lora(_make_tokenizer_manager())
+        rid = "lora_bad_request_rid"
+        state = _make_lora_req_state(rid)
+        tm.rid_to_state[rid] = state
+        out = {
+            "meta_info": {
+                "finish_reason": {
+                    "type": "abort",
+                    "status_code": HTTPStatus.BAD_REQUEST,
+                    "message": "rejected by the scheduler",
+                }
+            }
+        }
+
+        async def drive():
+            with self.assertRaises(ValueError):
+                await tm._handle_abort_finish_reason(out, state, is_stream=False)
+            await asyncio.sleep(0)
+
+        asyncio.run(drive())
+
+        self.assertNotIn(rid, tm.rid_to_state)
+        self.assertEqual(tm.lora_registry.released, ["lora-1"])
+
+    def test_dispatch_failure_releases_reference(self):
+        tm = _enable_lora(_make_tm_for_generate())
+        rid = "lora_overlen_rid"
+        obj = _make_generate_obj(rid, is_single=True)
+        obj.lora_path = "adapter"
+        obj.lora_id = "lora-1"
+        tm._tokenize_one_request = AsyncMock(side_effect=ValueError("input too long"))
+        tm._send_one_request = Mock()
+
+        async def drive():
+            with self.assertRaises(ValueError):
+                await tm.generate_request(obj).__anext__()
+            await asyncio.sleep(0)
+
+        asyncio.run(drive())
+
+        self.assertNotIn(rid, tm.rid_to_state)
+        self.assertEqual(tm.lora_registry.released, ["lora-1"])
+
+    def test_reference_released_once_when_cleanup_paths_race(self):
+        """A finish and a late abort echo for the same rid release only once."""
+        tm = _enable_lora(_make_tokenizer_manager())
+        rid = "lora_race_rid"
+        tm.rid_to_state[rid] = _make_lora_req_state(rid)
+
+        async def drive():
+            await tm._handle_batch_output(_make_batch_str_output(rid))
+            tm._handle_abort_req(_make_abort_req(rid))
+            await asyncio.sleep(0)
+
+        asyncio.run(drive())
+
+        self.assertEqual(tm.lora_registry.released, ["lora-1"])
+
+    def test_no_release_without_lora(self):
+        tm = _enable_lora(_make_tokenizer_manager())
+        rid = "no_lora_rid"
+        tm.rid_to_state[rid] = _make_req_state(rid)
+
+        async def drive():
+            tm._handle_abort_req(_make_abort_req(rid))
+            await asyncio.sleep(0)
+
+        asyncio.run(drive())
+
+        self.assertEqual(tm.lora_registry.released, [])
+
+
 if __name__ == "__main__":
     unittest.main(verbosity=2)

From ffc5748645125b1ce1760b029a180f955d9b5e28 Mon Sep 17 00:00:00 2001
From: "Ethan (Yusheng) Su" <yushengsu.thu@gmail.com>
Date: Mon, 3 Aug 2026 01:12:05 -0700
Subject: [PATCH 39/43] fix(lora): make request accounting lifecycle-safe

Port the exactly-once usage-counter lifecycle from #31808 onto sglang-miles, preserve #30913 stable upsert IDs, and balance native parallel multi-LoRA sampling across success, abort, and pre-dispatch failure paths.

Co-authored-by: Shafeeq <shafeeq@vmax.ai>
---
 .../sglang/srt/managers/tokenizer_manager.py  | 158 ++++++++++++++--
 .../test_tokenizer_manager_rid_cleanup.py     | 179 +++++++++++++++++-
 2 files changed, 314 insertions(+), 23 deletions(-)

diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py
index 74bb013b5877..2f31708ace98 100644
--- a/python/sglang/srt/managers/tokenizer_manager.py
+++ b/python/sglang/srt/managers/tokenizer_manager.py
@@ -182,6 +182,18 @@ class ReqState:
     last_completion_tokens: int = 1
     ttft_observed: bool = False
 
+    # A request may be observed as terminal by more than one path (normal
+    # output, an abort echo, an error finish reason, or dispatch cleanup).
+    # The LoRA registry reference must nevertheless be handed back once.
+    lora_released: bool = False
+
+    # Usually a state owns the single reference in ``obj.lora_id``. When an
+    # explicit rid list is used with parallel sampling, however, normalization
+    # keeps one parent rid per prompt while LoRARegistry.acquire has already
+    # acquired ``parallel_sample_num`` references per prompt. Keep that bundle
+    # on the parent until ownership is transferred to regenerated child states.
+    lora_ids_to_release: Optional[Union[str, List[Optional[str]]]] = None
+
     # An abort matched this rid while the request was still tokenizer-held
     # (parked at the pause gate / model-update lock / tokenization). The
     # scheduler never saw the rid, so the dispatch path must resolve the
@@ -453,6 +465,9 @@ def init_running_status(self):
         self.rid_to_state: Dict[str, ReqState] = {}
         self.event_loop = None
         self.asyncio_tasks = set()
+        # asyncio only keeps weak references to tasks. Keep release tasks alive
+        # until their registry decrement has completed.
+        self._lora_release_tasks: set[asyncio.Task] = set()
 
         # Health check
         self.server_status = ServerStatus.Starting
@@ -1490,7 +1505,10 @@ async def _handle_abort_finish_reason(
             finish_reason.get("type") == "abort"
             and finish_reason.get("status_code") == HTTPStatus.BAD_REQUEST
         ):
-            self._release_req_state(state.obj.rid)
+            self.rid_to_state.pop(state.obj.rid, None)
+            release_task = self._release_lora_once(state)
+            if release_task is not None:
+                await release_task
             if not is_stream:
                 raise ValueError(finish_reason["message"])
             return out
@@ -1503,7 +1521,11 @@ async def _handle_abort_finish_reason(
         ):
             # Delete the key to prevent resending abort request to the scheduler and
             # to ensure aborted request state is cleaned up.
-            self._release_req_state(state.obj.rid)
+            self.rid_to_state.pop(state.obj.rid, None)
+
+            release_task = self._release_lora_once(state)
+            if release_task is not None:
+                await release_task
             if not is_stream:
                 raise fastapi.HTTPException(
                     status_code=finish_reason["status_code"],
@@ -1513,6 +1535,38 @@ async def _handle_abort_finish_reason(
 
         return None
 
+    def _release_lora_once(self, state: ReqState) -> Optional[asyncio.Task]:
+        """Release a request's LoRA registry reference at most once.
+
+        A scheduler output, an abort echo, an error finish reason, and failed
+        dispatch cleanup may all observe the same terminal request. The
+        per-state guard makes those paths idempotent, while the task set keeps
+        the asynchronous registry decrement alive until it completes.
+
+        ``lora_id`` is the ownership signal. A request can carry ``lora_path``
+        before registry acquisition succeeds, so checking the path alone can
+        incorrectly schedule ``release(None)`` on validation failures.
+        """
+        if state.lora_released:
+            return None
+
+        lora_ids = state.lora_ids_to_release
+        if lora_ids is None:
+            lora_ids = getattr(state.obj, "lora_id", None)
+        if not getattr(self, "enable_lora", False) or lora_ids is None:
+            return None
+
+        if isinstance(lora_ids, list):
+            lora_ids = [lora_id for lora_id in lora_ids if lora_id is not None]
+            if not lora_ids:
+                return None
+
+        state.lora_released = True
+        task = asyncio.create_task(self.lora_registry.release(lora_ids))
+        self._lora_release_tasks.add(task)
+        task.add_done_callback(self._lora_release_tasks.discard)
+        return task
+
     async def _wait_one_response(
         self,
         obj: Union[GenerateReqInput, EmbeddingReqInput],
@@ -1676,9 +1730,27 @@ async def _handle_batch_request(
                 *(self._tokenize_one_request(obj) for obj in objs)
             )
 
+            # Every real sample below is sent under a regenerated rid. The
+            # original states are bookkeeping parents, not additional LoRA
+            # users: the acquire count belongs to the regenerated children.
+            # Remove all parents (including the otherwise orphaned expanded
+            # entries) and mark their accounting as transferred.
+            parent_states = {}
+            for parent_rid in obj.rid if isinstance(obj.rid, list) else [obj.rid]:
+                parent_state = self.rid_to_state.pop(parent_rid, None)
+                if parent_state is not None:
+                    parent_state.lora_released = True
+                    parent_states[parent_rid] = parent_state
+
             # Cache the common prefix for parallel sampling
             for i in range(batch_size):
                 tmp_obj = copy.copy(objs[i])
+                # This warm-up is an extra request beyond the n references
+                # acquired for the real samples. Keep LoRA on tokenized_obj so
+                # the scheduler caches the prefix under the right adapter, but
+                # make its tokenizer-side state a non-owner.
+                tmp_obj.lora_path = None
+                tmp_obj.lora_id = None
                 tokenized_obj = copy.copy(tokenized_objs[i])
                 # Ensure independent mm_items so wrap_shm_features won't mutate the original
                 if hasattr(tokenized_obj, "mm_inputs") and tokenized_obj.mm_inputs:
@@ -1715,8 +1787,8 @@ async def _handle_batch_request(
                     generators.append(self._wait_one_response(tmp_obj, request))
                     rids.append(tmp_obj.rid)
 
-                self.rid_to_state[objs[i].rid].time_stats.set_finished_time()
-                del self.rid_to_state[objs[i].rid]
+                if objs[i].rid in parent_states:
+                    parent_states[objs[i].rid].time_stats.set_finished_time()
 
         # Wait for all requests
         is_stream = hasattr(obj, "stream") and obj.stream
@@ -2213,7 +2285,8 @@ async def _handle_batch_output(
                         )
                     )
 
-                self._release_req_state(rid)
+                del self.rid_to_state[rid]
+                self._release_lora_once(state)
 
             if out_dict is not None:
                 state.out_list.append(out_dict)
@@ -2871,7 +2944,8 @@ def _handle_abort_req(self, recv_obj: AbortReq):
             "output_ids": output_ids,
             "meta_info": meta_info,
         }
-        self._release_req_state(recv_obj.rid)
+        del self.rid_to_state[recv_obj.rid]
+        self._release_lora_once(state)
 
         state.out_list.append(out)
         state.event.set()
@@ -3027,6 +3101,60 @@ async def _resolve_lora_path(self, obj: Union[GenerateReqInput, EmbeddingReqInpu
                 obj.lora_id[i] if isinstance(obj.lora_id, list) else obj.lora_id
             )
 
+        missing_rids, orphaned_lora_ids = self._assign_lora_release_ownership(obj)
+        if missing_rids:
+            # An abort ack can remove a tokenizer-held state while acquire() is
+            # awaiting the registry. In that race the abort cannot release yet
+            # (lora_id is still None), so return the newly acquired references
+            # here instead of leaking them.
+            if orphaned_lora_ids:
+                await self.lora_registry.release(orphaned_lora_ids)
+            raise ValueError(
+                "Request was aborted while resolving its LoRA adapter: "
+                + ", ".join(missing_rids)
+            )
+
+    def _assign_lora_release_ownership(
+        self, obj: Union[GenerateReqInput, EmbeddingReqInput]
+    ) -> Tuple[List[str], List[str]]:
+        """Assign every acquired LoRA reference to an initial request state.
+
+        Returns missing rids and the references that belonged to them. A rid
+        can disappear while ``LoRARegistry.acquire`` awaits; its references
+        must be released directly because no state remains to own them.
+        """
+        if not isinstance(obj.lora_id, list):
+            ownership = [(obj.rid, obj.lora_id)]
+        else:
+            rids = obj.rid if isinstance(obj.rid, list) else [obj.rid]
+            if len(rids) == len(obj.lora_id):
+                ownership = list(zip(rids, obj.lora_id))
+            else:
+                # Batch fields are expanded by repeating the whole batch, so
+                # prompt i owns i, i + batch_size, ... until fan-out. Using
+                # this partition for malformed input too keeps every acquired
+                # reference owned until later validation rejects the request.
+                ownership = [
+                    (rid, obj.lora_id[i :: len(rids)]) for i, rid in enumerate(rids)
+                ]
+
+        missing_rids = []
+        orphaned_lora_ids = []
+        for rid, lora_ids in ownership:
+            state = self.rid_to_state.get(rid)
+            if state is None:
+                missing_rids.append(rid)
+                if isinstance(lora_ids, list):
+                    orphaned_lora_ids.extend(
+                        lora_id for lora_id in lora_ids if lora_id is not None
+                    )
+                elif lora_ids is not None:
+                    orphaned_lora_ids.append(lora_ids)
+                continue
+            state.lora_ids_to_release = lora_ids
+
+        return missing_rids, orphaned_lora_ids
+
     def _init_req_state(
         self,
         obj: Union[GenerateReqInput, EmbeddingReqInput],
@@ -3072,20 +3200,6 @@ def _init_req_state(
                 time_stats.init_trace_ctx(rid, bootstrap_room, external_trace_header)
             time_stats.set_created_time(created_time)
 
-    def _release_req_state(self, rid: str) -> None:
-        """Drop the request's rid_to_state entry and release its LoRA reference.
-
-        Removal and release belong together: popping first keeps the release at
-        exactly one per request when several cleanup paths race for the same
-        rid, and a reference that is never handed back leaves
-        ``/unload_lora_adapter`` waiting on the registry counter forever.
-        """
-        state = self.rid_to_state.pop(rid, None)
-        if state is None:
-            return
-        if self.enable_lora and state.obj.lora_path:
-            asyncio.create_task(self.lora_registry.release(state.obj.lora_id))
-
     def _discard_pending_req_states(self, obj):
         """Drop rid_to_state entries created by _init_req_state for *obj*.
 
@@ -3098,7 +3212,9 @@ def _discard_pending_req_states(self, obj):
         else:
             rids = obj.rid
         for rid in rids:
-            self._release_req_state(rid)
+            state = self.rid_to_state.pop(rid, None)
+            if state is not None:
+                self._release_lora_once(state)
 
     def _should_dispatch_to_encoder(
         self, obj: Union[GenerateReqInput, EmbeddingReqInput]
diff --git a/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py b/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py
index c5b228032cce..5333a1b9a78d 100644
--- a/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py
+++ b/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py
@@ -16,6 +16,7 @@
 import asyncio
 import unittest
 from http import HTTPStatus
+from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock, Mock
 
 import msgspec
@@ -120,6 +121,7 @@ def _make_tokenizer_manager() -> TokenizerManager:
     tm.dump_requests_folder = ""
     tm.crash_dump_folder = ""
     tm.send_to_scheduler = MagicMock()
+    tm._lora_release_tasks = set()
     return tm
 
 
@@ -130,6 +132,7 @@ def _make_req_state(rid: str = "test_rid") -> ReqState:
     obj.stream = False
     obj.return_logprob = False
     obj.lora_path = None
+    obj.lora_id = None
     obj.log_metrics = False
     return ReqState(
         out_list=[],
@@ -524,13 +527,32 @@ async def drive():
 
 
 class _RecordingLoRARegistry:
-    """Stands in for LoRARegistry, recording which adapter ids were released."""
+    """Stands in for LoRARegistry, recording acquired and released IDs."""
 
     def __init__(self):
+        self.acquired = []
         self.released = []
 
+    async def get_unregistered_loras(self, lora_paths):
+        return []
+
+    async def acquire(self, lora_paths):
+        if isinstance(lora_paths, list):
+            lora_ids = [
+                f"id:{path}" if path is not None else None for path in lora_paths
+            ]
+            self.acquired.extend(lora_id for lora_id in lora_ids if lora_id is not None)
+            return lora_ids
+
+        lora_id = f"id:{lora_paths}"
+        self.acquired.append(lora_id)
+        return lora_id
+
     async def release(self, lora_id):
-        self.released.append(lora_id)
+        if isinstance(lora_id, list):
+            self.released.extend(item for item in lora_id if item is not None)
+        else:
+            self.released.append(lora_id)
 
 
 def _enable_lora(tm: TokenizerManager) -> TokenizerManager:
@@ -640,6 +662,159 @@ async def drive():
 
         self.assertEqual(tm.lora_registry.released, [])
 
+    def test_no_release_when_lora_acquire_never_succeeded(self):
+        tm = _enable_lora(_make_tokenizer_manager())
+        rid = "unacquired_lora_rid"
+        state = _make_req_state(rid)
+        state.obj.lora_path = "missing-adapter"
+        state.obj.lora_id = None
+        tm.rid_to_state[rid] = state
+
+        async def drive():
+            tm._handle_abort_req(_make_abort_req(rid))
+            await asyncio.sleep(0)
+
+        asyncio.run(drive())
+
+        self.assertEqual(tm.lora_registry.released, [])
+
+    def test_explicit_rids_release_every_parallel_acquire_on_failure(self):
+        """One parent rid can own n references until parallel fan-out."""
+
+        async def drive():
+            tm = _enable_lora(_make_tokenizer_manager())
+            tm.server_args.max_loaded_loras = None
+            obj = GenerateReqInput(
+                text=["hello", "world"],
+                rid=["parent-a", "parent-b"],
+                lora_path=["adapter-a", "adapter-b"],
+                sampling_params={"n": 3, "max_new_tokens": 8},
+            )
+            obj.received_time = 0.0
+            obj.normalize_batch_and_arguments()
+            tm._init_req_state(obj)
+            await tm._resolve_lora_path(obj)
+
+            # Simulate tokenization/validation failing before child states are
+            # created. Cleanup must balance all batch_size * n acquisitions.
+            tm._discard_pending_req_states(obj)
+            await asyncio.gather(*list(tm._lora_release_tasks))
+
+            self.assertCountEqual(
+                tm.lora_registry.released,
+                tm.lora_registry.acquired,
+            )
+            self.assertEqual(len(tm.lora_registry.released), 6)
+
+        asyncio.run(drive())
+
+    def test_abort_during_lora_acquire_releases_orphaned_reference(self):
+        """An abort racing acquire cannot strand the newly acquired ID."""
+
+        async def drive():
+            tm = _enable_lora(_make_tokenizer_manager())
+            tm.server_args.max_loaded_loras = None
+            obj = GenerateReqInput(
+                text="hello",
+                rid="acquire-race",
+                lora_path="adapter",
+            )
+            obj.received_time = 0.0
+            obj.normalize_batch_and_arguments()
+            tm._init_req_state(obj)
+
+            acquire_started = asyncio.Event()
+            finish_acquire = asyncio.Event()
+            original_acquire = tm.lora_registry.acquire
+
+            async def blocked_acquire(lora_path):
+                acquire_started.set()
+                await finish_acquire.wait()
+                return await original_acquire(lora_path)
+
+            tm.lora_registry.acquire = blocked_acquire
+            resolve_task = asyncio.create_task(tm._resolve_lora_path(obj))
+            await acquire_started.wait()
+            tm._handle_abort_req(_make_abort_req(obj.rid))
+            finish_acquire.set()
+
+            with self.assertRaisesRegex(ValueError, "aborted while resolving"):
+                await resolve_task
+
+            self.assertEqual(tm.lora_registry.acquired, ["id:adapter"])
+            self.assertEqual(tm.lora_registry.released, ["id:adapter"])
+
+        asyncio.run(drive())
+
+    def test_parallel_sampling_releases_only_real_samples(self):
+        """The prefix-cache warm-up and parent states do not own references."""
+
+        async def drive(text, rid, lora_path, lora_ids):
+            tm = _enable_lora(_make_tokenizer_manager())
+            obj = GenerateReqInput(
+                text=text,
+                rid=rid,
+                lora_path=lora_path,
+                sampling_params={"n": 4, "max_new_tokens": 8},
+            )
+            obj.received_time = 0.0
+            obj.normalize_batch_and_arguments()
+            tm._init_req_state(obj)
+
+            # Mirror _resolve_lora_path after acquiring four references.
+            obj.lora_id = lora_ids
+            for i, sub_obj in obj.__dict__.get("_sub_obj_cache", {}).items():
+                sub_obj.lora_id = obj.lora_id[i]
+            missing_rids, orphaned_lora_ids = tm._assign_lora_release_ownership(obj)
+            self.assertEqual(missing_rids, [])
+            self.assertEqual(orphaned_lora_ids, [])
+
+            async def tokenize(req):
+                return SimpleNamespace(
+                    input_ids=[1],
+                    mm_inputs=None,
+                    rid=req.rid,
+                    sampling_params=SimpleNamespace(max_new_tokens=8),
+                    stream=False,
+                )
+
+            async def finish(req, request=None):
+                state = tm.rid_to_state.pop(req.rid)
+                tm._release_lora_once(state)
+                yield {"meta_info": {"id": req.rid}}
+
+            tm._tokenize_one_request = tokenize
+            tm._send_one_request = Mock()
+            tm._wait_one_response = finish
+
+            outputs = [out async for out in tm._handle_batch_request(obj)]
+            await asyncio.sleep(0)
+
+            self.assertEqual(len(outputs), 1)
+            self.assertCountEqual(
+                tm.lora_registry.released,
+                lora_ids,
+            )
+            self.assertEqual(tm.rid_to_state, {})
+
+        # Both accepted rid shapes exercise different parent-state counts.
+        asyncio.run(
+            drive(
+                "hello",
+                None,
+                "adapter",
+                ["stable-upsert-id"] * 4,
+            )
+        )
+        asyncio.run(
+            drive(
+                ["hello", "world"],
+                ["parallel-a", "parallel-b"],
+                ["adapter-a", "adapter-b"],
+                ["stable-a", "stable-b"] * 4,
+            )
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

From e22f56c7fcddef5ffd82b5bba01b341cd694840d Mon Sep 17 00:00:00 2001
From: "Ethan (Yusheng) Su" <yushengsu@radixark.ai>
Date: Mon, 3 Aug 2026 17:33:46 -0700
Subject: [PATCH 40/43] Redesign multi-LoRA (#33299)

---
 python/sglang/srt/entrypoints/http_server.py  |   2 +
 python/sglang/srt/lora/lora_registry.py       |  47 ++--
 .../srt/managers/tokenizer_control_mixin.py   |  11 +-
 .../sglang/srt/managers/tokenizer_manager.py  |  56 ++++-
 test/registered/unit/lora/test_lora_lease.py  | 233 ++++++++++++++++++
 .../managers/test_abort_request_prefix.py     |  12 +-
 6 files changed, 334 insertions(+), 27 deletions(-)
 create mode 100644 test/registered/unit/lora/test_lora_lease.py

diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py
index 5b7de28d519c..4b4a34674ef4 100644
--- a/python/sglang/srt/entrypoints/http_server.py
+++ b/python/sglang/srt/entrypoints/http_server.py
@@ -1558,6 +1558,7 @@ async def load_lora_adapter(
 
 
 @app.api_route("/load_lora_adapter_from_tensors", methods=["POST"])
+@auth_level(AuthLevel.ADMIN_OPTIONAL)
 async def load_lora_adapter_from_tensors(
     obj: Annotated[LoadLoRAAdapterFromTensorsReqInput, Body()], request: Request
 ):
@@ -1570,6 +1571,7 @@ async def load_lora_adapter_from_tensors(
 
 
 @app.api_route("/load_lora_adapter_from_distributed", methods=["POST"])
+@auth_level(AuthLevel.ADMIN_OPTIONAL)
 async def load_lora_adapter_from_distributed(
     obj: Annotated[LoadLoRAAdapterFromDistributedReqInput, Body()], request: Request
 ):
diff --git a/python/sglang/srt/lora/lora_registry.py b/python/sglang/srt/lora/lora_registry.py
index 8ee933a8c402..ae6242b1ecbf 100644
--- a/python/sglang/srt/lora/lora_registry.py
+++ b/python/sglang/srt/lora/lora_registry.py
@@ -38,6 +38,12 @@ class LoRARef(msgspec.Struct, frozen=True, array_like=True):
     lora_name: Optional[str] = None
     lora_path: Optional[str] = None
     pinned: Optional[bool] = None
+    # False for adapters whose weights arrived over the wire (lora_path
+    # "__distributed__" / "__tensor__"): there is no disk artifact to reload
+    # from, so they must never be LRU-evicted nor implicitly reloaded.
+    # Trailing field with a default keeps the array_like wire format
+    # compatible with refs encoded before this field existed.
+    reloadable: bool = True
 
     def __post_init__(self):
         if self.lora_id is None:
@@ -184,24 +190,28 @@ def _lookup(name: str) -> str:
             self._registry.move_to_end(name)
             return lora_ref.lora_id
 
+        # Lookup and increment must be one atomic admission step under the
+        # writer lock: with the increment outside, an unload can slip between
+        # them — unregister the adapter, observe the counter at zero, delete
+        # the counter — leaving this request to KeyError or to run on an
+        # adapter that is already gone. wait_for_unload never awaits inside
+        # this lock, so holding it across the increment cannot deadlock.
         if isinstance(lora_name, str):
             async with self._registry_lock.writer_lock:
                 lora_id = _lookup(lora_name)
-
-            await self._counters[lora_id].increment(notify_all=False)
+                await self._counters[lora_id].increment(notify_all=False)
             return lora_id
         elif isinstance(lora_name, list):
             async with self._registry_lock.writer_lock:
                 lora_ids = [_lookup(name) for name in lora_name]
-
-            # Increment the counters only after all IDs are looked up.
-            await asyncio.gather(
-                *[
-                    self._counters[id].increment(notify_all=False)
-                    for id in lora_ids
-                    if id is not None
-                ]
-            )
+                # Increment the counters only after all IDs are looked up.
+                await asyncio.gather(
+                    *[
+                        self._counters[id].increment(notify_all=False)
+                        for id in lora_ids
+                        if id is not None
+                    ]
+                )
             return lora_ids
         else:
             raise TypeError("lora_name must be either a string or a list of strings.")
@@ -266,14 +276,15 @@ async def lru_lora_name(self, exclude_pinned=False):
         If exclude_pinned is True, then return the LRU LoRA adapter that isn't pinned.
         """
         async with self._registry_lock.reader_lock:
-            if not exclude_pinned:
-                return next(iter(self._registry), None)
-
             for lora_name, lora_ref in self._registry.items():
-                if not lora_ref.pinned:
-                    return lora_name
-            else:
-                return None
+                # Evicting a non-reloadable (wire-loaded) adapter is always
+                # destructive: there is no artifact to reload it from.
+                if not lora_ref.reloadable:
+                    continue
+                if exclude_pinned and lora_ref.pinned:
+                    continue
+                return lora_name
+            return None
 
     def _register_adapter(self, lora_ref: LoRARef):
         """
diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py
index 1d3ff8d6ddd2..9894a6605ecb 100644
--- a/python/sglang/srt/managers/tokenizer_control_mixin.py
+++ b/python/sglang/srt/managers/tokenizer_control_mixin.py
@@ -735,6 +735,7 @@ async def load_lora_adapter_from_tensors(
                     lora_name=obj.lora_name,
                     lora_path="__tensor__",
                     pinned=obj.pinned,
+                    reloadable=False,
                 )
                 obj.lora_id = new_adapter.lora_id
                 result = (await self.update_lora_adapter_communicator(obj))[0]
@@ -811,6 +812,7 @@ async def load_lora_adapter_from_distributed(
                         lora_name=obj.lora_name,
                         lora_path="__distributed__",
                         pinned=obj.pinned,
+                        reloadable=False,
                     ),
                     upsert=obj.upsert,
                 )
@@ -886,7 +888,14 @@ async def unload_lora_adapter(
             )
 
             async with self.lora_update_lock:
-                return await self._unload_lora_adapter_locked(obj)
+                result = await self._unload_lora_adapter_locked(obj)
+                # Explicit unload is a DELETE: drop the reload-catalog entry too.
+                # The max_loaded_loras LRU loop calls _unload_lora_adapter_locked
+                # directly — an EVICT — and must keep the entry so disk-backed
+                # adapters can be implicitly reloaded later.
+                if result.success:
+                    self.lora_ref_cache.pop(obj.lora_name, None)
+                return result
         except ValueError as e:
             return UnloadLoRAAdapterReqOutput(success=False, error_message=str(e))
 
diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py
index 0b435574b208..a6b83a7b07e8 100644
--- a/python/sglang/srt/managers/tokenizer_manager.py
+++ b/python/sglang/srt/managers/tokenizer_manager.py
@@ -188,6 +188,13 @@ class ReqState:
     # request as aborted instead of sending it.
     abort_before_dispatch: bool = False
 
+    # The LoRA lease for this request has already been released. Every terminal
+    # path funnels through TokenizerManager._finalize_lora_lease, which uses
+    # this flag to stay idempotent: a leaked release hangs unload's
+    # wait_for_zero forever, and a double release drives the ConcurrentCounter
+    # negative, which hangs it just the same (it waits for exactly zero).
+    lora_lease_released: bool = False
+
     # For streaming output
     last_output_offset: int = 0
 
@@ -1473,6 +1480,25 @@ def _coalesce_streaming_chunks(
             out["meta_info"] = meta_info
         return out
 
+    def _finalize_lora_lease(self, state: Optional[ReqState]) -> None:
+        """Release the request's LoRA lease exactly once, however it terminates:
+        normal finish, scheduler abort echo (queued / tokenizer-held / disagg),
+        status-code abort, or a failed dispatch. A request whose LoRA was never
+        acquired (pre-acquire failure leaves lora_id unset) has no lease to
+        release."""
+        if state is None or state.lora_lease_released:
+            return
+        # lora_path / lora_id are declared on both input structs with default
+        # None (= base model / lease never acquired), so plain None checks are
+        # the contract. State-derived checks come first: a request without a
+        # lease has nothing to release, whatever the server config says.
+        if state.obj.lora_path is None or state.obj.lora_id is None:
+            return
+        if not self.enable_lora:
+            return
+        state.lora_lease_released = True
+        asyncio.create_task(self.lora_registry.release(state.obj.lora_id))
+
     async def _handle_abort_finish_reason(
         self,
         out: dict,
@@ -1505,9 +1531,10 @@ async def _handle_abort_finish_reason(
             if state.obj.rid in self.rid_to_state:
                 del self.rid_to_state[state.obj.rid]
 
-            # Mark ongoing LoRA request as finished.
-            if self.enable_lora and state.obj.lora_path:
-                await self.lora_registry.release(state.obj.lora_id)
+            # Status-code aborts also arrive through _handle_batch_output, which
+            # already released the lease and deleted the state — the finalizer's
+            # idempotency is what prevents the counter from going negative here.
+            self._finalize_lora_lease(state)
             if not is_stream:
                 raise fastapi.HTTPException(
                     status_code=finish_reason["status_code"],
@@ -2220,8 +2247,7 @@ async def _handle_batch_output(
                 del self.rid_to_state[rid]
 
                 # Mark ongoing LoRA request as finished.
-                if self.enable_lora and state.obj.lora_path:
-                    asyncio.create_task(self.lora_registry.release(state.obj.lora_id))
+                self._finalize_lora_lease(state)
 
             if out_dict is not None:
                 state.out_list.append(out_dict)
@@ -2879,6 +2905,10 @@ def _handle_abort_req(self, recv_obj: AbortReq):
             "output_ids": output_ids,
             "meta_info": meta_info,
         }
+        # This abort echo is the scheduler-side terminal ACK for queued,
+        # tokenizer-held, and disagg-retracted requests — none of them reach
+        # _handle_batch_output, so this is their only lease-release point.
+        self._finalize_lora_lease(state)
         del self.rid_to_state[recv_obj.rid]
 
         state.out_list.append(out)
@@ -3010,8 +3040,15 @@ async def _resolve_lora_path(self, obj: Union[GenerateReqInput, EmbeddingReqInpu
                     f"All loaded adapters: {self.lora_ref_cache.keys()}."
                 )
 
-            logger.info(f"Reloading evicted adapter: {lora_path}")
             new_lora_ref = self.lora_ref_cache[lora_path]
+            if not new_lora_ref.reloadable:
+                raise ValueError(
+                    f"LoRA adapter '{lora_path}' was loaded over the wire "
+                    f"(lora_path={new_lora_ref.lora_path!r}) and has no disk "
+                    "artifact to reload from; it must be re-pushed by the "
+                    "trainer instead of implicitly reloaded."
+                )
+            logger.info(f"Reloading evicted adapter: {lora_path}")
             load_result = await self.load_lora_adapter(
                 LoadLoRAAdapterReqInput(
                     lora_name=new_lora_ref.lora_name,
@@ -3092,7 +3129,12 @@ def _discard_pending_req_states(self, obj):
         else:
             rids = obj.rid
         for rid in rids:
-            self.rid_to_state.pop(rid, None)
+            state = self.rid_to_state.pop(rid, None)
+            # A failed/partial dispatch may never produce a scheduler terminal
+            # for this rid, and once the state is dropped a late terminal has no
+            # release point either — so release here. Pre-acquire failures are
+            # covered by the finalizer's lora_id check.
+            self._finalize_lora_lease(state)
 
     def _should_dispatch_to_encoder(
         self, obj: Union[GenerateReqInput, EmbeddingReqInput]
diff --git a/test/registered/unit/lora/test_lora_lease.py b/test/registered/unit/lora/test_lora_lease.py
new file mode 100644
index 000000000000..8ab54f161b0b
--- /dev/null
+++ b/test/registered/unit/lora/test_lora_lease.py
@@ -0,0 +1,233 @@
+"""Unit tests for the exactly-once LoRA lease lifecycle.
+
+A leaked lease leaves the adapter's in-flight counter above zero forever, so
+unload's wait_for_zero never returns; a double release drives the counter
+negative, which hangs it just the same (it waits for exactly zero). Covers:
+
+  * TokenizerManager._finalize_lora_lease is the single, idempotent release
+    owner for every terminal path
+  * pre-acquire failures (lora_id never set) release nothing
+  * LoRARegistry.acquire: lookup + counter increment are one atomic admission
+    step under the writer lock, so acquire racing an unload fails cleanly
+    instead of touching a deleted counter
+  * LoRARegistry.lru_lora_name never picks non-reloadable (wire-loaded) refs
+  * explicit unload drops the lora_ref_cache entry (DELETE) while the locked
+    helper alone (the max_loaded_loras EVICT path) keeps it
+"""
+
+import asyncio
+import unittest
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock, Mock
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
+
+maybe_stub_sgl_kernel()
+
+from sglang.srt.lora.lora_registry import LoRARef, LoRARegistry
+from sglang.srt.managers.io_struct import UnloadLoRAAdapterReqInput
+from sglang.srt.managers.tokenizer_manager import ReqState, TokenizerManager
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+def _make_tm(enable_lora: bool = True) -> TokenizerManager:
+    """TokenizerManager with only the fields the lease path reads."""
+    tm = TokenizerManager.__new__(TokenizerManager)
+    tm.enable_lora = enable_lora
+    tm.lora_registry = MagicMock()
+    tm.lora_registry.release = AsyncMock()
+    return tm
+
+
+def _make_state(lora_path="adapter-a", lora_id="id-1") -> ReqState:
+    obj = SimpleNamespace(lora_path=lora_path, lora_id=lora_id, rid="rid-1")
+    return ReqState([], False, asyncio.Event(), obj, Mock())
+
+
+class TestFinalizeLoraLease(CustomTestCase):
+    def _finalize(self, tm, state, times=1):
+        async def run():
+            for _ in range(times):
+                tm._finalize_lora_lease(state)
+            # Let the create_task'd release coroutines run.
+            await asyncio.sleep(0)
+
+        asyncio.run(run())
+
+    def test_release_exactly_once_on_double_finalize(self):
+        tm = _make_tm()
+        state = _make_state()
+        self._finalize(tm, state, times=3)
+        tm.lora_registry.release.assert_awaited_once_with("id-1")
+        self.assertTrue(state.lora_lease_released)
+
+    def test_pre_acquire_failure_releases_nothing(self):
+        # _init_req_state runs before the LoRA acquire; a failure in between
+        # leaves lora_id unset and there is no lease to release.
+        tm = _make_tm()
+        state = _make_state(lora_id=None)
+        self._finalize(tm, state)
+        tm.lora_registry.release.assert_not_awaited()
+        self.assertFalse(state.lora_lease_released)
+
+    def test_no_lora_path_is_noop(self):
+        tm = _make_tm()
+        state = _make_state(lora_path=None)
+        self._finalize(tm, state)
+        tm.lora_registry.release.assert_not_awaited()
+
+    def test_lora_disabled_is_noop(self):
+        tm = _make_tm(enable_lora=False)
+        state = _make_state()
+        self._finalize(tm, state)
+        tm.lora_registry.release.assert_not_awaited()
+
+    def test_none_state_is_noop(self):
+        tm = _make_tm()
+
+        async def run():
+            tm._finalize_lora_lease(None)
+
+        asyncio.run(run())
+        tm.lora_registry.release.assert_not_awaited()
+
+
+class TestAcquireAtomicity(CustomTestCase):
+    def test_acquire_after_unregister_raises_value_error(self):
+        async def run():
+            registry = LoRARegistry()
+            ref = LoRARef(lora_name="a", lora_path="/x")
+            await registry.register(ref)
+            await registry.unregister("a")
+            with self.assertRaises(ValueError):
+                await registry.acquire("a")
+
+        asyncio.run(run())
+
+    def test_unload_waits_for_inflight_lease_then_completes(self):
+        async def run():
+            registry = LoRARegistry()
+            ref = LoRARef(lora_name="a", lora_path="/x")
+            await registry.register(ref)
+            lora_id = await registry.acquire("a")
+            await registry.unregister("a")
+
+            wait_task = asyncio.create_task(registry.wait_for_unload(lora_id))
+            await asyncio.sleep(0.01)
+            self.assertFalse(wait_task.done(), "unload must wait for the lease")
+
+            await registry.release(lora_id)
+            await asyncio.wait_for(wait_task, timeout=2)
+
+        asyncio.run(run())
+
+    def test_batch_acquire_release_roundtrip(self):
+        async def run():
+            registry = LoRARegistry()
+            for name in ("a", "b"):
+                await registry.register(LoRARef(lora_name=name, lora_path=f"/{name}"))
+            lora_ids = await registry.acquire(["a", "b", None])
+            self.assertEqual(len(lora_ids), 3)
+            self.assertIsNone(lora_ids[2])
+            await registry.release([i for i in lora_ids if i is not None])
+            for name in ("a", "b"):
+                await registry.unregister(name)
+
+        asyncio.run(run())
+
+
+class TestLruSkipsNonReloadable(CustomTestCase):
+    def test_lru_prefers_reloadable_even_if_older_entry_is_wire_loaded(self):
+        async def run():
+            registry = LoRARegistry()
+            # Registration order = LRU order; the wire-loaded ref is oldest.
+            await registry.register(
+                LoRARef(lora_name="wire", lora_path="__distributed__", reloadable=False)
+            )
+            await registry.register(LoRARef(lora_name="disk", lora_path="/d"))
+            self.assertEqual(await registry.lru_lora_name(), "disk")
+
+        asyncio.run(run())
+
+    def test_all_non_reloadable_yields_no_victim(self):
+        async def run():
+            registry = LoRARegistry()
+            await registry.register(
+                LoRARef(lora_name="w1", lora_path="__tensor__", reloadable=False)
+            )
+            await registry.register(
+                LoRARef(lora_name="w2", lora_path="__distributed__", reloadable=False)
+            )
+            self.assertIsNone(await registry.lru_lora_name())
+            self.assertIsNone(await registry.lru_lora_name(exclude_pinned=True))
+
+        asyncio.run(run())
+
+    def test_exclude_pinned_still_respected(self):
+        async def run():
+            registry = LoRARegistry()
+            await registry.register(
+                LoRARef(lora_name="pinned-disk", lora_path="/p", pinned=True)
+            )
+            await registry.register(LoRARef(lora_name="plain-disk", lora_path="/q"))
+            self.assertEqual(
+                await registry.lru_lora_name(exclude_pinned=True), "plain-disk"
+            )
+            self.assertEqual(await registry.lru_lora_name(), "pinned-disk")
+
+        asyncio.run(run())
+
+
+class TestUnloadRefCacheCleanup(CustomTestCase):
+    def _make_unload_tm(self, success: bool) -> TokenizerManager:
+        tm = TokenizerManager.__new__(TokenizerManager)
+        tm.auto_create_handle_loop = Mock()
+        tm.server_args = MagicMock()
+        tm.server_args.enable_lora = True
+        tm.server_args.dp_size = 1
+        tm.lora_update_lock = asyncio.Lock()
+        tm._unload_lora_adapter_locked = AsyncMock(
+            return_value=SimpleNamespace(success=success)
+        )
+        tm.lora_ref_cache = {"a": LoRARef(lora_name="a", lora_path="/x")}
+        return tm
+
+    def test_explicit_unload_drops_ref_cache_entry(self):
+        tm = self._make_unload_tm(success=True)
+        asyncio.run(tm.unload_lora_adapter(UnloadLoRAAdapterReqInput(lora_name="a")))
+        self.assertNotIn("a", tm.lora_ref_cache)
+
+    def test_failed_unload_keeps_ref_cache_entry(self):
+        tm = self._make_unload_tm(success=False)
+        asyncio.run(tm.unload_lora_adapter(UnloadLoRAAdapterReqInput(lora_name="a")))
+        self.assertIn("a", tm.lora_ref_cache)
+
+    def test_locked_helper_alone_keeps_ref_cache_entry(self):
+        # The max_loaded_loras LRU loop calls _unload_lora_adapter_locked
+        # directly — an EVICT — and a disk-backed adapter must keep its
+        # reload-catalog entry or it can never be implicitly reloaded.
+        tm = TokenizerManager.__new__(TokenizerManager)
+        tm.lora_update_lock = asyncio.Lock()
+        tm.lora_registry = MagicMock()
+        tm.lora_registry.unregister = AsyncMock(return_value="id-a")
+        tm.lora_registry.wait_for_unload = AsyncMock()
+        tm.update_lora_adapter_communicator = AsyncMock(
+            return_value=[SimpleNamespace(success=True)]
+        )
+        tm.lora_ref_cache = {"a": LoRARef(lora_name="a", lora_path="/x")}
+
+        async def run():
+            async with tm.lora_update_lock:
+                return await tm._unload_lora_adapter_locked(
+                    UnloadLoRAAdapterReqInput(lora_name="a")
+                )
+
+        result = asyncio.run(run())
+        self.assertTrue(result.success)
+        self.assertIn("a", tm.lora_ref_cache)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/registered/unit/managers/test_abort_request_prefix.py b/test/registered/unit/managers/test_abort_request_prefix.py
index 5734193e845e..36e83d046150 100644
--- a/test/registered/unit/managers/test_abort_request_prefix.py
+++ b/test/registered/unit/managers/test_abort_request_prefix.py
@@ -34,6 +34,7 @@ def _make_tokenizer_manager(rids=(), tokenizer_worker_num=1) -> TokenizerManager
     tm.server_args = MagicMock()
     tm.server_args.tokenizer_worker_num = tokenizer_worker_num
     tm.enable_metrics = False
+    tm.enable_lora = False
     tm.rid_to_state = {rid: Mock() for rid in rids}
     tm.send_to_scheduler = MagicMock()
     tm.tokenizer_ipc_name = None
@@ -104,7 +105,11 @@ def test_multi_tokenizer_worker_skips_local_check(self):
 
 
 def _make_state(rid: str) -> ReqState:
-    obj = SimpleNamespace(rid=rid, stream=False, return_logprob=False)
+    # lora_path / lora_id model the real input-struct contract: always present,
+    # None = base model (the lease finalizer does plain None checks on them).
+    obj = SimpleNamespace(
+        rid=rid, stream=False, return_logprob=False, lora_path=None, lora_id=None
+    )
     return ReqState([], False, asyncio.Event(), obj, MagicMock())
 
 
@@ -209,10 +214,15 @@ def _make_scheduler(waiting_rids=(), running_rids=(), chunked_rid=None):
     sched.chunked_req = FakeReq(chunked_rid) if chunked_rid is not None else None
     sched.waiting_queue = [FakeReq(rid) for rid in waiting_rids]
     sched.enable_hicache_storage = False
+    sched.dllm_config = None  # abort_request reads it since the dLLM rework (#27877)
     sched.disaggregation_mode = DisaggregationMode.NULL
     sched.grammar_manager = MagicMock()
     sched.running_batch = SimpleNamespace(reqs=[FakeReq(rid) for rid in running_rids])
     sched.cur_batch = None
+    # abort_request collects in-flight reqs from (running_batch, last_batch) on
+    # the pp_size == 1 path since the v0.5.16 rebase.
+    sched.ps = SimpleNamespace(pp_size=1)
+    sched.last_batch = None
     sched.ipc_channels = MagicMock()
     return sched
 

From 31ad05a54d8074072236f0fd6785eb36030b8a3e Mon Sep 17 00:00:00 2001
From: Yueming Yuan <yym022502@gmail.com>
Date: Mon, 3 Aug 2026 21:28:09 -0700
Subject: [PATCH 41/43] Fix NextN weight loading by sharing the unified-loader
 mappings (#33478)

---
 python/sglang/srt/models/deepseek_nextn.py |  1 +
 python/sglang/srt/models/deepseek_v2.py    | 70 ++++++++++++----------
 2 files changed, 41 insertions(+), 30 deletions(-)

diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py
index 222effb03b07..56286d919cf8 100644
--- a/python/sglang/srt/models/deepseek_nextn.py
+++ b/python/sglang/srt/models/deepseek_nextn.py
@@ -357,6 +357,7 @@ def __init__(
         # if not set, model load will be broken in DeepseekV3ForCausalLM load_weights()
         self.pp_group = get_pp_group()
         self.determine_num_fused_shared_experts("DeepseekV3ForCausalLMNextN")
+        self.init_unified_loader_mappings()
         self.use_dsa = is_deepseek_dsa(config)
         self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
         self.mla_enable_prefill_cp = is_mla_prefill_cp_enabled() and not self.use_dsa
diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py
index 80671b39b783..1ead4db52f60 100644
--- a/python/sglang/srt/models/deepseek_v2.py
+++ b/python/sglang/srt/models/deepseek_v2.py
@@ -2736,36 +2736,7 @@ def __init__(
             self.lm_head = PPMissingLayer()
         self.logits_processor = LogitsProcessor(config)
 
-        # Stacked params mapping for unified weight loading API
-        self.stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("gate_up_proj", "gate_proj", 0),
-            ("gate_up_proj", "up_proj", 1),
-        ]
-        # Add A-proj fusion mapping when q_lora_rank is enabled
-        # q_a_proj + kv_a_proj_with_mqa -> fused_qkv_a_proj_with_mqa
-        if self.fuse_qkv_a_proj:
-            self.stacked_params_mapping.extend(
-                [
-                    ("fused_qkv_a_proj_with_mqa", "q_a_proj", 0),
-                    ("fused_qkv_a_proj_with_mqa", "kv_a_proj_with_mqa", 1),
-                ]
-            )
-        self.expert_params_mapping = FusedMoE.make_expert_params_mapping(
-            ckpt_gate_proj_name="gate_proj",
-            ckpt_down_proj_name="down_proj",
-            ckpt_up_proj_name="up_proj",
-            num_experts=self.config.n_routed_experts + self.num_fused_shared_experts,
-        )
-        # Params for special naming rules in mixed-precision models, for example:
-        # model.layers.xx.mlp.experts.xx.w1.input_scale. For details,
-        # see https://huggingface.co/Barrrrry/DeepSeek-R1-W4AFP8/blob/main.
-        if is_wint4afp8_or_wint4a16_config(self.quant_config):
-            self.expert_params_mapping += (
-                FusedMoE.make_expert_input_scale_params_mapping(
-                    num_experts=self.config.n_routed_experts
-                )
-            )
+        self.init_unified_loader_mappings()
 
         self._routed_experts_weights_of_layer = LazyValue(
             lambda: {
@@ -2809,6 +2780,45 @@ def custom_scale_remap(self, name: str) -> str:
     def routed_experts_weights_of_layer(self):
         return self._routed_experts_weights_of_layer.value
 
+    def init_unified_loader_mappings(self):
+        """Mappings the unified weight loading API reads off the model.
+
+        Call after determine_num_fused_shared_experts(); subclasses that build
+        their own __init__ must call this or weight loading fails on the
+        missing attributes.
+        """
+        self.fuse_qkv_a_proj = (
+            hasattr(self.config, "q_lora_rank") and self.config.q_lora_rank is not None
+        )
+        self.stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+        # q_a_proj + kv_a_proj_with_mqa -> fused_qkv_a_proj_with_mqa
+        if self.fuse_qkv_a_proj:
+            self.stacked_params_mapping.extend(
+                [
+                    ("fused_qkv_a_proj_with_mqa", "q_a_proj", 0),
+                    ("fused_qkv_a_proj_with_mqa", "kv_a_proj_with_mqa", 1),
+                ]
+            )
+        self.expert_params_mapping = FusedMoE.make_expert_params_mapping(
+            ckpt_gate_proj_name="gate_proj",
+            ckpt_down_proj_name="down_proj",
+            ckpt_up_proj_name="up_proj",
+            num_experts=self.config.n_routed_experts + self.num_fused_shared_experts,
+        )
+        # Params for special naming rules in mixed-precision models, for example:
+        # model.layers.xx.mlp.experts.xx.w1.input_scale. For details,
+        # see https://huggingface.co/Barrrrry/DeepSeek-R1-W4AFP8/blob/main.
+        if is_wint4afp8_or_wint4a16_config(self.quant_config):
+            self.expert_params_mapping += (
+                FusedMoE.make_expert_input_scale_params_mapping(
+                    num_experts=self.config.n_routed_experts
+                )
+            )
+
     def determine_num_fused_shared_experts(
         self, architecture: str = "DeepseekV3ForCausalLM"
     ):

From c447264eba460045e6e961fd7cf1bfbaf038db6b Mon Sep 17 00:00:00 2001
From: Yueming Yuan <yym022502@gmail.com>
Date: Mon, 3 Aug 2026 21:28:19 -0700
Subject: [PATCH 42/43] Put the DSA cuda-graph page table in the pausable
 memory region (#33479)

---
 .../srt/layers/attention/dsa_backend.py       | 61 +++++++++++++------
 1 file changed, 41 insertions(+), 20 deletions(-)

diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py
index f8bfa0eadb3f..dd58570596c3 100644
--- a/python/sglang/srt/layers/attention/dsa_backend.py
+++ b/python/sglang/srt/layers/attention/dsa_backend.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+from contextlib import nullcontext
 from dataclasses import dataclass
 from typing import (
     TYPE_CHECKING,
@@ -15,6 +16,7 @@
 import torch
 
 from sglang.srt.configs.model_config import get_dsa_index_topk, is_deepseek_dsa
+from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH
 from sglang.srt.runtime_context import get_parallel
 
 logger = logging.getLogger(__name__)
@@ -70,6 +72,7 @@
     is_sm100_supported,
     print_warning_once,
 )
+from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
 
 # Opt-in (default off): route the fp8 sparse-MLA prefill path through the Triton
 # per-query flash kernel instead of TileLang. Validated on gfx950 (GLM-5.1 @
@@ -365,6 +368,7 @@ def __init__(
         )
         self.dsa_index_topk = get_dsa_index_topk(model_runner.model_config.hf_config)
         self.max_context_len = model_runner.model_config.context_len
+        self.enable_memory_saver = model_runner.server_args.enable_memory_saver
         self.num_q_heads = (
             model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
         )
@@ -1129,6 +1133,16 @@ def _cal_indexer_k_start_end(
             token_to_batch_idx = dsa_cp_round_robin_split_data(token_to_batch_idx)
         return (ks, ke), token_to_batch_idx
 
+    def _cuda_graph_memory_region(self):
+        """Region whose tensors are released while the engine is paused."""
+        adapter = TorchMemorySaverAdapter.create(
+            enable=self.enable_memory_saver
+            and envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get()
+        )
+        if not adapter.enabled:
+            return nullcontext()
+        return adapter.region(tag=GPU_MEMORY_TYPE_CUDA_GRAPH)
+
     def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
         """Initialize CUDA graph state for the attention backend.
 
@@ -1164,6 +1178,31 @@ def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
         )
 
         max_ctx_len = self.req_to_token.shape[1]
+        # The wide page_table (max_num_tokens x max_ctx_len int32) and the
+        # flashmla metadata are rewritten in full before every replay, so they
+        # can live in the pausable cuda-graph region. The others are written
+        # once at init and would come back zeroed after a resume.
+        with self._cuda_graph_memory_region():
+            page_table = (
+                None
+                if self.dsa_drop_wide_page_table
+                else torch.zeros(
+                    max_num_tokens,
+                    max_ctx_len,
+                    dtype=torch.int32,
+                    device=self.device,
+                )
+            )
+            flashmla_metadata = (
+                self._compute_flashmla_metadata(
+                    cache_seqlens=torch.ones(
+                        max_num_tokens, dtype=torch.int32, device=self.device
+                    ),
+                    seq_len_q=1,
+                )
+                if self.dsa_decode_impl == "flashmla_kv"
+                else None
+            )
         self.decode_cuda_graph_metadata: Dict = {
             "cache_seqlens": torch.ones(
                 max_num_tokens, dtype=torch.int32, device=self.device
@@ -1190,26 +1229,8 @@ def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
                 if self.dsa_drop_wide_page_table
                 else None
             ),
-            "page_table": (
-                None
-                if self.dsa_drop_wide_page_table
-                else torch.zeros(
-                    max_num_tokens,
-                    max_ctx_len,
-                    dtype=torch.int32,
-                    device=self.device,
-                )
-            ),
-            "flashmla_metadata": (
-                self._compute_flashmla_metadata(
-                    cache_seqlens=torch.ones(
-                        max_num_tokens, dtype=torch.int32, device=self.device
-                    ),
-                    seq_len_q=1,
-                )
-                if self.dsa_decode_impl == "flashmla_kv"
-                else None
-            ),
+            "page_table": page_table,
+            "flashmla_metadata": flashmla_metadata,
         }
 
     def _build_forward_metadata_cuda_graph(

From cb05a44f35a7c9e27e46d74112cc841ca674ef43 Mon Sep 17 00:00:00 2001
From: Zhichen Zeng <zczeng@uw.edu>
Date: Tue, 4 Aug 2026 22:49:06 -0700
Subject: [PATCH 43/43] [sglang-miles] Nemotron support on sglang-miles
 (#27110)

Co-authored-by: yueming-yuan <yym022502@gmail.com>
---
 .../srt/layers/attention/mamba/mamba2_metadata.py  | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py
index bffe8d100e8a..df52f7fe857f 100644
--- a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py
+++ b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py
@@ -234,10 +234,16 @@ def prepare_mixed(
             num_prefill_tokens = int(sum(extend_seq_lens_cpu))
         else:
             num_prefill_tokens = int(forward_batch.extend_num_tokens)
-        batch_size = getattr(forward_batch, "_original_batch_size", None)
-        if batch_size is None:
-            batch_size = len(forward_batch.seq_lens)
-        num_decodes = batch_size - num_prefills
+        if forward_batch._original_forward_mode is not None:
+            # mlp-sync converted the whole batch to EXTEND, so there are no decode
+            # rows; on an idle rank _original_batch_size is 0 and subtracting here
+            # would go negative.
+            num_decodes = 0
+        else:
+            batch_size = forward_batch._original_batch_size
+            if batch_size is None:
+                batch_size = len(forward_batch.seq_lens)
+            num_decodes = batch_size - num_prefills
         context_lens_tensor = forward_batch.extend_prefix_lens
         assert context_lens_tensor is not None
         has_initial_states = context_lens_tensor > 0