Introduce ModelRunner.ps ParallelState (#31161)

test/manual/test_forward_split_prefill.py

@@ -15,6 +15,7 @@ import torch
 
 from sglang.benchmark.one_batch import TreeCacheNamespace
 from sglang.srt.configs.model_config import ModelConfig
+from sglang.srt.distributed.parallel_state_wrapper import ParallelState
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.model_executor.model_runner import ModelRunner
@@ -57,14 +58,9 @@ class TestForwardSplitPrefill(CustomTestCase):
             model_config=cls.model_config,
             mem_fraction_static=cls.server_args.mem_fraction_static,
             gpu_id=0,
-            tp_rank=0,
-            tp_size=cls.tp_size,
-            pp_rank=0,
-            pp_size=1,
+            ps=ParallelState.trivial(tp_size=cls.tp_size),
             nccl_port=cls.port_args.nccl_port,
             server_args=cls.server_args,
-            moe_ep_rank=0,
-            moe_ep_size=1,
         )
 
         cls.tokenizer = get_tokenizer(