diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index bcddb9f6a..49159c999 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -377,6 +377,9 @@ class Envs: # even when is_cuda() is False (e.g. AMD/ROCm). On CUDA the barrier is # already enabled regardless of this flag (see start_event_loop). SGLANG_ENABLE_WAR_BARRIER = EnvBool(False) + # Force the WAR barrier to wait for the whole forward instead of the + # read-done fastpath event. + SGLANG_FORCE_COARSE_WAR_BARRIER = EnvBool(False) # PP: skip output send/recv when the entire batch consists of non-final chunked prefill requests, # since process_batch_result_prefill discards next_token_ids for those anyway. SGLANG_PP_SKIP_PURE_CHUNKED_OUTPUT_COMM = EnvBool(False) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index b2305a99b..3d509e90f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1487,14 +1487,15 @@ class Scheduler( # (result processing, next iteration's writes) behind the forward's # shared-buffer reads. Fast path: wait on the read-done event the # forward published after its snapshot (non-spec: decode graph; spec: - # draft_extend), then clear it. Else whole-forward wait_stream. + # draft_extend), then clear it. Else whole-forward wait_stream + # (forceable via SGLANG_FORCE_COARSE_WAR_BARRIER). if not self._war_barrier_enabled: return runner = self.model_worker.war_fastpath_runner ev = runner.war_fastpath_read_done_event - if ev is not None: + runner.war_fastpath_read_done_event = None + if ev is not None and not envs.SGLANG_FORCE_COARSE_WAR_BARRIER.get(): self.schedule_stream.wait_event(ev) - runner.war_fastpath_read_done_event = None else: self.schedule_stream.wait_stream(self.forward_stream) diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 6d3c750b3..f513c299e 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -881,6 +881,8 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): req.cache_protected_len = len(new_indices) req.last_node = new_last_node req.swa_uuid_for_lock = lock_result.swa_uuid_for_lock + # The rematch acquired a new SWA prefix lock. + req.swa_prefix_lock_released = False # cleanup for comp in self._components_tuple: diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 9a1407bf6..e432c46cb 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -932,6 +932,7 @@ class UnifiedRadixCacheSuite: req.last_node = cache.root_node req.cache_protected_len = 0 req.swa_uuid_for_lock = None + req.swa_prefix_lock_released = True req.extra_key = None req.full_untruncated_fill_ids = array("q", tokens) req.set_extend_range( @@ -975,6 +976,7 @@ class UnifiedRadixCacheSuite: self.assertGreater(len(req.prefix_indices), 0) self.assertEqual(req.cache_protected_len, len(req.prefix_indices)) self.assertIsNotNone(req.last_node) + self.assertFalse(req.swa_prefix_lock_released) cache.dec_lock_ref( req.last_node,