TP/PP Consensus checker (#34406)

docs/docs/references/environment_variables.mdx

@@ -1068,6 +1068,11 @@ SGLang supports various environment variables that can be used to configure its
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Detect and report ranks that fall behind during collective ops.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>false</code></td>
     </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_ENABLE_RANK_CONSENSUS_CHECKER</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Check for PP/TP-rank divergence.  Kill the server when divergence occurs.  Helpful for trouble-shooting server hangs issues.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>False</code></td>
+    </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_FORCE_SHUTDOWN</code></td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Force an immediate process-group shutdown on exit.</td>

python/sglang/srt/environ.py

@@ -329,6 +329,7 @@ class Envs:
     SGLANG_LOG_REQUEST_HEADERS = EnvTuple(tuple())
     SGLANG_LOG_SCHEDULER_STATUS_TARGET = EnvStr("")
     SGLANG_LOG_SCHEDULER_STATUS_INTERVAL = EnvFloat(60.0)
+    SGLANG_ENABLE_RANK_CONSENSUS_CHECKER = EnvBool(False)
 
     # ===================================================================
     # IPC, broadcasters, and ports

python/sglang/srt/managers/scheduler.py

@@ -309,6 +309,7 @@ from sglang.srt.utils import (
     is_hip,
     is_mps,
     kill_itself_when_parent_died,
+    rank_consensus_checker,
     require_mlp_sync,
     set_gpu_proc_affinity,
     set_random_seed,
@@ -655,6 +656,8 @@ class Scheduler(
 
         self.init_batch_result_processor()
 
+        self.init_rank_consensus_checker()
+
         self.is_initializing = False
         self.init_startup_timing_summary()
 
@@ -1688,6 +1691,8 @@ class Scheduler(
         if self.decode_offload_manager is not None:
             self.decode_offload_manager.release_host_resources()
 
+        rank_consensus_checker.shutdown()
+
     def run_event_loop(self) -> None:
         """Run the scheduler's event loop.
 
@@ -2114,6 +2119,16 @@ class Scheduler(
             get_running_batch=lambda: self.running_batch,
         )
 
+    def init_rank_consensus_checker(self) -> None:
+        groups = []
+        if self.attn_cp_group is not None and self.attn_tp_group is not None:
+            groups += [self.attn_cp_group, self.attn_tp_group]
+        else:
+            groups += [self.tp_group]
+        if self.pp_group is not None:
+            groups += [self.pp_group]
+        rank_consensus_checker.configure(groups)
+
     def init_kv_events_publisher(self) -> None:
         self.kv_events_publisher = SchedulerKvEventsPublisher(
             kv_events_config=get_observability().kv_events_config,

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -95,6 +95,7 @@ if TYPE_CHECKING:
     from sglang.srt.mem_cache.memory_pool_host import PoolEntry
     from sglang.srt.server_args import ServerArgs
 
+from sglang.srt.utils.rank_consensus_checker import rank_consensus
 
 T = TypeVar("T")
 
@@ -491,6 +492,10 @@ class UnifiedRadixCache(BasePrefixCache):
         if self.host_pool_group is not None:
             self.host_pool_group.destroy()
 
+    @rank_consensus(
+        same_params=["params"],
+        same_results=["result.full_kv_hit_length", "result.swa_host_hit_length"],
+    )
     def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
         result = self.session.try_match_prefix(params)
         if result is not None:
@@ -1700,6 +1705,7 @@ class UnifiedRadixCache(BasePrefixCache):
         operation_terminated = states[1].item() == 1
         return can_terminate or operation_terminated
 
+    @rank_consensus(same_params=True, same_results=True)
     def check_prefetch_progress(self, req_id: str) -> bool:
         if req_id not in self.ongoing_prefetch:
             return True
@@ -1917,6 +1923,7 @@ class UnifiedRadixCache(BasePrefixCache):
             return 0
         return self.buffer_pipeline.staged_prefetch_swa_tokens(req_id)
 
+    @rank_consensus(same_params=True)
     def release_aborted_request(self, rid: str) -> None:
         self.prefetch_loaded_tokens_by_reqid.pop(rid, None)
         if (

python/sglang/srt/utils/rank_consensus_checker.py

@@ -0,0 +1,432 @@
+from __future__ import annotations
+
+import functools
+import hashlib
+import inspect
+import logging
+import os
+import queue
+import threading
+from typing import TYPE_CHECKING, Any, Callable, List, Optional
+
+import torch
+import torch.distributed as dist
+
+from sglang.srt.environ import envs
+
+if TYPE_CHECKING:
+    from sglang.srt.distributed.parallel_state import GroupCoordinator
+
+logger = logging.getLogger(__name__)
+
+_sync_groups: List[dist.ProcessGroup] = []  # Dedicated gloo groups (one per rank-set).
+_q: Optional[queue.Queue[str]] = None
+_worker_thread: Optional[threading.Thread] = None
+_scheduler_thread: Optional[threading.Thread] = None
+
+
+def rank_consensus(func=None, *, same_params=None, same_results=None, **kwargs):
+    """
+    Mark a function that should be consensus in PP and TP ranks.  Here consensus means,
+    the same order of calling, same parameters and return values optionally.
+
+    The function must be called in the scheduler thread.
+
+    Usages:
+
+    * Assert that the function is called by all ranks.  The parameters or results may not be same.
+    @rank_consensus
+    def foo():
+        pass
+
+    * Assert that all parameters are same in all ranks.
+    @rank_consensus(same_params = True)
+    def foo(a, b):
+        pass
+
+    * Assert that some parameters are same in all ranks.
+    @rank_consensus(same_params = ["a", "c"])
+    def foo(a, b, c):
+        pass
+
+    * Assert that part of the parameters are same in all ranks.
+    @rank_consensus(same_params = ["a.req_id"])
+    def foo(a):
+        pass
+
+    * Assert that results are same in all ranks.
+    @rank_consensus(same_results = True)
+    def foo():
+        return 1
+
+    * Assert for part of the results are same.
+    @rank_consensus(same_results = ["result.some_field"]
+    def foo():
+        return SomeObject()
+
+    @rank_consensus(same_results = ["result.field", "len(result.field2)"]
+    def foo():
+        return SomeObject()
+
+    * Assert the function is called by all ranks and all parameters and results are the same.
+    @rank_consensus(same_params = True, same_results = True)
+    def foo():
+        return 1
+    """
+    if kwargs:
+        raise TypeError(
+            f"rank_consensus() got unexpected keyword argument(s): " f"{list(kwargs)}"
+        )
+
+    params_selector = _normalize_selector(same_params, "same_params")
+    results_selector = _normalize_selector(same_results, "same_results")
+
+    def decorator(func: Callable) -> Callable:
+        # This decorator function called at import time.  So it should be zero runtime overhead
+        # when the consensus checker is disabled.
+        if not envs.SGLANG_ENABLE_RANK_CONSENSUS_CHECKER.get():
+            return func
+
+        # Unwrap static/class-method descriptors so we always operate on the
+        # raw function. We remember the descriptor type so we can re-wrap the
+        # result and the class-body descriptor protocol keeps working.
+        if isinstance(func, (classmethod, staticmethod)):
+            raw_func = func.__func__
+            descriptor_type = type(func)
+        else:
+            raw_func = func
+            descriptor_type = None
+        sig = inspect.signature(raw_func)
+
+        # When calling class method or object method with "same_params=True",
+        # skip the first "cls" or "self", as the text format for that
+        # may include memory addresses, which are considered divergence.
+        skip_name: Optional[str] = None
+        if _is_method_with_receiver(func) and len(sig.parameters) > 0:
+            skip_name = next(iter(sig.parameters))
+
+        @functools.wraps(raw_func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            params_payload = "<no check>"
+            if params_selector is not None:
+                # Bind once and apply defaults so that name-based selectors work
+                # regardless of whether the caller passed positionally or by kw.
+                bound = sig.bind(*args, **kwargs)
+                bound.apply_defaults()
+                arguments = dict(bound.arguments)
+                params_payload = _build_payload(
+                    "call", params_selector, arguments, skip_name
+                )
+            assert_same("%s called params=%s", raw_func.__name__, params_payload)
+
+            result = raw_func(*args, **kwargs)
+
+            result_payload = "<no check>"
+            if results_selector is not None:
+                result_scope = {"result": result}
+                result_payload = _build_payload(
+                    "return",
+                    results_selector,
+                    result_scope,
+                )
+            assert_same("%s returns result=%s", raw_func.__name__, result_payload)
+            return result
+
+        # Re-wrap into the original descriptor type so class-body access
+        # (C.method / instance.method) still binds correctly.
+        if descriptor_type is staticmethod:
+            return staticmethod(wrapper)
+        if descriptor_type is classmethod:
+            return classmethod(wrapper)
+        return wrapper
+
+    if func is not None:
+        # Bare `@rank_consensus` form.
+        return decorator(func)
+    else:
+        # `@rank_consensus(same_params=True, same_results=True)` form.
+        return decorator
+
+
+def _normalize_selector(
+    value: None | bool | str | list[str], name: str
+) -> None | bool | list[str]:
+    """Normalize a selector argument to one of:
+    ``None`` (skip), ``True`` (compare everything), or ``list[str]`` (the
+    expressions to evaluate). ``False`` is treated as ``None``.
+    """
+    if value is None or value is False:
+        return None
+    if value is True:
+        return True
+    if isinstance(value, str):
+        return [value]
+    if isinstance(value, list) and all(isinstance(s, str) for s in value):
+        return list(value)
+    raise TypeError(f"{name} must be True / False / str / list[str], got {value!r}")
+
+
+def _is_method_with_receiver(func: Any) -> bool:
+    """Return True iff ``func`` is a method whose first parameter is a
+    receiver (instance for instance-methods, class for class-methods) that
+    should be dropped from the ``same_params=True`` payload.
+
+    Distinguishes:
+      * ``staticmethod`` object  -> False (no receiver)
+      * ``classmethod``  object  -> True  (receiver is the class)
+      * plain ``def`` defined inside a class body (``__qualname__`` has a
+        dot before the final segment and is not a ``<locals>`` closure) ->
+        True  (instance method)
+      * anything else (module-level function, nested function, lambda) ->
+        False
+    """
+    if isinstance(func, staticmethod):
+        return False
+    if isinstance(func, classmethod):
+        return True
+    if inspect.isfunction(func):
+        qualname = getattr(func, "__qualname__", "")
+        # ``C.m`` -> True; ``m`` -> False; ``outer.<locals>.m`` -> False
+        # (closures aren't class-body methods).
+        if "." in qualname and "<locals>" not in qualname:
+            return True
+    return False
+
+
+def _build_payload(
+    tag: str,
+    selector: bool | list[str],
+    scope: dict[str, Any],
+    skip_name: Optional[str] = None,
+) -> str:
+    """Serialize the selected values into a single comparable string.
+
+    ``skip_name`` only applies to the ``True`` (whole-scope) form and is used
+    to drop the receiver (``self`` / ``cls``) from method payloads; explicit
+    ``list[str]`` selectors honor exactly what the user listed.
+    """
+    if selector is True:
+        # Whole scope is the payload. For the call checkpoint, the scope is
+        # the arguments dict; for the return checkpoint, the caller wrapped
+        # result into the scope, so we repr ``result`` directly.
+        if tag == "call":
+            if skip_name is not None:
+                scope = {k: v for k, v in scope.items() if k != skip_name}
+            return repr(scope)
+        return repr(scope["result"])
+    parts: list[str] = []
+    for expr in selector:
+        value = _eval_selector(expr, scope)
+        parts.append(f"{expr}={value!r}")
+    return " | ".join(parts)
+
+
+def _eval_selector(expr: str, scope: dict[str, Any]) -> Any:
+    """Evaluate a selector expression in a restricted scope.
+
+    Errors (unknown parameter name, missing attribute, bad syntax) propagate
+    -- they are caller bugs and must not be silently swallowed or confused
+    with cross-rank divergence.
+    """
+    safe_builtins = {
+        "len": len,
+        "int": int,
+        "str": str,
+        "bool": bool,
+        "float": float,
+        "tuple": tuple,
+        "list": list,
+        "dict": dict,
+        "set": set,
+        "sorted": sorted,
+        "min": min,
+        "max": max,
+        "sum": sum,
+    }
+    return eval(expr, {"__builtins__": safe_builtins}, dict(scope))
+
+
+def enabled() -> bool:
+    """Test that the checker has been enabled and configure() is called."""
+    return _q is not None
+
+
+def assert_same(msg_fmt: str, *args: Any) -> None:
+    """Record a decision that every TP/PP rank must make identically.
+
+    Must be called from the scheduler thread. If the env var is set and the
+    checker is configured, an assertion guards that the caller is on the
+    scheduler thread recorded at configure() time — events from other threads
+    would interleave out of order with peer ranks and corrupt the lock-step
+    drain.
+
+    When the divergence checker is disabled, this is a zero-overhead no-op.
+
+    Example:
+    assert_same("my decision: %s %d", "foo", 100)
+
+    Prefer `@rank_consensus` over this function for code-cleanliness.
+    """
+    if not enabled():
+        return
+    # Sanity check: only the scheduler thread is allowed to enqueue. Other
+    # callers would race with the worker's min-length drain and desynchronize
+    # ranks, since their events would not exist on peer ranks.
+    if threading.current_thread() is not _scheduler_thread:
+        raise RuntimeError("rdc.assert_same must be called from the scheduler thread")
+    # Format eagerly: args may reference mutable state that mutates
+    # between now and when the worker thread drains the queue.
+    _q.put(msg_fmt % args)
+
+
+def configure(groups: List[GroupCoordinator]) -> None:
+    """Initialize the checker. No-op if SGLANG_ENABLE_RANK_CONSENSUS_CHECKER is not set."""
+    global _sync_groups, _q, _worker_thread, _scheduler_thread
+
+    if not envs.SGLANG_ENABLE_RANK_CONSENSUS_CHECKER.get():
+        return
+
+    logger.warning(
+        "Rank consensus checker is enabled. The server will suicide if rank divergence detected."
+    )
+    # Build a dedicated sync group.  So our synchronization work will not affect
+    # the scheduler thread at all.
+    _sync_groups = _create_sync_groups(groups)
+    _q = queue.Queue()
+    # Assume the calling thread is the schedule thread.
+    # We will check assert_same() must be called by the scheduler thread.
+    _scheduler_thread = threading.current_thread()
+    _worker_thread = threading.Thread(
+        target=_worker_loop, name="rank_consensus_checker", daemon=True
+    )
+    _worker_thread.start()
+
+
+def _create_sync_groups(
+    groups: List[GroupCoordinator],
+) -> List[dist.ProcessGroup]:
+    """Create duplicated groups, used for background thread"""
+    from sglang.srt.distributed.parallel_state import create_custom_parallel_group
+
+    dedicated: List[dist.ProcessGroup] = []
+    seen_rank_sets: set[tuple[int, ...]] = set()
+    for group in groups:
+        if group is None:
+            continue
+        # Skip single-rank groups: nothing to compare against.
+        if torch.distributed.get_world_size(group=group.cpu_group) == 1:
+            continue
+        group_ranks = tuple(torch.distributed.get_process_group_ranks(group.cpu_group))
+        if group_ranks in seen_rank_sets:
+            continue
+        seen_rank_sets.add(group_ranks)
+        pg = create_custom_parallel_group(group_ranks=list(group_ranks), backend="gloo")
+        if pg is not None:
+            dedicated.append(pg)
+    return dedicated
+
+
+def _destroy_dedicated_groups() -> None:
+    for pg in _sync_groups:
+        try:
+            torch.distributed.destroy_process_group(pg)
+        except Exception:
+            pass
+
+
+def shutdown() -> None:
+    """Flush the queue, stop the worker thread, and disable assert_same."""
+    global _q, _worker_thread, _sync_groups, _scheduler_thread
+
+    q = _q
+    if q is None:
+        return
+
+    # Put a sentinel value to wake the worker if it is blocked on _q.get().
+    q.put(None)
+    if _worker_thread is not None:
+        _worker_thread.join()
+        _worker_thread = None
+    # Tear down the dedicated gloo groups BEFORE clearing _groups so the
+    # destroy helper can see them. Worker thread is already joined, so there
+    # is no concurrent all_reduce on these groups.
+    _destroy_dedicated_groups()
+    _q = None
+    _sync_groups = []
+    _scheduler_thread = None
+
+
+def _worker_loop() -> None:
+    """Consume events in lock-step with peer ranks via gloo all-reduce.
+
+    Each iteration:
+      1. Determine the items available in _q.
+      2. Drain exactly the minimum number of items in all ranks.
+      3. Compare all events are identical across ranks.
+    """
+    while _q is not None:
+        # Drain first.  Block waiting for the first event.
+        first = _q.get()
+        if first is None:
+            # shutdown() is called.
+            return
+
+        # Drain more whenever available.
+        # Every rank should drain the same number.
+        count = _all_reduce_min_int(_q.qsize())
+        events: List[str] = [first]
+        shutdown_signaled = False
+        for _ in range(count):
+            event = _q.get()
+            if event is None:
+                # shutdown() sentinel arrived mid-batch: stop draining but
+                # still check the events we already hold — they are real
+                # decisions every rank must agree on. Then exit, since the
+                # sentinel means shutdown() is waiting on worker_thread.join().
+                shutdown_signaled = True
+                break
+            events.append(event)
+
+        # Cross-rank check.
+        _check_for_consensus(events)
+
+        if shutdown_signaled:
+            return
+
+
+def _all_reduce_min_int(value: int) -> int:
+    """Reduce `value` to its global minimum across every configured group."""
+    tensor = torch.tensor([value], dtype=torch.int64)
+    for group in _sync_groups:
+        dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=group)
+    return int(tensor.item())
+
+
+def _check_for_consensus(events: list[str]) -> None:
+    # Compute sha1 of concatenation of all msgs.
+    hasher = hashlib.sha1()
+    for msg in events:
+        hasher.update(msg.encode("utf-8"))
+
+    # Determine if some rank has a different value.
+    value_bytes = hasher.digest()
+    min_value = torch.tensor(list(hasher.digest()), dtype=torch.uint8)
+    max_value = min_value.clone()
+    for group in _sync_groups:
+        dist.all_reduce(min_value, op=dist.ReduceOp.MIN, group=group)
+        dist.all_reduce(max_value, op=dist.ReduceOp.MAX, group=group)
+    if not torch.equal(min_value, max_value):
+        # When divergence, all rank should output the following log.
+        logger.critical(
+            f"Found rank divergence for {len(events)} events(s)! local hash: {value_bytes.hex()}, events = {events}"
+        )
+        for handler in logger.handlers:
+            handler.flush()
+
+        # os._exit instead of sys.exit: this runs in a background thread, where
+        # SystemExit would only kill the thread, not the process. os._exit tears
+        # down the whole scheduler process so a TP/PP mismatch can never
+        # silently keep serving.
+        os._exit(1)
+
+    logger.debug(f"Consensus check passed for {len(events)} event(s).")

test/registered/cpu/test_rank_consensus_checker.py

@@ -0,0 +1,586 @@
+import os
+import queue
+import threading
+import traceback
+import unittest
+from multiprocessing import Process
+from unittest.mock import patch
+
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+from sglang.srt.distributed import parallel_state as ps
+from sglang.srt.distributed.parallel_state import (
+    get_pp_group,
+    get_tp_group,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from sglang.srt.utils.rank_consensus_checker import (
+    assert_same,
+    configure,
+    rank_consensus,
+    shutdown,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase, find_available_port
+
+register_cpu_ci(est_time=30, suite="base-b-test-cpu")
+
+
+def run_distributed_test(
+    rank: int,
+    world_size: int,
+    pp_size: int,
+    tp_size: int,
+    master_port: int,
+    fn,
+) -> None:
+    """Child-process entry point: set up gloo, then run fn.
+
+    Exit codes:
+      * 0 -> fn finished cleanly
+      * 1 -> rdc detected divergence and called os._exit(1) from its worker
+      * 2 -> fn raised (test setup/scenario bug)
+    """
+    # CUDA_VISIBLE_DEVICES is set to "99" (a non-existent device) by the parent
+    # in _spawn() before this process starts, so by the time the test module
+    # (and torch) is re-imported here, is_cuda_alike() returns False and
+    # GroupCoordinator picks device="cpu". That keeps this test CPU-only and
+    # lets world_size exceed the host's physical GPU count.
+
+    # The CUDA-only communicators (pynccl, custom allreduce) cannot be built
+    # without a GPU -- PyNcclCommunicator calls torch.cuda.device(device).
+    # initialize_model_parallel has no flag to disable pynccl, so patch
+    # init_model_parallel_group to force use_pynccl=False (and clear the
+    # module-level custom-allreduce default via its public setter). patch.object
+    # auto-restores on exit, including the os._exit(2) path below.
+    ps.set_custom_all_reduce(False)
+
+    def _cpu_init_model_parallel_group(
+        *args, _orig=ps.init_model_parallel_group, **kwargs
+    ):
+        kwargs.setdefault("use_pynccl", False)
+        kwargs.setdefault("use_custom_allreduce", False)
+        return _orig(*args, **kwargs)
+
+    with patch.object(ps, "init_model_parallel_group", _cpu_init_model_parallel_group):
+        try:
+            os.environ["RANK"] = str(rank)
+            os.environ["WORLD_SIZE"] = str(world_size)
+            os.environ["MASTER_ADDR"] = "localhost"
+            os.environ["MASTER_PORT"] = str(master_port)
+            os.environ["LOCAL_SIZE"] = str(world_size)
+
+            init_distributed_environment(
+                world_size=world_size,
+                rank=rank,
+                distributed_init_method="env://",
+                local_rank=rank,
+                backend="gloo",
+            )
+
+            initialize_model_parallel(
+                tensor_model_parallel_size=tp_size,
+                pipeline_model_parallel_size=pp_size,
+                backend="gloo",
+            )
+
+            fn()
+        except Exception as e:
+            print(f"subprocess[{rank=}] has error: {e}", flush=True)
+            traceback.print_exc()
+            os._exit(2)
+        finally:
+            try:
+                if dist.is_initialized():
+                    dist.destroy_process_group()
+            except Exception:
+                pass
+
+
+class _DummyClass:
+    def __init__(self, a: int = None, b: int = None):
+        self.a = a
+        self.b = b
+
+    def __repr__(self) -> str:
+        return f"DummyClass(a={self.a}, b={self.b})"
+
+
+class _MethodHost:
+    @rank_consensus(same_params=True)
+    def instance_method(obj, a, b):
+        return a + b
+
+    @rank_consensus(same_params=True)
+    @classmethod
+    def class_method(klass, a):
+        return a + 1
+
+    @rank_consensus(same_params=True)
+    @staticmethod
+    def static_method(a, b):
+        return a * b
+
+
+class RankConsensusCheckerTestCase(CustomTestCase):
+    def _spawn(self, fn, pp_size: int = 1, tp_size: int = 1, enable_env: bool = True):
+        """Run fn in world_size spawned gloo children. Returns True iff every
+        child exited with code 0. A detected divergence makes rdc call
+        os._exit(1) from its worker thread; an exception inside fn makes
+        run_distributed_test call os._exit(2). Either way _spawn returns
+        False for that child."""
+        mp.set_start_method("spawn", force=True)
+        master_port = find_available_port(23456)
+
+        old_env = os.getenv("SGLANG_ENABLE_RANK_CONSENSUS_CHECKER")
+        os.environ["SGLANG_ENABLE_RANK_CONSENSUS_CHECKER"] = str(enable_env)
+
+        world_size = pp_size * tp_size
+        processes = []
+        for rank in range(world_size):
+            p = Process(
+                target=run_distributed_test,
+                kwargs=dict(
+                    rank=rank,
+                    world_size=world_size,
+                    pp_size=pp_size,
+                    tp_size=tp_size,
+                    master_port=master_port,
+                    fn=fn,
+                ),
+            )
+            p.start()
+            processes.append(p)
+
+        for p in processes:
+            p.join()
+
+        if old_env is None:
+            os.environ.pop("SGLANG_ENABLE_RANK_CONSENSUS_CHECKER")
+        else:
+            os.environ["SGLANG_ENABLE_RANK_CONSENSUS_CHECKER"] = old_env
+
+        return all(p.exitcode == 0 for p in processes)
+
+
+class TestAssertSame(RankConsensusCheckerTestCase):
+    @staticmethod
+    def same_fn():
+        configure([get_tp_group()])
+        assert_same("same %d", 10)
+        shutdown()
+
+    def test_same(self):
+        """Same args on every rank -> no divergence, clean exit."""
+        self.assertTrue(self._spawn(TestAssertSame.same_fn, tp_size=2))
+
+    @staticmethod
+    def divergence_fn():
+        tp_group = get_tp_group()
+        configure([tp_group])
+        assert_same("diverge %d", tp_group.rank_in_group)
+        shutdown()
+
+    def test_divergence(self):
+        """Different args on different ranks -> rdc calls os._exit(1) -> child
+        exit code is 1 -> _spawn returns False."""
+        self.assertFalse(self._spawn(TestAssertSame.divergence_fn, tp_size=2))
+
+    @staticmethod
+    def divergent_multi_group_fn():
+        tp_group = get_tp_group()
+        pp_group = get_pp_group()
+        configure([tp_group, pp_group])
+        assert_same("diverge %d", tp_group.rank_in_group)
+        shutdown()
+
+    def test_divergence_detected_multi_group(self):
+        """Passing the same group twice must still surface the divergence."""
+        self.assertFalse(
+            self._spawn(
+                TestAssertSame.divergent_multi_group_fn,
+                pp_size=2,
+                tp_size=2,
+            )
+        )
+
+    @staticmethod
+    def wrong_thread_fn():
+        tp_group = get_tp_group()
+        configure([tp_group])
+
+        err_box: queue.Queue = queue.Queue()
+
+        def _other_thread():
+            try:
+                assert_same("from other thread")
+                err_box.put(None)
+            except Exception as e:  # noqa: BLE001
+                err_box.put(e)
+
+        t = threading.Thread(target=_other_thread)
+        t.start()
+        t.join()
+
+        err = err_box.get()
+        shutdown()
+        assert isinstance(
+            err, RuntimeError
+        ), f"Expected RuntimeError from stray-thread assert_same, got {err!r}"
+
+    def test_assert_same_rejects_non_scheduler_thread(self):
+        """Check that assert_same() must be called in the scheduler thread.  Otherwise report error."""
+        self.assertTrue(self._spawn(TestAssertSame.wrong_thread_fn, tp_size=2))
+
+    @staticmethod
+    def disabled_fn():
+        tp_group = get_tp_group()
+        configure([tp_group])
+        assert_same("diverge %d", tp_group.rank_in_group)
+
+    def test_disabled_is_noop(self):
+        """Test that when SGLANG_ENABLE_RANK_CONSENSUS_CHECKER=false, assert_same is no-op."""
+        self.assertTrue(
+            self._spawn(TestAssertSame.disabled_fn, tp_size=2, enable_env=False)
+        )
+
+
+class TestRankConsensusDecorator(RankConsensusCheckerTestCase):
+    @staticmethod
+    def consensus_bare_diverge_fn():
+        @rank_consensus
+        def foo(a: int) -> int:
+            return a
+
+        # Bare decorator only checks "was called", not args; even with diverging
+        # args this must exit clean (no rank divergence).
+        tp_group = get_tp_group()
+        configure([tp_group])
+        foo(tp_group.rank_in_group)
+        shutdown()
+
+    def test_bare_decorator_clean_with_diverging_args(self):
+        """Bare decorator only checks that every rank calls the function;
+        diverging args must NOT be flagged."""
+        self.assertTrue(
+            self._spawn(TestRankConsensusDecorator.consensus_bare_diverge_fn, tp_size=2)
+        )
+
+    @staticmethod
+    def consensus_all_params_same_fn():
+        @rank_consensus(same_params=True)
+        def foo(a: int, b: int) -> int:
+            return a + b
+
+        configure([get_tp_group()])
+        foo(1, 2)
+        shutdown()
+
+    def test_all_params_same(self):
+        self.assertTrue(
+            self._spawn(
+                TestRankConsensusDecorator.consensus_all_params_same_fn, tp_size=2
+            )
+        )
+
+    @staticmethod
+    def consensus_all_params_diverge_fn():
+        @rank_consensus(same_params=True)
+        def foo(a, b):
+            return a + b
+
+        tp_group = get_tp_group()
+        configure([tp_group])
+        # The second argument differs on rank.  Expect divergence.
+        foo(1, tp_group.rank_in_group)
+        shutdown()
+
+    def test_all_params_diverge(self):
+        self.assertFalse(
+            self._spawn(
+                TestRankConsensusDecorator.consensus_all_params_diverge_fn, tp_size=2
+            )
+        )
+
+    @staticmethod
+    def consensus_named_params_same_fn():
+        @rank_consensus(same_params=["a", "c"])
+        def foo(a: int, b: int, c: int) -> int:
+            return a + b + c
+
+        tp_group = get_tp_group()
+        configure([tp_group])
+        # b diverges but is NOT in the selector list.  Expect good.
+        foo(1, tp_group.rank_in_group, 3)
+        shutdown()
+
+    def test_named_params_ignores_unselected_divergence(self):
+        self.assertTrue(
+            self._spawn(
+                TestRankConsensusDecorator.consensus_named_params_same_fn, tp_size=2
+            )
+        )
+
+    @staticmethod
+    def consensus_named_params_diverge_fn():
+        @rank_consensus(same_params=["a", "c"])
+        def foo(a: int, b: int, c: int) -> int:
+            return a + b + c
+
+        # c diverges and IS in the selector list.  Expect divergence.
+        tp_group = get_tp_group()
+        configure([tp_group])
+        foo(1, 2, tp_group.rank_in_group)
+        shutdown()
+
+    def test_named_params_flags_selected_divergence(self):
+        self.assertFalse(
+            self._spawn(
+                TestRankConsensusDecorator.consensus_named_params_diverge_fn,
+                tp_size=2,
+            )
+        )
+
+    @staticmethod
+    def consensus_dotted_param_same_fn():
+        @rank_consensus(same_params=["_a.a"])
+        def foo(_a: _DummyClass) -> None:
+            pass
+
+        tp_group = get_tp_group()
+        configure([tp_group])
+        dummy = _DummyClass(a=10, b=tp_group.rank_in_group)
+        foo(dummy)
+        shutdown()
+
+    def test_dotted_param_same(self):
+        self.assertTrue(
+            self._spawn(
+                TestRankConsensusDecorator.consensus_dotted_param_same_fn, tp_size=2
+            )
+        )
+
+    @staticmethod
+    def consensus_dotted_param_diverge_fn():
+        @rank_consensus(same_params=["_a.a"])
+        def foo(_a: _DummyClass) -> None:
+            pass
+
+        tp_group = get_tp_group()
+        configure([tp_group])
+        dummy = _DummyClass(a=tp_group.rank_in_group, b=10)
+        foo(dummy)
+        shutdown()
+
+    def test_dotted_param_diverge(self):
+        self.assertFalse(
+            self._spawn(
+                TestRankConsensusDecorator.consensus_dotted_param_diverge_fn,
+                tp_size=2,
+            )
+        )
+
+    @staticmethod
+    def consensus_full_result_same_fn():
+        @rank_consensus(same_results=True)
+        def foo(value: int) -> _DummyClass:
+            return _DummyClass(a=value, b=value * 2)
+
+        configure([get_tp_group()])
+        foo(5)
+        shutdown()
+
+    def test_full_result_same(self):
+        self.assertTrue(
+            self._spawn(
+                TestRankConsensusDecorator.consensus_full_result_same_fn, tp_size=2
+            )
+        )
+
+    @staticmethod
+    def consensus_full_result_diverge_fn():
+        @rank_consensus(same_results=True)
+        def foo(value: int) -> _DummyClass:
+            return _DummyClass(a=value, b=value * 2)
+
+        tp_group = get_tp_group()
+        configure([tp_group])
+        foo(tp_group.rank_in_group)
+        shutdown()
+
+    def test_full_result_diverge(self):
+        self.assertFalse(
+            self._spawn(
+                TestRankConsensusDecorator.consensus_full_result_diverge_fn,
+                tp_size=2,
+            )
+        )
+
+    @staticmethod
+    def consensus_partial_result_same_fn():
+        @rank_consensus(same_results=["result.x", "len(result.y)"])
+        def foo(x, y_list):
+            class _R:
+                pass
+
+            r = _R()
+            r.x = x
+            r.y = y_list
+            return r
+
+        tp_group = get_tp_group()
+        configure([tp_group])
+        # x and len(y) both equal across ranks; y contents differ but are not selected.  Expect good.
+        foo(x=3, y_list=[tp_group.rank_in_group] * 4)
+        shutdown()
+
+    def test_partial_result_same(self):
+        self.assertTrue(
+            self._spawn(
+                TestRankConsensusDecorator.consensus_partial_result_same_fn,
+                tp_size=2,
+            )
+        )
+
+    @staticmethod
+    def consensus_partial_result_diverge_fn():
+        @rank_consensus(same_results=["result.x", "len(result.y)"])
+        def foo(x, y_list):
+            class _R:
+                pass
+
+            r = _R()
+            r.x = x
+            r.y = y_list
+            return r
+
+        tp_group = get_tp_group()
+        configure([tp_group])
+        # x diverges and IS selected.  Expect divergence.
+        foo(x=tp_group.rank_in_group, y_list=[1, 2, 3])
+        shutdown()
+
+    def test_partial_result_diverge(self):
+        self.assertFalse(
+            self._spawn(
+                TestRankConsensusDecorator.consensus_partial_result_diverge_fn,
+                tp_size=2,
+            )
+        )
+
+    @staticmethod
+    def consensus_both_same_fn():
+        @rank_consensus(same_params=True, same_results=True)
+        def foo(a: int) -> int:
+            return a * 2
+
+        configure([get_tp_group()])
+        foo(7)
+        shutdown()
+
+    def test_both_same(self):
+        self.assertTrue(
+            self._spawn(TestRankConsensusDecorator.consensus_both_same_fn, tp_size=2)
+        )
+
+    @staticmethod
+    def consensus_both_diverge_fn():
+        @rank_consensus(same_params=True, same_results=True)
+        def foo(a: int) -> int:
+            return a * 2
+
+        tp_group = get_tp_group()
+        configure([tp_group])
+        foo(tp_group.rank_in_group)
+        shutdown()
+
+    def test_both_diverge(self):
+        self.assertFalse(
+            self._spawn(TestRankConsensusDecorator.consensus_both_diverge_fn, tp_size=2)
+        )
+
+    @staticmethod
+    def consensus_instance_method_same_fn():
+        configure([get_tp_group()])
+        _MethodHost().instance_method(1, 2)
+        shutdown()
+
+    def test_instance_method_receiver_dropped(self):
+        # Two ranks build two different _MethodHost instances; without the
+        # receiver-skip the per-rank address would diverge. Clean exit
+        # confirms the receiver is dropped.
+        self.assertTrue(
+            self._spawn(
+                TestRankConsensusDecorator.consensus_instance_method_same_fn,
+                tp_size=2,
+            )
+        )
+
+    @staticmethod
+    def consensus_class_method_same_fn():
+        configure([get_tp_group()])
+        _MethodHost.class_method(5)
+        shutdown()
+
+    def test_class_method_receiver_dropped(self):
+        # First param is named ``klass`` (not cls); detection must still work.
+        self.assertTrue(
+            self._spawn(
+                TestRankConsensusDecorator.consensus_class_method_same_fn, tp_size=2
+            )
+        )
+
+    @staticmethod
+    def consensus_class_method_via_instance_same_fn():
+        configure([get_tp_group()])
+        _MethodHost().class_method(5)
+        shutdown()
+
+    def test_class_method_via_instance_receiver_dropped(self):
+        # Accessing the classmethod through an instance still binds the class
+        # as the receiver; verify it is still dropped.
+        self.assertTrue(
+            self._spawn(
+                TestRankConsensusDecorator.consensus_class_method_via_instance_same_fn,
+                tp_size=2,
+            )
+        )
+
+    @staticmethod
+    def consensus_static_method_same_fn():
+        configure([get_tp_group()])
+        _MethodHost.static_method(3, 4)
+        shutdown()
+
+    def test_static_method_no_receiver(self):
+        # Static method: no receiver, equal args -> clean.
+        self.assertTrue(
+            self._spawn(
+                TestRankConsensusDecorator.consensus_static_method_same_fn,
+                tp_size=2,
+            )
+        )
+
+    @staticmethod
+    def consensus_static_method_diverge_fn():
+        tp_group = get_tp_group()
+        configure([tp_group])
+        # Static method: no receiver to drop, so a rank-dependent arg diverges.
+        _MethodHost.static_method(tp_group.rank_in_group, 4)
+        shutdown()
+
+    def test_static_method_flags_diverging_arg(self):
+        # Static method: no receiver to drop, so a rank-dependent arg must
+        # still be flagged. Confirms we did not over-skip for static methods.
+        self.assertFalse(
+            self._spawn(
+                TestRankConsensusDecorator.consensus_static_method_diverge_fn,
+                tp_size=2,
+            )
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()