diff --git a/python/sglang/srt/layers/attention/qsa/mqa.py b/python/sglang/srt/layers/attention/qsa/mqa.py index 6446ed145cbe..5d4580529751 100644 --- a/python/sglang/srt/layers/attention/qsa/mqa.py +++ b/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, diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index c0dc373a4bc9..00caa048b9b4 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 7c3908f7977a..5b5ef7a38adc 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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 diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py index e6b78709e82d..18cd9d3820ef 100644 --- a/python/sglang/srt/managers/scheduler_components/weight_updater.py +++ b/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( diff --git a/test/registered/unit/managers/test_release_memory_flush_order.py b/test/registered/unit/managers/test_release_memory_flush_order.py new file mode 100644 index 000000000000..e0da06a7b426 --- /dev/null +++ b/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()