Read process groups through the runtime context (#40068)

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)