Enable runtime busy memory check for speculation topk>1 (#27228)

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

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import logging
-import warnings
 from collections import deque
 from dataclasses import dataclass, field
 from typing import (
@@ -164,12 +163,17 @@ class SchedulerInvariantChecker:
         """
         # After decode: running_batch IS last_batch (same object), count once.
         # After prefill: they differ, both hold uncached tokens.
-        batches = [self.get_last_batch()]
+        # Use identity (is / is not), not membership or ==: ScheduleBatch's
+        # dataclass __eq__ compares tensor fields and raises on ambiguous bools.
+        last_batch = self.get_last_batch()
+        running_batch = self.get_running_batch()
+        batches = [last_batch]
         if (
-            self.get_running_batch() not in (None, self.get_last_batch())
-            and not self.get_running_batch().is_empty()
+            running_batch is not None
+            and running_batch is not last_batch
+            and not running_batch.is_empty()
         ):
-            batches.append(self.get_running_batch())
+            batches.append(running_batch)
 
         full_uncached = 0
         swa_uncached = 0
@@ -196,13 +200,6 @@ class SchedulerInvariantChecker:
         if self.get_last_batch() is None:
             return
 
-        spec_topk = self.server_args.speculative_eagle_topk or 1
-        if spec_topk > 1:
-            warnings.warn(
-                "Runtime memory check (busy) is not supported when speculation topk > 1."
-            )
-            return
-
         ps = self.pool_stats_observer.get_pool_stats()
         full_uncached, swa_uncached = self._get_total_uncached_sizes()
 

python/sglang/test/server_fixtures/spec_eagle_fixture.py

@@ -66,7 +66,9 @@ class SpecEagleServerBase(CustomTestCase):
     trust_remote_code = True
 
     # -- extras --
-    # env_overrides: iterable of (env_var_obj, value) applied only around launch.
+    # env_overrides: (env_var_obj, value) pairs applied only around launch.
+    # Declare ONLY this class's own; _merged_env_overrides() unions them down the
+    # MRO (base first), so never restate a base's. Derived wins on a repeated env.
     env_overrides = ()
     extra_args = ()
 
@@ -105,6 +107,14 @@ class SpecEagleServerBase(CustomTestCase):
         args += [str(a) for a in cls.extra_args]
         return args
 
+    @classmethod
+    def _merged_env_overrides(cls):
+        # Base first so a derived class wins for a repeated env var.
+        merged = []
+        for klass in reversed(cls.__mro__):
+            merged.extend(klass.__dict__.get("env_overrides", ()))
+        return merged
+
     @classmethod
     def setUpClass(cls):
         cls.base_url = DEFAULT_URL_FOR_TEST
@@ -116,7 +126,7 @@ class SpecEagleServerBase(CustomTestCase):
             stack.enter_context(
                 envs.SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN.override(True)
             )
-            for env_var, value in cls.env_overrides:
+            for env_var, value in cls._merged_env_overrides():
                 stack.enter_context(env_var.override(value))
             cls.process = popen_launch_server(
                 cls.model,