[Scheduler] Add SGLANG_FORCE_COARSE_WAR_BARRIER opt-in for a whole-forward WAR barrier (#29353)

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)

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:
-            self.schedule_stream.wait_event(ev)
         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)
         else:
             self.schedule_stream.wait_stream(self.forward_stream)
 

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:

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,