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

python/sglang/srt/layers/attention/qsa/mqa.py

@@ -120,7 +120,6 @@ def _tilelang_qsa_mqa_prefill_kernel(
         head_dim: int,
         block_n: int = 64,
         block_q: int = 32,
-        num_stages: int = 3,
         threads: int = 512,
     ):
         rows = T.dynamic("rows")
@@ -150,9 +149,8 @@ def kernel(
                     end_max = T.max(end_max, T.min(Ends[row_base + qi], keys))
 
                 T.copy(Q[row_base * heads, 0], q_shared)
-                for ni in T.Pipelined(
-                    T.ceildiv(end_max - start_min, block_n), num_stages=num_stages
-                ):
+                # a serial loop: tilelang >= 0.1.14 only pipelines constant trip counts
+                for ni in T.serial(T.ceildiv(end_max - start_min, block_n)):
                     T.copy(K[start_min + ni * block_n, 0], k_shared)
                     T.gemm(
                         k_shared,

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

@@ -516,9 +516,6 @@ def release_memory_occupation(self, recv_req: ReleaseMemoryOccupationReqInput):
         if tags is None or len(tags) == 0:
             tags = GPU_MEMORY_ALL_TYPES
 
-        for tag in tags:
-            self.offload_tags.add(tag)
-
         if GPU_MEMORY_TYPE_KV_CACHE in tags:
             if scheduler is not None:
                 if scheduler.disaggregation_mode == DisaggregationMode.DECODE:
@@ -533,9 +530,15 @@ 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)
+            # flush_cache is a no-op once KV is tagged as offloaded, so flush before tagging
             self.flush_cache()
 
+        for tag in tags:
+            self.offload_tags.add(tag)
+
+        if GPU_MEMORY_TYPE_KV_CACHE in tags:
+            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")
             self.stashed_model_static_state = _export_static_state(

test/registered/unit/managers/test_release_memory_flush_order.py

@@ -0,0 +1,83 @@
+import unittest
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.managers.io_struct import (
+    ReleaseMemoryOccupationReqInput,
+    ResumeMemoryOccupationReqInput,
+)
+from sglang.srt.managers.scheduler import Scheduler
+from sglang.srt.managers.scheduler_components.weight_updater import (
+    SchedulerWeightUpdaterManager,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=11, suite="base-a-test-cpu")
+
+
+class TestReleaseMemoryFlushOrder(CustomTestCase):
+    """Releasing the KV region must flush the pools exactly once, before the region is tagged as offloaded."""
+
+    def _harness(self):
+        scheduler = SimpleNamespace(
+            _engine_paused=False,
+            is_fully_idle=lambda ignore_waiting=False: True,
+            cur_batch_for_debug=None,
+            last_batch=None,
+            tree_cache=Mock(),
+            req_to_token_pool=Mock(),
+            token_to_kv_pool_allocator=Mock(),
+            grammar_manager=Mock(),
+            metrics_reporter=SimpleNamespace(
+                reset_metrics=lambda: None, is_stats_logging_rank=False
+            ),
+            draft_worker=None,
+            waiting_queue=[],
+            running_batch=SimpleNamespace(reqs=[]),
+            disaggregation_mode=DisaggregationMode.NULL,
+        )
+        manager = SchedulerWeightUpdaterManager(
+            tp_worker=None,
+            draft_worker=None,
+            tp_cpu_group=None,
+            memory_saver_adapter=Mock(),
+            flush_cache=lambda empty_cache=False: Scheduler.flush_cache(
+                scheduler, empty_cache=False
+            ),
+            is_fully_idle=scheduler.is_fully_idle,
+            scheduler=scheduler,
+        )
+        scheduler.weight_updater = manager
+        return scheduler, manager
+
+    def test_release_flushes_once_before_the_pause_and_later_flushes_are_skipped(self):
+        scheduler, manager = self._harness()
+        with patch(
+            "sglang.srt.managers.scheduler_components.weight_updater.torch.get_device_module",
+            return_value=SimpleNamespace(synchronize=lambda: None),
+        ):
+            manager.release_memory_occupation(
+                ReleaseMemoryOccupationReqInput(tags=[GPU_MEMORY_TYPE_KV_CACHE])
+            )
+
+        self.assertEqual(scheduler.tree_cache.reset.call_count, 1)
+        self.assertEqual(manager.memory_saver_adapter.pause.call_count, 1)
+        self.assertIn(GPU_MEMORY_TYPE_KV_CACHE, manager.offload_tags)
+
+        # the trainer's /flush_cache while the KV region is unmapped must not touch the pools
+        self.assertTrue(Scheduler.flush_cache(scheduler, empty_cache=False))
+        self.assertEqual(scheduler.tree_cache.reset.call_count, 1)
+
+        manager.resume_memory_occupation(
+            ResumeMemoryOccupationReqInput(tags=[GPU_MEMORY_TYPE_KV_CACHE])
+        )
+        self.assertNotIn(GPU_MEMORY_TYPE_KV_CACHE, manager.offload_tags)
+        self.assertTrue(Scheduler.flush_cache(scheduler, empty_cache=False))
+        self.assertEqual(scheduler.tree_cache.reset.call_count, 2)
+
+
+if __name__ == "__main__":
+    unittest.main()