Read process groups through the runtime context (#40068)

test/registered/kernels/ops/layernorm/test_mhc_kernels.py

@@ -10,11 +10,25 @@ from sglang.test.ci.ci_register import register_cuda_ci
 register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-large")
 
 
+@pytest.fixture
+def stated_tp_group():
+    """A TP group for a test that runs in a process without one.
+
+    The production call passes the group *into* `use_symmetric_memory`, so
+    stubbing that context manager does not stop the read -- the argument is
+    evaluated first. Stating it on the context answers every spelling.
+    """
+    from sglang.srt.runtime_context import get_parallel
+
+    with get_parallel().override(tp_group=None):
+        yield
+
+
 @pytest.mark.parametrize("hidden_size", [4096, 7168])
 @pytest.mark.parametrize("num_tokens", [0, 1, 8, 17, 32, 64])
 @pytest.mark.parametrize("use_norm", [False, True])
 def test_mhc_fused_post_pre_matches_unfused(
-    monkeypatch, hidden_size, num_tokens, use_norm
+    monkeypatch, hidden_size, num_tokens, use_norm, stated_tp_group
 ):
     if not torch.cuda.is_available():
         pytest.skip("CUDA is required for TileLang mHC kernels")
@@ -22,12 +36,10 @@ def test_mhc_fused_post_pre_matches_unfused(
     monkeypatch.setattr(mhc, "is_dsa_prefill_cp_interleave", lambda: False)
     # This is a single-process kernel unit test with no TP group initialized.
     # mhc_pre / mhc_fused_post_pre allocate the MoE input in the symmetric-memory
-    # pool via use_symmetric_memory(get_tp_group(), ...); bypass that path so the
-    # kernel runs with a plain torch.empty allocation. Mirrors the workaround in
-    # test_mxfp4_sm90_cutlass.py for the same TP-group-not-initialized case.
+    # pool, which asks for the TP group; bypassing the allocation is enough, and
+    # then nothing asks. Mirrors the workaround in test_mxfp4_sm90_cutlass.py.
     monkeypatch.setattr(mhc, "use_symmetric_memory", lambda *a, **kw: nullcontext())
     monkeypatch.setattr(mhc, "is_allocation_symmetric", lambda: False)
-    monkeypatch.setattr(mhc, "get_tp_group", lambda: None)
     torch.manual_seed(0)
     device = torch.device("cuda")
     hc_mult = 4

test/registered/kernels/ops/moe/test_minimax_quant_scatter.py

@@ -35,6 +35,20 @@ register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="4-gpu-b
 dev = "cuda"
 
 
+@pytest.fixture
+def stated_tp_group():
+    """A TP group for a test that runs in a process without one.
+
+    The production call passes the group *into* `use_symmetric_memory`, so
+    stubbing that context manager does not stop the read -- the argument is
+    evaluated first. Stating it on the context answers every spelling.
+    """
+    from sglang.srt.runtime_context import get_parallel
+
+    with get_parallel().override(tp_group=None):
+        yield
+
+
 def test_sm120_mxfp8_dispatch_preserves_activation_scale_recipe(monkeypatch):
     """SM120 group-128 activations must not use the MXFP8 weight-scale recipe."""
     from sglang.srt.layers import deep_gemm_wrapper
@@ -257,7 +271,9 @@ def test_standard_layout_auto_memory_policy(monkeypatch):
 
 
 @pytest.mark.parametrize("weight_dtype", ["fp8", "bf16"])
-def test_standard_masked_runner_matches_compact_end_to_end(monkeypatch, weight_dtype):
+def test_standard_masked_runner_matches_compact_end_to_end(
+    monkeypatch, weight_dtype, stated_tp_group
+):
     """Exercise both production grouped GEMMs through the standard path."""
     arch_major, _ = torch.cuda.get_device_capability(torch.cuda.current_device())
     if arch_major <= 9:
@@ -266,7 +282,6 @@ def test_standard_masked_runner_matches_compact_end_to_end(monkeypatch, weight_d
     # This kernel test runs outside a model-parallel process. Bypass only the
     # symmetric-allocation context; all pre-permute, DeepGEMM, activation,
     # quantization, down-GEMM, and post-permute kernels remain real.
-    monkeypatch.setattr(deep_gemm_runner, "get_tp_group", lambda: None)
     monkeypatch.setattr(
         deep_gemm_runner,
         "use_symmetric_memory",

test/registered/ops/test_aiter_greedy_sample_amd.py

@@ -22,7 +22,7 @@ register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd")
 
 
 def _mock_global_server_args(backend="pytorch"):
-    from sglang.srt.layers import sampler as sampler_mod
+    from sglang.srt.runtime_context import get_parallel
     from sglang.srt.server_args import (
         ServerArgs,
         set_global_server_args_for_scheduler,
@@ -37,7 +37,11 @@ def _mock_global_server_args(backend="pytorch"):
     class _DummyTPGroup:
         device_group = None
 
-    sampler_mod.get_tp_group = lambda: _DummyTPGroup()
+    # `Sampler.__init__` asks the context for the group; state one for the rest
+    # of the process, since this process has no distributed init. Not the scoped
+    # `override()`: its context manager would be collected here and take the
+    # value back down with it.
+    get_parallel().override_permanently(tp_group=_DummyTPGroup())
     from sglang.srt.runtime_context import get_flags
 
     get_flags().dp.enabled = False

test/registered/unit/disaggregation/test_encode_server.py

@@ -1012,7 +1012,7 @@ class TestEncoderDelivery(CustomTestCase):
 
             with (
                 patch(
-                    "sglang.srt.disaggregation.encoder.server.get_tp_group",
+                    "sglang.srt.distributed.parallel_state.get_tp_group",
                     return_value=TPGroup(),
                 ),
                 patch(
@@ -1052,7 +1052,7 @@ class TestEncoderDelivery(CustomTestCase):
 
             with (
                 patch(
-                    "sglang.srt.disaggregation.encoder.server.get_tp_group",
+                    "sglang.srt.distributed.parallel_state.get_tp_group",
                     return_value=TPGroup(),
                 ),
                 patch(
@@ -1096,7 +1096,7 @@ class TestEncoderDelivery(CustomTestCase):
 
             with (
                 patch(
-                    "sglang.srt.disaggregation.encoder.server.get_tp_group",
+                    "sglang.srt.distributed.parallel_state.get_tp_group",
                     return_value=TPGroup(),
                 ),
                 patch(

test/registered/unit/disaggregation/test_register_to_bootstrap.py

@@ -201,7 +201,9 @@ class TestRegisterToBootstrap(CustomTestCase):
         self.assertIn("10.0.0.1", url_used)
 
     @patch("sglang.srt.disaggregation.common.conn.requests.put")
-    @patch("sglang.srt.disaggregation.common.conn.get_world_group")
+    # The consumer reads the group through `get_parallel()`, which reads
+    # through to the canonical getter, so that is where the stub belongs.
+    @patch("sglang.srt.distributed.parallel_state.get_world_group")
     def test_rust_attention_dp_replicates_complete_topology_across_hosts(
         self, mock_world_group, mock_put
     ):

test/registered/unit/layers/quantization/test_mxfp4_sm100_trtllm_gen.py

@@ -31,6 +31,20 @@ if not is_sm100_supported():
 GROUP_SIZE = 32  # MXFP4 block size
 
 
+@pytest.fixture
+def stated_tp_group():
+    """A TP group for a test that runs in a process without one.
+
+    The production call passes the group *into* `use_symmetric_memory`, so
+    stubbing that context manager does not stop the read -- the argument is
+    evaluated first. Stating it on the context answers every spelling.
+    """
+    from sglang.srt.runtime_context import get_parallel
+
+    with get_parallel().override(tp_group=None):
+        yield
+
+
 class _MockLayer:
     """Hand-built ``FusedMoE`` stand-in (avoids distributed init)."""
 
@@ -283,7 +297,7 @@ def _ref_trtllm(x, layer, method, precision, top_k, router_logits):
     ],
 )
 def test_apply_trtllm_gen_matches_flashinfer_direct(
-    tokens, num_experts, hidden, inter, top_k, precision, monkeypatch
+    tokens, num_experts, hidden, inter, top_k, precision, monkeypatch, stated_tp_group
 ):
     """``Mxfp4MoEMethod.apply`` (SM100 branch) must produce the same output as a
     direct ``trtllm_fp4_block_scale_moe`` call fed the same inputs.
@@ -301,7 +315,6 @@ def test_apply_trtllm_gen_matches_flashinfer_direct(
         fi_trtllm_mod, "use_symmetric_memory", lambda *a, **kw: nullcontext()
     )
     monkeypatch.setattr(fi_trtllm_mod, "is_allocation_symmetric", lambda: False)
-    monkeypatch.setattr(fi_trtllm_mod, "get_tp_group", lambda: None)
 
     fixtures = _make_random_mxfp4(num_experts, hidden, inter)
     x = torch.randn(tokens, hidden, dtype=torch.bfloat16, device="cuda") * 0.1

test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py

@@ -17,6 +17,20 @@ from sglang.test.ci.ci_register import register_cuda_ci
 register_cuda_ci(est_time=14, stage="base-b", runner_config="1-gpu-small")
 
 
+@pytest.fixture
+def stated_tp_group():
+    """A TP group for a test that runs in a process without one.
+
+    The production call passes the group *into* `use_symmetric_memory`, so
+    stubbing that context manager does not stop the read -- the argument is
+    evaluated first. Stating it on the context answers every spelling.
+    """
+    from sglang.srt.runtime_context import get_parallel
+
+    with get_parallel().override(tp_group=None):
+        yield
+
+
 def _random_weights(num_experts: int, hidden: int, intermediate: int):
     generator = torch.Generator(device="cuda").manual_seed(0)
     w13 = torch.randint(
@@ -110,7 +124,7 @@ def test_dsv4_sm120_load_contract(monkeypatch, request):
     assert captured["fp4_scale_dtype"] == torch.float8_e8m0fnu
 
 
-def test_dsv4_sm120_matches_direct_flashinfer(monkeypatch):
+def test_dsv4_sm120_matches_direct_flashinfer(monkeypatch, stated_tp_group):
     if not torch.cuda.is_available():
         pytest.skip("CUDA required")
     if torch.cuda.get_device_capability()[0] != 12:
@@ -134,7 +148,6 @@ def test_dsv4_sm120_matches_direct_flashinfer(monkeypatch):
         runner_module, "use_symmetric_memory", lambda *args, **kwargs: nullcontext()
     )
     monkeypatch.setattr(runner_module, "is_allocation_symmetric", lambda: False)
-    monkeypatch.setattr(runner_module, "get_tp_group", lambda: None)
 
     num_experts, hidden, intermediate = 4, 256, 256
     w13, w2, w13_scale, w2_scale = _random_weights(num_experts, hidden, intermediate)
@@ -254,7 +267,7 @@ def test_dsv4_sm120_matches_direct_flashinfer(monkeypatch):
     assert torch.equal(actual, expected)
 
 
-def test_gpt_oss_sm120_padding_layout_and_kernel(monkeypatch):
+def test_gpt_oss_sm120_padding_layout_and_kernel(monkeypatch, stated_tp_group):
     if not torch.cuda.is_available():
         pytest.skip("CUDA required")
     if torch.cuda.get_device_capability() != (12, 0):
@@ -277,7 +290,6 @@ def test_gpt_oss_sm120_padding_layout_and_kernel(monkeypatch):
         runner_module, "use_symmetric_memory", lambda *args, **kwargs: nullcontext()
     )
     monkeypatch.setattr(runner_module, "is_allocation_symmetric", lambda: False)
-    monkeypatch.setattr(runner_module, "get_tp_group", lambda: None)
 
     num_experts, hidden, intermediate = 4, 160, 160
     padded_hidden = padded_intermediate = 256

test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py

@@ -63,6 +63,20 @@ from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
 GROUP_SIZE = 32  # MXFP4 block size
 
 
+@pytest.fixture
+def stated_tp_group():
+    """A TP group for a test that runs in a process without one.
+
+    The production call passes the group *into* `use_symmetric_memory`, so
+    stubbing that context manager does not stop the read -- the argument is
+    evaluated first. Stating it on the context answers every spelling.
+    """
+    from sglang.srt.runtime_context import get_parallel
+
+    with get_parallel().override(tp_group=None):
+        yield
+
+
 class _MockLayer:
     """Stand-in for ``FusedMoE`` carrying the attributes the SM90 helpers read.
 
@@ -336,7 +350,7 @@ def test_process_weights_matches_direct_interleave(num_experts, hidden, inter):
     ],
 )
 def test_apply_sm90_cutlass_matches_flashinfer_direct(
-    tokens, num_experts, hidden, inter, top_k, monkeypatch
+    tokens, num_experts, hidden, inter, top_k, monkeypatch, stated_tp_group
 ):
     """End-to-end: SGLang's ``_apply_sm90_cutlass`` must produce the same
     output as a direct FlashInfer ``cutlass_fused_moe`` call fed with the
@@ -352,7 +366,6 @@ def test_apply_sm90_cutlass_matches_flashinfer_direct(
         fi_cutlass_mod, "use_symmetric_memory", lambda *a, **kw: nullcontext()
     )
     monkeypatch.setattr(fi_cutlass_mod, "is_allocation_symmetric", lambda: False)
-    monkeypatch.setattr(fi_cutlass_mod, "get_tp_group", lambda: None)
     monkeypatch.setattr(
         fi_cutlass_mod.envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE,
         "get",
@@ -593,7 +606,7 @@ def test_humming_range_ignores_prerounded_hidden_tail():
     [(8, 256, 256, 1, 0), (8, 192, 192, 1, 0), (8, 256, 256, 2, 1)],
 )
 def test_apply_sm90_humming_matches_flashinfer_direct(
-    tokens, hidden, inter, ep_size, ep_rank, monkeypatch
+    tokens, hidden, inter, ep_size, ep_rank, monkeypatch, stated_tp_group
 ):
     """SGLang must forward the five Humming scales and enable the new kernel."""
     import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass as fi_cutlass_mod
@@ -602,7 +615,6 @@ def test_apply_sm90_humming_matches_flashinfer_direct(
         fi_cutlass_mod, "use_symmetric_memory", lambda *a, **kw: nullcontext()
     )
     monkeypatch.setattr(fi_cutlass_mod, "is_allocation_symmetric", lambda: False)
-    monkeypatch.setattr(fi_cutlass_mod, "get_tp_group", lambda: None)
     monkeypatch.setattr(
         fi_cutlass_mod.envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE,
         "get",
@@ -725,7 +737,7 @@ def _make_random_dsv4_mxfp4(num_experts, hidden, inter, seed=0):
     ],
 )
 def test_dsv4_apply_matches_flashinfer_direct(
-    tokens, num_experts, hidden, inter, top_k, monkeypatch
+    tokens, num_experts, hidden, inter, top_k, monkeypatch, stated_tp_group
 ):
     """End-to-end: SGLang's DSv4 ``Mxfp4FlashinferCutlassMoEMethod.apply``
     output must match a direct FlashInfer ``cutlass_fused_moe`` call with
@@ -741,7 +753,6 @@ def test_dsv4_apply_matches_flashinfer_direct(
         fi_cutlass_mod, "use_symmetric_memory", lambda *a, **kw: nullcontext()
     )
     monkeypatch.setattr(fi_cutlass_mod, "is_allocation_symmetric", lambda: False)
-    monkeypatch.setattr(fi_cutlass_mod, "get_tp_group", lambda: None)
 
     w13, w2, w13_s, w2_s = _make_random_dsv4_mxfp4(num_experts, hidden, inter)
     w1, w3 = w13.chunk(2, dim=1)

test/registered/unit/mem_cache/test_mem_pool_host.py

@@ -270,8 +270,9 @@ class TestHostMemoryBudget(CustomTestCase):
             unittest.mock.patch.object(
                 torch.distributed, "is_initialized", return_value=True
             ),
-            unittest.mock.patch.object(
-                base, "get_world_group", return_value=fake_group
+            unittest.mock.patch(
+                "sglang.srt.distributed.parallel_state.get_world_group",
+                return_value=fake_group,
             ),
         ):
             self.assertEqual(base.ranks_per_host(), 8)

test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py

@@ -217,7 +217,12 @@ class TestGraphPoolBorrow(CustomTestCase):
                 "sglang.srt.layers.dp_attention.is_dp_attention_enabled",
                 return_value=False,
             ),
-            patch("sglang.srt.distributed.get_tp_group", return_value=tp_group),
+            # `parallel_state`, not the package re-export: a stub on the
+            # re-export is never consulted.
+            patch(
+                "sglang.srt.distributed.parallel_state.get_tp_group",
+                return_value=tp_group,
+            ),
             patch(
                 "sglang.kernels.ops.speculative.sampling.tree_speculative_sampling_target_only",
                 side_effect=fake_sampling,

test/registered/unit/model_executor/test_kv_canary_headroom.py

@@ -60,9 +60,8 @@ class TestCanaryHeadroom(CustomTestCase):
                 ),
             ),
             patch.object(kv_pool_runtime.torch.cuda, "synchronize"),
-            patch.object(
-                kv_pool_runtime,
-                "get_world_group",
+            patch(
+                "sglang.srt.distributed.parallel_state.get_world_group",
                 return_value=SimpleNamespace(world_size=1, cpu_group=None),
             ),
             patch.object(kv_pool_runtime, "get_available_gpu_memory", return_value=20),

test/registered/unit/model_loader/test_prefetch_checkpoints.py

@@ -238,7 +238,7 @@ class TestPrefetchCheckpoints(CustomTestCase):
             patch("concurrent.futures.ThreadPoolExecutor", _InlineExecutor),
             patch("concurrent.futures.wait", side_effect=_wait_all),
             patch(
-                "sglang.srt.model_loader.weight_utils.get_world_group",
+                "sglang.srt.distributed.parallel_state.get_world_group",
                 return_value=FakeWorldGroup(),
             ),
             patch(

test/registered/unit/model_loader/test_presharded_loader.py

@@ -653,7 +653,8 @@ class TestStructuralSignature(unittest.TestCase):
         fake_group.all_gather_object.side_effect = lambda local: ["sig-pp0", "sig-pp1"]
 
         with mock.patch(
-            "sglang.srt.distributed.get_world_group", return_value=fake_group
+            "sglang.srt.distributed.parallel_state.get_world_group",
+            return_value=fake_group,
         ):
             agg_from_rank0 = (
                 PreshardedModelLoader._make_rank_invariant_structural_signature(
@@ -673,7 +674,8 @@ class TestStructuralSignature(unittest.TestCase):
             "sig-pp1-changed",
         ]
         with mock.patch(
-            "sglang.srt.distributed.get_world_group", return_value=fake_group
+            "sglang.srt.distributed.parallel_state.get_world_group",
+            return_value=fake_group,
         ):
             agg_changed = (
                 PreshardedModelLoader._make_rank_invariant_structural_signature(

test/registered/unit/model_loader/test_transformers_fallback.py

@@ -55,7 +55,7 @@ class TestTransformersFallbackSkipSubstrs(CustomTestCase):
 
         with (
             patch(
-                "sglang.srt.models.transformers.get_pp_group",
+                "sglang.srt.distributed.parallel_state.get_pp_group",
                 return_value=SimpleNamespace(),
             ),
             patch(

test/registered/unit/multimodal/test_vit_cuda_graph_runner.py

@@ -33,7 +33,7 @@ def _runner(*, use_data_parallel: bool) -> ViTCudaGraphRunner:
 def test_dp_vit_graph_capture_does_not_enter_tp_communication_capture():
     runner = _runner(use_data_parallel=True)
     with patch(
-        "sglang.srt.multimodal.vit_cuda_graph_runner.get_tp_group",
+        "sglang.srt.distributed.parallel_state.get_tp_group",
         side_effect=AssertionError("DP capture must be rank-local"),
     ):
         with runner._capture_context():
@@ -53,7 +53,7 @@ def test_non_dp_vit_graph_capture_uses_tp_communication_capture():
     group = SimpleNamespace(ca_comm=SimpleNamespace(capture=lambda: Capture()))
     runner = _runner(use_data_parallel=False)
     with patch(
-        "sglang.srt.multimodal.vit_cuda_graph_runner.get_tp_group", return_value=group
+        "sglang.srt.distributed.parallel_state.get_tp_group", return_value=group
     ):
         with runner._capture_context():
             pass