Add get_parallel(): a structured accessor for parallel-topology state (#28567)

python/sglang/srt/batch_overlap/two_batch_overlap.py

@@ -19,7 +19,6 @@ from sglang.srt.layers.communicator import (
     CommunicateSummableTensorPairFn,
     ScatterMode,
 )
-from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.layers.moe import (
     get_deepep_mode,
     get_moe_a2a_backend,
@@ -40,6 +39,7 @@ from sglang.srt.model_executor.forward_batch_info import (
     compute_position,
 )
 from sglang.srt.model_executor.forward_context import get_attn_backend
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.speculative.spec_info import SpecInput
 from sglang.srt.utils import BumpAllocator, empty_context, get_bool_env_var, is_hip
@@ -649,7 +649,7 @@ class TboForwardBatchPreparer:
             ), f"{key=} {old_value=} {num_tokens=} {batch=}"
             output_dict[key] = old_value[start_token_index:end_token_index]
 
-        attention_tp_size = get_attention_tp_size()
+        attention_tp_size = get_parallel().attn_tp_size
         output_dict["tbo_padded_len"] = (
             (end_token_index - start_token_index - 1) // attention_tp_size + 1
         ) * attention_tp_size

python/sglang/srt/layers/activation.py

@@ -24,8 +24,6 @@ from transformers import PretrainedConfig
 
 from sglang.srt.distributed import (
     divide,
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
 )
 from sglang.srt.environ import envs
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
@@ -35,6 +33,7 @@ from sglang.srt.model_executor.cuda_graph_config import (
     Phase,
     check_cuda_graph_backend,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
     cpu_has_amx_support,
@@ -356,7 +355,7 @@ class ScaledActivation(nn.Module):
         self.act = act_module
         self.input_is_parallel = input_is_parallel
         if input_is_parallel:
-            tp_size = get_tensor_model_parallel_world_size()
+            tp_size = get_parallel().tp_size
             intermediate_size_per_partition = divide(intermediate_size, tp_size)
         else:
             intermediate_size_per_partition = intermediate_size
@@ -373,7 +372,7 @@ class ScaledActivation(nn.Module):
     def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
         param_data = param.data
         if self.input_is_parallel:
-            tp_rank = get_tensor_model_parallel_rank()
+            tp_rank = get_parallel().tp_rank
             shard_size = param_data.shape[0]
             start_idx = tp_rank * shard_size
             loaded_weight = loaded_weight.narrow(0, start_idx, shard_size)

python/sglang/srt/layers/attention/aiter_backend.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from sglang.srt.runtime_context import get_parallel
+
 """
 end to end attention solution with aiter kernels
 """
@@ -24,7 +26,6 @@ from sglang.srt.layers.attention.utils import (
     get_num_kv_index_blocks_flashmla,
 )
 from sglang.srt.layers.dp_attention import (
-    get_attention_tp_size,
     is_dp_attention_enabled,
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
@@ -154,11 +155,11 @@ class AiterAttnBackend(AttentionBackend):
         self.speculative_num_steps = model_runner.server_args.speculative_num_steps
         self.topk = topk
         self.num_head = (
-            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+            model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
         )
         self.head_dim = model_runner.model_config.head_dim
         self.num_kv_head = model_runner.model_config.get_num_kv_heads(
-            get_attention_tp_size()
+            get_parallel().attn_tp_size
         )
         self.kv_cache_dtype = model_runner.kv_cache_dtype
 
@@ -2565,10 +2566,10 @@ class AiterIndicesUpdaterPrefill:
     def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend):
         # Parse Constants
         self.num_qo_heads = (
-            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+            model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
         )
         self.num_kv_heads = model_runner.model_config.get_num_kv_heads(
-            get_attention_tp_size()
+            get_parallel().attn_tp_size
         )
         self.head_dim = model_runner.model_config.head_dim
         self.data_type = model_runner.kv_cache_dtype
@@ -2781,7 +2782,7 @@ class AiterMultiStepDraftBackend:
             )
         self.max_context_len = self.attn_backends[0].max_context_len
         self.num_head = (
-            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+            model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
         )
         self.device = model_runner.device
         # Cached variables for generate_draft_decode_kv_indices

python/sglang/srt/layers/attention/cutlass_mla_backend.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from sglang.srt.runtime_context import get_parallel
+
 """
 Support attention backend for Cutlass MLA.
 
@@ -16,7 +18,6 @@ from sglang.srt.layers.attention.utils import (
     create_flashmla_kv_indices_triton,
     get_num_kv_index_blocks_flashmla,
 )
-from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.utils import is_cuda
 
@@ -62,14 +63,14 @@ class CutlassMLABackend(FlashInferMLAAttnBackend):
         )
 
         self.num_q_heads = (
-            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+            model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
         )
         self.num_kv_heads = model_runner.model_config.get_num_kv_heads(
-            get_attention_tp_size()
+            get_parallel().attn_tp_size
         )
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
         self.num_local_heads = (
-            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+            model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
         )
         self.forward_metadata: Union[CutlassMLADecodeMetadata] = None
         self.kv_lora_rank = model_runner.model_config.kv_lora_rank

python/sglang/srt/layers/attention/deepseek_v4_backend.py

@@ -20,6 +20,7 @@ import torch.nn.functional as F
 
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.runtime_context import get_parallel
 
 if envs.SGLANG_OPT_USE_COMPRESSOR_V2.get():
     # NOTE: should eventually be the only compressor backend
@@ -55,10 +56,6 @@ from sglang.srt.layers.attention.dsv4.quant_k_cache import (
 from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import (
     SparsePrefillChunkCache,
 )
-from sglang.srt.layers.dp_attention import (
-    get_attention_cp_rank,
-    get_attention_cp_size,
-)
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.speculative.eagle_utils import per_step_draft_out_cache_loc
@@ -311,8 +308,8 @@ class DSV4AttnMetadata:
     ]
 
     def apply_cp_reindex(self) -> None:
-        cp_rank = get_attention_cp_rank()
+        cp_rank = get_parallel().attn_cp_rank
-        cp_size = get_attention_cp_size()
+        cp_size = get_parallel().attn_cp_size
         idx = slice(cp_rank, None, cp_size)
         pre_global_len = self.seq_lens_casual.shape[0]
         assert pre_global_len % cp_size == 0, (

python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py

@@ -20,6 +20,7 @@ import torch.nn.functional as F
 
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.runtime_context import get_parallel
 
 if envs.SGLANG_OPT_USE_COMPRESSOR_V2.get():
     from sglang.srt.layers.attention.dsv4.compressor_v2 import (
@@ -33,6 +34,7 @@ else:
         FusedCompressMetadata,
         create_paged_compressor_data,
     )
+
 from sglang.srt.layers.attention.dsv4.indexer import C4IndexerBackendMixin
 from sglang.srt.layers.attention.dsv4.metadata import (
     PagedIndexerMetadata,
@@ -45,10 +47,6 @@ from sglang.srt.layers.attention.dsv4.metadata_kernel import (
 from sglang.srt.layers.attention.dsv4.quant_k_cache import (
     quant_to_nope_fp8_rope_bf16_pack_triton,
 )
-from sglang.srt.layers.dp_attention import (
-    get_attention_cp_rank,
-    get_attention_cp_size,
-)
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.speculative.eagle_utils import per_step_draft_out_cache_loc
@@ -289,8 +287,8 @@ class DSV4AttnMetadata:
     ]
 
     def apply_cp_reindex(self) -> None:
-        cp_rank = get_attention_cp_rank()
+        cp_rank = get_parallel().attn_cp_rank
-        cp_size = get_attention_cp_size()
+        cp_size = get_parallel().attn_cp_size
         idx = slice(cp_rank, None, cp_size)
         pre_global_len = self.seq_lens_casual.shape[0]
         assert pre_global_len % cp_size == 0, (
@@ -1203,7 +1201,7 @@ class DeepseekV4HipRadixBackend(
         # HIP backend (DeepseekV4HipRadixBackend, selected only when is_hip()).
         # The NVIDIA path uses DeepseekV4AttnBackend and never reaches here, so
         # these CP changes do not affect B200/H200 execution.
-        _cp_size = get_attention_cp_size()
+        _cp_size = get_parallel().attn_cp_size
         _cp_active = (
             _cp_size > 1
             and is_dsa_prefill_cp_round_robin_split()
@@ -1214,7 +1212,7 @@ class DeepseekV4HipRadixBackend(
         final_pos_full = final_pos
         positions_full = positions
         if _cp_active:
-            _sl = slice(get_attention_cp_rank(), None, _cp_size)
+            _sl = slice(get_parallel().attn_cp_rank, None, _cp_size)
             state_slot = state_slot[_sl].contiguous()
             chunk_start = chunk_start[_sl].contiguous()
             cu_q = cu_q[_sl].contiguous()

python/sglang/srt/layers/attention/dsa/dsa_indexer.py

@@ -26,6 +26,7 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo
     get_tc_piecewise_forward_context,
     is_in_tc_piecewise_cuda_graph,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.state_capturer.indexer_topk import (
     maybe_capture_indexer_topk,
 )
@@ -72,8 +73,6 @@ if is_npu():
     from sglang.srt.hardware_backend.npu.utils import get_indexer_weight_stream
 
 from sglang.srt.distributed import (
-    get_attn_context_model_parallel_rank,
-    get_attn_context_model_parallel_world_size,
     get_attn_tp_group,
 )
 from sglang.srt.distributed.parallel_state import get_pp_group
@@ -330,8 +329,8 @@ class Indexer(MultiPlatformOp):
         self.alt_stream = alt_stream
         self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
         if self.dsa_enable_prefill_cp:
-            self.cp_size = get_attn_context_model_parallel_world_size()
+            self.cp_size = get_parallel().attn_cp_size
-            self.cp_rank = get_attn_context_model_parallel_rank()
+            self.cp_rank = get_parallel().attn_cp_rank
         else:
             self.cp_size = None
             self.cp_rank = None

python/sglang/srt/layers/attention/dsa/utils.py

@@ -8,10 +8,8 @@ import triton.language as tl
 from sglang.srt.environ import envs
 from sglang.srt.layers.dp_attention import (
     DpPaddingMode,
-    get_attention_cp_rank,
-    get_attention_cp_size,
-    get_attention_dp_rank,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import get_bool_env_var, is_hip
 from sglang.srt.utils.common import ceil_align, ceil_div
@@ -85,7 +83,7 @@ def is_dsa_prefill_cp_round_robin_split():
 def can_dsa_prefill_cp_round_robin_split(forward_batch: "ForwardBatch"):
     if not forward_batch.forward_mode.is_context_parallel_extend():
         return False
-    cp_size = get_attention_cp_size()
+    cp_size = get_parallel().attn_cp_size
     seq_len = sum(forward_batch.extend_seq_lens_cpu)
     return (
         is_dsa_prefill_cp_round_robin_split()
@@ -108,8 +106,8 @@ def dsa_cp_round_robin_split_data(input_: Union[torch.Tensor, List]):
     | dp_atten_tp3: token3, token7, token11, token15, token19, ... |
     |   +-------------------------+
     """
-    cp_size = get_attention_cp_size()
+    cp_size = get_parallel().attn_cp_size
-    cp_rank = get_attention_cp_rank()
+    cp_rank = get_parallel().attn_cp_rank
     if isinstance(input_, (tuple, list)):
         indices = range(cp_rank, len(input_), cp_size)
         return input_[indices]
@@ -133,7 +131,7 @@ def cal_padded_tokens(forward_batch: "ForwardBatch"):
 
     global_num_tokens = forward_batch.global_num_tokens_cpu.copy()
     sync_group_size = len(global_num_tokens)
-    attn_cp_size = get_attention_cp_size()
+    attn_cp_size = get_parallel().attn_cp_size
     # Must match the CP padding in ForwardBatch.prepare_mlp_sync_batch.
     cp_align_size = get_cp_padding_align_size()
     for i in range(sync_group_size):
@@ -144,7 +142,7 @@ def cal_padded_tokens(forward_batch: "ForwardBatch"):
     if dp_padding_mode.is_max_len():
         tokens = max(global_num_tokens)
     elif len(global_num_tokens) > 1:
-        tokens = global_num_tokens[get_attention_dp_rank()]
+        tokens = global_num_tokens[get_parallel().attn_dp_rank]
     else:
         tokens = global_num_tokens[0]
     if can_dsa_prefill_cp_round_robin_split(forward_batch):
@@ -153,7 +151,7 @@ def cal_padded_tokens(forward_batch: "ForwardBatch"):
 
 
 def pad_dsa_cache_seqlens(forward_batch: "ForwardBatch", dsa_cache_seqlens):
-    attn_cp_size = get_attention_cp_size()
+    attn_cp_size = get_parallel().attn_cp_size
     needs_cp_pad = attn_cp_size > 1 and can_dsa_prefill_cp_round_robin_split(
         forward_batch
     )
@@ -219,8 +217,8 @@ def dsa_cp_round_robin_split_q_seqs_kernel(
 
 
 def dsa_cp_round_robin_split_q_seqs_cpu(extend_seqs):
-    cp_size = get_attention_cp_size()
+    cp_size = get_parallel().attn_cp_size
-    cp_rank = get_attention_cp_rank()
+    cp_rank = get_parallel().attn_cp_rank
     extra_seq = 0
     q_seqs = []
     for bs, cur_len in enumerate(extend_seqs):
@@ -245,8 +243,8 @@ def dsa_cp_round_robin_split_q_seqs(
     bs_idx_cpu(List) and bs_idx(torch.Tensor): marks which sequences are ultimately selected,
         i.e., those with a partitioned length greater than zero.
     """
-    cp_size = get_attention_cp_size()
+    cp_size = get_parallel().attn_cp_size
-    cp_rank = get_attention_cp_rank()
+    cp_rank = get_parallel().attn_cp_rank
     # len(ret_q_lens_cpu) == len(bs_idx_cpu)
     ret_q_lens_cpu, bs_idx_cpu = dsa_cp_round_robin_split_q_seqs_cpu(extend_seqs_cpu)
     ret_q_lens = torch.empty(

python/sglang/srt/layers/attention/dsa_backend.py

@@ -15,6 +15,7 @@ from typing import (
 import torch
 
 from sglang.srt.configs.model_config import get_dsa_index_topk, is_deepseek_dsa
+from sglang.srt.runtime_context import get_parallel
 
 logger = logging.getLogger(__name__)
 from sglang.srt.environ import envs
@@ -48,7 +49,6 @@ from sglang.srt.layers.attention.utils import (
     mla_quantize_and_rope_for_fp8,
     seqlens_expand_triton,
 )
-from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.utils import is_cuda, is_hip, is_sm100_supported
 
@@ -311,7 +311,7 @@ class DeepseekSparseAttnBackend(
         self.dsa_index_topk = get_dsa_index_topk(model_runner.model_config.hf_config)
         self.max_context_len = model_runner.model_config.context_len
         self.num_q_heads = (
-            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+            model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
         )
         self.kv_cache_dim = model_runner.token_to_kv_pool.kv_cache_dim
         self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim

python/sglang/srt/layers/attention/dsv4/compressor.py

@@ -19,7 +19,6 @@ from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
 from sglang.srt.layers.attention.dsv4.quant_k_cache import (
     quant_to_nope_fp8_rope_bf16_pack_triton,
 )
-from sglang.srt.layers.dp_attention import get_attention_cp_size
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import ReplicatedLinear
 from sglang.srt.layers.utils.cp_utils import cp_all_gather_rerange_output
@@ -28,6 +27,7 @@ from sglang.srt.mem_cache.deepseek_v4_compress_state import (
 )
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.srt.models.deepseek_v2 import _is_hip
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, get_bool_env_var, set_weight_attrs
 
 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
@@ -427,7 +427,7 @@ class Compressor(nn.Module):
         if dsa_use_prefill_cp(forward_batch):
             kv_score = cp_all_gather_rerange_output(
                 kv_score,
-                get_attention_cp_size(),
+                get_parallel().attn_cp_size,
                 forward_batch,
                 torch.cuda.current_stream(),
             )

python/sglang/srt/layers/attention/dual_chunk_flashattention_backend.py

@@ -19,12 +19,12 @@ from sglang.jit_kernel.flash_attention import (
     flash_attn_varlen_func,
     flash_attn_with_kvcache,
 )
-from sglang.srt.distributed.parallel_state import get_tensor_model_parallel_rank
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.attention.flashattention_backend import (
     FlashAttentionMetadata,
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.runtime_context import get_parallel
 
 if TYPE_CHECKING:
     from sglang.srt.layers.radix_attention import RadixAttention
@@ -170,7 +170,7 @@ class DualChunkFlashAttentionBackend(AttentionBackend):
         layer_sparse_attention_config = {
             int(i): j for i, j in self.sparse_attention_config[layer_idx].items()
         }
-        start_head = self.num_heads * get_tensor_model_parallel_rank()
+        start_head = self.num_heads * get_parallel().tp_rank
         end_head = start_head + self.num_heads
         return [layer_sparse_attention_config[i] for i in range(start_head, end_head)]
 

python/sglang/srt/layers/attention/flashinfer_backend.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from sglang.srt.runtime_context import get_parallel
+
 """
 Support different attention backends.
 Now there are two backends: FlashInfer and Triton.
@@ -24,10 +26,6 @@ from sglang.srt.layers.attention.utils import (
     assert_buffer_fits,
     create_flashinfer_kv_indices_triton,
 )
-from sglang.srt.layers.dp_attention import (
-    get_attention_cp_size,
-    get_attention_tp_size,
-)
 from sglang.srt.layers.radix_attention import AttentionType
 from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
 from sglang.srt.mem_cache.memory_pool import KVWriteLoc
@@ -67,9 +65,9 @@ def _cuda_graph_capture_max_bs(server_args, max_bs: int) -> int:
     if server_args.enable_two_batch_overlap:
         mul_base *= 2
     if require_gathered_buffer(server_args):
-        mul_base *= get_attention_tp_size()
+        mul_base *= get_parallel().attn_tp_size
-    if mul_base % get_attention_cp_size() != 0:
+    if mul_base % get_parallel().attn_cp_size != 0:
-        mul_base *= get_attention_cp_size()
+        mul_base *= get_parallel().attn_cp_size
     return (max_bs + mul_base - 1) // mul_base * mul_base
 
 
@@ -208,9 +206,9 @@ class FlashInferAttnBackend(AttentionBackend):
         self.decode_use_tensor_cores = should_use_tensor_core(
             kv_cache_dtype=model_runner.kv_cache_dtype,
             num_attention_heads=model_runner.model_config.num_attention_heads
-            // get_attention_tp_size(),
+            // get_parallel().attn_tp_size,
             num_kv_heads=model_runner.model_config.get_num_kv_heads(
-                get_attention_tp_size()
+                get_parallel().attn_tp_size
             ),
         )
         self.max_context_len = model_runner.model_config.context_len
@@ -1005,10 +1003,10 @@ class FlashInferIndicesUpdaterDecode:
     def __init__(self, model_runner: ModelRunner, attn_backend: FlashInferAttnBackend):
         # Parse Constants
         self.num_qo_heads = (
-            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+            model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
         )
         self.num_kv_heads = model_runner.model_config.get_num_kv_heads(
-            get_attention_tp_size()
+            get_parallel().attn_tp_size
         )
         self.head_dim = model_runner.model_config.head_dim
         self.data_type = model_runner.kv_cache_dtype
@@ -1273,10 +1271,10 @@ class FlashInferIndicesUpdaterPrefill:
     def __init__(self, model_runner: ModelRunner, attn_backend: FlashInferAttnBackend):
         # Parse Constants
         self.num_qo_heads = (
-            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+            model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
         )
         self.num_kv_heads = model_runner.model_config.get_num_kv_heads(
-            get_attention_tp_size()
+            get_parallel().attn_tp_size
         )
         self.head_dim = model_runner.model_config.head_dim
         self.data_type = model_runner.kv_cache_dtype

python/sglang/srt/layers/attention/flashinfer_mla_backend.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from sglang.srt.runtime_context import get_parallel
+
 """
 Support attention backend for flashinfer MLA.
 The flashinfer_mla_disable_ragged flag controls whether to use ragged prefill wrapper and defaults to be false.
@@ -21,7 +23,6 @@ from sglang.srt.layers.attention.flashinfer_backend import (
     create_flashinfer_kv_indices_triton,
 )
 from sglang.srt.layers.attention.utils import assert_buffer_fits
-from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
     is_in_tc_piecewise_cuda_graph,
@@ -81,7 +82,7 @@ class FlashInferMhaChunkKVRunner:
     ):
         # Parse Constants
         self.num_local_heads = (
-            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+            model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
         )
         self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim
         self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim
@@ -639,7 +640,7 @@ class FlashInferMLAIndicesUpdaterDecode:
     def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend):
         # Parse Constants
         self.num_local_heads = (
-            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+            model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
         )
         self.kv_lora_rank = model_runner.model_config.kv_lora_rank
         self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim
@@ -748,7 +749,7 @@ class FlashInferMLAIndicesUpdaterPrefill:
     def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend):
         # Parse Constants
         self.num_local_heads = (
-            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+            model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
         )
         self.kv_lora_rank = model_runner.model_config.kv_lora_rank
         self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim

python/sglang/srt/layers/attention/flashmla_backend.py

@@ -17,9 +17,9 @@ from sglang.srt.layers.attention.utils import (
     create_flashmla_kv_indices_triton,
     get_num_kv_index_blocks_flashmla,
 )
-from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.runtime_context import get_parallel
 
 if TYPE_CHECKING:
     from sglang.srt.layers.radix_attention import RadixAttention
@@ -60,11 +60,11 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
         )
 
         self.num_q_heads = (
-            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+            model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
         )
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
         self.num_local_heads = (
-            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+            model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
         )
         self.forward_metadata: Union[FlashMLADecodeMetadata] = None
         self.kv_lora_rank = model_runner.model_config.kv_lora_rank

python/sglang/srt/layers/attention/linear/lightning_backend.py

@@ -15,6 +15,7 @@ from sglang.srt.layers.attention.linear.seg_la import SegLaMeta, seg_la_fwd
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.model_runner import ModelRunner
+from sglang.srt.runtime_context import get_parallel
 
 logger = logging.getLogger(__name__)
 
@@ -133,13 +134,9 @@ class LightningAttentionBackend(MambaAttnBackendBase):
         slopes = torch.tensor(
             get_slopes(n_attention_heads), dtype=torch.float32
         ).reshape(n_attention_heads, 1, 1)
-        from sglang.srt.layers.dp_attention import (
-            get_attention_tp_rank,
-            get_attention_tp_size,
-        )
 
-        tp_heads = n_attention_heads // get_attention_tp_size()
+        tp_heads = n_attention_heads // get_parallel().attn_tp_size
-        tp_rank = get_attention_tp_rank()
+        tp_rank = get_parallel().attn_tp_rank
         if num_hidden_layers <= 1:
             slope_rate_list = [slopes * (1 + 1e-5)]
         else:

python/sglang/srt/layers/attention/mamba/mamba.py

@@ -10,8 +10,6 @@ from sglang.srt.configs.mamba_utils import (
 )
 from sglang.srt.distributed import (
     divide,
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
 )
 from sglang.srt.layers.attention.mamba.mamba2_metadata import Mamba2Metadata
 from sglang.srt.layers.attention.mamba.mixer2_rms_norm_gated import Mixer2RMSNormGated
@@ -20,8 +18,6 @@ from sglang.srt.layers.attention.mamba.ops import (
     selective_state_update,
 )
 from sglang.srt.layers.dp_attention import (
-    get_attention_tp_rank,
-    get_attention_tp_size,
     is_dp_attention_enabled,
 )
 from sglang.srt.layers.linear import (
@@ -36,6 +32,7 @@ from sglang.srt.model_loader.weight_utils import (
     composed_weight_loader,
     sharded_weight_loader,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import (
     is_cpu,
     is_cuda,
@@ -232,11 +229,11 @@ class MambaMixer2(torch.nn.Module):
         # - NOTE: currently for the world size DOES NOT divide groups
         #   case, we only support the case when n_groups == 1
         if is_dp_attention_enabled():
-            self.tp_size = get_attention_tp_size()
+            self.tp_size = get_parallel().attn_tp_size
-            self.tp_rank = get_attention_tp_rank()
+            self.tp_rank = get_parallel().attn_tp_rank
         else:
-            self.tp_size = get_tensor_model_parallel_world_size()
+            self.tp_size = get_parallel().tp_size
-            self.tp_rank = get_tensor_model_parallel_rank()
+            self.tp_rank = get_parallel().tp_rank
 
         self.num_heads = num_heads = cache_params.shape.num_heads
         self.head_dim = cache_params.shape.head_dim

python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py

@@ -6,20 +6,15 @@ from sglang.srt.distributed.communication_op import (
     tensor_model_parallel_all_gather,
     tensor_model_parallel_all_reduce,
 )
-from sglang.srt.distributed.parallel_state import (
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
-)
 from sglang.srt.layers.attention.fla.layernorm_gated import rms_norm_gated
 from sglang.srt.layers.dp_attention import (
     attn_tp_all_reduce,
     get_attention_tp_group,
-    get_attention_tp_rank,
-    get_attention_tp_size,
     is_dp_attention_enabled,
 )
 from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.model_loader.weight_utils import sharded_weight_loader
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils.common import set_weight_attrs
 
 
@@ -34,11 +29,11 @@ class Mixer2RMSNormGated(MultiPlatformOp):
         super().__init__()
         self.use_attn_tp_group = is_dp_attention_enabled()
         if self.use_attn_tp_group:
-            self.tp_size = get_attention_tp_size()
+            self.tp_size = get_parallel().attn_tp_size
-            self.tp_rank = get_attention_tp_rank()
+            self.tp_rank = get_parallel().attn_tp_rank
         else:
-            self.tp_size = get_tensor_model_parallel_world_size()
+            self.tp_size = get_parallel().tp_size
-            self.tp_rank = get_tensor_model_parallel_rank()
+            self.tp_rank = get_parallel().tp_rank
         self.full_hidden_size = full_hidden_size
         self.group_size = full_hidden_size // full_n_groups
         self.per_rank_hidden_size = full_hidden_size // self.tp_size

python/sglang/srt/layers/attention/triton_backend.py

@@ -12,12 +12,12 @@ from sglang.srt.layers.attention.triton_ops.kv_indices import (
     create_flashinfer_kv_indices_triton,
 )
 from sglang.srt.layers.attention.triton_ops.metadata import get_num_kv_splits_triton
-from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.layers.radix_attention import AttentionType
 from sglang.srt.mem_cache.memory_pool import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.speculative.spec_utils import (
     draft_kv_indices_buffer_width,
     draft_kv_indices_used_len,
@@ -132,10 +132,10 @@ class TritonAttnBackend(AttentionBackend):
         self.speculative_num_steps = model_runner.server_args.speculative_num_steps
         self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA
         self.num_head = (
-            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+            model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
         )
         self.num_kv_head = model_runner.model_config.get_num_kv_heads(
-            get_attention_tp_size()
+            get_parallel().attn_tp_size
         )
         # The decode triton kernel derives attn_lse offsets from attn_logits
         # strides via integer division by v_head_dim (the "// Lv" trick in
@@ -1386,7 +1386,7 @@ class TritonMultiStepDraftBackend:
             )
         self.max_context_len = self.attn_backends[0].max_context_len
         self.num_head = (
-            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+            model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
         )
         self.device = model_runner.device
         # Cached variables for generate_draft_decode_kv_indices

python/sglang/srt/layers/attention/trtllm_mla_backend.py

@@ -33,12 +33,12 @@ from sglang.srt.layers.attention.utils import (
     concat_mla_absorb_q_general,
     mla_quantize_and_rope_for_fp8,
 )
-from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
     is_in_tc_piecewise_cuda_graph,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import is_flashinfer_available, is_float4_e2m1fn_x2
 
@@ -149,9 +149,9 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         config = model_runner.model_config
 
         # Model parameters
-        self.num_q_heads = config.num_attention_heads // get_attention_tp_size()
+        self.num_q_heads = config.num_attention_heads // get_parallel().attn_tp_size
-        self.num_kv_heads = config.get_num_kv_heads(get_attention_tp_size())
+        self.num_kv_heads = config.get_num_kv_heads(get_parallel().attn_tp_size)
-        self.num_local_heads = config.num_attention_heads // get_attention_tp_size()
+        self.num_local_heads = config.num_attention_heads // get_parallel().attn_tp_size
 
         # MLA-specific dimensions
         self.kv_lora_rank = config.kv_lora_rank

python/sglang/srt/layers/attention/vision.py

@@ -14,8 +14,8 @@ from einops import rearrange
 
 from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm as can_use_jit_qk_norm
 from sglang.srt.environ import envs
-from sglang.srt.layers.dp_attention import get_attention_tp_rank, get_attention_tp_size
 from sglang.srt.models.utils import apply_qk_norm
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import (
     cpu_has_amx_support,
     get_bool_env_var,
@@ -347,7 +347,7 @@ class VisionTritonAttention(nn.Module):
         use_data_parallel = (
             kwargs["use_data_parallel"] if "use_data_parallel" in kwargs else False
         )
-        self.tp_size = 1 if use_data_parallel else get_attention_tp_size()
+        self.tp_size = 1 if use_data_parallel else get_parallel().attn_tp_size
 
     def forward(
         self,
@@ -420,7 +420,7 @@ class VisionFlash3Attention(nn.Module):
         use_data_parallel = (
             kwargs["use_data_parallel"] if "use_data_parallel" in kwargs else False
         )
-        self.tp_size = 1 if use_data_parallel else get_attention_tp_size()
+        self.tp_size = 1 if use_data_parallel else get_parallel().attn_tp_size
 
     def forward(
         self,
@@ -925,8 +925,8 @@ class VisionAttention(nn.Module):
                 DeprecationWarning,
                 stacklevel=2,
             )
-        self.tp_size = 1 if use_data_parallel else get_attention_tp_size()
+        self.tp_size = 1 if use_data_parallel else get_parallel().attn_tp_size
-        self.tp_rank = 0 if use_data_parallel else get_attention_tp_rank()
+        self.tp_rank = 0 if use_data_parallel else get_parallel().attn_tp_rank
         self.dropout = dropout
         num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
         self.head_size = head_dim if head_dim is not None else embed_dim // num_heads

python/sglang/srt/layers/attention/vision_utils.py

@@ -2,12 +2,12 @@
 
 import torch
 
-from sglang.srt.layers.dp_attention import get_attention_tp_size
+from sglang.srt.runtime_context import get_parallel
 
 
 def update_vit_attn_dummy_heads_config(config):
     """Update HF config to ensure vision attention num_attention_heads is divisible by tp_size"""
-    tp_size = get_attention_tp_size()
+    tp_size = get_parallel().attn_tp_size
     num_heads = getattr(
         config.vision_config,
         "num_heads",

python/sglang/srt/layers/attention/wave_backend.py

@@ -12,8 +12,8 @@ from sglang.srt.layers.attention.triton_ops.kv_indices import (
     create_flashinfer_kv_indices_triton,
 )
 from sglang.srt.layers.attention.triton_ops.metadata import get_num_kv_splits_triton
-from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import get_bool_env_var, get_device_core_count
 
 if TYPE_CHECKING:
@@ -95,10 +95,10 @@ class WaveAttnBackend(AttentionBackend):
         self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens
 
         self.num_head = (
-            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+            model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
         )
         self.num_kv_head = model_runner.model_config.get_num_kv_heads(
-            get_attention_tp_size()
+            get_parallel().attn_tp_size
         )
 
         self.static_kv_splits = get_bool_env_var(

python/sglang/srt/layers/clippable_linear.py

@@ -27,7 +27,6 @@ from typing import Optional, Tuple
 import torch
 import torch.nn as nn
 
-from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.layers.linear import (
     ColumnParallelLinear,
     MergedColumnParallelLinear,
@@ -35,6 +34,7 @@ from sglang.srt.layers.linear import (
     RowParallelLinear,
 )
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix
 
 _INF = float("inf")
@@ -128,7 +128,7 @@ class ClippableQKVParallelLinear(nn.Module):
         prefix: str = "",
     ):
         super().__init__()
-        tp_size = get_attention_tp_size()
+        tp_size = get_parallel().attn_tp_size
         self.q_size = (total_num_heads // tp_size) * head_size
         self.kv_size = (total_num_kv_heads // tp_size) * head_size
 
@@ -192,7 +192,7 @@ class ClippableGLUParallelLinear(nn.Module):
         prefix: str = "",
     ):
         super().__init__()
-        tp_size = get_attention_tp_size()
+        tp_size = get_parallel().attn_tp_size
         self.proj_size = hidden_size // tp_size
 
         self.linear = MergedColumnParallelLinear(
@@ -255,7 +255,7 @@ class ClippableGateUpParallelLinear(nn.Module):
         prefix: str = "",
     ):
         super().__init__()
-        tp_size = get_attention_tp_size()
+        tp_size = get_parallel().attn_tp_size
         self.proj_size = intermediate_size // tp_size
 
         self.gate_up_proj = MergedColumnParallelLinear(

python/sglang/srt/layers/communicator.py

@@ -23,8 +23,6 @@ import torch
 from sglang.srt.distributed import (
     attention_tensor_model_parallel_all_reduce,
     attention_tensor_model_parallel_quant_all_reduce,
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
     get_tp_group,
     moe_tensor_model_parallel_all_reduce,
     tensor_model_parallel_all_reduce,
@@ -44,12 +42,7 @@ from sglang.srt.layers.dp_attention import (
     dp_gather_replicate,
     dp_reduce_scatter_tensor,
     dp_scatter,
-    get_attention_cp_rank,
-    get_attention_cp_size,
-    get_attention_dp_size,
     get_attention_tp_group,
-    get_attention_tp_rank,
-    get_attention_tp_size,
     get_dp_global_num_tokens,
     get_global_dp_buffer,
     get_local_dp_buffer,
@@ -77,6 +70,7 @@ from sglang.srt.model_executor.cuda_graph_config import (
     check_cuda_graph_backend,
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.utils import (
@@ -189,7 +183,7 @@ def apply_aiter_all_reduce_fusion(input_tensor: torch.Tensor):
         and total_bytes > 0
         and n <= 16384
         and total_bytes <= 8 * 1024 * 8192
-        and get_tensor_model_parallel_world_size() != 6
+        and get_parallel().tp_size != 6
         and not is_dp_attention_enabled()
         and get_global_server_args().enable_aiter_allreduce_fusion
     )
@@ -276,7 +270,7 @@ class AttnTpContext:
             and (_is_cuda or _is_npu)
             and q_lora_rank is not None
             and not is_dsa
-            and get_tensor_model_parallel_world_size() > 1
+            and get_parallel().tp_size > 1
             and not is_dp_attention_enabled()
             and get_moe_a2a_backend().is_none()
             and not enable_moe_dense_fully_dp()
@@ -801,7 +795,7 @@ class LayerCommunicator:
                 or (
                     _use_aiter
                     and batch_size > 0
-                    and get_tensor_model_parallel_world_size() != 6
+                    and get_parallel().tp_size != 6
                     and get_global_server_args().enable_aiter_allreduce_fusion
                 )
             )
@@ -828,13 +822,13 @@ class CommunicateContext:
 
     @classmethod
     def init_new(cls):
-        attn_tp_rank = get_attention_tp_rank()
+        attn_tp_rank = get_parallel().attn_tp_rank
-        attn_tp_size = get_attention_tp_size()
+        attn_tp_size = get_parallel().attn_tp_size
-        attn_dp_size = get_attention_dp_size()
+        attn_dp_size = get_parallel().attn_dp_size
-        attn_cp_size = get_attention_cp_size()
+        attn_cp_size = get_parallel().attn_cp_size
-        attn_cp_rank = get_attention_cp_rank()
+        attn_cp_rank = get_parallel().attn_cp_rank
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
-        tp_rank = get_tensor_model_parallel_rank()
+        tp_rank = get_parallel().tp_rank
         moe_cp_size = get_moe_cp_size()
         process_group_sizes = {
             ScatterMode.SCATTERED: 1,
@@ -1312,7 +1306,7 @@ class CommunicateSummableTensorPairFn:
         context: CommunicateContext,
         allow_reduce_scatter: bool = False,
     ):
-        if get_tensor_model_parallel_world_size() == get_attention_dp_size():
+        if get_parallel().tp_size == get_parallel().attn_dp_size:
             group = get_tp_group()
         else:
             group = get_attention_tp_group()
@@ -1402,7 +1396,7 @@ class CommunicateSummableTensorPairFn:
 
         # DP scatter (if DP attention is enabled)
         if context.attn_dp_size > 1:
-            if get_tensor_model_parallel_world_size() == get_attention_dp_size():
+            if get_parallel().tp_size == get_parallel().attn_dp_size:
                 group = get_tp_group()
             else:
                 group = get_attention_tp_group()

python/sglang/srt/layers/communicator_dsa_cp.py

@@ -35,14 +35,11 @@ from sglang.srt.layers.dp_attention import (
     attn_cp_all_gather_into_tensor,
     attn_cp_reduce_scatter_tensor,
     get_attention_cp_group,
-    get_attention_cp_rank,
-    get_attention_cp_size,
-    get_attention_dp_size,
-    get_attention_tp_size,
     get_local_dp_buffer,
 )
 from sglang.srt.layers.utils.cp_utils import mla_use_prefill_cp
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.runtime_context import get_parallel
 
 
 def dsa_enable_prefill_cp():
@@ -53,8 +50,8 @@ def dsa_enable_prefill_cp():
 
 
 def dsa_cp_gather_hidden_states(hidden_states: torch.Tensor):
-    attn_dp_size = get_attention_dp_size()
+    attn_dp_size = get_parallel().attn_dp_size
-    attn_tp_size = get_attention_tp_size()
+    attn_tp_size = get_parallel().attn_tp_size
     assert attn_dp_size == 1 and attn_tp_size == 1
     hidden_states, local_hidden_states = (
         get_local_dp_buffer(get_attention_cp_group()),
@@ -65,11 +62,11 @@ def dsa_cp_gather_hidden_states(hidden_states: torch.Tensor):
 
 
 def dsa_cp_reduce_scatter_hidden_states(hidden_states: torch.Tensor):
-    attn_dp_size = get_attention_dp_size()
+    attn_dp_size = get_parallel().attn_dp_size
-    attn_tp_size = get_attention_tp_size()
+    attn_tp_size = get_parallel().attn_tp_size
     assert attn_dp_size == 1 and attn_tp_size == 1
-    cp_size = get_attention_cp_size()
+    cp_size = get_parallel().attn_cp_size
-    cp_rank = get_attention_cp_rank()
+    cp_rank = get_parallel().attn_cp_rank
     input_hidden_states = hidden_states
     hidden_states = hidden_states.tensor_split(cp_size)[cp_rank]
     attn_cp_reduce_scatter_tensor(hidden_states, input_hidden_states)

python/sglang/srt/layers/cp/base.py

@@ -29,6 +29,8 @@ from dataclasses import dataclass
 from enum import IntEnum
 from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple
 
+from sglang.srt.runtime_context import get_parallel
+
 if TYPE_CHECKING:
     from sglang.srt.model_executor.forward_batch_info import ForwardBatch
     from sglang.srt.server_args import ServerArgs
@@ -93,9 +95,8 @@ class ContextParallelStrategy(ABC):
 
     @property
     def cp_rank(self) -> int:
-        from sglang.srt.layers.dp_attention import get_attention_cp_rank
 
-        return get_attention_cp_rank()
+        return get_parallel().attn_cp_rank
 
     @property
     def per_layer_attn_cp_comm(self) -> bool:

python/sglang/srt/layers/flashinfer_comm_fusion.py

@@ -7,18 +7,13 @@ import torch.distributed as dist
 from torch.distributed import ProcessGroup
 
 from sglang.srt.distributed import (
-    get_attn_tensor_model_parallel_rank,
-    get_attn_tensor_model_parallel_world_size,
     get_attn_tp_group,
     get_moe_ep_group,
-    get_moe_expert_parallel_rank,
-    get_moe_expert_parallel_world_size,
-    get_moe_tensor_parallel_rank,
-    get_moe_tensor_parallel_world_size,
     get_moe_tp_group,
     get_tp_group,
 )
 from sglang.srt.distributed.parallel_state import in_the_same_node_as
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
     ceil_align,
@@ -637,17 +632,17 @@ def ensure_workspace_initialized(
         return False
 
     if use_attn_tp_group:
-        world_size = get_attn_tensor_model_parallel_world_size()
+        world_size = get_parallel().attn_tp_size
-        rank = get_attn_tensor_model_parallel_rank()
+        rank = get_parallel().attn_tp_rank
         coordinator = get_attn_tp_group()
     else:
-        if get_moe_expert_parallel_world_size() > 1:
+        if get_parallel().moe_ep_size > 1:
-            world_size = get_moe_expert_parallel_world_size()
+            world_size = get_parallel().moe_ep_size
-            rank = get_moe_expert_parallel_rank()
+            rank = get_parallel().moe_ep_rank
             coordinator = get_moe_ep_group()
         else:
-            world_size = get_moe_tensor_parallel_world_size()
+            world_size = get_parallel().moe_tp_size
-            rank = get_moe_tensor_parallel_rank()
+            rank = get_parallel().moe_tp_rank
             coordinator = get_moe_tp_group()
 
     # Always pass the coordinator's groups: flashinfer >=0.6.10 reads the
@@ -757,12 +752,12 @@ def flashinfer_allreduce_residual_rmsnorm(
         return None, None
 
     if use_attn_tp_group:
-        world_size = get_attn_tensor_model_parallel_world_size()
+        world_size = get_parallel().attn_tp_size
     else:
-        if get_moe_expert_parallel_world_size() > 1:
+        if get_parallel().moe_ep_size > 1:
-            world_size = get_moe_expert_parallel_world_size()
+            world_size = get_parallel().moe_ep_size
         else:
-            world_size = get_moe_tensor_parallel_world_size()
+            world_size = get_parallel().moe_tp_size
 
     if world_size <= 1:
         logger.debug("Single GPU, no need for allreduce fusion")

python/sglang/srt/layers/layernorm.py

@@ -31,6 +31,7 @@ from sglang.srt.model_executor.cuda_graph_config import (
     Phase,
     check_cuda_graph_backend,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
     cpu_has_amx_support,
@@ -153,9 +154,6 @@ def _forward_with_allreduce_fusion(
     """Shared allreduce-fused RMSNorm logic usable by any norm."""
     if residual is not None:
         from sglang.srt.distributed import (
-            get_attn_tensor_model_parallel_world_size,
-            get_moe_expert_parallel_world_size,
-            get_moe_tensor_parallel_world_size,
             tensor_model_parallel_all_reduce,
             tensor_model_parallel_fused_allreduce_rmsnorm,
         )
@@ -164,12 +162,12 @@ def _forward_with_allreduce_fusion(
         )
 
         if use_attn_tp_group:
-            world_size = get_attn_tensor_model_parallel_world_size()
+            world_size = get_parallel().attn_tp_size
         else:
-            if get_moe_expert_parallel_world_size() > 1:
+            if get_parallel().moe_ep_size > 1:
-                world_size = get_moe_expert_parallel_world_size()
+                world_size = get_parallel().moe_ep_size
             else:
-                world_size = get_moe_tensor_parallel_world_size()
+                world_size = get_parallel().moe_tp_size
 
         if world_size > 1:
             if post_residual_addition is not None:

python/sglang/srt/layers/linear.py

@@ -15,8 +15,6 @@ from torch.nn.parameter import Parameter, UninitializedParameter
 from sglang.kernel_api_logging import wrap_method_with_debug_kernel_once
 from sglang.srt.distributed import (
     divide,
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
     get_tp_group,
     split_tensor_along_last_dim,
     tensor_model_parallel_all_gather,
@@ -40,6 +38,7 @@ from sglang.srt.layers.parameter import (
     _ColumnvLLMParameter,
 )
 from sglang.srt.layers.utils import pad_or_narrow_weight
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import get_bool_env_var, is_cpu, is_hip, is_npu, set_weight_attrs
 
@@ -338,9 +337,9 @@ class ColumnParallelLinear(LinearBase):
 
         # Divide the weight matrix along the last dimension.
         if tp_rank is None:
-            tp_rank = get_tensor_model_parallel_rank()
+            tp_rank = get_parallel().tp_rank
         if tp_size is None:
-            tp_size = get_tensor_model_parallel_world_size()
+            tp_size = get_parallel().tp_size
         self.tp_rank, self.tp_size = tp_rank, tp_size
         assert self.quant_method is not None
         self.output_size_per_partition = divide(self.output_size, tp_size)
@@ -526,9 +525,9 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
     ):
         self.output_sizes = output_sizes
         if tp_rank is None:
-            tp_rank = get_tensor_model_parallel_rank()
+            tp_rank = get_parallel().tp_rank
         if tp_size is None:
-            tp_size = get_tensor_model_parallel_world_size()
+            tp_size = get_parallel().tp_size
         self.tp_rank, self.tp_size = tp_rank, tp_size
         assert all(output_size % tp_size == 0 for output_size in output_sizes)
         self.use_presharded_weights = use_presharded_weights
@@ -943,9 +942,9 @@ class QKVParallelLinear(ColumnParallelLinear):
         self.total_num_kv_heads = total_num_kv_heads
         # Divide the weight matrix along the last dimension.
         if tp_rank is None:
-            tp_rank = get_tensor_model_parallel_rank()
+            tp_rank = get_parallel().tp_rank
         if tp_size is None:
-            tp_size = get_tensor_model_parallel_world_size()
+            tp_size = get_parallel().tp_size
         self.tp_rank, self.tp_size = tp_rank, tp_size
         self.num_heads = divide(self.total_num_heads, tp_size)
         if tp_size >= self.total_num_kv_heads:
@@ -1390,9 +1389,9 @@ class RowParallelLinear(LinearBase):
 
         # Divide the weight matrix along the last dimension.
         if tp_rank is None:
-            tp_rank = get_tensor_model_parallel_rank()
+            tp_rank = get_parallel().tp_rank
         if tp_size is None:
-            tp_size = get_tensor_model_parallel_world_size()
+            tp_size = get_parallel().tp_size
         self.tp_rank, self.tp_size = tp_rank, tp_size
         self.input_size_per_partition = divide(input_size, self.tp_size)
         assert self.quant_method is not None
@@ -1605,8 +1604,8 @@ class MergedColumnParallelRepeatedLinear(LinearBase):
             prefix=prefix,
         )
         self.num_column_parallel = len(column_output_sizes)
-        self.tp_rank = get_tensor_model_parallel_rank()
+        self.tp_rank = get_parallel().tp_rank
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
 
         self.output_partition_sizes = [
             divide(x, self.tp_size) for x in column_output_sizes
@@ -1657,8 +1656,8 @@ class ColumnParallelBatchedLinear(nn.Module):
         self, batch: int, input_size: int, output_size: int, dtype: torch.dtype
     ):
         super().__init__()
-        self.tp_rank = get_tensor_model_parallel_rank()
+        self.tp_rank = get_parallel().tp_rank
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.weight = nn.Parameter(
             torch.empty(batch, output_size // self.tp_size, input_size, dtype=dtype),
             requires_grad=False,

python/sglang/srt/layers/logits_processor.py

@@ -22,7 +22,6 @@ import torch
 from torch import nn
 
 from sglang.srt.distributed import (
-    get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_gather,
 )
 from sglang.srt.environ import envs
@@ -32,9 +31,6 @@ from sglang.srt.layers.dp_attention import (
     attn_tp_all_gather_into_tensor,
     dp_gather_replicate,
     dp_scatter,
-    get_attention_dp_rank,
-    get_attention_dp_size,
-    get_attention_tp_size,
     get_dp_device,
     get_dp_dtype,
     get_dp_hidden_size,
@@ -53,6 +49,7 @@ from sglang.srt.model_executor.forward_batch_info import (
     ForwardBatch,
     ForwardMode,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils.common import (
     is_cpu,
@@ -229,7 +226,7 @@ class LogitsMetadata:
 
     def compute_dp_attention_metadata(self):
         cumtokens = torch.cumsum(self.global_num_tokens_for_logprob_gpu, dim=0)
-        dp_rank = get_attention_dp_rank()
+        dp_rank = get_parallel().attn_dp_rank
         if dp_rank == 0:
             dp_local_start_pos = torch.zeros_like(
                 self.global_num_tokens_for_logprob_gpu[0]
@@ -275,17 +272,17 @@ class LogitsProcessor(nn.Module):
         self.use_attn_tp_group = get_global_server_args().enable_dp_lm_head
         self.use_fp32_lm_head = get_global_server_args().enable_fp32_lm_head
         if self.use_attn_tp_group:
-            self.attn_tp_size = get_attention_tp_size()
+            self.attn_tp_size = get_parallel().attn_tp_size
             self.do_tensor_parallel_all_gather = (
                 not skip_all_gather and self.attn_tp_size > 1
             )
             self.do_tensor_parallel_all_gather_dp_attn = False
         else:
             self.do_tensor_parallel_all_gather = (
-                not skip_all_gather and get_tensor_model_parallel_world_size() > 1
+                not skip_all_gather and get_parallel().tp_size > 1
             )
             self.do_tensor_parallel_all_gather_dp_attn = (
-                self.do_tensor_parallel_all_gather and get_attention_dp_size() != 1
+                self.do_tensor_parallel_all_gather and get_parallel().attn_dp_size != 1
             )
         self.final_logit_softcapping = getattr(
             self.config, "final_logit_softcapping", None

python/sglang/srt/layers/moe/cutlass_w4a8_moe.py

@@ -5,6 +5,7 @@ from typing import Optional
 
 import torch
 
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import is_cuda, is_cuda_alike
 
 _is_cuda = is_cuda()
@@ -22,7 +23,6 @@ else:
     from sgl_kernel import silu_and_mul
 
 from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
-from sglang.srt.distributed import get_moe_expert_parallel_world_size
 from sglang.srt.layers.moe.ep_moe.kernels import (
     cutlass_w4_run_moe_ep_preproess,
     deepep_ll_get_cutlass_w4a8_moe_mm_data,
@@ -124,7 +124,7 @@ def cutlass_w4a8_moe(
         assert topk == 1, "apply_router_weight_on_input is only implemented for topk=1"
 
     device = a.device
-    if get_moe_expert_parallel_world_size() > 1:
+    if get_parallel().moe_ep_size > 1:
         topk_ids = torch.where(topk_ids == -1, num_local_experts, topk_ids)
 
     src2dst = cutlass_w4_run_moe_ep_preproess(

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -13,10 +13,6 @@ from torch.nn.parameter import UninitializedParameter
 from sglang.srt.batch_overlap.single_batch_overlap import DownGemmOverlapArgs
 from sglang.srt.batch_overlap.two_batch_overlap import MaybeTboDeepEPDispatcher
 from sglang.srt.distributed import (
-    get_moe_expert_parallel_rank,
-    get_moe_expert_parallel_world_size,
-    get_moe_tensor_parallel_rank,
-    get_moe_tensor_parallel_world_size,
     get_tp_group,
     tensor_model_parallel_all_reduce,
 )
@@ -65,6 +61,7 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo
     is_in_tc_piecewise_cuda_graph,
 )
 from sglang.srt.model_loader.weight_utils import narrow_padded_param_and_loaded_weight
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
     cpu_has_amx_support,
@@ -196,10 +193,10 @@ class FusedMoE(torch.nn.Module):
         self.enable_flashinfer_cutlass_moe = (
             get_moe_runner_backend().is_flashinfer_cutlass()
         )
-        self.moe_ep_size = get_moe_expert_parallel_world_size()
+        self.moe_ep_size = get_parallel().moe_ep_size
-        self.moe_ep_rank = get_moe_expert_parallel_rank()
+        self.moe_ep_rank = get_parallel().moe_ep_rank
-        self.moe_tp_size = get_moe_tensor_parallel_world_size()
+        self.moe_tp_size = get_parallel().moe_tp_size
-        self.moe_tp_rank = get_moe_tensor_parallel_rank()
+        self.moe_tp_rank = get_parallel().moe_tp_rank
 
         # DeepEP: each rank has its own shared expert slot, so total shared
         # weight slots = num_fused_shared_experts * ep_size.

python/sglang/srt/layers/moe/kt_ep_wrapper.py

@@ -12,8 +12,8 @@ from typing import TYPE_CHECKING, Optional
 
 import torch
 
-from sglang.srt.distributed import get_tensor_model_parallel_rank
 from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import get_compiler_backend
 
 if TYPE_CHECKING:
@@ -154,7 +154,7 @@ class KTEPWrapperMethod(FusedMoEMethodBase):
         self.num_gpu_experts = kt_config.num_gpu_experts
         self.override_num_local_experts = True
         self.gpu_method.num_gpu_experts = self.num_gpu_experts
-        self.tp_rank = get_tensor_model_parallel_rank()
+        self.tp_rank = get_parallel().tp_rank
 
         # KT wrapper will be initialized in create_weights
         self.wrapper: Optional[KTMoEWrapper] = None

python/sglang/srt/layers/moe/token_dispatcher/moriep.py

@@ -20,6 +20,7 @@ from sglang.srt.layers.moe.utils import (
     DeepEPMode,
     is_tbo_enabled,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import (
     get_bool_env_var,
     get_int_env_var,
@@ -35,10 +36,6 @@ from functools import lru_cache
 
 import torch
 
-from sglang.srt.distributed import (
-    get_moe_expert_parallel_rank,
-    get_moe_expert_parallel_world_size,
-)
 from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype
 
 # Blockwise quantization group sizes: number of elements sharing one scale factor
@@ -217,8 +214,8 @@ def init_mori_op(
 
     import mori
 
-    world_size = get_moe_expert_parallel_world_size()
+    world_size = get_parallel().moe_ep_size
-    rank = get_moe_expert_parallel_rank()
+    rank = get_parallel().moe_ep_rank
 
     gpu_per_node = 8 if world_size >= 8 else world_size
 
@@ -1048,7 +1045,7 @@ class MoriEPDispatcher(BaseDispatcher):
         # experts that are not local to this rank.
         self.expert_mask_gpu = None
         if _use_aiter and num_experts is not None and num_local_experts is not None:
-            ep_rank = get_moe_expert_parallel_rank()
+            ep_rank = get_parallel().moe_ep_rank
             expert_mask = torch.zeros(
                 num_experts,
                 device=torch.cuda.current_device(),

python/sglang/srt/layers/moe/token_dispatcher/standard.py

@@ -5,8 +5,6 @@ from typing import TYPE_CHECKING, NamedTuple, Optional
 import torch
 
 from sglang.srt.distributed import (
-    get_moe_expert_parallel_rank,
-    get_moe_expert_parallel_world_size,
     get_tp_group,
 )
 from sglang.srt.distributed.device_communicators.pynccl_allocator import (
@@ -30,6 +28,7 @@ from sglang.srt.layers.moe.utils import (
     get_moe_runner_backend,
     should_use_flashinfer_cutlass_moe_fp4_allgather,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils.common import (
     get_bool_env_var,
     get_device,
@@ -88,7 +87,7 @@ class StandardDispatcher(BaseDispatcher):
 
     def __init__(self, moe_runner_config: MoeRunnerConfig):
         super().__init__()
-        self.moe_ep_size = get_moe_expert_parallel_world_size()
+        self.moe_ep_size = get_parallel().moe_ep_size
         backend = get_moe_runner_backend()
         self.enable_flashinfer_cutlass_moe = backend.is_flashinfer_cutlass()
         self.enable_flashinfer_mxfp4_moe = backend.is_flashinfer_mxfp4()
@@ -110,7 +109,7 @@ class StandardDispatcher(BaseDispatcher):
         self.num_local_routed_experts = (
             self.num_local_experts - self.num_local_shared_experts
         )
-        self.moe_ep_rank = get_moe_expert_parallel_rank()
+        self.moe_ep_rank = get_parallel().moe_ep_rank
         self.local_expert_mapping = None
         self.expert_mask_gpu = None
 

python/sglang/srt/layers/moe/topk.py

@@ -31,6 +31,8 @@ from typing import (
 import torch
 import torch.nn.functional as F
 
+from sglang.srt.runtime_context import get_parallel
+
 try:
     from triton_kernels.matmul_ogs import GatherIndx, RoutingData, ScatterIndx
     from triton_kernels.tensor import make_ragged_tensor_metadata
@@ -81,8 +83,6 @@ except ImportError:
 
 from sglang.jit_kernel.dsv4 import mask_topk_ids
 from sglang.srt.distributed import (
-    get_moe_expert_parallel_rank,
-    get_moe_expert_parallel_world_size,
     get_tp_group,
 )
 from sglang.srt.distributed.device_communicators.pynccl_allocator import (
@@ -1484,8 +1484,8 @@ def _remap_topk_for_deepep(
     if topk_ids.shape[0] == 0:
         return topk_ids, topk_weights
 
-    ep_size = get_moe_expert_parallel_world_size()
+    ep_size = get_parallel().moe_ep_size
-    ep_rank = get_moe_expert_parallel_rank()
+    ep_rank = get_parallel().moe_ep_rank
     # Static EPLB may add redundant physical experts. At this point routed
     # topk_ids have already been remapped from logical to physical ids, so the
     # DeepEP interleaved layout must use the physical routed count.

python/sglang/srt/layers/moe/utils.py

@@ -8,12 +8,11 @@ from typing import TYPE_CHECKING, Optional
 
 import torch
 
-from sglang.srt.distributed.parallel_state import get_moe_expert_parallel_world_size
 from sglang.srt.environ import envs
 from sglang.srt.layers.dp_attention import (
-    get_attention_dp_size,
     is_dp_attention_enabled,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import is_cuda, is_npu
 
 _is_npu = is_npu()
@@ -409,7 +408,7 @@ def should_use_flashinfer_cutlass_moe_fp4_allgather():
         and get_moe_runner_backend().is_flashinfer_cutlass()
         and is_dp_attention_enabled()
         and MOE_QUANTIZATION == "modelopt_fp4"
-        and get_moe_expert_parallel_world_size() == get_attention_dp_size()
+        and get_parallel().moe_ep_size == get_parallel().attn_dp_size
     )
 
 
@@ -423,8 +422,8 @@ def should_use_dp_reduce_scatterv():
         not should_use_flashinfer_cutlass_moe_fp4_allgather()
         and get_moe_a2a_backend().is_none()
         and is_dp_attention_enabled()
-        and get_attention_dp_size() > 1
+        and get_parallel().attn_dp_size > 1
-        and get_moe_expert_parallel_world_size() == get_attention_dp_size()
+        and get_parallel().moe_ep_size == get_parallel().attn_dp_size
     )
 
 

python/sglang/srt/layers/quantization/blockwise_int8.py

@@ -10,7 +10,6 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional
 import torch
 from torch.nn import Module
 
-from sglang.srt.distributed import get_tensor_model_parallel_world_size
 from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
 from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
 from sglang.srt.layers.parameter import BlockQuantScaleParameter, ModelWeightParameter
@@ -23,6 +22,7 @@ from sglang.srt.layers.quantization.base_config import (
 from sglang.srt.layers.quantization.int8_utils import apply_w8a8_block_int8_linear
 from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
 from sglang.srt.layers.quantization.utils import is_layer_skipped
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import set_weight_attrs
 
 if TYPE_CHECKING:
@@ -149,7 +149,7 @@ class BlockInt8LinearMethod(LinearMethodBase):
         output_size_per_partition = sum(output_partition_sizes)
         weight_loader = extra_weight_attrs.get("weight_loader")
 
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
 
         block_n, block_k = (
             self.quant_config.weight_block_size[0],
@@ -271,7 +271,7 @@ class BlockInt8MoEMethod(FusedMoEMethodBase):
 
         if self.quant_config.is_checkpoint_int8_serialized:
             params_dtype = torch.int8
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
 
         block_n, block_k = (
             self.quant_config.weight_block_size[0],

python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxint4_moe.py

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING
 import torch
 from compressed_tensors import CompressionFormat
 
-from sglang.srt.distributed import get_moe_expert_parallel_rank, get_tp_group
+from sglang.srt.distributed import get_tp_group
 from sglang.srt.distributed.device_communicators.pynccl_allocator import (
     use_symmetric_memory,
 )
@@ -17,6 +17,7 @@ from sglang.srt.layers.quantization.compressed_tensors.schemes import (
     CompressedTensorsMoEScheme,
 )
 from sglang.srt.layers.quantization.utils import replace_parameter
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import is_flashinfer_available, next_power_of_2, set_weight_attrs
 
 logger = logging.getLogger(__name__)
@@ -65,7 +66,7 @@ class CompressedTensorsMxInt4MoE(CompressedTensorsMoEScheme):
         assert (
             not config.actorder
         ), "Actorder is not supported by flashinfer_trtllm backend"
-        self.moe_ep_rank = get_moe_expert_parallel_rank()
+        self.moe_ep_rank = get_parallel().moe_ep_rank
 
         if self.quant_config.quant_format != CompressionFormat.pack_quantized.value:
             raise ValueError(

python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py

@@ -6,7 +6,6 @@ from typing import TYPE_CHECKING
 import torch
 from compressed_tensors.quantization import QuantizationStrategy
 
-from sglang.srt.distributed import get_tensor_model_parallel_world_size
 from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
 from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (
     FlashInferTrtllmFp8MoeQuantInfo,
@@ -27,6 +26,7 @@ from sglang.srt.layers.quantization.utils import (
     per_tensor_dequantize,
     swap_w13_to_w31,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import get_bool_env_var, is_hip, set_weight_attrs
 
 if TYPE_CHECKING:
@@ -99,7 +99,7 @@ class CompressedTensorsW8A8Fp8MoE(CompressedTensorsMoEScheme):
         if self.block_quant:
             assert self.weight_block_size is not None
             layer.weight_block_size = self.weight_block_size
-            tp_size = get_tensor_model_parallel_world_size()
+            tp_size = get_parallel().tp_size
             block_n, block_k = (
                 self.weight_block_size[0],
                 self.weight_block_size[1],

python/sglang/srt/layers/quantization/fp8.py

@@ -12,7 +12,7 @@ import torch.nn.functional as F
 from torch.nn import Module
 from torch.nn.parameter import Parameter
 
-from sglang.srt.distributed import get_tensor_model_parallel_world_size, get_tp_group
+from sglang.srt.distributed import get_tp_group
 from sglang.srt.distributed.device_communicators.pynccl_allocator import (
     use_symmetric_memory,
 )
@@ -79,6 +79,7 @@ from sglang.srt.layers.quantization.utils import (
     requantize_with_max_scale,
 )
 from sglang.srt.layers.utils import copy_or_rebind_param
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import (
     cpu_has_amx_support,
     get_bool_env_var,
@@ -374,7 +375,7 @@ class Fp8LinearMethod(LinearMethodBase):
         output_partition_sizes: List[int],
         skip_block_quant_check: bool = False,
     ):
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         block_n, block_k = (
             self.quant_config.weight_block_size[0],
             self.quant_config.weight_block_size[1],
@@ -916,7 +917,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
 
         if self.quant_config.is_checkpoint_fp8_serialized:
             params_dtype = torch.uint32 if _use_hip_int4 else torch.float8_e4m3fn
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
 
         w13_up_dim, w2_up_dim, weight_padded = get_moe_weight_sizes(
             intermediate_size_per_partition,

python/sglang/srt/layers/quantization/fp8_utils.py

@@ -13,6 +13,7 @@ from sglang.srt.layers.quantization.fp8_kernel import (
     sglang_per_token_group_quant_fp8_row_padded,
 )
 from sglang.srt.layers.quantization.mxfp4_tensor import MXFP4QuantizeUtil
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils.common import torch_release
 
 if TYPE_CHECKING:
@@ -1779,9 +1780,8 @@ def validate_fp8_block_shape(
     block_size: list[int],
 ) -> None:
     """Validate block quantization shapes for tensor parallelism."""
-    from sglang.srt.distributed import get_tensor_model_parallel_world_size
 
-    tp_size = getattr(layer, "tp_size", get_tensor_model_parallel_world_size())
+    tp_size = getattr(layer, "tp_size", get_parallel().tp_size)
     block_n, block_k = block_size[0], block_size[1]
 
     # Required by row parallel

python/sglang/srt/layers/quantization/moe_wna16.py

@@ -9,7 +9,6 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional
 import numpy as np
 import torch
 
-from sglang.srt.distributed import get_tensor_model_parallel_rank
 from sglang.srt.distributed.parallel_state import get_tp_group
 from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
 from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
@@ -24,6 +23,7 @@ from sglang.srt.layers.quantization.unquant import (
     UnquantizedFusedMoEMethod,
     UnquantizedLinearMethod,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import get_device_capability, set_weight_attrs
 
 logger = logging.getLogger(__name__)
@@ -455,7 +455,7 @@ class MoeWNA16Method(FusedMoEMethodBase):
                 return
 
             device = get_tp_group().device
-            tp_rank = get_tensor_model_parallel_rank()
+            tp_rank = get_parallel().tp_rank
             loaded_weight = loaded_weight.to(device)
             shard_size = layer.intermediate_size_per_partition
 

python/sglang/srt/layers/quantization/quark_int4fp8_moe.py

@@ -5,7 +5,6 @@ import torch
 from tqdm import tqdm
 from tqdm.std import EMA
 
-from sglang.srt.distributed import get_tensor_model_parallel_rank
 from sglang.srt.layers.int4fp8_utils import (
     pack_int4_to_int32,
     quantize_fp8_scale_tensorwise,
@@ -18,6 +17,7 @@ from sglang.srt.layers.quantization.base_config import (
     QuantizeMethodBase,
 )
 from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import BAR_FORMAT, is_hip, set_weight_attrs
 
 if TYPE_CHECKING:
@@ -71,7 +71,7 @@ class QuarkInt4Fp8Config(QuantizationConfig):
 
         self.num_quant_layers = 0
 
-        tp_rank = get_tensor_model_parallel_rank()
+        tp_rank = get_parallel().tp_rank
 
         # The weight iterator already has a progress bar on rank=0, account for that.
         position = 1 + tqdm._get_free_pos()
@@ -138,7 +138,7 @@ class QuarkInt4Fp8MoEMethod(FusedMoEMethodBase):
 
         self.online_quant_progress_bar = self.quant_config.online_quant_progress_bar
 
-        self.tp_rank = get_tensor_model_parallel_rank()
+        self.tp_rank = get_parallel().tp_rank
 
         if not _is_hip:
             raise NotImplementedError(

python/sglang/srt/layers/quantization/w8a8_int8.py

@@ -7,7 +7,6 @@ from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, cast
 import torch
 from torch.nn.parameter import Parameter
 
-from sglang.srt.distributed import get_tensor_model_parallel_world_size
 from sglang.srt.layers.amx_utils import (
     CPUQuantMethod,
     _amx_process_weight_after_loading,
@@ -24,6 +23,7 @@ from sglang.srt.layers.quantization.base_config import (
 from sglang.srt.layers.quantization.compressed_tensors.utils import should_ignore_layer
 from sglang.srt.layers.quantization.int8_kernel import per_token_quant_int8
 from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import (
     cpu_has_amx_support,
     is_cpu,
@@ -260,7 +260,7 @@ class W8A8Int8MoEMethod(FusedMoEMethodBase):
     ):
         from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
 
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
 
         # WEIGHTS
         w13_weight = torch.nn.Parameter(

python/sglang/srt/layers/utils/cp_utils.py

@@ -11,13 +11,12 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
 from sglang.srt.layers.dp_attention import (
     attn_cp_all_gather_into_tensor,
     get_attention_cp_group,
-    get_attention_cp_rank,
-    get_attention_cp_size,
     is_allocation_symmetric,
 )
 from sglang.srt.layers.moe import get_moe_a2a_backend
 from sglang.srt.mem_cache.memory_pool import KVWriteLoc
 from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 
 
@@ -80,7 +79,7 @@ def get_cp_padding_align_size() -> int:
     """
     from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_in_seq_split
 
-    attn_cp_size = get_attention_cp_size()
+    attn_cp_size = get_parallel().attn_cp_size
     if is_prefill_cp_in_seq_split() or is_dsa_prefill_cp_in_seq_split():
         return attn_cp_size * 2
     return attn_cp_size
@@ -150,7 +149,7 @@ def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
     )
 
     if is_dsa_prefill_cp_round_robin_split():
-        cp_size = get_attention_cp_size()
+        cp_size = get_parallel().attn_cp_size
         assert (
             input_.shape[0] % cp_size == 0
         ), f"Expect input shape 0 can divided by cp size, but got input shape {input_.shape}, cp size {cp_size}"
@@ -172,7 +171,7 @@ def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor):
     )
 
     if is_dsa_prefill_cp_round_robin_split():
-        cp_size = get_attention_cp_size()
+        cp_size = get_parallel().attn_cp_size
         assert positions.shape[0] % cp_size == 0, (
             f"Expect positions shape 0 can divided by cp size, but got positions shape {positions.shape}, "
             f"cp size {cp_size}"
@@ -204,8 +203,8 @@ def cp_round_robin_input_ids(input_ids):
     rank2: 2,10,18,...
     ...
     """
-    cp_size = get_attention_cp_size()
+    cp_size = get_parallel().attn_cp_size
-    cp_rank = get_attention_cp_rank()
+    cp_rank = get_parallel().attn_cp_rank
     if get_moe_a2a_backend().is_none():
         input_ids = input_ids.reshape(-1, cp_size).T.flatten()
     else:

python/sglang/srt/layers/vocab_parallel_embedding.py

@@ -11,8 +11,6 @@ from torch.nn.parameter import Parameter, UninitializedParameter
 
 from sglang.srt.distributed import (
     divide,
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
     get_tp_group,
     tensor_model_parallel_all_reduce,
 )
@@ -24,8 +22,6 @@ from sglang.srt.layers.amx_utils import PackWeightMethod
 from sglang.srt.layers.communicator import get_attn_tp_context
 from sglang.srt.layers.dp_attention import (
     attn_tp_all_reduce,
-    get_attention_tp_rank,
-    get_attention_tp_size,
     is_allocation_symmetric,
     is_dp_attention_enabled,
 )
@@ -36,6 +32,7 @@ from sglang.srt.layers.quantization.base_config import (
     method_has_implemented_embedding,
 )
 from sglang.srt.layers.quantization.unquant import UnquantizedEmbeddingMethod
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import (
     cpu_has_amx_support,
     get_compiler_backend,
@@ -245,11 +242,11 @@ class VocabParallelEmbedding(torch.nn.Module):
         self.use_attn_tp_group = use_attn_tp_group
         if self.enable_tp:
             if use_attn_tp_group:
-                tp_rank = get_attention_tp_rank()
+                tp_rank = get_parallel().attn_tp_rank
-                self.tp_size = get_attention_tp_size()
+                self.tp_size = get_parallel().attn_tp_size
             else:
-                tp_rank = get_tensor_model_parallel_rank()
+                tp_rank = get_parallel().tp_rank
-                self.tp_size = get_tensor_model_parallel_world_size()
+                self.tp_size = get_parallel().tp_size
         else:
             assert use_attn_tp_group is False
             tp_rank = 0

python/sglang/srt/model_executor/forward_batch_info.py

@@ -36,19 +36,12 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
 
 import torch
 
-from sglang.srt.distributed.parallel_state import (
-    get_moe_expert_parallel_world_size,
-    get_tensor_model_parallel_world_size,
-)
 from sglang.srt.environ import envs
 from sglang.srt.kv_canary.req_to_expected_token_ids_manager import (
     compute_req_all_ids_info,
 )
 from sglang.srt.layers.dp_attention import (
     DpPaddingMode,
-    get_attention_dp_rank,
-    get_attention_tp_rank,
-    get_attention_tp_size,
     set_dp_buffer_len,
     set_is_extend_in_batch,
 )
@@ -56,6 +49,7 @@ from sglang.srt.model_executor.forward_batch_deepseek_mha_mixin import (
     ForwardBatchDeepSeekMHAMixin,
 )
 from sglang.srt.model_executor.triton_ops.position import compute_position_triton
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
     is_cuda,
@@ -212,8 +206,8 @@ def compute_local_num_token_non_padded(
     Converts a global count (across all TP ranks) to a local count for this rank.
     The "global" scope is within the current DP rank; DP is handled via num_tokens_per_dp.
     """
-    attn_tp_rank = get_attention_tp_rank()
+    attn_tp_rank = get_parallel().attn_tp_rank
-    attn_tp_size = get_attention_tp_size()
+    attn_tp_size = get_parallel().attn_tp_size
     tokens_per_rank = num_tokens_per_dp // attn_tp_size
 
     return torch.clamp(
@@ -849,7 +843,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
         """Make num_token_non_padded local to this attention-TP rank."""
         from sglang.srt.utils.common import require_mlp_tp_gather
 
-        dp_rank = get_attention_dp_rank()
+        dp_rank = get_parallel().attn_dp_rank
         assert self.global_num_tokens_cpu is not None
 
         if require_mlp_tp_gather(server_args):
@@ -1063,7 +1057,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
         self._original_batch_size = self.batch_size
         global_num_tokens = self.global_num_tokens_cpu
         sync_group_size = len(global_num_tokens)
-        attn_tp_size = get_attention_tp_size()
+        attn_tp_size = get_parallel().attn_tp_size
 
         for i in range(sync_group_size):
             # make sure that the padded length is divisible by attn_tp_size because we may need reduce-scatter across attn_tp dim.
@@ -1096,7 +1090,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             buffer_len = sum(global_num_tokens)
 
         if len(global_num_tokens) > 1:
-            num_tokens = global_num_tokens[get_attention_dp_rank()]
+            num_tokens = global_num_tokens[get_parallel().attn_dp_rank]
         else:
             num_tokens = global_num_tokens[0]
 
@@ -1299,7 +1293,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             return
         assert self.forward_mode.is_extend()
         tokens = self.input_ids.shape[0]
-        rank_size = get_tensor_model_parallel_world_size()
+        rank_size = get_parallel().tp_size
         tokens_padded = (tokens + rank_size - 1) // rank_size * rank_size
         self._pad_inputs_to_size(model_runner, tokens_padded, self.batch_size)
 
@@ -1359,7 +1353,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
 
 
 def enable_num_token_non_padded():
-    return get_moe_expert_parallel_world_size() > 1
+    return get_parallel().moe_ep_size > 1
 
 
 def build_inner_fb_view(

python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py

@@ -25,11 +25,7 @@ from typing import TYPE_CHECKING, Any, List, Sequence, Tuple
 import torch
 
 from sglang.srt.batch_overlap.two_batch_overlap import TboCudaGraphRunnerPlugin
-from sglang.srt.layers.dp_attention import (
+from sglang.srt.runtime_context import get_parallel
-    get_attention_cp_size,
-    get_attention_tp_rank,
-    get_attention_tp_size,
-)
 from sglang.srt.utils import require_gathered_buffer
 
 if TYPE_CHECKING:
@@ -81,10 +77,10 @@ def get_batch_sizes_to_capture(
         num_tokens_per_bs = 1
 
     if require_gathered_buffer(server_args):
-        mul_base *= get_attention_tp_size()
+        mul_base *= get_parallel().attn_tp_size
 
-    if mul_base % get_attention_cp_size() != 0:
+    if mul_base % get_parallel().attn_cp_size != 0:
-        mul_base *= get_attention_cp_size()
+        mul_base *= get_parallel().attn_cp_size
 
     # pad `num_max_requests` to avoid being filtered out
     num_max_requests = (num_max_requests + mul_base - 1) // mul_base * mul_base
@@ -146,8 +142,8 @@ class BaseCudaGraphRunner(ABC):
         self.tp_size = model_runner.server_args.tp_size
         self.dp_size = model_runner.server_args.dp_size
         self.pp_size = model_runner.server_args.pp_size
-        self.attn_tp_size = get_attention_tp_size()
+        self.attn_tp_size = get_parallel().attn_tp_size
-        self.attn_tp_rank = get_attention_tp_rank()
+        self.attn_tp_rank = get_parallel().attn_tp_rank
         self.tbo_plugin = TboCudaGraphRunnerPlugin()
 
     @staticmethod

python/sglang/srt/model_loader/weight_utils.py

@@ -42,11 +42,8 @@ from tqdm.auto import tqdm
 from sglang.srt.configs.load_config import LoadConfig
 from sglang.srt.configs.model_config import ModelConfig
 from sglang.srt.distributed import (
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
     get_world_group,
 )
-from sglang.srt.layers.dp_attention import get_attention_tp_rank
 from sglang.srt.layers.quantization import QuantizationConfig, get_quantization_config
 from sglang.srt.layers.quantization.fp8 import Fp8Config
 from sglang.srt.layers.quantization.modelopt_quant import (
@@ -57,6 +54,7 @@ from sglang.srt.model_loader.ci_weight_validation import (
     ci_download_with_validation_and_retry,
     ci_validate_and_cleanup_local_snapshot,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import (
     BAR_FORMAT,
     find_local_repo_dir,
@@ -1327,7 +1325,7 @@ def row_parallel_weight_loader(
     param: torch.Tensor, loaded_weight: torch.Tensor
 ) -> None:
     """Load weights that are row-parallelized."""
-    tp_rank = get_tensor_model_parallel_rank()
+    tp_rank = get_parallel().tp_rank
     shard_dim = 0 if param.dim() != 1 else None
 
     if shard_dim is not None:
@@ -1345,7 +1343,7 @@ def sharded_weight_loader(shard_axis: int) -> LoaderFunction:
     """Create a weight loader that shards the weights along the given axis"""
 
     def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
-        tp_rank = get_attention_tp_rank()
+        tp_rank = get_parallel().attn_tp_rank
 
         shard_size = param.data.shape[shard_axis]
         start_idx = tp_rank * shard_size
@@ -1353,9 +1351,8 @@ def sharded_weight_loader(shard_axis: int) -> LoaderFunction:
         if (
             is_cpu()
             and (
-                loaded_weight.size(0) % get_tensor_model_parallel_world_size() != 0
+                loaded_weight.size(0) % get_parallel().tp_size != 0
-                or loaded_weight.size(0)
+                or loaded_weight.size(0) < get_parallel().tp_size * shard_size
-                < get_tensor_model_parallel_world_size() * shard_size
             )
             and loaded_weight.dim() == 1
         ):

python/sglang/srt/models/afmoe.py

@@ -32,8 +32,6 @@ from torch import nn
 from transformers import PretrainedConfig
 
 from sglang.srt.distributed import (
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
 )
 from sglang.srt.layers.activation import SiluAndMul
@@ -58,6 +56,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, is_npu
 
 _is_npu = is_npu()
@@ -160,8 +159,8 @@ class AfmoeMoE(nn.Module):
     ):
         super().__init__()
         self.config = config
-        self.rank = get_tensor_model_parallel_rank()
+        self.rank = get_parallel().tp_rank
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
 
         self.n_routed_experts = getattr(config, "num_experts", None)
         if self.n_routed_experts is None:
@@ -309,7 +308,7 @@ class AfmoeAttention(nn.Module):
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         self.total_num_heads = num_heads
         assert self.total_num_heads % tp_size == 0
         self.num_heads = self.total_num_heads // tp_size

python/sglang/srt/models/apertus.py

@@ -28,8 +28,6 @@ from transformers import ApertusConfig
 
 from sglang.srt.distributed import (
     get_pp_group,
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
 )
 from sglang.srt.layers.activation import XIELU
 from sglang.srt.layers.layernorm import RMSNorm
@@ -54,6 +52,7 @@ from sglang.srt.model_loader.weight_utils import (
     kv_cache_scales_loader,
     maybe_remap_kv_scale_name,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix, make_layers
 
@@ -130,7 +129,7 @@ class ApertusAttention(nn.Module):
         super().__init__()
         self.layer_id = layer_id
         self.hidden_size = hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         self.total_num_heads = num_heads
         assert self.total_num_heads % tp_size == 0
         self.num_heads = self.total_num_heads // tp_size
@@ -383,8 +382,8 @@ class ApertusModel(nn.Module):
     # factors (or else raise an exception). Thus, handled exceptions should
     # make sure to leave KV cache scale factors in a known good (dummy) state
     def load_kv_cache_scales(self, quantization_param_path: str) -> None:
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
-        tp_rank = get_tensor_model_parallel_rank()
+        tp_rank = get_parallel().tp_rank
         for layer_idx, scaling_factor in kv_cache_scales_loader(
             quantization_param_path,
             tp_rank,

python/sglang/srt/models/arcee.py

@@ -22,8 +22,6 @@ from transformers import LlamaConfig
 
 from sglang.srt.distributed import (
     get_pp_group,
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
 )
 from sglang.srt.layers.activation import get_act_fn
 from sglang.srt.layers.layernorm import RMSNorm
@@ -48,6 +46,7 @@ from sglang.srt.model_loader.weight_utils import (
     kv_cache_scales_loader,
     maybe_remap_kv_scale_name,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix, make_layers
 
@@ -119,7 +118,7 @@ class ArceeAttention(nn.Module):
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         self.total_num_heads = num_heads
         assert self.total_num_heads % tp_size == 0
         self.num_heads = self.total_num_heads // tp_size
@@ -348,8 +347,8 @@ class ArceeModel(nn.Module):
         return hidden_states, aux_hidden_states
 
     def load_kv_cache_scales(self, quantization_param_path: str) -> None:
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
-        tp_rank = get_tensor_model_parallel_rank()
+        tp_rank = get_parallel().tp_rank
         for layer_idx, scaling_factor in kv_cache_scales_loader(
             quantization_param_path,
             tp_rank,

python/sglang/srt/models/baichuan.py

@@ -30,10 +30,6 @@ import torch
 from torch import nn
 from transformers import PretrainedConfig
 
-from sglang.srt.distributed import (
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
-)
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
@@ -51,6 +47,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, is_npu
 from sglang.srt.utils.hf_transformers_utils import get_rope_config
 
@@ -137,7 +134,7 @@ class BaiChuanAttention(nn.Module):
     ):
         super().__init__()
         self.hidden_size = hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         self.total_num_heads = num_heads
         self.total_num_kv_heads = self.total_num_heads
         assert self.total_num_heads % tp_size == 0
@@ -187,7 +184,7 @@ class BaiChuanAttention(nn.Module):
 
         # Create the alibi slopes and slice them.
         if self.position_embedding == "ALIBI":
-            tp_rank = get_tensor_model_parallel_rank()
+            tp_rank = get_parallel().tp_rank
             head_start = tp_rank * self.num_heads
             head_end = (tp_rank + 1) * self.num_heads
             alibi_slopes = _get_alibi_slopes(self.total_num_heads)

python/sglang/srt/models/bailing_moe.py

@@ -29,7 +29,6 @@ from transformers import PretrainedConfig
 
 from sglang.srt.distributed import (
     get_pp_group,
-    get_tensor_model_parallel_world_size,
     parallel_state,
     tensor_model_parallel_all_reduce,
 )
@@ -43,9 +42,6 @@ from sglang.srt.layers.communicator import (
     enable_moe_dense_fully_dp,
 )
 from sglang.srt.layers.dp_attention import (
-    get_attention_dp_size,
-    get_attention_tp_rank,
-    get_attention_tp_size,
     is_dp_attention_enabled,
 )
 from sglang.srt.layers.layernorm import RMSNorm
@@ -81,6 +77,7 @@ from sglang.srt.models.utils import (
     create_fused_set_kv_buffer_arg,
     enable_fused_set_kv_buffer,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix, is_cuda, is_non_idle_and_non_empty, make_layers
 
@@ -188,7 +185,7 @@ class BailingMoESparseMoeBlock(nn.Module):
         super().__init__()
         self.layer_id = layer_id
         self.alt_stream = alt_stream
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.top_k = config.num_experts_per_tok
         self.norm_topk_prob = config.norm_topk_prob
         self.hidden_size = config.hidden_size
@@ -290,7 +287,7 @@ class BailingMoESparseMoeBlock(nn.Module):
         # dispatcher
         if get_moe_a2a_backend().is_deepep():
             # TODO: we will support tp < ep in the future
-            self.ep_size = get_tensor_model_parallel_world_size()
+            self.ep_size = get_parallel().tp_size
 
             self.deepep_dispatcher = DeepEPDispatcher(
                 group=parallel_state.get_tp_group().device_group,
@@ -434,9 +431,9 @@ class BailingMoEAttention(nn.Module):
         self.hidden_size = config.hidden_size
         self.total_num_heads = config.num_attention_heads
         self.total_kv_heads = config.num_key_value_heads
-        self.dp_size = get_attention_dp_size()
+        self.dp_size = get_parallel().attn_dp_size
-        attn_tp_rank = get_attention_tp_rank()
+        attn_tp_rank = get_parallel().attn_tp_rank
-        attn_tp_size = get_attention_tp_size()
+        attn_tp_size = get_parallel().attn_tp_size
 
         assert self.total_num_heads % attn_tp_size == 0
         if self.total_kv_heads >= attn_tp_size:
@@ -574,7 +571,7 @@ class BailingMoEBlock(nn.Module):
         hidden_size = config.hidden_size
 
         self.input_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps)
-        self.dp_size = get_attention_dp_size()
+        self.dp_size = get_parallel().attn_dp_size
         self.attention = BailingMoEAttention(
             config,
             layer_id,
@@ -584,8 +581,8 @@ class BailingMoEBlock(nn.Module):
             alt_stream=alt_stream,
         )
         self.layer_id = layer_id
-        self.attn_tp_size = get_attention_tp_size()
+        self.attn_tp_size = get_parallel().attn_tp_size
-        self.attn_tp_rank = get_attention_tp_rank()
+        self.attn_tp_rank = get_parallel().attn_tp_rank
 
         self.is_layer_sparse = self._is_layer_sparse(
             config, layer_id=layer_id, is_nextn=False

python/sglang/srt/models/bailing_moe_linear.py

@@ -11,8 +11,6 @@ from transformers import PretrainedConfig
 
 from sglang.srt.distributed import (
     get_pp_group,
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
 )
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
@@ -22,8 +20,6 @@ from sglang.srt.layers.attention.fla.layernorm_gated import RMSNorm as RMSNormGa
 from sglang.srt.layers.attention.fla.layernorm_gated import layernorm_fn
 from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes
 from sglang.srt.layers.dp_attention import (
-    get_attention_tp_rank,
-    get_attention_tp_size,
     is_dp_attention_enabled,
 )
 from sglang.srt.layers.layernorm import RMSNorm
@@ -62,6 +58,7 @@ from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA, DeepseekV2MLP, _is_hip
 from sglang.srt.models.utils import WeightsMapper
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
     BumpAllocator,
@@ -251,8 +248,8 @@ class BailingMoE(nn.Module):
         self.alt_stream = alt_stream
         self.layer_id = layer_id
 
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
-        self.tp_rank = get_tensor_model_parallel_rank()
+        self.tp_rank = get_parallel().tp_rank
 
         self.top_k = config.num_experts_per_tok
         self.norm_expert_prob = getattr(config, "norm_topk_prob", False)
@@ -406,8 +403,8 @@ class BailingGroupRMSNormGate(RMSNormGated):
         param: torch.nn.Parameter,
         loaded_weight: torch.Tensor,
     ) -> None:
-        tp_size = get_attention_tp_size()
+        tp_size = get_parallel().attn_tp_size
-        tp_rank = get_attention_tp_rank()
+        tp_rank = get_parallel().attn_tp_rank
         shard_size = loaded_weight.shape[0] // tp_size
         shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)
         param.data.copy_(loaded_weight[shard].contiguous())
@@ -437,8 +434,8 @@ class BailingMoELinearAttention(nn.Module):
 
         self.hidden_inner_size = self.head_dim * self.total_num_heads
         self.scaling = self.head_dim**-0.5
-        self.tp_size = get_attention_tp_size()
+        self.tp_size = get_parallel().attn_tp_size
-        self.tp_rank = get_attention_tp_rank()
+        self.tp_rank = get_parallel().attn_tp_rank
 
         assert self.total_num_heads % self.tp_size == 0
         self.tp_heads = self.total_num_heads // self.tp_size
@@ -642,7 +639,7 @@ class BailingMoEAttention(nn.Module):
         self.layer_id = layer_id
 
         self.hidden_size = config.hidden_size
-        tp_size = get_attention_tp_size()
+        tp_size = get_parallel().attn_tp_size
         self.total_num_heads = config.num_attention_heads
         assert self.total_num_heads % tp_size == 0
         self.num_heads = self.total_num_heads // tp_size

python/sglang/srt/models/bailing_moe_nextn.py

@@ -26,7 +26,6 @@ import torch
 from torch import nn
 from transformers import PretrainedConfig
 
-from sglang.srt.distributed import get_tensor_model_parallel_world_size
 from sglang.srt.layers.dp_attention import is_dp_attention_enabled
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import ReplicatedLinear
@@ -43,6 +42,7 @@ from sglang.srt.models.bailing_moe_linear import (
     BailingMoeV2_5ForCausalLM,
 )
 from sglang.srt.models.utils import WeightsMapper
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import BumpAllocator, add_prefix
 
@@ -195,7 +195,7 @@ class BailingMoeForCausalLMNextN(nn.Module):
     ) -> None:
         nn.Module.__init__(self)
         self.config = config
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.quant_config = quant_config
         if hasattr(self, "determine_num_fused_shared_experts"):
             # Asystem has determine_num_fused_shared_experts but theta does not.

python/sglang/srt/models/bert.py

@@ -4,7 +4,6 @@ from typing import Iterable, Optional, Set, Tuple
 import torch
 from torch import nn
 
-from sglang.srt.distributed import get_tensor_model_parallel_world_size
 from sglang.srt.layers.activation import get_act_fn
 from sglang.srt.layers.linear import (
     ColumnParallelLinear,
@@ -17,6 +16,7 @@ from sglang.srt.layers.radix_attention import AttentionType, RadixAttention
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix
 
@@ -220,7 +220,7 @@ class BertSelfAttention(nn.Module):
     ):
         super().__init__()
         self.hidden_size = hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
 
         self.total_num_heads = num_attention_heads
         assert self.total_num_heads % tp_size == 0

python/sglang/srt/models/chatglm.py

@@ -23,7 +23,6 @@ from torch import nn
 from torch.nn import LayerNorm
 
 from sglang.srt.configs import ChatGLMConfig
-from sglang.srt.distributed import get_tensor_model_parallel_world_size
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
@@ -41,6 +40,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix
 
 LoraConfig = None
@@ -56,7 +56,7 @@ class GLMAttention(nn.Module):
     ):
         super().__init__()
         self.hidden_size = config.hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         self.total_num_heads = config.num_attention_heads
         assert self.total_num_heads % tp_size == 0
         self.num_heads = self.total_num_heads // tp_size

python/sglang/srt/models/cohere2_moe.py

@@ -11,7 +11,6 @@ from torch import nn
 from transformers import PretrainedConfig
 
 from sglang.srt.distributed import (
-    get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
 )
 from sglang.srt.layers.activation import SiluAndMul
@@ -32,6 +31,7 @@ from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, get_compiler_backend, is_cuda, make_layers
 
 
@@ -117,7 +117,7 @@ class Cohere2MoeAttention(nn.Module):
         prefix: str = "",
     ):
         super().__init__()
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         self.config = config
         self.layer_id = layer_id
         self.hidden_size = config.hidden_size
@@ -232,7 +232,7 @@ class Cohere2MoeSparseMoeBlock(nn.Module):
         prefix: str = "",
     ):
         super().__init__()
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.hidden_size = config.hidden_size
         self.num_experts = config.num_experts
         self.top_k = config.num_experts_per_tok
@@ -403,7 +403,7 @@ class Cohere2MoeDecoderLayer(nn.Module):
 
         norm_eps = getattr(config, "layer_norm_eps", 1e-5)
         self.input_layernorm = Cohere2MoeLayerNorm(config.hidden_size, eps=norm_eps)
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
 
     def forward(
         self,

python/sglang/srt/models/commandr.py

@@ -49,10 +49,6 @@ from torch import nn
 from torch.nn.parameter import Parameter
 from transformers import Cohere2Config, CohereConfig, PretrainedConfig
 
-from sglang.srt.distributed import (
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
-)
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.linear import (
     MergedColumnParallelLinear,
@@ -69,6 +65,7 @@ from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, get_compiler_backend, set_weight_attrs
 
 
@@ -97,7 +94,7 @@ class LayerNorm(nn.Module):
         return hidden_states, residuals
 
     def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
-        tp_rank = get_tensor_model_parallel_rank()
+        tp_rank = get_parallel().tp_rank
         shard_dim = 0 if param.dim() != 1 else None
         param_data = param.data
         if shard_dim is not None:
@@ -152,7 +149,7 @@ class CohereAttention(nn.Module):
         prefix: str = "",
     ):
         super().__init__()
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         self.config = config
         self.attention_dropout = config.attention_dropout
         self.hidden_size = config.hidden_size

python/sglang/srt/models/dbrx.py

@@ -25,8 +25,6 @@ import torch.nn as nn
 
 from sglang.srt.configs import DbrxConfig
 from sglang.srt.distributed import (
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
 )
 from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
@@ -54,6 +52,7 @@ from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, is_npu, set_weight_attrs
 
 _is_npu = is_npu()
@@ -71,7 +70,7 @@ class DbrxRouter(nn.Module):
         prefix: str = "",
     ):
         super().__init__()
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.num_total_experts = config.ffn_config.moe_num_experts
         self.d_model = config.d_model
         self.layer = ReplicatedLinear(
@@ -103,7 +102,7 @@ class DbrxExperts(nn.Module):
         prefix: str = "",
     ):
         super().__init__()
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.num_total_experts = config.ffn_config.moe_num_experts
         self.top_k = config.ffn_config.moe_top_k
         self.d_model = config.d_model
@@ -155,7 +154,7 @@ class DbrxExperts(nn.Module):
     def weight_loader(
         self, param: nn.Parameter, loaded_weight: torch.Tensor, weight_name: str
     ):
-        tp_rank = get_tensor_model_parallel_rank()
+        tp_rank = get_parallel().tp_rank
         param_data = param.data
         shard_size = self.intermediate_size
         shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)
@@ -242,7 +241,7 @@ class DbrxAttention(nn.Module):
             is_neox_style=True,
         )
 
-        tp_world_size = get_tensor_model_parallel_world_size()
+        tp_world_size = get_parallel().tp_size
         self.tp_size = tp_world_size
         assert self.total_num_heads % tp_world_size == 0
         self.num_heads = self.total_num_heads // tp_world_size

python/sglang/srt/models/deepseek.py

@@ -25,8 +25,6 @@ from torch import nn
 from transformers import PretrainedConfig
 
 from sglang.srt.distributed import (
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
 )
 from sglang.srt.layers.activation import SiluAndMul
@@ -49,6 +47,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, cpu_has_amx_support, is_cpu, is_npu
 from sglang.srt.utils.hf_transformers_utils import get_rope_config
 
@@ -118,8 +117,8 @@ class DeepseekMoE(nn.Module):
     ):
         super().__init__()
         self.config = config
-        self.rank = get_tensor_model_parallel_rank()
+        self.rank = get_parallel().tp_rank
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.n_routed_experts = config.n_routed_experts
         self.top_k = config.num_experts_per_tok
         if self.tp_size > self.n_routed_experts:
@@ -244,7 +243,7 @@ class DeepseekAttention(nn.Module):
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         self.total_num_heads = num_heads
         assert self.total_num_heads % tp_size == 0
         self.num_heads = self.total_num_heads // tp_size

python/sglang/srt/models/deepseek_nextn.py

@@ -25,7 +25,7 @@ from torch import nn
 from transformers import PretrainedConfig
 
 from sglang.srt.configs.model_config import is_deepseek_dsa
-from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size
+from sglang.srt.distributed import get_pp_group
 from sglang.srt.environ import envs
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
 from sglang.srt.layers.attention.dsa.utils import (
@@ -33,10 +33,6 @@ from sglang.srt.layers.attention.dsa.utils import (
     dsa_use_prefill_cp,
     is_dsa_enable_prefill_cp,
 )
-from sglang.srt.layers.dp_attention import (
-    get_attention_cp_rank,
-    get_attention_cp_size,
-)
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import ReplicatedLinear
 from sglang.srt.layers.logits_processor import LogitsProcessor
@@ -60,6 +56,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.models.deepseek_common.utils import enable_nextn_moe_bf16_cast_to_fp8
 from sglang.srt.models.deepseek_v2 import DeepseekV2DecoderLayer, DeepseekV3ForCausalLM
 from sglang.srt.models.utils import WeightsMapper
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import BumpAllocator, add_prefix, is_cuda, is_npu
 
@@ -145,7 +142,7 @@ class DeepseekModelNextN(nn.Module):
             is_mla_prefill_cp_enabled() and not is_deepseek_dsa(config)
         )
         if self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp:
-            self.cp_size = get_attention_cp_size()
+            self.cp_size = get_parallel().attn_cp_size
         else:
             self.cp_size = None
         self.decoder = DeepseekV2DecoderLayer(
@@ -280,7 +277,7 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
     ) -> None:
         nn.Module.__init__(self)
         self.config = config
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.quant_config = quant_config
         # if not set, model load will be broken in DeepseekV3ForCausalLM load_weights()
         self.pp_group = get_pp_group()
@@ -289,8 +286,8 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
         self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
         self.mla_enable_prefill_cp = is_mla_prefill_cp_enabled() and not self.use_dsa
         if self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp:
-            self.cp_rank = get_attention_cp_rank()
+            self.cp_rank = get_parallel().attn_cp_rank
-            self.cp_size = get_attention_cp_size()
+            self.cp_size = get_parallel().attn_cp_size
         else:
             self.cp_rank = None
             self.cp_size = None

python/sglang/srt/models/deepseek_v2.py

@@ -47,9 +47,7 @@ from sglang.srt.configs.model_config import (
 )
 from sglang.srt.distributed import (
     divide,
-    get_moe_expert_parallel_world_size,
     get_pp_group,
-    get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
 )
 from sglang.srt.environ import envs
@@ -72,12 +70,6 @@ from sglang.srt.layers.communicator import (
     get_attn_tp_context,
 )
 from sglang.srt.layers.communicator_dsa_cp import DSACPLayerCommunicator
-from sglang.srt.layers.dp_attention import (
-    get_attention_cp_rank,
-    get_attention_cp_size,
-    get_attention_tp_rank,
-    get_attention_tp_size,
-)
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
     ColumnParallelLinear,
@@ -169,6 +161,7 @@ from sglang.srt.models.deepseek_common.utils import (
     _use_aiter_bpreshuffle_gfx95,
     _use_aiter_gfx95,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.utils import (
@@ -527,8 +520,8 @@ class DeepseekV2MoE(nn.Module):
         mla_enable_prefill_cp: bool = False,
     ):
         super().__init__()
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
-        self.moe_ep_size = get_moe_expert_parallel_world_size()
+        self.moe_ep_size = get_parallel().moe_ep_size
         self.routed_scaling_factor = config.routed_scaling_factor
         self.n_shared_experts = config.n_shared_experts
 
@@ -776,7 +769,7 @@ class DeepseekV2MoE(nn.Module):
             or get_moe_a2a_backend().is_ascend_fuseep()
         ):
             # TODO: we will support tp < ep in the future
-            self.ep_size = get_moe_expert_parallel_world_size()
+            self.ep_size = get_parallel().moe_ep_size
             self.num_experts = (
                 config.n_routed_experts
                 + get_global_server_args().ep_num_redundant_experts
@@ -1510,8 +1503,8 @@ class DeepseekV2AttentionMLA(
         self.kv_lora_rank = kv_lora_rank
         self.quant_config = quant_config
         self.is_nextn = is_nextn
-        attn_tp_rank = get_attention_tp_rank()
+        attn_tp_rank = get_parallel().attn_tp_rank
-        attn_tp_size = get_attention_tp_size()
+        attn_tp_size = get_parallel().attn_tp_size
         self.use_dsa = is_deepseek_dsa(config)
         self.dsa_enable_prefill_cp = dsa_enable_prefill_cp
         self.mla_enable_prefill_cp = mla_enable_prefill_cp
@@ -1521,7 +1514,7 @@ class DeepseekV2AttentionMLA(
         # store cp_size whenever either CP flavor is active so rebuild_cp_kv_cache
         # and the FA3 MLA wrapper can reach it on the dense MLA path too.
         if self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp:
-            self.cp_size = get_attention_cp_size()
+            self.cp_size = get_parallel().attn_cp_size
         self.num_heads = num_heads
         assert num_heads % attn_tp_size == 0
         self.num_local_heads = num_heads // attn_tp_size
@@ -2287,7 +2280,7 @@ class DeepseekV2Model(nn.Module):
             is_prefill_context_parallel_enabled() and not is_deepseek_dsa(config)
         )
         if self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp:
-            self.cp_size = get_attention_cp_size()
+            self.cp_size = get_parallel().attn_cp_size
         else:
             self.cp_size = None
 
@@ -2372,11 +2365,9 @@ class DeepseekV2Model(nn.Module):
             allocate_size = 0
             for i in range(len(self.layers)):
                 if isinstance(self.layers[i].mlp, DeepseekV2MoE):
-                    # tp_size = get_tensor_model_parallel_world_size()
+                    # tp_size = get_parallel().tp_size
                     is_a2a_moe = is_deepep_class_backend()
-                    tp_size = (
+                    tp_size = 1 if is_a2a_moe else get_parallel().tp_size
-                        1 if is_a2a_moe else get_tensor_model_parallel_world_size()
-                    )
                     intermediate_size = (
                         config.moe_intermediate_size * config.n_shared_experts
                     )
@@ -2576,7 +2567,7 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
 
         self.pp_group = get_pp_group()
         self.config = config
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.quant_config = quant_config
         self.determine_num_fused_shared_experts()
         self.use_dsa = is_deepseek_dsa(config)
@@ -2614,8 +2605,8 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
             is_prefill_context_parallel_enabled() and not is_deepseek_dsa(config)
         )
         if self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp:
-            self.cp_rank = get_attention_cp_rank()
+            self.cp_rank = get_parallel().attn_cp_rank
-            self.cp_size = get_attention_cp_size()
+            self.cp_size = get_parallel().attn_cp_size
         else:
             self.cp_rank = self.cp_size = None
 
@@ -2672,7 +2663,7 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
                 "or AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization."
                 "or MT-platform with capability >= 31 can use shared experts fusion optimization."
             )
-        elif get_moe_expert_parallel_world_size() > 1 and (
+        elif get_parallel().moe_ep_size > 1 and (
             not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4)
         ):
             disable_reason = (

python/sglang/srt/models/deepseek_v4.py

@@ -29,7 +29,6 @@ from sglang.srt.compilation.compilation_config import register_split_op
 from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config
 from sglang.srt.distributed import (
     get_pp_group,
-    get_tensor_model_parallel_world_size,
     get_tp_group,
 )
 from sglang.srt.environ import envs
@@ -53,11 +52,6 @@ from sglang.srt.layers.dp_attention import (
     attn_tp_all_gather,
     dp_gather_partial,
     dp_scatter,
-    get_attention_cp_rank,
-    get_attention_cp_size,
-    get_attention_dp_size,
-    get_attention_tp_rank,
-    get_attention_tp_size,
     get_dp_global_num_tokens,
     get_global_dp_buffer,
     get_local_dp_buffer,
@@ -116,6 +110,7 @@ from sglang.srt.models.deepseek_v2 import ParallelLMHead, _is_cuda, _is_hip, _is
 from sglang.srt.models.triton_ops.deepseek_v4 import (
     rms_normalize_triton as rms_normalize_triton,
 )
+from sglang.srt.runtime_context import get_parallel
 
 if not _is_hip:
     from sglang.srt.layers.utils.cp_utils import (
@@ -271,11 +266,11 @@ class MQALayer(nn.Module):
         compress_ratio_override: Optional[int] = None,
     ) -> None:
         super().__init__()
-        self.tp_rank = attn_tp_rank = get_attention_tp_rank()
+        self.tp_rank = attn_tp_rank = get_parallel().attn_tp_rank
-        self.tp_size = attn_tp_size = get_attention_tp_size()
+        self.tp_size = attn_tp_size = get_parallel().attn_tp_size
         self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
         if self.dsa_enable_prefill_cp:
-            self.cp_size = get_attention_cp_size()
+            self.cp_size = get_parallel().attn_cp_size
             self.tp_rank = attn_tp_rank = 0
             self.tp_size = attn_tp_size = 1
         self.layer_id = layer_id
@@ -1494,13 +1489,13 @@ class DeepseekV4DecoderLayer(nn.Module):
         _use_cp = self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch)
         _use_tp_moe_gather = (
             not _use_cp
-            and get_attention_dp_size() > 1
+            and get_parallel().attn_dp_size > 1
             and get_moe_a2a_backend().is_none()
         )
         _use_tp_attn_a2a_scatter = (
             not _use_cp
             and envs.SGLANG_DSV4_FIX_TP_ATTN_A2A_SCATTER.get()
-            and get_attention_tp_size() > 1
+            and get_parallel().attn_tp_size > 1
             and not get_moe_a2a_backend().is_none()
         )
         # symmetric gather+scatter for the no-EP TP-MoE dp-attn path:
@@ -1532,7 +1527,7 @@ class DeepseekV4DecoderLayer(nn.Module):
             dp_gather_partial(hidden_states, local_hidden_states, forward_batch)
         _a2a_scatter_chunks: Optional[List[torch.Tensor]] = None
         if _use_tp_attn_a2a_scatter:
-            s, r = get_attention_tp_size(), get_attention_tp_rank()
+            s, r = get_parallel().attn_tp_size, get_parallel().attn_tp_rank
             _a2a_scatter_chunks = list(hidden_states.tensor_split(s))
             hidden_states = _a2a_scatter_chunks[r].contiguous()
             input_ids = input_ids.tensor_split(s)[r].contiguous()
@@ -1646,7 +1641,7 @@ class DeepseekV4Model(nn.Module):
         self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
         self.use_fused_mhc_post_pre = _is_fused_mhc_post_pre_enabled()
         if self.dsa_enable_prefill_cp:
-            self.cp_size = get_attention_cp_size()
+            self.cp_size = get_parallel().attn_cp_size
 
     def hc_head(
         self,
@@ -1694,7 +1689,7 @@ class DeepseekV4Model(nn.Module):
                     hidden_states.shape[0], self.hc_mult, self.hidden_size
                 )
 
-        if get_attention_dp_size() > 1 and get_moe_a2a_backend().is_none():
+        if get_parallel().attn_dp_size > 1 and get_moe_a2a_backend().is_none():
             input_ids_global = torch.empty(
                 (_DpGatheredBufferWrapper._global_dp_buffer_len, 1),
                 dtype=input_ids.dtype,
@@ -1776,7 +1771,7 @@ class DeepseekV4ForCausalLM(nn.Module):
     ) -> None:
         super().__init__()
         self.config = config
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.quant_config = quant_config
         self.determine_num_fused_shared_experts()
         self.model = DeepseekV4Model(
@@ -1816,8 +1811,8 @@ class DeepseekV4ForCausalLM(nn.Module):
 
         self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
         if self.dsa_enable_prefill_cp:
-            self.cp_rank = get_attention_cp_rank()
+            self.cp_rank = get_parallel().attn_cp_rank
-            self.cp_size = get_attention_cp_size()
+            self.cp_size = get_parallel().attn_cp_size
 
     @property
     def routed_experts_weights_of_layer(self):

python/sglang/srt/models/deepseek_v4_nextn.py

@@ -6,7 +6,7 @@ import torch.nn.functional as F
 from torch import nn
 from transformers import PretrainedConfig
 
-from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size
+from sglang.srt.distributed import get_pp_group
 from sglang.srt.layers.attention.dsa.utils import (
     can_dsa_cp_split,
     dsa_use_prefill_cp,
@@ -16,9 +16,6 @@ from sglang.srt.layers.attention.dsa.utils import (
 from sglang.srt.layers.dp_attention import (
     _DpGatheredBufferWrapper,
     dp_gather_partial,
-    get_attention_cp_rank,
-    get_attention_cp_size,
-    get_attention_dp_size,
     is_dp_attention_enabled,
 )
 from sglang.srt.layers.layernorm import RMSNorm
@@ -40,6 +37,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.forward_context import get_attn_backend
 from sglang.srt.models.deepseek_v4 import DeepseekV4DecoderLayer, DeepseekV4ForCausalLM
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix
 
@@ -108,7 +106,7 @@ class DeepseekV4ModelNextN(nn.Module):
 
         self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
         if self.dsa_enable_prefill_cp:
-            self.cp_size = get_attention_cp_size()
+            self.cp_size = get_parallel().attn_cp_size
         else:
             self.cp_size = None
 
@@ -156,7 +154,7 @@ class DeepseekV4ModelNextN(nn.Module):
         else:
             hidden_states = hidden_states.unsqueeze(1).repeat(1, self.hc_mult, 1)
 
-        if get_attention_dp_size() > 1 and get_moe_a2a_backend().is_none():
+        if get_parallel().attn_dp_size > 1 and get_moe_a2a_backend().is_none():
             input_ids_global = torch.empty(
                 (_DpGatheredBufferWrapper._global_dp_buffer_len, 1),
                 dtype=input_ids.dtype,
@@ -213,14 +211,14 @@ class DeepseekV4ForCausalLMNextN(DeepseekV4ForCausalLM):
     ) -> None:
         nn.Module.__init__(self)
         self.config = config
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.pp_group = get_pp_group()
         self.quant_config = quant_config
         self.determine_num_fused_shared_experts()
         self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
         if self.dsa_enable_prefill_cp:
-            self.cp_rank = get_attention_cp_rank()
+            self.cp_rank = get_parallel().attn_cp_rank
-            self.cp_size = get_attention_cp_size()
+            self.cp_size = get_parallel().attn_cp_size
         else:
             self.cp_rank = None
             self.cp_size = None

python/sglang/srt/models/dflash.py

@@ -12,7 +12,6 @@ import torch
 import torch.nn.functional as F
 from torch import nn
 
-from sglang.srt.distributed import get_tensor_model_parallel_world_size
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
@@ -26,6 +25,7 @@ from sglang.srt.layers.rotary_embedding import get_rope
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.utils import apply_qk_norm
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.speculative.dflash_utils import (
     can_dflash_slice_qkv_weight,
     get_dflash_attention_sliding_window_size,
@@ -70,7 +70,7 @@ class DFlashAttention(nn.Module):
     def __init__(self, config, layer_id: int) -> None:
         super().__init__()
         hidden_size = int(config.hidden_size)
-        tp_size = int(get_tensor_model_parallel_world_size())
+        tp_size = int(get_parallel().tp_size)
         total_num_heads = int(config.num_attention_heads)
         total_num_kv_heads = int(
             getattr(config, "num_key_value_heads", total_num_heads)

python/sglang/srt/models/dots_vlm_vit.py

@@ -9,10 +9,10 @@ from torch.nn import LayerNorm
 from transformers.modeling_utils import PreTrainedModel
 
 from sglang.srt.configs.dots_vlm import DotsVisionConfig
-from sglang.srt.distributed import parallel_state
 from sglang.srt.layers.attention.vision import VisionAttention
 from sglang.srt.layers.conv import Conv2dLayer
 from sglang.srt.layers.quantization import QuantizationConfig
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, is_npu
 
 logger = logging.getLogger(__name__)
@@ -224,7 +224,7 @@ class DotsVisionTransformer(PreTrainedModel):
 
     def _update_vision_config(self):
         """update vision config to support tp"""
-        world_size = parallel_state.get_tensor_model_parallel_world_size()
+        world_size = get_parallel().tp_size
         num_heads = self.config.num_attention_heads
         head_dim = self.config.embed_dim // num_heads
         num_dummy_heads = 0

python/sglang/srt/models/ernie4.py

@@ -24,7 +24,6 @@ from transformers.models.ernie4_5_moe.configuration_ernie4_5_moe import (
 )
 
 from sglang.srt.distributed import (
-    get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
 )
 from sglang.srt.layers.communicator import enable_moe_dense_fully_dp
@@ -42,6 +41,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.deepseek_v2 import DeepseekV2MLP as Ernie4MLP
 from sglang.srt.models.llama import LlamaAttention as Ernie4Attention
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, is_npu, make_layers
 from sglang.srt.utils.hf_transformers_utils import get_rope_config
 
@@ -77,7 +77,7 @@ class Ernie4Moe(nn.Module):
     ):
         super().__init__()
         self.layer_id = layer_id
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.moe_num_shared_experts = getattr(config, "moe_num_shared_experts", 0)
 
         if config.hidden_act != "silu":

python/sglang/srt/models/ernie45_moe_vl.py

@@ -24,7 +24,6 @@ from transformers import PretrainedConfig
 
 from sglang.srt.distributed import (
     get_pp_group,
-    get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
 )
 from sglang.srt.layers.dp_attention import is_dp_attention_enabled
@@ -43,6 +42,7 @@ from sglang.srt.layers.utils import PPMissingLayer
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
 from sglang.srt.models.deepseek_v2 import DeepseekV2MLP as Ernie4_5_VLMoeMLP
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, make_layers
 
 logger = logging.getLogger(__name__)
@@ -67,7 +67,7 @@ class Ernie4_5_VLMoeAttention(nn.Module):
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         self.total_num_heads = num_heads
         assert self.total_num_heads % tp_size == 0
         self.num_heads = self.total_num_heads // tp_size
@@ -158,7 +158,7 @@ class Ernie4_5_VLMoeMoE(nn.Module):
     ):
         super().__init__()
         self.layer_id = layer_id
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.moe_num_shared_experts = getattr(config, "moe_num_shared_experts", 0)
         self.hidden_size = config.hidden_size
 

python/sglang/srt/models/exaone.py

@@ -21,7 +21,6 @@ from typing import Any, Dict, Iterable, Optional, Tuple
 import torch
 from torch import nn
 
-from sglang.srt.distributed import get_tensor_model_parallel_world_size
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
@@ -39,6 +38,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix
 from sglang.srt.utils.hf_transformers_utils import get_rope_config
 
@@ -98,7 +98,7 @@ class ExaoneAttention(nn.Module):
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         self.total_num_heads = num_heads
         assert self.total_num_heads % tp_size == 0
         self.num_heads = self.total_num_heads // tp_size

python/sglang/srt/models/exaone4.py

@@ -5,11 +5,9 @@ import torch
 from torch import nn
 from transformers import Exaone4Config
 
-from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size
+from sglang.srt.distributed import get_pp_group
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.dp_attention import (
-    get_attention_tp_rank,
-    get_attention_tp_size,
     get_local_attention_dp_size,
 )
 from sglang.srt.layers.layernorm import RMSNorm
@@ -33,6 +31,7 @@ from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix, make_layers
 from sglang.utils import get_exception_traceback, logger
@@ -106,10 +105,10 @@ class Exaone4Attention(nn.Module):
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
 
-        attn_tp_rank = get_attention_tp_rank()
+        attn_tp_rank = get_parallel().attn_tp_rank
-        attn_tp_size = get_attention_tp_size()
+        attn_tp_size = get_parallel().attn_tp_size
 
         self.total_num_heads = num_heads
         assert self.total_num_heads % tp_size == 0
@@ -241,8 +240,8 @@ class Exaone4DecoderLayer(nn.Module):
         max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
 
         self.local_dp_size = get_local_attention_dp_size()
-        self.attn_tp_size = get_attention_tp_size()
+        self.attn_tp_size = get_parallel().attn_tp_size
-        self.attn_tp_rank = get_attention_tp_rank()
+        self.attn_tp_rank = get_parallel().attn_tp_rank
 
         self.self_attn = Exaone4Attention(
             config=config,

python/sglang/srt/models/exaone_moe.py

@@ -25,9 +25,7 @@ from torch import nn
 from transformers import PretrainedConfig
 
 from sglang.srt.distributed import (
-    get_moe_expert_parallel_world_size,
     get_pp_group,
-    get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
 )
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
@@ -35,8 +33,6 @@ from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
 from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.dp_attention import (
-    get_attention_tp_rank,
-    get_attention_tp_size,
     is_dp_attention_enabled,
 )
 from sglang.srt.layers.layernorm import RMSNorm
@@ -66,6 +62,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
 from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import LazyValue, add_prefix, is_cuda, make_layers
 
@@ -147,8 +144,8 @@ class ExaoneMoESparseMoEBlock(nn.Module):
         prefix: str = "",
     ):
         super().__init__()
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
-        self.moe_ep_size = get_moe_expert_parallel_world_size()
+        self.moe_ep_size = get_parallel().moe_ep_size
         self.layer_id = layer_id
         self.routed_scaling_factor = config.routed_scaling_factor
         self.alt_stream = alt_stream
@@ -214,7 +211,7 @@ class ExaoneMoESparseMoEBlock(nn.Module):
             )
 
         if get_moe_a2a_backend().is_deepep():
-            self.ep_size = get_moe_expert_parallel_world_size()
+            self.ep_size = get_parallel().moe_ep_size
             self.num_experts = (
                 config.num_experts + get_global_server_args().ep_num_redundant_experts
             )
@@ -330,8 +327,8 @@ class ExaoneMoEAttention(nn.Module):
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
-        attn_tp_rank = get_attention_tp_rank()
+        attn_tp_rank = get_parallel().attn_tp_rank
-        attn_tp_size = get_attention_tp_size()
+        attn_tp_size = get_parallel().attn_tp_size
 
         self.total_num_heads = num_heads
         assert self.total_num_heads % attn_tp_size == 0
@@ -469,8 +466,8 @@ class ExaoneMoEDecoderLayer(nn.Module):
         attention_bias = getattr(config, "attention_bias", False) or getattr(
             config, "bias", False
         )
-        self.attn_tp_size = get_attention_tp_size()
+        self.attn_tp_size = get_parallel().attn_tp_size
-        self.attn_tp_rank = get_attention_tp_rank()
+        self.attn_tp_rank = get_parallel().attn_tp_rank
 
         self.self_attn = ExaoneMoEAttention(
             config=config,

python/sglang/srt/models/exaone_moe_mtp.py

@@ -23,13 +23,14 @@ import torch
 from torch import nn
 from transformers import PretrainedConfig
 
-from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size
+from sglang.srt.distributed import get_pp_group
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.models.exaone_moe import ExaoneMoEForCausalLM, ExaoneMoEModel
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix
 
@@ -46,7 +47,7 @@ class ExaoneMoEForCausalLMMTP(ExaoneMoEForCausalLM):
         nn.Module.__init__(self)
         self.config = config
         config.num_hidden_layers = 1
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.quant_config = quant_config
         self.pp_group = get_pp_group()
 

python/sglang/srt/models/falcon_h1.py

@@ -5,7 +5,7 @@ import torch
 from torch import nn
 
 from sglang.srt.configs.falcon_h1 import FalconH1Config
-from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size
+from sglang.srt.distributed import get_pp_group
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
     HybridLinearAttnBackend,
@@ -14,8 +14,6 @@ from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
 from sglang.srt.layers.attention.mamba.mamba import MambaMixer2
 from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes
 from sglang.srt.layers.dp_attention import (
-    get_attention_tp_rank,
-    get_attention_tp_size,
     is_dp_attention_enabled,
 )
 from sglang.srt.layers.layernorm import RMSNorm
@@ -35,6 +33,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.forward_context import get_attn_backend
 from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix, is_cuda, make_layers
 
@@ -79,7 +78,7 @@ class FalconH1MLP(nn.Module):
         self.layer_id = layer_id
 
         self.intermediate_size = intermediate_size
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
 
         self.gate_multiplier, self.down_multiplier = mlp_multipliers
 
@@ -114,9 +113,9 @@ class FalconH1HybridAttentionDecoderLayer(nn.Module):
         super().__init__()
         self.config = config
         self.hidden_size = config.hidden_size
-        self.attn_tp_rank = get_attention_tp_rank()
+        self.attn_tp_rank = get_parallel().attn_tp_rank
-        self.attn_tp_size = get_attention_tp_size()
+        self.attn_tp_size = get_parallel().attn_tp_size
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.total_num_heads = config.num_attention_heads
         assert self.total_num_heads % self.attn_tp_size == 0
         self.num_heads = self.total_num_heads // self.attn_tp_size

python/sglang/srt/models/gemma.py

@@ -25,7 +25,6 @@ import torch
 from torch import nn
 from transformers import PretrainedConfig
 
-from sglang.srt.distributed import get_tensor_model_parallel_world_size
 from sglang.srt.layers.activation import GeluAndMul
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
@@ -40,6 +39,7 @@ from sglang.srt.layers.rotary_embedding import get_rope
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix
 
 
@@ -90,7 +90,7 @@ class GemmaAttention(nn.Module):
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         self.total_num_heads = num_heads
         assert self.total_num_heads % tp_size == 0
         self.num_heads = self.total_num_heads // tp_size

python/sglang/srt/models/gemma2.py

@@ -24,7 +24,6 @@ import torch
 from torch import nn
 from transformers import PretrainedConfig
 
-from sglang.srt.distributed import get_tensor_model_parallel_world_size
 from sglang.srt.layers.activation import GeluAndMul
 from sglang.srt.layers.layernorm import GemmaRMSNorm
 from sglang.srt.layers.linear import (
@@ -42,6 +41,7 @@ from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, is_npu, make_layers
 
 _is_npu = is_npu()
@@ -111,7 +111,7 @@ class Gemma2Attention(nn.Module):
         self.layer_id = layer_id
         self.config = config
         self.hidden_size = hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         self.total_num_heads = num_heads
         assert self.total_num_heads % tp_size == 0
         self.num_heads = self.total_num_heads // tp_size

python/sglang/srt/models/gemma3_causal.py

@@ -26,10 +26,6 @@ from transformers import (
     PreTrainedModel,
 )
 
-from sglang.srt.distributed import (
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
-)
 from sglang.srt.layers.activation import GeluAndMul
 from sglang.srt.layers.layernorm import Gemma3RMSNorm
 from sglang.srt.layers.linear import (
@@ -47,6 +43,7 @@ from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, cpu_has_amx_support, is_cpu, make_layers
 
 _is_cpu = is_cpu()
@@ -126,7 +123,7 @@ class Gemma3Attention(nn.Module):
         super().__init__()
         self.layer_id = layer_id
         self.config = config
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
 
         self.total_num_heads = config.num_attention_heads
         assert self.total_num_heads % tp_size == 0
@@ -922,10 +919,10 @@ class Gemma3ForCausalLM(PreTrainedModel):
         VocabParallelEmbedding (sharded). This method extracts the correct
         shard so the weights can be shared.
         """
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         if tp_size <= 1:
             return weight
-        tp_rank = get_tensor_model_parallel_rank()
+        tp_rank = get_parallel().tp_rank
         shard_size = (weight.shape[0] + tp_size - 1) // tp_size
         return weight[tp_rank * shard_size : (tp_rank + 1) * shard_size]
 

python/sglang/srt/models/gemma3n_causal.py

@@ -5,7 +5,6 @@ import torch.nn.functional as F
 from torch import nn
 from transformers import AutoModel, Gemma3nTextConfig, PretrainedConfig, PreTrainedModel
 
-from sglang.srt.distributed import get_tensor_model_parallel_world_size
 from sglang.srt.layers.activation import GeluAndMul
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
@@ -26,6 +25,7 @@ from sglang.srt.model_loader.weight_utils import (
     maybe_remap_kv_scale_name,
 )
 from sglang.srt.models.gemma3_causal import Gemma3TextScaledWordEmbedding
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, make_layers
 
 
@@ -325,7 +325,7 @@ class Gemma3nAttention(nn.Module):
         super().__init__()
         self.layer_id = layer_id
         self.config = config
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
 
         self.total_num_heads = config.num_attention_heads
         assert self.total_num_heads % tp_size == 0

python/sglang/srt/models/gemma4_audio.py

@@ -37,16 +37,13 @@ from sglang.srt.layers.clippable_linear import (
     ClippableQKVParallelLinear,
     ClippableRowParallelLinear,
 )
-from sglang.srt.layers.dp_attention import (
-    get_attention_tp_rank,
-    get_attention_tp_size,
-)
 from sglang.srt.layers.layernorm import Gemma4RMSNorm
 from sglang.srt.layers.linear import (
     ColumnParallelLinear,
     RowParallelLinear,
 )
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, make_layers, set_weight_attrs
 
 # SSCP convolution constants (no longer in config.json, never varied across models)
@@ -69,7 +66,7 @@ class Gemma4AudioRelativePositionEmbedding(nn.Module):
         super().__init__()
         self.config = config
 
-        tp_size = get_attention_tp_size()
+        tp_size = get_parallel().attn_tp_size
         total_num_heads = config.num_attention_heads
         self.channels = config.hidden_size
         self.head_dim = self.channels // total_num_heads
@@ -219,7 +216,7 @@ class Gemma4AudioAttention(nn.Module):
         super().__init__()
         self.config = config
 
-        tp_size = get_attention_tp_size()
+        tp_size = get_parallel().attn_tp_size
         total_num_heads = config.num_attention_heads
         self.hidden_size = config.hidden_size
         self.head_dim = self.hidden_size // total_num_heads
@@ -641,7 +638,7 @@ class Gemma4AudioConformerLightConv1d(nn.Module):
         super().__init__()
         self.config = config
         self.causal_padding = config.conv_kernel_size - 1
-        tp_size = get_attention_tp_size()
+        tp_size = get_parallel().attn_tp_size
         hidden_per_tp = config.hidden_size // tp_size
 
         self.register_buffer(
@@ -673,7 +670,7 @@ class Gemma4AudioConformerLightConv1d(nn.Module):
             hidden_per_tp, eps=config.rms_norm_eps, scale_shift=0.0
         )
 
-        tp_rank = get_attention_tp_rank()
+        tp_rank = get_parallel().attn_tp_rank
 
         def _shard_dim0(param, loaded_weight, _rank=tp_rank, _tp=tp_size):
             shard = param.shape[0]

python/sglang/srt/models/gemma4_causal.py

@@ -26,8 +26,6 @@ from transformers import (
 
 from sglang.srt.distributed import (
     get_pp_group,
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
 )
 from sglang.srt.layers.gemma4_fused_ops import (
     gemma4_fused_routing,
@@ -60,6 +58,7 @@ from sglang.srt.models.gemma3_causal import Gemma3MLP, Gemma3TextScaledWordEmbed
 from sglang.srt.models.utils import (
     create_fused_set_kv_buffer_arg,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix, make_layers
 
@@ -210,7 +209,7 @@ class Gemma4MoE(nn.Module):
         self.layer_id = layer_id
         self.hidden_size = hidden_size
         self.num_experts = config.num_experts
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
 
         # Per-expert output scale folded into routing weights so that
         # MoE's fused kernel computes: Σ_e (expert_e * w_e * scale_e)
@@ -291,7 +290,7 @@ class Gemma4Attention(nn.Module):
 
         self.layer_id = layer_id
         self.config = config
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
 
         layer_type = config.layer_types[layer_id]
         self.sliding_window = (
@@ -1379,10 +1378,10 @@ class Gemma4ForCausalLM(PreTrainedModel):
         VocabParallelEmbedding (sharded). This method extracts the correct
         shard so the weights can be shared.
         """
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         if tp_size <= 1:
             return weight
-        tp_rank = get_tensor_model_parallel_rank()
+        tp_rank = get_parallel().tp_rank
         shard_size = (weight.shape[0] + tp_size - 1) // tp_size
         return weight[tp_rank * shard_size : (tp_rank + 1) * shard_size]
 

python/sglang/srt/models/gemma4_vision.py

@@ -27,9 +27,9 @@ from sglang.srt.layers.clippable_linear import (
     ClippableQKVParallelLinear,
     ClippableRowParallelLinear,
 )
-from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.layers.layernorm import Gemma4RMSNorm
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, get_device_capability, is_cuda, is_hip
 
 # ---------------------------------------------------------------------------
@@ -140,7 +140,7 @@ class Gemma4VisionAttention(nn.Module):
         super().__init__()
         self.head_dim = config.head_dim
 
-        tp_size = get_attention_tp_size()
+        tp_size = get_parallel().attn_tp_size
         self.num_heads_per_partition = config.num_attention_heads // tp_size
         self.num_kv_heads_per_partition = config.num_key_value_heads // tp_size
 

python/sglang/srt/models/glm4.py

@@ -25,8 +25,6 @@ from torch import nn
 
 from sglang.srt.distributed import (
     get_pp_group,
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
 )
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.dp_attention import is_dp_attention_enabled
@@ -51,6 +49,7 @@ from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
     kv_cache_scales_loader,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, make_layers
 from sglang.srt.utils.hf_transformers_utils import get_rope_config
 
@@ -125,7 +124,7 @@ class Glm4Attention(nn.Module):
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         self.total_num_heads = num_heads
         assert self.total_num_heads % tp_size == 0
         self.num_heads = self.total_num_heads // tp_size
@@ -397,8 +396,8 @@ class Glm4Model(nn.Module):
     # factors (or else raise an exception). Thus, handled exceptions should
     # make sure to leave KV cache scale factors in a known good (dummy) state
     def load_kv_cache_scales(self, quantization_param_path: str) -> None:
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
-        tp_rank = get_tensor_model_parallel_rank()
+        tp_rank = get_parallel().tp_rank
         for layer_idx, scaling_factor in kv_cache_scales_loader(
             quantization_param_path,
             tp_rank,

python/sglang/srt/models/glm4_moe.py

@@ -26,11 +26,8 @@ from transformers import PretrainedConfig
 from sglang.srt.batch_overlap.single_batch_overlap import SboFlags
 from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo
 from sglang.srt.distributed import (
-    get_moe_expert_parallel_world_size,
     get_pp_group,
     get_pp_indices,
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
     parallel_state,
     tensor_model_parallel_all_reduce,
 )
@@ -48,8 +45,6 @@ from sglang.srt.layers.communicator import (
     enable_moe_dense_fully_dp,
 )
 from sglang.srt.layers.dp_attention import (
-    get_attention_tp_rank,
-    get_attention_tp_size,
     is_allocation_symmetric,
     is_dp_attention_enabled,
 )
@@ -87,6 +82,7 @@ from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM
 from sglang.srt.models.utils import apply_qk_norm
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
     add_prefix,
@@ -205,8 +201,8 @@ class Glm4MoeAttention(nn.Module):
         self.hidden_size = hidden_size
         self.start_layer = start_layer
 
-        attn_tp_rank = get_attention_tp_rank()
+        attn_tp_rank = get_parallel().attn_tp_rank
-        attn_tp_size = get_attention_tp_size()
+        attn_tp_size = get_parallel().attn_tp_size
 
         self.total_num_heads = num_heads
         assert self.total_num_heads % attn_tp_size == 0
@@ -228,7 +224,7 @@ class Glm4MoeAttention(nn.Module):
         self.rope_theta = rope_theta
         self.use_qk_norm = use_qk_norm
         self.max_position_embeddings = max_position_embeddings
-        self.tp_rank = get_tensor_model_parallel_rank()
+        self.tp_rank = get_parallel().tp_rank
 
         self.qkv_proj = QKVParallelLinear(
             hidden_size,
@@ -403,8 +399,8 @@ class Glm4MoeSparseMoeBlock(nn.Module):
     ):
         nn.Module.__init__(self)
         self.top_k = config.num_experts_per_tok
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
-        self.moe_ep_size = get_moe_expert_parallel_world_size()
+        self.moe_ep_size = get_parallel().moe_ep_size
         self.routed_scaling_factor = config.routed_scaling_factor
         self.n_shared_experts = config.n_shared_experts
         self.num_fused_shared_experts = (
@@ -527,7 +523,7 @@ class Glm4MoeSparseMoeBlock(nn.Module):
             or get_moe_a2a_backend().is_ascend_fuseep()
         ):
             # TODO: we will support tp < ep in the future
-            self.ep_size = get_moe_expert_parallel_world_size()
+            self.ep_size = get_parallel().moe_ep_size
             self.num_experts = (
                 config.n_routed_experts
                 + get_global_server_args().ep_num_redundant_experts
@@ -1178,7 +1174,7 @@ class Glm4MoeForCausalLM(nn.Module):
         nn.Module.__init__(self)
         self.pp_group = get_pp_group()
         self.config = config
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.quant_config = quant_config
         self.num_fused_shared_experts = 0
         self.determine_num_fused_shared_experts()
@@ -1209,7 +1205,7 @@ class Glm4MoeForCausalLM(nn.Module):
                 "Only GLM-4.5 on NV-platform with capability >= 80 "
                 "or AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization."
             )
-        elif get_moe_expert_parallel_world_size() > 1 and (
+        elif get_parallel().moe_ep_size > 1 and (
             not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4)
         ):
             disable_reason = "Only GLM-4.5 on AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization under expert parallelism."

python/sglang/srt/models/glm4_moe_lite.py

@@ -26,9 +26,7 @@ from transformers import PretrainedConfig
 from sglang.srt.batch_overlap.single_batch_overlap import SboFlags
 from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo
 from sglang.srt.distributed import (
-    get_moe_expert_parallel_world_size,
     get_pp_group,
-    get_tensor_model_parallel_world_size,
     parallel_state,
     tensor_model_parallel_all_reduce,
 )
@@ -76,6 +74,7 @@ from sglang.srt.models.deepseek_common.deepseek_weight_loader import (
 )
 from sglang.srt.models.deepseek_common.utils import _is_cuda, _use_aiter
 from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
     BumpAllocator,
@@ -185,7 +184,7 @@ class Glm4MoeLiteSparseMoeBlock(nn.Module):
         is_nextn: bool = False,
     ):
         super().__init__()
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.routed_scaling_factor = config.routed_scaling_factor
         self.n_shared_experts = config.n_shared_experts
         self.num_fused_shared_experts = (
@@ -283,7 +282,7 @@ class Glm4MoeLiteSparseMoeBlock(nn.Module):
 
         if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake():
             # TODO: we will support tp < ep in the future
-            self.ep_size = get_moe_expert_parallel_world_size()
+            self.ep_size = get_parallel().moe_ep_size
             self.num_experts = (
                 config.n_routed_experts
                 + get_global_server_args().ep_num_redundant_experts
@@ -907,7 +906,7 @@ class Glm4MoeLiteForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
         super().__init__()
         config.moe_layer_freq = 1
         self.config = config
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.quant_config = quant_config
         self.pp_group = get_pp_group()
         self.determine_num_fused_shared_experts("Glm4MoeLiteForCausalLM")
@@ -951,7 +950,7 @@ class Glm4MoeLiteForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
             or self.config.n_shared_experts != 1
         ):
             disable_reason = "Only GLM-4.5 or GLM-4.6 on NV-platform with capability >= 80 can use shared experts fusion optimization."
-        elif get_moe_expert_parallel_world_size() > 1:
+        elif get_parallel().moe_ep_size > 1:
             disable_reason = "GLM-4.5 or GLM-4.6 cannot use shared experts fusion optimization under expert parallelism."
 
         if disable_reason is not None:

python/sglang/srt/models/glm4_moe_lite_nextn.py

@@ -21,7 +21,6 @@ import torch
 from torch import nn
 from transformers import PretrainedConfig
 
-from sglang.srt.distributed import get_tensor_model_parallel_world_size
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
 from sglang.srt.layers.dp_attention import is_dp_attention_enabled
 from sglang.srt.layers.layernorm import RMSNorm
@@ -36,6 +35,7 @@ from sglang.srt.models.glm4_moe_lite import (
     Glm4MoeLiteDecoderLayer,
     Glm4MoeLiteForCausalLM,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import BumpAllocator, add_prefix, is_npu
 
@@ -139,7 +139,7 @@ class Glm4MoeLiteForCausalLMNextN(Glm4MoeLiteForCausalLM):
     ) -> None:
         nn.Module.__init__(self)
         self.config = config
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         if (
             is_npu()
             and get_global_server_args().speculative_draft_model_quantization is None

python/sglang/srt/models/glm4_moe_nextn.py

@@ -21,7 +21,6 @@ import torch
 from torch import nn
 from transformers import PretrainedConfig
 
-from sglang.srt.distributed import get_tensor_model_parallel_world_size
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
 from sglang.srt.layers.dp_attention import is_dp_attention_enabled
 from sglang.srt.layers.layernorm import RMSNorm
@@ -33,6 +32,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.models.glm4_moe import Glm4MoeDecoderLayer, Glm4MoeForCausalLM
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix, is_npu
 
@@ -125,7 +125,7 @@ class Glm4MoeForCausalLMNextN(Glm4MoeForCausalLM):
     ) -> None:
         nn.Module.__init__(self)
         self.config = config
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         if (
             is_npu()
             and get_global_server_args().speculative_draft_model_quantization is None

python/sglang/srt/models/glm4v.py

@@ -27,10 +27,6 @@ import torch.nn.functional as F
 from einops import rearrange
 from transformers.models.glm4v.configuration_glm4v import Glm4vConfig, Glm4vVisionConfig
 
-from sglang.srt.distributed import (
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
-)
 from sglang.srt.distributed.parallel_state import get_pp_group
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.attention import vision_utils
@@ -57,6 +53,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTe
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.glm4 import Glm4Model
 from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix, is_npu
 from sglang.srt.utils.hf_transformers_utils import get_processor
@@ -86,10 +83,8 @@ class Glm4vVisionMLP(nn.Module):
         use_data_parallel: bool = False,
     ):
         super().__init__()
-        self.tp_size = (
+        self.tp_size = 1 if use_data_parallel else get_parallel().tp_size
-            1 if use_data_parallel else get_tensor_model_parallel_world_size()
+        self.tp_rank = 0 if use_data_parallel else get_parallel().tp_rank
-        )
-        self.tp_rank = 0 if use_data_parallel else get_tensor_model_parallel_rank()
         self.gate_up_proj = MergedColumnParallelLinear(
             input_size=in_features,
             output_sizes=[hidden_features] * 2,  # [gate_proj, up_proj]
@@ -237,8 +232,8 @@ class Glm4vPatchMerger(nn.Module):
     ) -> None:
         super().__init__()
         self.hidden_size = d_model
-        tp_size = 1 if use_data_parallel else get_tensor_model_parallel_world_size()
+        tp_size = 1 if use_data_parallel else get_parallel().tp_size
-        tp_rank = 0 if use_data_parallel else get_tensor_model_parallel_rank()
+        tp_rank = 0 if use_data_parallel else get_parallel().tp_rank
         self.proj = ReplicatedLinear(
             self.hidden_size,
             self.hidden_size,

python/sglang/srt/models/glm4v_moe.py

@@ -6,10 +6,6 @@ import torch
 import torch.nn as nn
 from transformers.models.glm4v_moe.configuration_glm4v_moe import Glm4vMoeConfig
 
-from sglang.srt.distributed import (
-    get_moe_expert_parallel_world_size,
-    get_tensor_model_parallel_world_size,
-)
 from sglang.srt.distributed.parallel_state import get_pp_group
 from sglang.srt.layers.attention import vision_utils
 from sglang.srt.layers.logits_processor import LogitsProcessor
@@ -22,6 +18,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.glm4_moe import Glm4MoeModel
 from sglang.srt.models.glm4v import Glm4vForConditionalGeneration, Glm4vVisionModel
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix, get_device_sm, is_cuda, log_info_on_rank0
 from sglang.srt.utils.hf_transformers_utils import get_processor
@@ -47,7 +44,7 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration):
         self.config = config
         self.use_data_parallel = get_global_server_args().mm_enable_dp_encoder
         vision_utils.update_vit_attn_dummy_heads_config(self.config)
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.quant_config = quant_config
         self.num_fused_shared_experts = 0
         self.determine_num_fused_shared_experts()
@@ -97,7 +94,7 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration):
             disable_reason = "Shared experts fusion currently requires CUDA devices."
         elif _is_cuda and (_device_sm is not None) and (_device_sm < 80):
             disable_reason = "Shared experts fusion requires SM80 or newer GPUs."
-        elif get_moe_expert_parallel_world_size() > 1:
+        elif get_parallel().moe_ep_size > 1:
             disable_reason = "Shared experts fusion is not supported together with expert parallelism yet."
         elif get_moe_a2a_backend().is_deepep():
             disable_reason = "Shared experts fusion is not supported when Deepep MoE backend is enabled."

python/sglang/srt/models/glm_ocr_nextn.py

@@ -21,7 +21,6 @@ import torch
 from torch import nn
 from transformers import PretrainedConfig
 
-from sglang.srt.distributed import get_tensor_model_parallel_world_size
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
 from sglang.srt.layers.dp_attention import is_dp_attention_enabled
 from sglang.srt.layers.layernorm import RMSNorm
@@ -34,6 +33,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.models.glm4 import Glm4DecoderLayer
 from sglang.srt.models.glm_ocr import GlmOcrForConditionalGeneration
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix
 
@@ -125,7 +125,7 @@ class GlmOcrForConditionalGenerationNextN(GlmOcrForConditionalGeneration):
     ) -> None:
         nn.Module.__init__(self)
         self.config = config
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.quant_config = quant_config
         self.model = GlmOcrModelNextN(
             config, quant_config, prefix=add_prefix("model", prefix)

python/sglang/srt/models/gpt2.py

@@ -24,7 +24,6 @@ import torch
 from torch import nn
 from transformers import GPT2Config
 
-from sglang.srt.distributed.parallel_state import get_tensor_model_parallel_world_size
 from sglang.srt.layers.activation import NewGELU
 from sglang.srt.layers.linear import (
     ColumnParallelLinear,
@@ -37,6 +36,7 @@ from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix
 
 
@@ -52,7 +52,7 @@ class GPT2Attention(nn.Module):
         super().__init__()
         self.hidden_size = config.hidden_size
         total_num_heads = config.num_attention_heads
-        tensor_model_parallel_world_size = get_tensor_model_parallel_world_size()
+        tensor_model_parallel_world_size = get_parallel().tp_size
         assert total_num_heads % tensor_model_parallel_world_size == 0
         self.num_heads = total_num_heads // tensor_model_parallel_world_size
         self.head_dim = self.hidden_size // total_num_heads

python/sglang/srt/models/gpt_bigcode.py

@@ -25,7 +25,6 @@ import torch
 from torch import nn
 from transformers import GPTBigCodeConfig
 
-from sglang.srt.distributed import get_tensor_model_parallel_world_size
 from sglang.srt.layers.activation import get_act_fn
 from sglang.srt.layers.linear import (
     ColumnParallelLinear,
@@ -38,6 +37,7 @@ from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix
 
 
@@ -53,7 +53,7 @@ class GPTBigCodeAttention(nn.Module):
         super().__init__()
         self.hidden_size = config.hidden_size
         total_num_heads = config.num_attention_heads
-        self.tensor_model_parallel_world_size = get_tensor_model_parallel_world_size()
+        self.tensor_model_parallel_world_size = get_parallel().tp_size
         assert total_num_heads % self.tensor_model_parallel_world_size == 0
         self.num_heads = total_num_heads // self.tensor_model_parallel_world_size
         self.head_dim = self.hidden_size // total_num_heads

python/sglang/srt/models/gpt_j.py

@@ -25,7 +25,6 @@ import torch
 from torch import nn
 from transformers import GPTJConfig
 
-from sglang.srt.distributed.parallel_state import get_tensor_model_parallel_world_size
 from sglang.srt.layers.activation import get_act_fn
 from sglang.srt.layers.linear import (
     ColumnParallelLinear,
@@ -45,6 +44,7 @@ from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix
 
 
@@ -78,7 +78,7 @@ class GPTJAttention(nn.Module):
             prefix=add_prefix("out_proj", prefix),
         )
 
-        tensor_model_parallel_world_size = get_tensor_model_parallel_world_size()
+        tensor_model_parallel_world_size = get_parallel().tp_size
         assert total_num_heads % tensor_model_parallel_world_size == 0
         num_heads = total_num_heads // tensor_model_parallel_world_size
 

python/sglang/srt/models/gpt_oss.py

@@ -28,21 +28,13 @@ from transformers import PretrainedConfig
 
 from sglang.jit_kernel.utils import is_arch_support_pdl
 from sglang.srt.distributed import (
-    get_moe_expert_parallel_rank,
-    get_moe_expert_parallel_world_size,
-    get_moe_tensor_parallel_rank,
-    get_moe_tensor_parallel_world_size,
     get_pp_group,
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
 )
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
 from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
 from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes
 from sglang.srt.layers.dp_attention import (
-    get_attention_tp_rank,
-    get_attention_tp_size,
     is_dp_attention_enabled,
 )
 from sglang.srt.layers.layernorm import RMSNorm
@@ -76,6 +68,7 @@ from sglang.srt.models.utils import (
     create_fused_set_kv_buffer_arg,
     enable_fused_set_kv_buffer,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
     LazyValue,
@@ -188,7 +181,7 @@ def _resolve_moe_input_pad_multiple(
     # output directly.
     if quant_config.get_name() != "mxfp4":
         return 0
-    if get_tensor_model_parallel_world_size() != 1:
+    if get_parallel().tp_size != 1:
         # Mid-layer hidden_states still flow through CommunicateWith...
         # AllReduceAndLayerNormFn helpers other than `_simple` when
         # attn_tp_size > 1; those helpers haven't been updated to handle
@@ -207,7 +200,7 @@ class GptOssSparseMoeBlock(nn.Module):
         prefix: str = "",
     ):
         super().__init__()
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
         self.layer_id = layer_id
         self.hidden_size = config.hidden_size
         self.activation = config.hidden_act
@@ -358,8 +351,8 @@ class GptOssAttention(nn.Module):
         self.hidden_size = hidden_size
         self.sliding_window_size = sliding_window_size
 
-        attn_tp_rank = get_attention_tp_rank()
+        attn_tp_rank = get_parallel().attn_tp_rank
-        attn_tp_size = get_attention_tp_size()
+        attn_tp_size = get_parallel().attn_tp_size
 
         self.total_num_heads = num_heads
         assert self.total_num_heads % attn_tp_size == 0
@@ -380,7 +373,7 @@ class GptOssAttention(nn.Module):
         self.scaling = self.head_dim**-0.5
         self.rope_theta = rope_theta
         self.max_position_embeddings = max_position_embeddings
-        self.tp_rank = get_tensor_model_parallel_rank()
+        self.tp_rank = get_parallel().tp_rank
 
         self.qkv_proj = QKVParallelLinear(
             hidden_size,
@@ -535,8 +528,8 @@ class GptOssDecoderLayer(nn.Module):
 
         self.layer_id = layer_id
 
-        self.attn_tp_size = get_attention_tp_size()
+        self.attn_tp_size = get_parallel().attn_tp_size
-        self.attn_tp_rank = get_attention_tp_rank()
+        self.attn_tp_rank = get_parallel().attn_tp_rank
 
         # GptOss all layers are sparse and have no nextn now
         self.is_layer_sparse = True
@@ -923,10 +916,10 @@ class GptOssForCausalLM(nn.Module):
         loaded_params: set[str] = set()
         mxfp4_block = 32
 
-        moe_tp_rank = get_moe_tensor_parallel_rank()
+        moe_tp_rank = get_parallel().moe_tp_rank
-        moe_tp_size = get_moe_tensor_parallel_world_size()
+        moe_tp_size = get_parallel().moe_tp_size
-        moe_ep_rank = get_moe_expert_parallel_rank()
+        moe_ep_rank = get_parallel().moe_ep_rank
-        moe_ep_size = get_moe_expert_parallel_world_size()
+        moe_ep_size = get_parallel().moe_ep_size
 
         intermediate_size = self.config.intermediate_size
         assert (
@@ -1217,7 +1210,7 @@ class GptOssForCausalLM(nn.Module):
                     weight_loader = param.weight_loader
                     if "bias" not in name:
                         loaded_weight = loaded_weight.transpose(-2, -1)
-                    if "w2_weight_bias" in name and get_moe_tensor_parallel_rank() != 0:
+                    if "w2_weight_bias" in name and get_parallel().moe_tp_rank != 0:
                         loaded_weight = loaded_weight.zero_()
 
                     weight_loader(
@@ -1235,8 +1228,8 @@ class GptOssForCausalLM(nn.Module):
                     if name in params_dict.keys():
                         param = params_dict[name]
                         if "sinks" in name:
-                            start = get_attention_tp_rank() * param.numel()
+                            start = get_parallel().attn_tp_rank * param.numel()
-                            tp_size = get_tensor_model_parallel_world_size()
+                            tp_size = get_parallel().tp_size
                             full_shard_size = param.numel() * tp_size
                             # This handles TP padding: if the checkpoint dim is not divisible by tp_size,
                             # the last TP shard extends beyond `loaded_weight`, pad with zeros before slicing.
@@ -1337,7 +1330,7 @@ def _canonicalize_weights(config, weights_in: Iterable[Tuple[str, torch.Tensor]]
 
 
 def _dequant_mlp_weight(debug_name, w_blocks, w_scales):
-    if get_tensor_model_parallel_rank() == 0:
+    if get_parallel().tp_rank == 0:
         logger.info(f"Dequantize {debug_name} start")
 
     original_device = w_blocks.device
@@ -1348,7 +1341,7 @@ def _dequant_mlp_weight(debug_name, w_blocks, w_scales):
     w_bf16 = dequant_mxfp4(w_block=w_blocks, w_scale=w_scales, out_dtype=torch.bfloat16)
     w_bf16 = w_bf16.transpose(-2, -1).contiguous()
 
-    if get_tensor_model_parallel_rank() == 0:
+    if get_parallel().tp_rank == 0:
         logger.info(
             f"Dequantize {debug_name} end {w_blocks.shape=} {w_scales.shape=} {w_bf16.shape=}"
         )

python/sglang/srt/models/granite.py

@@ -26,7 +26,6 @@ import torch
 from torch import nn
 from transformers import GraniteConfig
 
-from sglang.srt.distributed import get_tensor_model_parallel_world_size
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
@@ -45,6 +44,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix
 from sglang.utils import get_exception_traceback
 
@@ -106,7 +106,7 @@ class GraniteAttention(nn.Module):
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         self.total_num_heads = num_heads
         assert self.total_num_heads % tp_size == 0
         self.num_heads = self.total_num_heads // tp_size

python/sglang/srt/models/granitemoe.py

@@ -6,7 +6,6 @@ import torch
 from torch import nn
 from transformers import GraniteConfig
 
-from sglang.srt.distributed import get_tensor_model_parallel_world_size
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
     QKVParallelLinear,
@@ -26,6 +25,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.models import mixtral
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix
 
 
@@ -105,7 +105,7 @@ class GraniteMoeAttention(nn.Module):
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         self.total_num_heads = num_heads
         assert self.total_num_heads % tp_size == 0
         self.num_heads = self.total_num_heads // tp_size

python/sglang/srt/models/granitemoehybrid.py

@@ -5,7 +5,7 @@ from torch import nn
 from transformers.models.granitemoeshared import GraniteMoeSharedConfig
 
 from sglang.srt.configs.granitemoehybrid import GraniteMoeHybridConfig
-from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size
+from sglang.srt.distributed import get_pp_group
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
     HybridLinearAttnBackend,
@@ -32,6 +32,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTe
 from sglang.srt.model_executor.forward_context import get_attn_backend
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.transformers import maybe_prefix
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import make_layers
 
 from .granitemoe import GraniteMoeMoE
@@ -112,7 +113,7 @@ class GraniteMoeHybridMambaDecoderLayer(nn.Module):
                 intermediate_size=config.intermediate_size,
                 layer_id=layer_idx,
                 quant_config=quant_config,
-                tp_size=get_tensor_model_parallel_world_size(),
+                tp_size=get_parallel().tp_size,
                 prefix=f"{prefix}.block_sparse_moe",
             )
 
@@ -192,7 +193,7 @@ class GraniteMoeHybridAttention(nn.Module):
         self.total_num_kv_heads = config.num_key_value_heads
 
         # TensorParallel logic
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         assert self.total_num_heads % tp_size == 0
         self.num_heads = self.total_num_heads // tp_size
         if self.total_num_kv_heads >= tp_size:
@@ -299,7 +300,7 @@ class GraniteMoeHybridAttentionDecoderLayer(nn.Module):
                 intermediate_size=config.intermediate_size,
                 layer_id=layer_idx,
                 quant_config=quant_config,
-                tp_size=get_tensor_model_parallel_world_size(),
+                tp_size=get_parallel().tp_size,
                 prefix=f"{prefix}.block_sparse_moe",
             )
 

python/sglang/srt/models/grok.py

@@ -23,8 +23,6 @@ from torch import nn
 from transformers import PretrainedConfig
 
 from sglang.srt.distributed import (
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
 )
 from sglang.srt.layers.activation import GeluAndMul
@@ -60,6 +58,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.loader import DefaultModelLoader
 from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, is_npu
 
 _is_npu = is_npu()
@@ -333,8 +332,8 @@ class Grok1Attention(nn.Module):
         self.config = config
         self.layer_id = layer_id
         self.hidden_size = hidden_size
-        attn_tp_rank = get_tensor_model_parallel_rank()
+        attn_tp_rank = get_parallel().tp_rank
-        attn_tp_size = get_tensor_model_parallel_world_size()
+        attn_tp_size = get_parallel().tp_size
         self.total_num_heads = num_heads
         assert self.total_num_heads % attn_tp_size == 0
         self.num_heads = self.total_num_heads // attn_tp_size
@@ -542,7 +541,7 @@ class Grok1DecoderLayer(nn.Module):
             if self.residual_moe:
                 # NOTE: self.block_sparse_moe modifies the input in-place,
                 # so we have to call it later. Be aware of any possible related errors.
-                if get_tensor_model_parallel_world_size() > 1:
+                if get_parallel().tp_size > 1:
                     self.ffn = lambda x: tensor_model_parallel_all_reduce(
                         self.moe_with_rmoe(x)
                     )
@@ -593,7 +592,7 @@ class Grok1DecoderLayer(nn.Module):
             forward_batch=forward_batch,
         )
 
-        if get_tensor_model_parallel_world_size() > 1:
+        if get_parallel().tp_size > 1:
             hidden_states = tensor_model_parallel_all_reduce(hidden_states)
 
         hidden_states, residual = fused_dual_residual_rmsnorm(
@@ -710,7 +709,7 @@ class Grok1ForCausalLM(nn.Module):
         self.load_presharded_moe = (
             getattr(config, "load_presharded_moe", True)
             and self.config.num_local_experts > 0
-            and get_tensor_model_parallel_world_size() > 1
+            and get_parallel().tp_size > 1
         )
         self.load_presharded_attn = getattr(config, "load_presharded_attn", False)
         self.load_presharded_embedding = getattr(
@@ -722,7 +721,7 @@ class Grok1ForCausalLM(nn.Module):
             config, "replicate_lm_head", default_replicate_lm_head
         )
 
-        if get_tensor_model_parallel_world_size() > 1:
+        if get_parallel().tp_size > 1:
             setattr(DefaultModelLoader, "_prepare_weights", _prepare_presharded_weights)
 
         self.replicate_embedding = getattr(config, "replicate_embedding", False)
@@ -939,10 +938,7 @@ class Grok1ForCausalLM(nn.Module):
         return wq + wkv + out + ffn1 + ffn2 + embed
 
     def get_num_params_torch(self):
-        return (
+        return sum(p.numel() for p in self.parameters()) * get_parallel().tp_size
-            sum(p.numel() for p in self.parameters())
-            * get_tensor_model_parallel_world_size()
-        )
 
 
 old_prepare_weights = getattr(DefaultModelLoader, "_prepare_weights")
@@ -954,7 +950,7 @@ def _prepare_presharded_weights(
     import glob
     import os
 
-    if get_tensor_model_parallel_world_size() == 1:
+    if get_parallel().tp_size == 1:
         return old_prepare_weights(self, model_name_or_path, revision, fall_back_to_pt)
 
     if not os.path.isdir(model_name_or_path):
@@ -971,7 +967,7 @@ def _prepare_presharded_weights(
     else:
         hf_folder = model_name_or_path
 
-    tp_rank = get_tensor_model_parallel_rank()
+    tp_rank = get_parallel().tp_rank
 
     # The old format
     allow_patterns = [f"*-{tp_rank:03d}.bin"]

python/sglang/srt/models/hunyuan.py

@@ -21,8 +21,6 @@ from torch import nn
 from transformers import PretrainedConfig
 
 from sglang.srt.distributed import (
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
 )
 from sglang.srt.eplb.expert_distribution import ExpertDistributionRecorder
@@ -52,6 +50,7 @@ from sglang.srt.model_loader.weight_utils import (
     kv_cache_scales_loader,
     maybe_remap_kv_scale_name,
 )
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import is_hip
 from sglang.srt.utils.hf_transformers_utils import get_rope_config
 
@@ -125,7 +124,7 @@ class HunYuanSparseMoeBlock(nn.Module):
         layer_id: int = -1,
     ):
         super().__init__()
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_parallel().tp_size
 
         if self.tp_size > config.num_experts:
             raise ValueError(
@@ -263,7 +262,7 @@ class HunYuanAttention(nn.Module):
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
         self.total_num_heads = num_heads
         assert self.total_num_heads % tp_size == 0
         self.num_heads = self.total_num_heads // tp_size
@@ -783,8 +782,8 @@ class HunYuanMoEV1ForCausalLM(nn.Module):
     # factors (or else raise an exception). Thus, handled exceptions should
     # make sure to leave KV cache scale factors in a known good (dummy) state
     def load_kv_cache_scales(self, quantization_param_path: str) -> None:
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = get_parallel().tp_size
-        tp_rank = get_tensor_model_parallel_rank()
+        tp_rank = get_parallel().tp_rank
         for layer_idx, scaling_factor in kv_cache_scales_loader(
             quantization_param_path,
             tp_rank,