[sglang-miles] Colocated RL for hybrid-state models: flush order and bf16 MoE weight-update layout

python/sglang/srt/arg_groups/moe_hook.py

@@ -13,6 +13,7 @@
     _a2a_fusion_adjustments,
     _moe_runner_backend_quant_constraints,
     _moe_runner_fusion_disable,
+    _routed_experts_capture_backend_guard,
     cutedsl_moe_max_num_tokens,
     declare_resolution,
     max_prefill_buffer_tokens,
@@ -39,6 +40,7 @@ def handle_moe_kernel_config(server_args: Any):
     cfg = resolving_view(server_args)
 
     run_post_process_pass(server_args, _moe_runner_backend_quant_constraints)
+    run_post_process_pass(server_args, _routed_experts_capture_backend_guard)
 
     view = resolved_view(server_args)
     if view.moe_runner_backend == "flashinfer_cutlass":

python/sglang/srt/arg_groups/overrides.py

@@ -1578,6 +1578,36 @@ def _moe_runner_backend_quant_constraints(view: Any) -> dict:
     return {}
 
 
+_TOPK_BYPASSING_MOE_RUNNER_BACKENDS = frozenset(
+    {
+        "flashinfer_trtllm",
+        "experimental_sgl_trtllm",
+        "triton_kernel",
+        "flashinfer_mxfp4",
+    }
+)
+
+
+@register_post_process
+def _routed_experts_capture_backend_guard(view: Any) -> dict:
+    """Routed-experts capture needs a MoE runner that materializes topk ids."""
+    if not view.enable_return_routed_experts:
+        return {}
+    if view.moe_runner_backend not in _TOPK_BYPASSING_MOE_RUNNER_BACKENDS:
+        return {}
+    if view.quantization in (None, "fp8"):
+        logger.warning(
+            f"moe_runner_backend={view.moe_runner_backend!r} bypasses TopK, so "
+            "--enable-return-routed-experts falls back to moe_runner_backend='auto'."
+        )
+        return {"moe_runner_backend": "auto"}
+    raise ValueError(
+        f"--enable-return-routed-experts is incompatible with moe_runner_backend="
+        f"{view.moe_runner_backend!r} under quantization={view.quantization!r}; "
+        "pass a topk-materializing --moe-runner-backend such as triton."
+    )
+
+
 @register_post_process
 def _moe_runner_fusion_disable(view: Any) -> dict:
     """FlashInfer CuteDSL / TRT-LLM / TRT-LLM-routed MoE runners require the

python/sglang/srt/layers/quantization/unquant.py

@@ -906,7 +906,7 @@ def maybe_restore_flashinfer_trtllm_bf16_weight_shape_for_load(
         block layout. During weight update, checkpoint tensors are in
         canonical layout and need a temporary shape restore for copy.
         """
-        if not get_moe_runner_backend().is_flashinfer_trtllm_routed():
+        if not self.use_flashinfer_trtllm_moe:
             return
 
         expected_shape = None

python/sglang/srt/managers/scheduler.py

@@ -72,6 +72,7 @@
     is_deepseek_v4,
     is_minimax_sparse,
 )
+from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
 from sglang.srt.constrained.grammar_manager import GrammarManager
 from sglang.srt.debug_utils.pr_fix_toggle import maybe_revert_pr_fix
 from sglang.srt.disaggregation.decode import (
@@ -4998,6 +4999,9 @@ 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 GPU_MEMORY_TYPE_KV_CACHE in self.weight_updater.offload_tags:
+            logger.info("Cache flush skipped: KV cache memory is released.")
+            return True
         if self.is_fully_idle(ignore_waiting=self._engine_paused):
             self.cur_batch_for_debug = None
             self.last_batch = None

python/sglang/srt/managers/scheduler_components/weight_updater.py

@@ -533,8 +533,8 @@ def release_memory_occupation(self, recv_req: ReleaseMemoryOccupationReqInput):
                     queue = getattr(scheduler, "disagg_prefill_bootstrap_queue", None)
                     if queue is not None:
                         queue.release_memory_occupation()
-            self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_KV_CACHE)
             self.flush_cache()
+            self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_KV_CACHE)
 
         if GPU_MEMORY_TYPE_WEIGHTS in tags:
             self._assert_weight_cache_inactive("release_memory_occupation")