From 53318911cab4e721d62c4c7678caad79b251fcf3 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Wed, 17 Jun 2026 20:23:43 -0700 Subject: [PATCH] Add get_parallel(): a structured accessor for parallel-topology state (#28567) --- .../srt/batch_overlap/two_batch_overlap.py | 4 +- python/sglang/srt/layers/activation.py | 7 +- .../srt/layers/attention/aiter_backend.py | 13 +- .../layers/attention/cutlass_mla_backend.py | 9 +- .../layers/attention/deepseek_v4_backend.py | 9 +- .../deepseek_v4_backend_hip_radix.py | 14 +- .../srt/layers/attention/dsa/dsa_indexer.py | 7 +- .../sglang/srt/layers/attention/dsa/utils.py | 24 +- .../srt/layers/attention/dsa_backend.py | 4 +- .../srt/layers/attention/dsv4/compressor.py | 4 +- .../dual_chunk_flashattention_backend.py | 4 +- .../layers/attention/flashinfer_backend.py | 24 +- .../attention/flashinfer_mla_backend.py | 9 +- .../srt/layers/attention/flashmla_backend.py | 6 +- .../attention/linear/lightning_backend.py | 9 +- .../srt/layers/attention/mamba/mamba.py | 13 +- .../attention/mamba/mixer2_rms_norm_gated.py | 15 +- .../srt/layers/attention/triton_backend.py | 8 +- .../layers/attention/trtllm_mla_backend.py | 8 +- python/sglang/srt/layers/attention/vision.py | 10 +- .../srt/layers/attention/vision_utils.py | 4 +- .../srt/layers/attention/wave_backend.py | 6 +- python/sglang/srt/layers/clippable_linear.py | 8 +- python/sglang/srt/layers/communicator.py | 32 +- .../sglang/srt/layers/communicator_dsa_cp.py | 17 +- python/sglang/srt/layers/cp/base.py | 5 +- .../srt/layers/flashinfer_comm_fusion.py | 29 +- python/sglang/srt/layers/layernorm.py | 12 +- python/sglang/srt/layers/linear.py | 27 +- python/sglang/srt/layers/logits_processor.py | 13 +- .../sglang/srt/layers/moe/cutlass_w4a8_moe.py | 4 +- .../srt/layers/moe/fused_moe_triton/layer.py | 13 +- python/sglang/srt/layers/moe/kt_ep_wrapper.py | 4 +- .../srt/layers/moe/token_dispatcher/moriep.py | 11 +- .../layers/moe/token_dispatcher/standard.py | 7 +- python/sglang/srt/layers/moe/topk.py | 8 +- python/sglang/srt/layers/moe/utils.py | 9 +- .../srt/layers/quantization/blockwise_int8.py | 6 +- .../compressed_tensors_w4a4_mxint4_moe.py | 5 +- .../compressed_tensors_w8a8_fp8_moe.py | 4 +- python/sglang/srt/layers/quantization/fp8.py | 7 +- .../srt/layers/quantization/fp8_utils.py | 4 +- .../srt/layers/quantization/moe_wna16.py | 4 +- .../layers/quantization/quark_int4fp8_moe.py | 6 +- .../srt/layers/quantization/w8a8_int8.py | 4 +- python/sglang/srt/layers/utils/cp_utils.py | 13 +- .../srt/layers/vocab_parallel_embedding.py | 13 +- .../srt/model_executor/forward_batch_info.py | 22 +- .../runner/base_cuda_graph_runner.py | 16 +- .../sglang/srt/model_loader/weight_utils.py | 13 +- python/sglang/srt/models/afmoe.py | 9 +- python/sglang/srt/models/apertus.py | 1373 ++++++++--------- python/sglang/srt/models/arcee.py | 9 +- python/sglang/srt/models/baichuan.py | 9 +- python/sglang/srt/models/bailing_moe.py | 21 +- .../sglang/srt/models/bailing_moe_linear.py | 19 +- python/sglang/srt/models/bailing_moe_nextn.py | 4 +- python/sglang/srt/models/bert.py | 4 +- python/sglang/srt/models/chatglm.py | 4 +- python/sglang/srt/models/cohere2_moe.py | 8 +- python/sglang/srt/models/commandr.py | 9 +- python/sglang/srt/models/dbrx.py | 11 +- python/sglang/srt/models/deepseek.py | 9 +- python/sglang/srt/models/deepseek_nextn.py | 15 +- python/sglang/srt/models/deepseek_v2.py | 37 +- python/sglang/srt/models/deepseek_v4.py | 29 +- python/sglang/srt/models/deepseek_v4_nextn.py | 16 +- python/sglang/srt/models/dflash.py | 4 +- python/sglang/srt/models/dots_vlm_vit.py | 4 +- python/sglang/srt/models/ernie4.py | 4 +- python/sglang/srt/models/ernie45_moe_vl.py | 6 +- python/sglang/srt/models/exaone.py | 4 +- python/sglang/srt/models/exaone4.py | 15 +- python/sglang/srt/models/exaone_moe.py | 19 +- python/sglang/srt/models/exaone_moe_mtp.py | 5 +- python/sglang/srt/models/falcon_h1.py | 13 +- python/sglang/srt/models/gemma.py | 4 +- python/sglang/srt/models/gemma2.py | 4 +- python/sglang/srt/models/gemma3_causal.py | 11 +- python/sglang/srt/models/gemma3n_causal.py | 4 +- python/sglang/srt/models/gemma4_audio.py | 13 +- python/sglang/srt/models/gemma4_causal.py | 11 +- python/sglang/srt/models/gemma4_vision.py | 4 +- python/sglang/srt/models/glm4.py | 9 +- python/sglang/srt/models/glm4_moe.py | 22 +- python/sglang/srt/models/glm4_moe_lite.py | 11 +- .../sglang/srt/models/glm4_moe_lite_nextn.py | 4 +- python/sglang/srt/models/glm4_moe_nextn.py | 4 +- python/sglang/srt/models/glm4v.py | 15 +- python/sglang/srt/models/glm4v_moe.py | 9 +- python/sglang/srt/models/glm_ocr_nextn.py | 4 +- python/sglang/srt/models/gpt2.py | 4 +- python/sglang/srt/models/gpt_bigcode.py | 4 +- python/sglang/srt/models/gpt_j.py | 4 +- python/sglang/srt/models/gpt_oss.py | 41 +- python/sglang/srt/models/granite.py | 4 +- python/sglang/srt/models/granitemoe.py | 4 +- python/sglang/srt/models/granitemoehybrid.py | 9 +- python/sglang/srt/models/grok.py | 24 +- python/sglang/srt/models/hunyuan.py | 11 +- python/sglang/srt/models/hunyuan_v3.py | 10 +- python/sglang/srt/models/internlm2.py | 4 +- python/sglang/srt/models/interns1pro.py | 6 +- python/sglang/srt/models/internvl.py | 11 +- python/sglang/srt/models/iquest_loopcoder.py | 6 +- python/sglang/srt/models/kimi_linear.py | 13 +- python/sglang/srt/models/laguna.py | 10 +- python/sglang/srt/models/lfm2.py | 5 +- python/sglang/srt/models/lfm2_moe.py | 7 +- python/sglang/srt/models/llada2.py | 21 +- python/sglang/srt/models/llama.py | 9 +- python/sglang/srt/models/llama4.py | 14 +- python/sglang/srt/models/longcat_flash.py | 12 +- .../sglang/srt/models/longcat_flash_nextn.py | 7 +- python/sglang/srt/models/mimo_mtp.py | 4 +- python/sglang/srt/models/mimo_v2.py | 23 +- python/sglang/srt/models/mimo_v2_nextn.py | 7 +- python/sglang/srt/models/mindspore.py | 11 +- python/sglang/srt/models/minicpm.py | 4 +- python/sglang/srt/models/minicpm3.py | 4 +- python/sglang/srt/models/minimax_m2.py | 17 +- python/sglang/srt/models/mixtral.py | 6 +- python/sglang/srt/models/mixtral_quant.py | 9 +- python/sglang/srt/models/mllama.py | 4 +- python/sglang/srt/models/moss_vl.py | 13 +- python/sglang/srt/models/nemotron_h.py | 10 +- python/sglang/srt/models/nemotron_h_mtp.py | 11 +- python/sglang/srt/models/olmo.py | 4 +- python/sglang/srt/models/olmo2.py | 7 +- python/sglang/srt/models/olmoe.py | 4 +- python/sglang/srt/models/opt.py | 9 +- python/sglang/srt/models/orion.py | 4 +- python/sglang/srt/models/persimmon.py | 5 +- python/sglang/srt/models/phi.py | 5 +- python/sglang/srt/models/phi3_small.py | 5 +- python/sglang/srt/models/phimoe.py | 9 +- python/sglang/srt/models/qwen.py | 4 +- python/sglang/srt/models/qwen2.py | 9 +- python/sglang/srt/models/qwen2_5_vl.py | 19 +- python/sglang/srt/models/qwen2_moe.py | 27 +- python/sglang/srt/models/qwen3.py | 12 +- python/sglang/srt/models/qwen3_5.py | 11 +- python/sglang/srt/models/qwen3_5_mtp.py | 5 +- python/sglang/srt/models/qwen3_moe.py | 30 +- python/sglang/srt/models/qwen3_moe_mtp.py | 5 +- python/sglang/srt/models/qwen3_next.py | 11 +- python/sglang/srt/models/qwen3_next_mtp.py | 5 +- python/sglang/srt/models/qwen3_omni_moe.py | 6 +- python/sglang/srt/models/qwen3_vl.py | 16 +- python/sglang/srt/models/sarvam_moe.py | 12 +- python/sglang/srt/models/sdar.py | 13 +- python/sglang/srt/models/sdar_moe.py | 15 +- python/sglang/srt/models/solar.py | 69 +- python/sglang/srt/models/stablelm.py | 4 +- python/sglang/srt/models/starcoder2.py | 5 +- python/sglang/srt/models/step3_vl.py | 17 +- python/sglang/srt/models/step3p5.py | 18 +- python/sglang/srt/models/step3p5_mtp.py | 4 +- .../sglang/srt/models/torch_native_llama.py | 9 +- python/sglang/srt/models/transformers.py | 9 +- python/sglang/srt/models/whisper.py | 4 +- python/sglang/srt/models/xverse.py | 4 +- python/sglang/srt/models/xverse_moe.py | 9 +- python/sglang/srt/models/zaya.py | 16 +- python/sglang/srt/runtime_context.py | 226 +++ .../attention_methods/dense_attention.py | 5 +- .../attention_methods/dsa_attention.py | 6 +- .../attention_methods/dsv4_attention.py | 10 +- .../attention_methods/dual_chunk_attention.py | 4 +- .../attention_methods/gdn_attention.py | 5 +- .../attention_methods/kda_attention.py | 5 +- .../attention_methods/lightning_attention.py | 6 +- .../attention_methods/mamba2_attention.py | 19 +- .../attention_methods/mla_attention.py | 5 +- .../speculative_draft_extend_runner.py | 6 +- .../runner_modes/speculative_draft_runner.py | 11 +- .../attention/test_trtllm_mla_backend.py | 7 +- test/registered/kernels/test_dsa_indexer.py | 6 +- .../layers/mamba/test_mamba2_mixer.py | 15 +- .../test_tbo_filter_batch_marker.py | 3 +- .../unit/eplb/test_deepep_waterfill_eplb.py | 6 +- .../layers/test_flashinfer_comm_fusion.py | 7 +- .../test_cuda_graph_buffer_registry.py | 7 +- test/registered/unit/test_runtime_context.py | 146 ++ 184 files changed, 1871 insertions(+), 1733 deletions(-) create mode 100644 python/sglang/srt/runtime_context.py create mode 100644 test/registered/unit/test_runtime_context.py diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index 64998cbeb..ba924e918 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/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 diff --git a/python/sglang/srt/layers/activation.py b/python/sglang/srt/layers/activation.py index ffd4896c1..292be3c01 100644 --- a/python/sglang/srt/layers/activation.py +++ b/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) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 855d3e3d0..5c2149728 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/cutlass_mla_backend.py b/python/sglang/srt/layers/attention/cutlass_mla_backend.py index 342298890..ca156d0e7 100644 --- a/python/sglang/srt/layers/attention/cutlass_mla_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 54f502442..8b504153f 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/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_size = get_attention_cp_size() + cp_rank = get_parallel().attn_cp_rank + 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, ( diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index 4857c9551..a5e55570b 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/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_size = get_attention_cp_size() + cp_rank = get_parallel().attn_cp_rank + 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() diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index 0f5b3c050..c242a1c99 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/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_rank = get_attn_context_model_parallel_rank() + self.cp_size = get_parallel().attn_cp_size + self.cp_rank = get_parallel().attn_cp_rank else: self.cp_size = None self.cp_rank = None diff --git a/python/sglang/srt/layers/attention/dsa/utils.py b/python/sglang/srt/layers/attention/dsa/utils.py index 04f542752..506a299ff 100644 --- a/python/sglang/srt/layers/attention/dsa/utils.py +++ b/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_rank = get_attention_cp_rank() + cp_size = get_parallel().attn_cp_size + 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_rank = get_attention_cp_rank() + cp_size = get_parallel().attn_cp_size + 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_rank = get_attention_cp_rank() + cp_size = get_parallel().attn_cp_size + 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( diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index cdc7b8ef2..169019f15 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/dsv4/compressor.py b/python/sglang/srt/layers/attention/dsv4/compressor.py index d3e87f906..f5822ddd0 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor.py +++ b/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(), ) diff --git a/python/sglang/srt/layers/attention/dual_chunk_flashattention_backend.py b/python/sglang/srt/layers/attention/dual_chunk_flashattention_backend.py index 96feb0c28..040199ff9 100644 --- a/python/sglang/srt/layers/attention/dual_chunk_flashattention_backend.py +++ b/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)] diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 11b917208..957c0c4ec 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/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() - if mul_base % get_attention_cp_size() != 0: - mul_base *= get_attention_cp_size() + mul_base *= get_parallel().attn_tp_size + if mul_base % get_parallel().attn_cp_size != 0: + 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 diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py index 17d4509c7..d2c8b9395 100644 --- a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/flashmla_backend.py b/python/sglang/srt/layers/attention/flashmla_backend.py index c4a6c4ecd..0c382790e 100644 --- a/python/sglang/srt/layers/attention/flashmla_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/linear/lightning_backend.py b/python/sglang/srt/layers/attention/linear/lightning_backend.py index d20e50a97..ce343ba25 100644 --- a/python/sglang/srt/layers/attention/linear/lightning_backend.py +++ b/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_rank = get_attention_tp_rank() + tp_heads = n_attention_heads // get_parallel().attn_tp_size + tp_rank = get_parallel().attn_tp_rank if num_hidden_layers <= 1: slope_rate_list = [slopes * (1 + 1e-5)] else: diff --git a/python/sglang/srt/layers/attention/mamba/mamba.py b/python/sglang/srt/layers/attention/mamba/mamba.py index c3617bcd7..80ec14481 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba.py +++ b/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_rank = get_attention_tp_rank() + self.tp_size = get_parallel().attn_tp_size + self.tp_rank = get_parallel().attn_tp_rank else: - self.tp_size = get_tensor_model_parallel_world_size() - self.tp_rank = get_tensor_model_parallel_rank() + self.tp_size = get_parallel().tp_size + 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 diff --git a/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py b/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py index 79155affe..753150054 100644 --- a/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py +++ b/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_rank = get_attention_tp_rank() + self.tp_size = get_parallel().attn_tp_size + self.tp_rank = get_parallel().attn_tp_rank else: - self.tp_size = get_tensor_model_parallel_world_size() - self.tp_rank = get_tensor_model_parallel_rank() + self.tp_size = get_parallel().tp_size + 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 diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index d2384656a..a44dbd2fc 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index d37859d58..418ed488b 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/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_kv_heads = config.get_num_kv_heads(get_attention_tp_size()) - self.num_local_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_parallel().attn_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 diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index a6fc032c3..9ffd181ab 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/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_rank = 0 if use_data_parallel else get_attention_tp_rank() + self.tp_size = 1 if use_data_parallel else get_parallel().attn_tp_size + 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 diff --git a/python/sglang/srt/layers/attention/vision_utils.py b/python/sglang/srt/layers/attention/vision_utils.py index ecccb1f85..32e5091f3 100644 --- a/python/sglang/srt/layers/attention/vision_utils.py +++ b/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", diff --git a/python/sglang/srt/layers/attention/wave_backend.py b/python/sglang/srt/layers/attention/wave_backend.py index 55871b32d..68aa893d2 100644 --- a/python/sglang/srt/layers/attention/wave_backend.py +++ b/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( diff --git a/python/sglang/srt/layers/clippable_linear.py b/python/sglang/srt/layers/clippable_linear.py index a253bb421..4f53ae8b6 100644 --- a/python/sglang/srt/layers/clippable_linear.py +++ b/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( diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 36cdb23c3..2a84ab161 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/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_size = get_attention_tp_size() - attn_dp_size = get_attention_dp_size() - attn_cp_size = get_attention_cp_size() - attn_cp_rank = get_attention_cp_rank() - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size + attn_dp_size = get_parallel().attn_dp_size + attn_cp_size = get_parallel().attn_cp_size + attn_cp_rank = get_parallel().attn_cp_rank + tp_size = get_parallel().tp_size + 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() diff --git a/python/sglang/srt/layers/communicator_dsa_cp.py b/python/sglang/srt/layers/communicator_dsa_cp.py index dc979a8ab..73845649e 100644 --- a/python/sglang/srt/layers/communicator_dsa_cp.py +++ b/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_tp_size = get_attention_tp_size() + attn_dp_size = get_parallel().attn_dp_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_tp_size = get_attention_tp_size() + attn_dp_size = get_parallel().attn_dp_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_rank = get_attention_cp_rank() + cp_size = get_parallel().attn_cp_size + 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) diff --git a/python/sglang/srt/layers/cp/base.py b/python/sglang/srt/layers/cp/base.py index 57a7dcd15..b63f46933 100644 --- a/python/sglang/srt/layers/cp/base.py +++ b/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: diff --git a/python/sglang/srt/layers/flashinfer_comm_fusion.py b/python/sglang/srt/layers/flashinfer_comm_fusion.py index 9f0e0b63c..2411e09c0 100644 --- a/python/sglang/srt/layers/flashinfer_comm_fusion.py +++ b/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() - rank = get_attn_tensor_model_parallel_rank() + world_size = get_parallel().attn_tp_size + rank = get_parallel().attn_tp_rank coordinator = get_attn_tp_group() else: - if get_moe_expert_parallel_world_size() > 1: - world_size = get_moe_expert_parallel_world_size() - rank = get_moe_expert_parallel_rank() + if get_parallel().moe_ep_size > 1: + world_size = get_parallel().moe_ep_size + rank = get_parallel().moe_ep_rank coordinator = get_moe_ep_group() else: - world_size = get_moe_tensor_parallel_world_size() - rank = get_moe_tensor_parallel_rank() + world_size = get_parallel().moe_tp_size + 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: - world_size = get_moe_expert_parallel_world_size() + if get_parallel().moe_ep_size > 1: + 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") diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index d416e937e..6504923a3 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/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: - world_size = get_moe_expert_parallel_world_size() + if get_parallel().moe_ep_size > 1: + 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: diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index fc197860d..50c6d3cc9 100644 --- a/python/sglang/srt/layers/linear.py +++ b/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_size = get_tensor_model_parallel_world_size() + self.tp_rank = get_parallel().tp_rank + 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_size = get_tensor_model_parallel_world_size() + self.tp_rank = get_parallel().tp_rank + 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, diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 15a3d6aac..a99d25267 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/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 diff --git a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py b/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py index 16d5428c7..82b1c239c 100644 --- a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py +++ b/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( diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 923036a91..8f74b354e 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/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_rank = get_moe_expert_parallel_rank() - self.moe_tp_size = get_moe_tensor_parallel_world_size() - self.moe_tp_rank = get_moe_tensor_parallel_rank() + self.moe_ep_size = get_parallel().moe_ep_size + self.moe_ep_rank = get_parallel().moe_ep_rank + self.moe_tp_size = get_parallel().moe_tp_size + 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. diff --git a/python/sglang/srt/layers/moe/kt_ep_wrapper.py b/python/sglang/srt/layers/moe/kt_ep_wrapper.py index 128853931..63b52e5e0 100644 --- a/python/sglang/srt/layers/moe/kt_ep_wrapper.py +++ b/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 diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py index 1e21ac701..6f486194f 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py +++ b/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() - rank = get_moe_expert_parallel_rank() + world_size = get_parallel().moe_ep_size + 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(), diff --git a/python/sglang/srt/layers/moe/token_dispatcher/standard.py b/python/sglang/srt/layers/moe/token_dispatcher/standard.py index 40f305712..650da2797 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/standard.py +++ b/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 diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 6e80ee794..37964b346 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/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_rank = get_moe_expert_parallel_rank() + ep_size = get_parallel().moe_ep_size + 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. diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index f7eda3c36..67dbe70e3 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/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_moe_expert_parallel_world_size() == get_attention_dp_size() + and get_parallel().attn_dp_size > 1 + and get_parallel().moe_ep_size == get_parallel().attn_dp_size ) diff --git a/python/sglang/srt/layers/quantization/blockwise_int8.py b/python/sglang/srt/layers/quantization/blockwise_int8.py index d39c66c9b..432374354 100644 --- a/python/sglang/srt/layers/quantization/blockwise_int8.py +++ b/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], diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxint4_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxint4_moe.py index 569e20454..7dc430e9c 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxint4_moe.py +++ b/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( diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py index e0f59c9f8..b36e28f21 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py +++ b/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], diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 08ba0cada..6d33b9121 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/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, diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index ca16262f1..82129b102 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/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 diff --git a/python/sglang/srt/layers/quantization/moe_wna16.py b/python/sglang/srt/layers/quantization/moe_wna16.py index a7d87c638..e5221c8a9 100644 --- a/python/sglang/srt/layers/quantization/moe_wna16.py +++ b/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 diff --git a/python/sglang/srt/layers/quantization/quark_int4fp8_moe.py b/python/sglang/srt/layers/quantization/quark_int4fp8_moe.py index 3972bca08..1fdb45780 100644 --- a/python/sglang/srt/layers/quantization/quark_int4fp8_moe.py +++ b/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( diff --git a/python/sglang/srt/layers/quantization/w8a8_int8.py b/python/sglang/srt/layers/quantization/w8a8_int8.py index f94c90b30..18fcb184e 100644 --- a/python/sglang/srt/layers/quantization/w8a8_int8.py +++ b/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( diff --git a/python/sglang/srt/layers/utils/cp_utils.py b/python/sglang/srt/layers/utils/cp_utils.py index e40db43bc..74ddc07db 100644 --- a/python/sglang/srt/layers/utils/cp_utils.py +++ b/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_rank = get_attention_cp_rank() + cp_size = get_parallel().attn_cp_size + 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: diff --git a/python/sglang/srt/layers/vocab_parallel_embedding.py b/python/sglang/srt/layers/vocab_parallel_embedding.py index cbba8cd06..46c711aa7 100644 --- a/python/sglang/srt/layers/vocab_parallel_embedding.py +++ b/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() - self.tp_size = get_attention_tp_size() + tp_rank = get_parallel().attn_tp_rank + self.tp_size = get_parallel().attn_tp_size else: - tp_rank = get_tensor_model_parallel_rank() - self.tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_parallel().tp_rank + self.tp_size = get_parallel().tp_size else: assert use_attn_tp_group is False tp_rank = 0 diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index dc4483edc..30c600bfa 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/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_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + 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( diff --git a/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py index b81fedebf..0cb801333 100644 --- a/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py +++ b/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 ( - get_attention_cp_size, - get_attention_tp_rank, - get_attention_tp_size, -) +from sglang.srt.runtime_context import get_parallel 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: - mul_base *= get_attention_cp_size() + if mul_base % get_parallel().attn_cp_size != 0: + 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_rank = get_attention_tp_rank() + self.attn_tp_size = get_parallel().attn_tp_size + self.attn_tp_rank = get_parallel().attn_tp_rank self.tbo_plugin = TboCudaGraphRunnerPlugin() @staticmethod diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index e53a3193b..afd17dc16 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/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 - or loaded_weight.size(0) - < get_tensor_model_parallel_world_size() * shard_size + loaded_weight.size(0) % get_parallel().tp_size != 0 + or loaded_weight.size(0) < get_parallel().tp_size * shard_size ) and loaded_weight.dim() == 1 ): diff --git a/python/sglang/srt/models/afmoe.py b/python/sglang/srt/models/afmoe.py index efdfaade7..4f0d3bbef 100644 --- a/python/sglang/srt/models/afmoe.py +++ b/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.tp_size = get_tensor_model_parallel_world_size() + self.rank = get_parallel().tp_rank + 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 diff --git a/python/sglang/srt/models/apertus.py b/python/sglang/srt/models/apertus.py index efdc5aa2d..f20380724 100644 --- a/python/sglang/srt/models/apertus.py +++ b/python/sglang/srt/models/apertus.py @@ -1,687 +1,686 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Copyright 2025 The SwissAI Initiative -# Copyright 2023-2024 SGLang Team -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== - -# Adapted from -# https://github.com/vllm-project/vllm/blob/c7f2cf2b7f67bce5842fedfdba508440fe257375/vllm/model_executor/models/llama.py#L1 -"""Inference-only Apertus model compatible with HuggingFace weights.""" - -import logging -from typing import Any, Dict, Iterable, List, Optional, Tuple, Union - -import torch -from torch import nn -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 -from sglang.srt.layers.linear import ( - ColumnParallelLinear, - QKVParallelLinear, - RowParallelLinear, -) -from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput -from sglang.srt.layers.pooler import Pooler, PoolingType -from sglang.srt.layers.quantization.base_config import QuantizationConfig -from sglang.srt.layers.radix_attention import RadixAttention -from sglang.srt.layers.rotary_embedding import get_rope -from sglang.srt.layers.utils import PPMissingLayer, get_layer_id -from sglang.srt.layers.vocab_parallel_embedding import ( - ParallelLMHead, - VocabParallelEmbedding, -) -from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors -from sglang.srt.model_loader.weight_utils import ( - default_weight_loader, - kv_cache_scales_loader, - maybe_remap_kv_scale_name, -) -from sglang.srt.server_args import get_global_server_args -from sglang.srt.utils import add_prefix, make_layers - -logger = logging.getLogger(__name__) - - -class ApertusMLP(nn.Module): - def __init__( - self, - hidden_size: int, - intermediate_size: int, - hidden_act: str, - quant_config: Optional[QuantizationConfig] = None, - bias: bool = False, - prefix: str = "", - reduce_results: bool = True, - ) -> None: - super().__init__() - self.up_proj = ColumnParallelLinear( - hidden_size, - intermediate_size, - bias=bias, - quant_config=quant_config, - prefix=add_prefix("up_proj", prefix), - ) - self.down_proj = RowParallelLinear( - intermediate_size, - hidden_size, - bias=bias, - quant_config=quant_config, - prefix=add_prefix("down_proj", prefix), - reduce_results=reduce_results, - ) - if hidden_act != "xielu": - raise ValueError( - f"Unsupported activation: {hidden_act}. " - "Only xIELU is supported for now." - ) - self.act_fn = XIELU() - - def forward( - self, - x, - forward_batch=None, - use_reduce_scatter: bool = False, - ): - # note: with xielu, there's no gate_proj - x, _ = self.up_proj(x) - x = self.act_fn(x) - x, _ = self.down_proj( - x, - skip_all_reduce=use_reduce_scatter, - ) - return x - - -class ApertusAttention(nn.Module): - def __init__( - self, - config: ApertusConfig, - hidden_size: int, - num_heads: int, - num_kv_heads: int, - layer_id: int = 0, - rope_theta: float = 10000, - rope_scaling: Optional[Dict[str, Any]] = None, - rope_is_neox_style: bool = True, - max_position_embeddings: int = 8192, - quant_config: Optional[QuantizationConfig] = None, - prefix: str = "", - bias: bool = False, - bias_o_proj: bool = False, - ) -> None: - super().__init__() - self.layer_id = layer_id - self.hidden_size = hidden_size - tp_size = get_tensor_model_parallel_world_size() - self.total_num_heads = num_heads - assert self.total_num_heads % tp_size == 0 - self.num_heads = self.total_num_heads // tp_size - self.total_num_kv_heads = num_kv_heads - if self.total_num_kv_heads >= tp_size: - # Number of KV heads is greater than TP size, so we partition - # the KV heads across multiple tensor parallel GPUs. - assert self.total_num_kv_heads % tp_size == 0 - else: - # Number of KV heads is less than TP size, so we replicate - # the KV heads across multiple tensor parallel GPUs. - assert tp_size % self.total_num_kv_heads == 0 - self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - # MistralConfig has an optional head_dim introduced by Mistral-Nemo - self.head_dim = getattr( - config, "head_dim", self.hidden_size // self.total_num_heads - ) - partial_rotary_factor = getattr(config, "partial_rotary_factor", 1) - self.rotary_dim = int(partial_rotary_factor * self.head_dim) - self.q_size = self.num_heads * self.head_dim - self.kv_size = self.num_kv_heads * self.head_dim - self.scaling = self.head_dim**-0.5 - self.rope_theta = rope_theta - self.max_position_embeddings = max_position_embeddings - - self.qkv_proj = QKVParallelLinear( - hidden_size, - self.head_dim, - self.total_num_heads, - self.total_num_kv_heads, - bias=bias, - quant_config=quant_config, - prefix=add_prefix("qkv_proj", prefix), - ) - self.o_proj = RowParallelLinear( - self.total_num_heads * self.head_dim, - hidden_size, - bias=bias_o_proj, - quant_config=quant_config, - prefix=add_prefix("o_proj", prefix), - ) - - self.rotary_emb = get_rope( - self.head_dim, - rotary_dim=self.rotary_dim, - max_position=max_position_embeddings, - base=rope_theta, - rope_scaling=rope_scaling, - is_neox_style=rope_is_neox_style, - ) - self.attn = RadixAttention( - self.num_heads, - self.head_dim, - self.scaling, - num_kv_heads=self.num_kv_heads, - layer_id=layer_id, - quant_config=quant_config, - prefix=add_prefix("attn", prefix), - ) - self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) - self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - forward_batch: ForwardBatch, - ) -> torch.Tensor: - qkv, _ = self.qkv_proj(hidden_states) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q = self.q_norm(q.contiguous().view(-1, self.head_dim)).view_as(q) - k = self.k_norm(k.contiguous().view(-1, self.head_dim)).view_as(k) - q, k = self.rotary_emb(positions, q, k) - attn_output = self.attn(q, k, v, forward_batch) - output, _ = self.o_proj(attn_output) - return output - - -class ApertusDecoderLayer(nn.Module): - def __init__( - self, - config: ApertusConfig, - layer_id: int = 0, - quant_config: Optional[QuantizationConfig] = None, - prefix: str = "", - ) -> None: - super().__init__() - self.hidden_size = config.hidden_size - rope_theta = config.rope_parameters["rope_theta"] - rope_scaling = config.rope_parameters - if rope_scaling is not None and getattr( - config, "original_max_position_embeddings", None - ): - rope_scaling["original_max_position_embeddings"] = ( - config.original_max_position_embeddings - ) - rope_is_neox_style = getattr(config, "rope_is_neox_style", True) - max_position_embeddings = getattr(config, "max_position_embeddings", 8192) - # Support llamafy/Qwen-Qwen2.5-7B-Instruct-llamafied with attention_bias - # Support internlm/internlm-7b with bias - attention_bias = getattr(config, "attention_bias", False) or getattr( - config, "bias", False - ) - bias_o_proj = attention_bias - # support internlm/internlm3-8b with qkv_bias - if hasattr(config, "qkv_bias"): - attention_bias = config.qkv_bias - self.self_attn = ApertusAttention( - config=config, - hidden_size=self.hidden_size, - num_heads=config.num_attention_heads, - num_kv_heads=config.num_key_value_heads, - layer_id=layer_id, - rope_theta=rope_theta, - rope_scaling=rope_scaling, - rope_is_neox_style=rope_is_neox_style, - max_position_embeddings=max_position_embeddings, - quant_config=quant_config, - prefix=add_prefix("self_attn", prefix), - bias=attention_bias, - bias_o_proj=bias_o_proj, - ) - self.mlp = ApertusMLP( - hidden_size=self.hidden_size, - intermediate_size=config.intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - bias=getattr(config, "mlp_bias", False), - prefix=add_prefix("mlp", prefix), - ) - self.attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.feedforward_layernorm = RMSNorm( - config.hidden_size, eps=config.rms_norm_eps - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - forward_batch: ForwardBatch, - residual: Optional[torch.Tensor], - ) -> Tuple[torch.Tensor, torch.Tensor]: - # Self Attention - if residual is None: - residual = hidden_states - hidden_states = self.attention_layernorm(hidden_states) - else: - hidden_states, residual = self.attention_layernorm(hidden_states, residual) - hidden_states = self.self_attn( - positions=positions, - hidden_states=hidden_states, - forward_batch=forward_batch, - ) - - # Fully Connected - hidden_states, residual = self.feedforward_layernorm(hidden_states, residual) - hidden_states = self.mlp(hidden_states) - return hidden_states, residual - - -class ApertusModel(nn.Module): - def __init__( - self, - config: ApertusConfig, - quant_config: Optional[QuantizationConfig] = None, - prefix: str = "", - ) -> None: - super().__init__() - self.quant_config = quant_config - self.config = config - self.padding_idx = config.pad_token_id - self.vocab_size = config.vocab_size - self.org_vocab_size = config.vocab_size - self.pp_group = get_pp_group() - if self.pp_group.is_first_rank: - self.embed_tokens = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=add_prefix("embed_tokens", prefix), - ) - else: - self.embed_tokens = PPMissingLayer() - - self.layers, self.start_layer, self.end_layer = make_layers( - config.num_hidden_layers, - lambda idx, prefix: ApertusDecoderLayer( - config=config, quant_config=quant_config, layer_id=idx, prefix=prefix - ), - pp_rank=self.pp_group.rank_in_group, - pp_size=self.pp_group.world_size, - prefix="model.layers", - ) - - if self.pp_group.is_last_rank: - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - else: - self.norm = PPMissingLayer(return_tuple=True) - self.layers_to_capture = [] - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - forward_batch: ForwardBatch, - input_embeds: torch.Tensor = None, - pp_proxy_tensors: Optional[PPProxyTensors] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[torch.Tensor]], PPProxyTensors]: - if self.pp_group.is_first_rank: - if input_embeds is None: - hidden_states = self.embed_tokens(input_ids) - else: - hidden_states = input_embeds - residual = None - else: - assert pp_proxy_tensors is not None - # FIXME(@ying): reduce the number of proxy tensors by not fusing layer norms - hidden_states = pp_proxy_tensors["hidden_states"] - residual = pp_proxy_tensors["residual"] - deferred_norm = None - - aux_hidden_states = [] - for i in range(self.start_layer, self.end_layer): - if i in self.layers_to_capture: - aux_hidden_states.append(hidden_states + residual) - layer = self.layers[i] - hidden_states, residual = layer( - positions, - hidden_states, - forward_batch, - residual, - ) - - if not self.pp_group.is_last_rank: - return PPProxyTensors( - { - "hidden_states": hidden_states, - "residual": residual, - } - ) - else: - hidden_states, _ = self.norm(hidden_states, residual) - - if len(aux_hidden_states) == 0: - return hidden_states - - return hidden_states, aux_hidden_states - - # If this function is called, it should always initialize KV cache scale - # 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_rank = get_tensor_model_parallel_rank() - for layer_idx, scaling_factor in kv_cache_scales_loader( - quantization_param_path, - tp_rank, - tp_size, - self.config.num_hidden_layers, - self.config.__class__.model_type, - ): - if not isinstance(self.layers[layer_idx], nn.Identity): - layer_self_attn = self.layers[layer_idx].self_attn - - if hasattr(layer_self_attn.attn, "k_scale"): - layer_self_attn.attn.k_scale = scaling_factor - layer_self_attn.attn.v_scale = scaling_factor - else: - raise RuntimeError( - "Self attention has no KV cache scaling " "factor attribute!" - ) - - -class ApertusForCausalLM(nn.Module): - # LoRA specific attributes - embedding_modules = { - "embed_tokens": "input_embeddings", - "lm_head": "output_embeddings", - } - embedding_padding_modules = ["lm_head"] - # BitandBytes specific attributes - default_bitsandbytes_target_modules = [ - ".down_proj.", - ".up_proj.", - ".q_proj.", - ".k_proj.", - ".v_proj.", - ".o_proj.", - ] - # in TP, these weights are partitioned along the column dimension (dim=-1) - column_parallel_weights_modules = [".down_proj.", ".o_proj."] - bitsandbytes_stacked_params_mapping = { - # shard_name, weight_name, index - ".q_proj": (".qkv_proj", 0), - ".k_proj": (".qkv_proj", 1), - ".v_proj": (".qkv_proj", 2), - } - - def __init__( - self, - config: ApertusConfig, - quant_config: Optional[QuantizationConfig] = None, - prefix: str = "", - ) -> None: - super().__init__() - self.pp_group = get_pp_group() - self.config = config - self.quant_config = quant_config - self.model = self._init_model(config, quant_config, add_prefix("model", prefix)) - if self.config.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=add_prefix("lm_head", prefix), - use_attn_tp_group=get_global_server_args().enable_dp_lm_head, - ) - self.logits_processor = LogitsProcessor(config) - self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True) - self.stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - (".qkv_proj", ".q_proj", "q"), - (".qkv_proj", ".k_proj", "k"), - (".qkv_proj", ".v_proj", "v"), - ] - - self.capture_aux_hidden_states = False - - def _init_model( - self, - config: ApertusConfig, - quant_config: Optional[QuantizationConfig] = None, - prefix: str = "", - ): - return ApertusModel(config, quant_config=quant_config, prefix=prefix) - - @torch.no_grad() - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - forward_batch: ForwardBatch, - input_embeds: torch.Tensor = None, - get_embedding: bool = False, - pp_proxy_tensors: Optional[PPProxyTensors] = None, - ) -> LogitsProcessorOutput: - hidden_states = self.model( - input_ids, - positions, - forward_batch, - input_embeds, - pp_proxy_tensors=pp_proxy_tensors, - ) - - aux_hidden_states = None - if self.capture_aux_hidden_states: - hidden_states, aux_hidden_states = hidden_states - - if self.pp_group.is_last_rank: - if not get_embedding: - return self.logits_processor( - input_ids, - hidden_states, - self.lm_head, - forward_batch, - aux_hidden_states, - ) - else: - return self.pooler(hidden_states, forward_batch) - else: - return hidden_states - - @torch.no_grad() - def forward_split_prefill( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - forward_batch: ForwardBatch, - split_interval: Tuple[int, int], # [start, end) 0-based - input_embeds: torch.Tensor = None, - ) -> Optional[LogitsProcessorOutput]: - start, end = split_interval - # embed - if start == 0: - if input_embeds is None: - forward_batch.hidden_states = self.model.embed_tokens(input_ids) - else: - forward_batch.hidden_states = input_embeds - # decoder layer - for i in range(start, end): - layer = self.model.layers[i] - forward_batch.hidden_states, forward_batch.residual = layer( - positions, - forward_batch.hidden_states, - forward_batch, - forward_batch.residual, - ) - - if end == self.model.config.num_hidden_layers: - # norm - hidden_states, _ = self.model.norm( - forward_batch.hidden_states, forward_batch.residual - ) - forward_batch.hidden_states = hidden_states - # logits process - result = self.logits_processor( - input_ids, forward_batch.hidden_states, self.lm_head, forward_batch - ) - else: - result = None - - return result - - @property - def start_layer(self): - return self.model.start_layer - - @property - def end_layer(self): - return self.model.end_layer - - def get_input_embeddings(self) -> nn.Embedding: - return self.model.embed_tokens - - def get_module_name_from_weight_name(self, name): - for param_name, weight_name, shard_id, num_shard in self.stacked_params_mapping: - if weight_name in name: - return ( - name.replace(weight_name, param_name)[: -len(".weight")], - num_shard, - ) - return name[: -len(".weight")], 1 - - def get_num_params(self): - params_dict = dict(self.named_parameters()) - return len(params_dict) - - def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - (".qkv_proj", ".q_proj", "q"), - (".qkv_proj", ".k_proj", "k"), - (".qkv_proj", ".v_proj", "v"), - ] - - params_dict = dict(self.named_parameters()) - - for name, buffer in self.named_buffers(): - if name.endswith(".beta") or name.endswith(".eps"): - params_dict[name] = buffer - - for name, loaded_weight in weights: - layer_id = get_layer_id(name) - if ( - layer_id is not None - and hasattr(self.model, "start_layer") - and ( - layer_id < self.model.start_layer - or layer_id >= self.model.end_layer - ) - ): - continue - if "rotary_emb.inv_freq" in name or "projector" in name: - continue - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: - # Models trained using ColossalAI may include these tensors in - # the checkpoint. Skip them. - continue - if name.startswith("model.vision_tower") and name not in params_dict: - continue - if self.config.tie_word_embeddings and "lm_head.weight" in name: - continue - # Handle FP8 kv-scale remapping - if "scale" in name: - name = maybe_remap_kv_scale_name(name, params_dict) - if name is None: - continue - - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if name not in params_dict: - continue - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - # Skip loading kv_scale from ckpts towards new design. - if name.endswith(".kv_scale") and name not in params_dict: - continue - if name in params_dict.keys(): - param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader - ) - weight_loader(param, loaded_weight) - else: - logger.warning(f"Parameter {name} not found in params_dict") - - def get_embed_and_head(self): - return self.model.embed_tokens.weight, self.lm_head.weight - - def set_embed_and_head(self, embed, head): - del self.model.embed_tokens.weight - del self.lm_head.weight - self.model.embed_tokens.weight = embed - self.lm_head.weight = head - torch.cuda.empty_cache() - torch.cuda.synchronize() - - def get_embed(self): - return self.model.embed_tokens.weight - - def set_embed(self, embed): - # NOTE: If draft hidden size != target hidden size, the embed weight cannot be shared for EAGLE3 - if ( - hasattr(self.config, "target_hidden_size") - and self.config.target_hidden_size != self.config.hidden_size - ): - return - del self.model.embed_tokens.weight - self.model.embed_tokens.weight = embed - torch.cuda.empty_cache() - torch.cuda.synchronize() - - def load_kv_cache_scales(self, quantization_param_path: str) -> None: - self.model.load_kv_cache_scales(quantization_param_path) - - def set_eagle3_layers_to_capture(self, layer_ids: Optional[List[int]] = None): - if not self.pp_group.is_last_rank: - return - - if layer_ids is None: - self.capture_aux_hidden_states = True - num_layers = self.config.num_hidden_layers - self.model.layers_to_capture = [2, num_layers // 2, num_layers - 3] - else: - self.capture_aux_hidden_states = True - # we plus 1 here because in sglang, for the ith layer, it takes the output - # of the (i-1)th layer as aux hidden state - self.model.layers_to_capture = [val + 1 for val in layer_ids] - - -EntryClass = [ApertusForCausalLM] +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright 2025 The SwissAI Initiative +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from +# https://github.com/vllm-project/vllm/blob/c7f2cf2b7f67bce5842fedfdba508440fe257375/vllm/model_executor/models/llama.py#L1 +"""Inference-only Apertus model compatible with HuggingFace weights.""" + +import logging +from typing import Any, Dict, Iterable, List, Optional, Tuple, Union + +import torch +from torch import nn +from transformers import ApertusConfig + +from sglang.srt.distributed import ( + get_pp_group, +) +from sglang.srt.layers.activation import XIELU +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + ColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput +from sglang.srt.layers.pooler import Pooler, PoolingType +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.utils import PPMissingLayer, get_layer_id +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.model_loader.weight_utils import ( + default_weight_loader, + 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 + +logger = logging.getLogger(__name__) + + +class ApertusMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + bias: bool = False, + prefix: str = "", + reduce_results: bool = True, + ) -> None: + super().__init__() + self.up_proj = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=bias, + quant_config=quant_config, + prefix=add_prefix("up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=bias, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + reduce_results=reduce_results, + ) + if hidden_act != "xielu": + raise ValueError( + f"Unsupported activation: {hidden_act}. " + "Only xIELU is supported for now." + ) + self.act_fn = XIELU() + + def forward( + self, + x, + forward_batch=None, + use_reduce_scatter: bool = False, + ): + # note: with xielu, there's no gate_proj + x, _ = self.up_proj(x) + x = self.act_fn(x) + x, _ = self.down_proj( + x, + skip_all_reduce=use_reduce_scatter, + ) + return x + + +class ApertusAttention(nn.Module): + def __init__( + self, + config: ApertusConfig, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + layer_id: int = 0, + rope_theta: float = 10000, + rope_scaling: Optional[Dict[str, Any]] = None, + rope_is_neox_style: bool = True, + max_position_embeddings: int = 8192, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + bias: bool = False, + bias_o_proj: bool = False, + ) -> None: + super().__init__() + self.layer_id = layer_id + self.hidden_size = hidden_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 + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + # MistralConfig has an optional head_dim introduced by Mistral-Nemo + self.head_dim = getattr( + config, "head_dim", self.hidden_size // self.total_num_heads + ) + partial_rotary_factor = getattr(config, "partial_rotary_factor", 1) + self.rotary_dim = int(partial_rotary_factor * self.head_dim) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=bias, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=bias_o_proj, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + ) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.rotary_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + is_neox_style=rope_is_neox_style, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q = self.q_norm(q.contiguous().view(-1, self.head_dim)).view_as(q) + k = self.k_norm(k.contiguous().view(-1, self.head_dim)).view_as(k) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.o_proj(attn_output) + return output + + +class ApertusDecoderLayer(nn.Module): + def __init__( + self, + config: ApertusConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + rope_theta = config.rope_parameters["rope_theta"] + rope_scaling = config.rope_parameters + if rope_scaling is not None and getattr( + config, "original_max_position_embeddings", None + ): + rope_scaling["original_max_position_embeddings"] = ( + config.original_max_position_embeddings + ) + rope_is_neox_style = getattr(config, "rope_is_neox_style", True) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + # Support llamafy/Qwen-Qwen2.5-7B-Instruct-llamafied with attention_bias + # Support internlm/internlm-7b with bias + attention_bias = getattr(config, "attention_bias", False) or getattr( + config, "bias", False + ) + bias_o_proj = attention_bias + # support internlm/internlm3-8b with qkv_bias + if hasattr(config, "qkv_bias"): + attention_bias = config.qkv_bias + self.self_attn = ApertusAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + layer_id=layer_id, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + rope_is_neox_style=rope_is_neox_style, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + bias=attention_bias, + bias_o_proj=bias_o_proj, + ) + self.mlp = ApertusMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + bias=getattr(config, "mlp_bias", False), + prefix=add_prefix("mlp", prefix), + ) + self.attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.feedforward_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.attention_layernorm(hidden_states) + else: + hidden_states, residual = self.attention_layernorm(hidden_states, residual) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + # Fully Connected + hidden_states, residual = self.feedforward_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class ApertusModel(nn.Module): + def __init__( + self, + config: ApertusConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.quant_config = quant_config + self.config = config + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.org_vocab_size = config.vocab_size + self.pp_group = get_pp_group() + if self.pp_group.is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("embed_tokens", prefix), + ) + else: + self.embed_tokens = PPMissingLayer() + + self.layers, self.start_layer, self.end_layer = make_layers( + config.num_hidden_layers, + lambda idx, prefix: ApertusDecoderLayer( + config=config, quant_config=quant_config, layer_id=idx, prefix=prefix + ), + pp_rank=self.pp_group.rank_in_group, + pp_size=self.pp_group.world_size, + prefix="model.layers", + ) + + if self.pp_group.is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer(return_tuple=True) + self.layers_to_capture = [] + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[torch.Tensor]], PPProxyTensors]: + if self.pp_group.is_first_rank: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + residual = None + else: + assert pp_proxy_tensors is not None + # FIXME(@ying): reduce the number of proxy tensors by not fusing layer norms + hidden_states = pp_proxy_tensors["hidden_states"] + residual = pp_proxy_tensors["residual"] + deferred_norm = None + + aux_hidden_states = [] + for i in range(self.start_layer, self.end_layer): + if i in self.layers_to_capture: + aux_hidden_states.append(hidden_states + residual) + layer = self.layers[i] + hidden_states, residual = layer( + positions, + hidden_states, + forward_batch, + residual, + ) + + if not self.pp_group.is_last_rank: + return PPProxyTensors( + { + "hidden_states": hidden_states, + "residual": residual, + } + ) + else: + hidden_states, _ = self.norm(hidden_states, residual) + + if len(aux_hidden_states) == 0: + return hidden_states + + return hidden_states, aux_hidden_states + + # If this function is called, it should always initialize KV cache scale + # 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_parallel().tp_size + tp_rank = get_parallel().tp_rank + for layer_idx, scaling_factor in kv_cache_scales_loader( + quantization_param_path, + tp_rank, + tp_size, + self.config.num_hidden_layers, + self.config.__class__.model_type, + ): + if not isinstance(self.layers[layer_idx], nn.Identity): + layer_self_attn = self.layers[layer_idx].self_attn + + if hasattr(layer_self_attn.attn, "k_scale"): + layer_self_attn.attn.k_scale = scaling_factor + layer_self_attn.attn.v_scale = scaling_factor + else: + raise RuntimeError( + "Self attention has no KV cache scaling " "factor attribute!" + ) + + +class ApertusForCausalLM(nn.Module): + # LoRA specific attributes + embedding_modules = { + "embed_tokens": "input_embeddings", + "lm_head": "output_embeddings", + } + embedding_padding_modules = ["lm_head"] + # BitandBytes specific attributes + default_bitsandbytes_target_modules = [ + ".down_proj.", + ".up_proj.", + ".q_proj.", + ".k_proj.", + ".v_proj.", + ".o_proj.", + ] + # in TP, these weights are partitioned along the column dimension (dim=-1) + column_parallel_weights_modules = [".down_proj.", ".o_proj."] + bitsandbytes_stacked_params_mapping = { + # shard_name, weight_name, index + ".q_proj": (".qkv_proj", 0), + ".k_proj": (".qkv_proj", 1), + ".v_proj": (".qkv_proj", 2), + } + + def __init__( + self, + config: ApertusConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.pp_group = get_pp_group() + self.config = config + self.quant_config = quant_config + self.model = self._init_model(config, quant_config, add_prefix("model", prefix)) + if self.config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + use_attn_tp_group=get_global_server_args().enable_dp_lm_head, + ) + self.logits_processor = LogitsProcessor(config) + self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True) + self.stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + ] + + self.capture_aux_hidden_states = False + + def _init_model( + self, + config: ApertusConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + return ApertusModel(config, quant_config=quant_config, prefix=prefix) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + get_embedding: bool = False, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> LogitsProcessorOutput: + hidden_states = self.model( + input_ids, + positions, + forward_batch, + input_embeds, + pp_proxy_tensors=pp_proxy_tensors, + ) + + aux_hidden_states = None + if self.capture_aux_hidden_states: + hidden_states, aux_hidden_states = hidden_states + + if self.pp_group.is_last_rank: + if not get_embedding: + return self.logits_processor( + input_ids, + hidden_states, + self.lm_head, + forward_batch, + aux_hidden_states, + ) + else: + return self.pooler(hidden_states, forward_batch) + else: + return hidden_states + + @torch.no_grad() + def forward_split_prefill( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + split_interval: Tuple[int, int], # [start, end) 0-based + input_embeds: torch.Tensor = None, + ) -> Optional[LogitsProcessorOutput]: + start, end = split_interval + # embed + if start == 0: + if input_embeds is None: + forward_batch.hidden_states = self.model.embed_tokens(input_ids) + else: + forward_batch.hidden_states = input_embeds + # decoder layer + for i in range(start, end): + layer = self.model.layers[i] + forward_batch.hidden_states, forward_batch.residual = layer( + positions, + forward_batch.hidden_states, + forward_batch, + forward_batch.residual, + ) + + if end == self.model.config.num_hidden_layers: + # norm + hidden_states, _ = self.model.norm( + forward_batch.hidden_states, forward_batch.residual + ) + forward_batch.hidden_states = hidden_states + # logits process + result = self.logits_processor( + input_ids, forward_batch.hidden_states, self.lm_head, forward_batch + ) + else: + result = None + + return result + + @property + def start_layer(self): + return self.model.start_layer + + @property + def end_layer(self): + return self.model.end_layer + + def get_input_embeddings(self) -> nn.Embedding: + return self.model.embed_tokens + + def get_module_name_from_weight_name(self, name): + for param_name, weight_name, shard_id, num_shard in self.stacked_params_mapping: + if weight_name in name: + return ( + name.replace(weight_name, param_name)[: -len(".weight")], + num_shard, + ) + return name[: -len(".weight")], 1 + + def get_num_params(self): + params_dict = dict(self.named_parameters()) + return len(params_dict) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + ] + + params_dict = dict(self.named_parameters()) + + for name, buffer in self.named_buffers(): + if name.endswith(".beta") or name.endswith(".eps"): + params_dict[name] = buffer + + for name, loaded_weight in weights: + layer_id = get_layer_id(name) + if ( + layer_id is not None + and hasattr(self.model, "start_layer") + and ( + layer_id < self.model.start_layer + or layer_id >= self.model.end_layer + ) + ): + continue + if "rotary_emb.inv_freq" in name or "projector" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + if name.startswith("model.vision_tower") and name not in params_dict: + continue + if self.config.tie_word_embeddings and "lm_head.weight" in name: + continue + # Handle FP8 kv-scale remapping + if "scale" in name: + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + # Skip loading kv_scale from ckpts towards new design. + if name.endswith(".kv_scale") and name not in params_dict: + continue + if name in params_dict.keys(): + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + else: + logger.warning(f"Parameter {name} not found in params_dict") + + def get_embed_and_head(self): + return self.model.embed_tokens.weight, self.lm_head.weight + + def set_embed_and_head(self, embed, head): + del self.model.embed_tokens.weight + del self.lm_head.weight + self.model.embed_tokens.weight = embed + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + def get_embed(self): + return self.model.embed_tokens.weight + + def set_embed(self, embed): + # NOTE: If draft hidden size != target hidden size, the embed weight cannot be shared for EAGLE3 + if ( + hasattr(self.config, "target_hidden_size") + and self.config.target_hidden_size != self.config.hidden_size + ): + return + del self.model.embed_tokens.weight + self.model.embed_tokens.weight = embed + torch.cuda.empty_cache() + torch.cuda.synchronize() + + def load_kv_cache_scales(self, quantization_param_path: str) -> None: + self.model.load_kv_cache_scales(quantization_param_path) + + def set_eagle3_layers_to_capture(self, layer_ids: Optional[List[int]] = None): + if not self.pp_group.is_last_rank: + return + + if layer_ids is None: + self.capture_aux_hidden_states = True + num_layers = self.config.num_hidden_layers + self.model.layers_to_capture = [2, num_layers // 2, num_layers - 3] + else: + self.capture_aux_hidden_states = True + # we plus 1 here because in sglang, for the ith layer, it takes the output + # of the (i-1)th layer as aux hidden state + self.model.layers_to_capture = [val + 1 for val in layer_ids] + + +EntryClass = [ApertusForCausalLM] diff --git a/python/sglang/srt/models/arcee.py b/python/sglang/srt/models/arcee.py index 9ee50f02c..7934b4418 100644 --- a/python/sglang/srt/models/arcee.py +++ b/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_rank = get_tensor_model_parallel_rank() + tp_size = get_parallel().tp_size + tp_rank = get_parallel().tp_rank for layer_idx, scaling_factor in kv_cache_scales_loader( quantization_param_path, tp_rank, diff --git a/python/sglang/srt/models/baichuan.py b/python/sglang/srt/models/baichuan.py index 456a5efc5..b1c7b7b99 100644 --- a/python/sglang/srt/models/baichuan.py +++ b/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) diff --git a/python/sglang/srt/models/bailing_moe.py b/python/sglang/srt/models/bailing_moe.py index d0d066647..093bccd60 100644 --- a/python/sglang/srt/models/bailing_moe.py +++ b/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() - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() + self.dp_size = get_parallel().attn_dp_size + attn_tp_rank = get_parallel().attn_tp_rank + 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_rank = get_attention_tp_rank() + self.attn_tp_size = get_parallel().attn_tp_size + 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 diff --git a/python/sglang/srt/models/bailing_moe_linear.py b/python/sglang/srt/models/bailing_moe_linear.py index c1bc85503..1e38b2916 100644 --- a/python/sglang/srt/models/bailing_moe_linear.py +++ b/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_rank = get_tensor_model_parallel_rank() + self.tp_size = get_parallel().tp_size + 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_rank = get_attention_tp_rank() + tp_size = get_parallel().attn_tp_size + 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_rank = get_attention_tp_rank() + self.tp_size = get_parallel().attn_tp_size + 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 diff --git a/python/sglang/srt/models/bailing_moe_nextn.py b/python/sglang/srt/models/bailing_moe_nextn.py index 2c392a5ae..648b304be 100644 --- a/python/sglang/srt/models/bailing_moe_nextn.py +++ b/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. diff --git a/python/sglang/srt/models/bert.py b/python/sglang/srt/models/bert.py index 976b69ab8..ed81a26c4 100644 --- a/python/sglang/srt/models/bert.py +++ b/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 diff --git a/python/sglang/srt/models/chatglm.py b/python/sglang/srt/models/chatglm.py index 9cf585a02..9b1b464c0 100644 --- a/python/sglang/srt/models/chatglm.py +++ b/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 diff --git a/python/sglang/srt/models/cohere2_moe.py b/python/sglang/srt/models/cohere2_moe.py index e1d45e741..aa3e39a29 100644 --- a/python/sglang/srt/models/cohere2_moe.py +++ b/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, diff --git a/python/sglang/srt/models/commandr.py b/python/sglang/srt/models/commandr.py index 2306cecc8..4b899c632 100644 --- a/python/sglang/srt/models/commandr.py +++ b/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 diff --git a/python/sglang/srt/models/dbrx.py b/python/sglang/srt/models/dbrx.py index e957081f3..6707fe7c2 100644 --- a/python/sglang/srt/models/dbrx.py +++ b/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 diff --git a/python/sglang/srt/models/deepseek.py b/python/sglang/srt/models/deepseek.py index 5e122a05c..133c0e6f9 100644 --- a/python/sglang/srt/models/deepseek.py +++ b/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.tp_size = get_tensor_model_parallel_world_size() + self.rank = get_parallel().tp_rank + 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 diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index 2dbe37b58..10ad9c3ea 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/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_size = get_attention_cp_size() + self.cp_rank = get_parallel().attn_cp_rank + self.cp_size = get_parallel().attn_cp_size else: self.cp_rank = None self.cp_size = None diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index fa60ae4ff..66f7e3fce 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/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.moe_ep_size = get_moe_expert_parallel_world_size() + self.tp_size = get_parallel().tp_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_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + 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 = ( - 1 if is_a2a_moe else get_tensor_model_parallel_world_size() - ) + tp_size = 1 if is_a2a_moe else get_parallel().tp_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_size = get_attention_cp_size() + self.cp_rank = get_parallel().attn_cp_rank + 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 = ( diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index c8c932d12..787a1e1dd 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/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_size = attn_tp_size = get_attention_tp_size() + self.tp_rank = attn_tp_rank = get_parallel().attn_tp_rank + 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_size = get_attention_cp_size() + self.cp_rank = get_parallel().attn_cp_rank + self.cp_size = get_parallel().attn_cp_size @property def routed_experts_weights_of_layer(self): diff --git a/python/sglang/srt/models/deepseek_v4_nextn.py b/python/sglang/srt/models/deepseek_v4_nextn.py index fd75783fd..a11880b22 100644 --- a/python/sglang/srt/models/deepseek_v4_nextn.py +++ b/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_size = get_attention_cp_size() + self.cp_rank = get_parallel().attn_cp_rank + self.cp_size = get_parallel().attn_cp_size else: self.cp_rank = None self.cp_size = None diff --git a/python/sglang/srt/models/dflash.py b/python/sglang/srt/models/dflash.py index 71b56e311..a8b2aeefd 100644 --- a/python/sglang/srt/models/dflash.py +++ b/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) diff --git a/python/sglang/srt/models/dots_vlm_vit.py b/python/sglang/srt/models/dots_vlm_vit.py index caf6e38b1..84e94a5c5 100644 --- a/python/sglang/srt/models/dots_vlm_vit.py +++ b/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 diff --git a/python/sglang/srt/models/ernie4.py b/python/sglang/srt/models/ernie4.py index 7c292a125..6ed8a15f3 100644 --- a/python/sglang/srt/models/ernie4.py +++ b/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": diff --git a/python/sglang/srt/models/ernie45_moe_vl.py b/python/sglang/srt/models/ernie45_moe_vl.py index 265cca20e..7321a59db 100644 --- a/python/sglang/srt/models/ernie45_moe_vl.py +++ b/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 diff --git a/python/sglang/srt/models/exaone.py b/python/sglang/srt/models/exaone.py index 27ed2a024..441cfafde 100644 --- a/python/sglang/srt/models/exaone.py +++ b/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 diff --git a/python/sglang/srt/models/exaone4.py b/python/sglang/srt/models/exaone4.py index 76d5998ec..e14bf6e2b 100644 --- a/python/sglang/srt/models/exaone4.py +++ b/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_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + 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_rank = get_attention_tp_rank() + self.attn_tp_size = get_parallel().attn_tp_size + self.attn_tp_rank = get_parallel().attn_tp_rank self.self_attn = Exaone4Attention( config=config, diff --git a/python/sglang/srt/models/exaone_moe.py b/python/sglang/srt/models/exaone_moe.py index 0eae269bc..1ee8feeb3 100755 --- a/python/sglang/srt/models/exaone_moe.py +++ b/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.moe_ep_size = get_moe_expert_parallel_world_size() + self.tp_size = get_parallel().tp_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_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + 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_rank = get_attention_tp_rank() + self.attn_tp_size = get_parallel().attn_tp_size + self.attn_tp_rank = get_parallel().attn_tp_rank self.self_attn = ExaoneMoEAttention( config=config, diff --git a/python/sglang/srt/models/exaone_moe_mtp.py b/python/sglang/srt/models/exaone_moe_mtp.py index 05e63dcae..ed7125e7b 100644 --- a/python/sglang/srt/models/exaone_moe_mtp.py +++ b/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() diff --git a/python/sglang/srt/models/falcon_h1.py b/python/sglang/srt/models/falcon_h1.py index 3be39824e..1bc2d8103 100644 --- a/python/sglang/srt/models/falcon_h1.py +++ b/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_size = get_attention_tp_size() - self.tp_size = get_tensor_model_parallel_world_size() + self.attn_tp_rank = get_parallel().attn_tp_rank + self.attn_tp_size = get_parallel().attn_tp_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 diff --git a/python/sglang/srt/models/gemma.py b/python/sglang/srt/models/gemma.py index 6fd0cf2cf..2249e48ab 100644 --- a/python/sglang/srt/models/gemma.py +++ b/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 diff --git a/python/sglang/srt/models/gemma2.py b/python/sglang/srt/models/gemma2.py index 4c0942839..8df231622 100644 --- a/python/sglang/srt/models/gemma2.py +++ b/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 diff --git a/python/sglang/srt/models/gemma3_causal.py b/python/sglang/srt/models/gemma3_causal.py index a2ae549e6..d83d54b3f 100644 --- a/python/sglang/srt/models/gemma3_causal.py +++ b/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] diff --git a/python/sglang/srt/models/gemma3n_causal.py b/python/sglang/srt/models/gemma3n_causal.py index 8351fc77d..1fb27a737 100644 --- a/python/sglang/srt/models/gemma3n_causal.py +++ b/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 diff --git a/python/sglang/srt/models/gemma4_audio.py b/python/sglang/srt/models/gemma4_audio.py index db825165f..621e4bf15 100644 --- a/python/sglang/srt/models/gemma4_audio.py +++ b/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] diff --git a/python/sglang/srt/models/gemma4_causal.py b/python/sglang/srt/models/gemma4_causal.py index 8c89aead6..339257ca1 100644 --- a/python/sglang/srt/models/gemma4_causal.py +++ b/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] diff --git a/python/sglang/srt/models/gemma4_vision.py b/python/sglang/srt/models/gemma4_vision.py index f0c49cbc6..c4c6fb1c2 100644 --- a/python/sglang/srt/models/gemma4_vision.py +++ b/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 diff --git a/python/sglang/srt/models/glm4.py b/python/sglang/srt/models/glm4.py index 016941b4b..f4614a159 100644 --- a/python/sglang/srt/models/glm4.py +++ b/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_rank = get_tensor_model_parallel_rank() + tp_size = get_parallel().tp_size + tp_rank = get_parallel().tp_rank for layer_idx, scaling_factor in kv_cache_scales_loader( quantization_param_path, tp_rank, diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py index ac6bffe8d..5bc6346c5 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/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_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + 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.moe_ep_size = get_moe_expert_parallel_world_size() + self.tp_size = get_parallel().tp_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." diff --git a/python/sglang/srt/models/glm4_moe_lite.py b/python/sglang/srt/models/glm4_moe_lite.py index a36d4ee86..7530b4b5d 100644 --- a/python/sglang/srt/models/glm4_moe_lite.py +++ b/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: diff --git a/python/sglang/srt/models/glm4_moe_lite_nextn.py b/python/sglang/srt/models/glm4_moe_lite_nextn.py index 103ce18ab..7682afadb 100644 --- a/python/sglang/srt/models/glm4_moe_lite_nextn.py +++ b/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 diff --git a/python/sglang/srt/models/glm4_moe_nextn.py b/python/sglang/srt/models/glm4_moe_nextn.py index 149fec102..3eeeecf6a 100644 --- a/python/sglang/srt/models/glm4_moe_nextn.py +++ b/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 diff --git a/python/sglang/srt/models/glm4v.py b/python/sglang/srt/models/glm4v.py index 9bb5a92b2..7fafa200d 100644 --- a/python/sglang/srt/models/glm4v.py +++ b/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 = ( - 1 if use_data_parallel else get_tensor_model_parallel_world_size() - ) - self.tp_rank = 0 if use_data_parallel else get_tensor_model_parallel_rank() + self.tp_size = 1 if use_data_parallel else get_parallel().tp_size + self.tp_rank = 0 if use_data_parallel else get_parallel().tp_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_rank = 0 if use_data_parallel else get_tensor_model_parallel_rank() + tp_size = 1 if use_data_parallel else get_parallel().tp_size + tp_rank = 0 if use_data_parallel else get_parallel().tp_rank self.proj = ReplicatedLinear( self.hidden_size, self.hidden_size, diff --git a/python/sglang/srt/models/glm4v_moe.py b/python/sglang/srt/models/glm4v_moe.py index 2f0074924..b76ee11c1 100644 --- a/python/sglang/srt/models/glm4v_moe.py +++ b/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." diff --git a/python/sglang/srt/models/glm_ocr_nextn.py b/python/sglang/srt/models/glm_ocr_nextn.py index ae771af53..af09d1af7 100644 --- a/python/sglang/srt/models/glm_ocr_nextn.py +++ b/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) diff --git a/python/sglang/srt/models/gpt2.py b/python/sglang/srt/models/gpt2.py index 6dac103e2..613690f2f 100644 --- a/python/sglang/srt/models/gpt2.py +++ b/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 diff --git a/python/sglang/srt/models/gpt_bigcode.py b/python/sglang/srt/models/gpt_bigcode.py index 38f52a40f..47ac77609 100644 --- a/python/sglang/srt/models/gpt_bigcode.py +++ b/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 diff --git a/python/sglang/srt/models/gpt_j.py b/python/sglang/srt/models/gpt_j.py index 6429e58de..72c736724 100644 --- a/python/sglang/srt/models/gpt_j.py +++ b/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 diff --git a/python/sglang/srt/models/gpt_oss.py b/python/sglang/srt/models/gpt_oss.py index 4b0174b75..7d411a7a0 100644 --- a/python/sglang/srt/models/gpt_oss.py +++ b/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_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + 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_rank = get_attention_tp_rank() + self.attn_tp_size = get_parallel().attn_tp_size + 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_size = get_moe_tensor_parallel_world_size() - moe_ep_rank = get_moe_expert_parallel_rank() - moe_ep_size = get_moe_expert_parallel_world_size() + moe_tp_rank = get_parallel().moe_tp_rank + moe_tp_size = get_parallel().moe_tp_size + moe_ep_rank = get_parallel().moe_ep_rank + 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() - tp_size = get_tensor_model_parallel_world_size() + start = get_parallel().attn_tp_rank * param.numel() + 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=}" ) diff --git a/python/sglang/srt/models/granite.py b/python/sglang/srt/models/granite.py index eb9d2aa7f..21fb338cb 100644 --- a/python/sglang/srt/models/granite.py +++ b/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 diff --git a/python/sglang/srt/models/granitemoe.py b/python/sglang/srt/models/granitemoe.py index ffeb13742..822c335dd 100644 --- a/python/sglang/srt/models/granitemoe.py +++ b/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 diff --git a/python/sglang/srt/models/granitemoehybrid.py b/python/sglang/srt/models/granitemoehybrid.py index 85385b4fd..a1c7fa7ec 100644 --- a/python/sglang/srt/models/granitemoehybrid.py +++ b/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", ) diff --git a/python/sglang/srt/models/grok.py b/python/sglang/srt/models/grok.py index a5eb64abc..ebfef7714 100644 --- a/python/sglang/srt/models/grok.py +++ b/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_size = get_tensor_model_parallel_world_size() + attn_tp_rank = get_parallel().tp_rank + 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 ( - sum(p.numel() for p in self.parameters()) - * get_tensor_model_parallel_world_size() - ) + return sum(p.numel() for p in self.parameters()) * get_parallel().tp_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"] diff --git a/python/sglang/srt/models/hunyuan.py b/python/sglang/srt/models/hunyuan.py index 9c01e5307..e5c10469e 100644 --- a/python/sglang/srt/models/hunyuan.py +++ b/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_rank = get_tensor_model_parallel_rank() + tp_size = get_parallel().tp_size + tp_rank = get_parallel().tp_rank for layer_idx, scaling_factor in kv_cache_scales_loader( quantization_param_path, tp_rank, diff --git a/python/sglang/srt/models/hunyuan_v3.py b/python/sglang/srt/models/hunyuan_v3.py index f15a18cff..a10e1e9b5 100644 --- a/python/sglang/srt/models/hunyuan_v3.py +++ b/python/sglang/srt/models/hunyuan_v3.py @@ -19,9 +19,6 @@ from torch import nn from transformers import PretrainedConfig from sglang.srt.distributed import ( - get_moe_expert_parallel_world_size, - get_moe_tensor_parallel_world_size, - get_tensor_model_parallel_world_size, moe_expert_parallel_all_reduce, moe_tensor_model_parallel_all_reduce, ) @@ -47,6 +44,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ( from sglang.srt.managers.schedule_batch 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 is_cuda from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -100,8 +98,8 @@ class HYV3MoEFused(nn.Module): alt_stream: Optional[torch.cuda.Stream] = None, ): super().__init__() - self.tp_size = get_moe_tensor_parallel_world_size() - self.ep_size = get_moe_expert_parallel_world_size() + self.tp_size = get_parallel().moe_tp_size + self.ep_size = get_parallel().moe_ep_size self.layer_id = layer_id self.alt_stream = alt_stream self.n_routed_experts = config.num_experts @@ -258,7 +256,7 @@ class HYV3Attention(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 diff --git a/python/sglang/srt/models/internlm2.py b/python/sglang/srt/models/internlm2.py index 6a276ef0a..ee9807ac6 100644 --- a/python/sglang/srt/models/internlm2.py +++ b/python/sglang/srt/models/internlm2.py @@ -23,7 +23,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 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 @@ -97,7 +97,7 @@ class InternLM2Attention(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 diff --git a/python/sglang/srt/models/interns1pro.py b/python/sglang/srt/models/interns1pro.py index b22ff20ec..df3b4d03b 100644 --- a/python/sglang/srt/models/interns1pro.py +++ b/python/sglang/srt/models/interns1pro.py @@ -5,7 +5,6 @@ from typing import Any, Dict, Iterable, Optional, Tuple import torch from transformers import PretrainedConfig -from sglang.srt.layers.dp_attention import get_attention_tp_rank, get_attention_tp_size from sglang.srt.layers.moe.topk import TopK from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.rotary_embedding import get_rope @@ -16,6 +15,7 @@ from sglang.srt.models.qwen3_vl_moe import ( Qwen3MoeLLMModel, Qwen3VLMoeForConditionalGeneration, ) +from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import add_prefix logger = logging.getLogger(__name__) @@ -214,8 +214,8 @@ class InternS1ProForConditionalGeneration(Qwen3VLMoeForConditionalGeneration): def _load_fope_weights(self, name: str, loaded_weight: torch.Tensor, params_dict): """load fope weights""" - attn_tp_size = get_attention_tp_size() - attn_tp_rank = get_attention_tp_rank() + attn_tp_size = get_parallel().attn_tp_size + attn_tp_rank = get_parallel().attn_tp_rank num_key_value_heads = loaded_weight.size(0) # replicate head if necessary diff --git a/python/sglang/srt/models/internvl.py b/python/sglang/srt/models/internvl.py index 9a804ffe4..a59fc1164 100644 --- a/python/sglang/srt/models/internvl.py +++ b/python/sglang/srt/models/internvl.py @@ -11,10 +11,6 @@ from torch import nn from transformers import PretrainedConfig, PreTrainedModel from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling -from sglang.srt.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) from sglang.srt.environ import envs from sglang.srt.layers.activation import get_act_fn from sglang.srt.layers.attention import vision_utils @@ -44,6 +40,7 @@ from sglang.srt.multimodal.internvl_vit_cuda_graph_runner import ( InternViTCudaGraphRunner, ) from sglang.srt.multimodal.mm_utils import run_dp_sharded_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 is_cuda from sglang.utils import logger @@ -191,10 +188,8 @@ class InternMLP(nn.Module): use_data_parallel: bool = False, ): super().__init__() - self.tp_size = ( - 1 if use_data_parallel else get_tensor_model_parallel_world_size() - ) - self.tp_rank = 0 if use_data_parallel else get_tensor_model_parallel_rank() + self.tp_size = 1 if use_data_parallel else get_parallel().tp_size + self.tp_rank = 0 if use_data_parallel else get_parallel().tp_rank self.config = config self.act = get_act_fn(config.hidden_act) self.fc1 = ColumnParallelLinear( diff --git a/python/sglang/srt/models/iquest_loopcoder.py b/python/sglang/srt/models/iquest_loopcoder.py index 286d4fb19..62488a432 100644 --- a/python/sglang/srt/models/iquest_loopcoder.py +++ b/python/sglang/srt/models/iquest_loopcoder.py @@ -20,7 +20,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.layernorm import RMSNorm from sglang.srt.layers.linear import ( ColumnParallelLinear, @@ -38,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_loader.weight_utils import default_weight_loader from sglang.srt.models.llama import LlamaMLP as LoopCoderMLP +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 @@ -64,7 +64,7 @@ class LoopGateProjection(nn.Module): super().__init__() self.total_num_heads = total_num_heads self.head_dim = head_dim - 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 @@ -131,7 +131,7 @@ class LoopCoderAttention(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 diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index 139f311cc..004ba0c49 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/python/sglang/srt/models/kimi_linear.py @@ -12,12 +12,10 @@ from sglang.srt.configs.kimi_linear import KimiLinearConfig from sglang.srt.distributed import ( divide, 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 from sglang.srt.layers.attention.fla.fused_norm_gate import FusedRMSNormGated -from sglang.srt.layers.dp_attention import get_attention_tp_rank, get_attention_tp_size from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( ColumnParallelBatchedLinear, @@ -49,6 +47,7 @@ from sglang.srt.model_loader.weight_utils import ( from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA as KimiMLAAttention from sglang.srt.models.llama import LlamaMLP as KimiMLP from sglang.srt.models.transformers import maybe_prefix +from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import make_layers from sglang.srt.utils.common import BumpAllocator, add_prefix, set_weight_attrs @@ -68,7 +67,7 @@ class KimiMoE(nn.Module): moe_intermediate_size = config.moe_intermediate_size num_experts = config.num_experts moe_renormalize = config.moe_renormalize - 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.num_shared_experts = config.num_shared_experts self.layer_idx = layer_idx @@ -177,8 +176,8 @@ class KimiDeltaAttention(nn.Module): **kwargs, ) -> None: super().__init__() - self.tp_size = get_tensor_model_parallel_world_size() - self.attn_tp_size = get_attention_tp_size() + self.tp_size = get_parallel().tp_size + self.attn_tp_size = get_parallel().attn_tp_size self.hidden_size = hidden_size self.config = config self.head_dim = config.linear_attn_config["head_dim"] @@ -225,7 +224,7 @@ class KimiDeltaAttention(nn.Module): ) else: # Unfused path: separate QKVParallelLinear - attn_tp_rank = get_attention_tp_rank() + attn_tp_rank = get_parallel().attn_tp_rank self.qkv_proj = QKVParallelLinear( self.hidden_size, self.head_dim, @@ -549,7 +548,7 @@ class KimiLinearModel(nn.Module): else: self.norm = PPMissingLayer() - world_size = get_tensor_model_parallel_world_size() + world_size = get_parallel().tp_size assert ( config.num_attention_heads % world_size == 0 ), "num_attention_heads must be divisible by world_size" diff --git a/python/sglang/srt/models/laguna.py b/python/sglang/srt/models/laguna.py index cc8c65894..57bded6ba 100644 --- a/python/sglang/srt/models/laguna.py +++ b/python/sglang/srt/models/laguna.py @@ -19,7 +19,6 @@ from torch import nn from sglang.srt.configs.laguna import LagunaConfig from sglang.srt.distributed import ( get_pp_group, - get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) from sglang.srt.layers.activation import SiluAndMul @@ -28,8 +27,6 @@ from sglang.srt.layers.communicator import ( 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 @@ -55,6 +52,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ( from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors 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.server_args import get_global_server_args from sglang.srt.utils import LazyValue, add_prefix, make_layers @@ -142,7 +140,7 @@ class LagunaMoE(nn.Module): prefix: str = "", ): super().__init__() - self.tp_size = get_tensor_model_parallel_world_size() + self.tp_size = get_parallel().tp_size self.routed_scaling_factor = config.moe_routed_scaling_factor self.router_logit_softcapping = getattr( config, "moe_router_logit_softcapping", 0.0 @@ -257,8 +255,8 @@ class LagunaAttention(nn.Module): self.gating = bool(gating) self.gate_per_head = gating is True or gating == "per-head" - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size self.total_num_heads = num_heads assert self.total_num_heads % attn_tp_size == 0 diff --git a/python/sglang/srt/models/lfm2.py b/python/sglang/srt/models/lfm2.py index fb2fde239..3a93ddf1d 100644 --- a/python/sglang/srt/models/lfm2.py +++ b/python/sglang/srt/models/lfm2.py @@ -19,7 +19,7 @@ import torch.nn.functional as F from torch import nn from sglang.srt.configs.lfm2 import Lfm2Config -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.mamba.causal_conv1d import ( causal_conv1d_fn, causal_conv1d_update, @@ -45,6 +45,7 @@ from sglang.srt.model_loader.weight_utils import ( default_weight_loader, sharded_weight_loader, ) +from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import add_prefix, make_layers, set_weight_attrs logger = logging.getLogger(__name__) @@ -223,7 +224,7 @@ class Lfm2ShortConv(nn.Module): self.use_bias = bool(config.conv_bias) self.hidden_size = config.hidden_size - tp_size = get_tensor_model_parallel_world_size() + tp_size = get_parallel().tp_size self.hidden_size_per_partition = self.hidden_size // tp_size # Use MergedColumnParallelLinear so each output (B, C, x) is sharded separately diff --git a/python/sglang/srt/models/lfm2_moe.py b/python/sglang/srt/models/lfm2_moe.py index c37167faa..b3c033a89 100644 --- a/python/sglang/srt/models/lfm2_moe.py +++ b/python/sglang/srt/models/lfm2_moe.py @@ -18,7 +18,7 @@ import torch from torch import nn from sglang.srt.configs.lfm2_moe import Lfm2MoeConfig -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.mamba.causal_conv1d import ( causal_conv1d_fn, @@ -47,6 +47,7 @@ from sglang.srt.model_loader.weight_utils import ( default_weight_loader, sharded_weight_loader, ) +from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import add_prefix, make_layers, set_weight_attrs @@ -103,7 +104,7 @@ class Lfm2MoeSparseMoeBlock(nn.Module): prefix: str = "", ): 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 if self.tp_size > config.num_experts: @@ -286,7 +287,7 @@ class Lfm2MoeShortConv(nn.Module): self.hidden_size = config.hidden_size # Get tensor parallel size for sharding - self.tp_size = get_tensor_model_parallel_world_size() + self.tp_size = get_parallel().tp_size self.hidden_size_per_partition = self.hidden_size // self.tp_size # Use MergedColumnParallelLinear so each output (B, C, x) is sharded separately diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index 17a47641a..fafde609d 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.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 @@ -80,6 +76,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, @@ -194,7 +191,7 @@ class LLaDA2MoeSparseMoeBlock(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 @@ -301,7 +298,7 @@ class LLaDA2MoeSparseMoeBlock(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, @@ -435,9 +432,9 @@ class LLaDA2MoeAttention(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() - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() + self.dp_size = get_parallel().attn_dp_size + attn_tp_rank = get_parallel().attn_tp_rank + 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: @@ -576,7 +573,7 @@ class LLaDA2MoeBlock(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 = LLaDA2MoeAttention( config, layer_id, @@ -586,8 +583,8 @@ class LLaDA2MoeBlock(nn.Module): alt_stream=alt_stream, ) self.layer_id = layer_id - self.attn_tp_size = get_attention_tp_size() - self.attn_tp_rank = get_attention_tp_rank() + self.attn_tp_size = get_parallel().attn_tp_size + self.attn_tp_rank = get_parallel().attn_tp_rank self.is_layer_sparse = self._is_layer_sparse(config, layer_id=layer_id) is_previous_layer_sparse = self._is_layer_sparse(config, layer_id=layer_id - 1) diff --git a/python/sglang/srt/models/llama.py b/python/sglang/srt/models/llama.py index 742f810ac..d88789099 100644 --- a/python/sglang/srt/models/llama.py +++ b/python/sglang/srt/models/llama.py @@ -28,8 +28,6 @@ from transformers import LlamaConfig from sglang.srt.distributed import ( get_pp_group, get_pp_indices, - 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 @@ -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, is_cuda, is_npu, is_xpu, make_layers from sglang.utils import get_exception_traceback @@ -144,7 +143,7 @@ class LlamaAttention(nn.Module): super().__init__() self.hidden_size = hidden_size self.start_layer = start_layer - 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 @@ -435,8 +434,8 @@ class LlamaModel(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_rank = get_tensor_model_parallel_rank() + tp_size = get_parallel().tp_size + tp_rank = get_parallel().tp_rank for layer_idx, scaling_factor in kv_cache_scales_loader( quantization_param_path, tp_rank, diff --git a/python/sglang/srt/models/llama4.py b/python/sglang/srt/models/llama4.py index 4b6ff13e4..bf322ef4e 100644 --- a/python/sglang/srt/models/llama4.py +++ b/python/sglang/srt/models/llama4.py @@ -26,13 +26,10 @@ from torch import nn from transformers import Llama4TextConfig from sglang.srt.distributed import ( - get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) 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 @@ -55,6 +52,7 @@ from sglang.srt.model_executor.forward_batch_info import ( ) from sglang.srt.models.llama import LlamaForCausalLM, LlamaMLP from sglang.srt.models.utils import apply_qk_norm +from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import ( add_prefix, fast_topk, @@ -98,7 +96,7 @@ class Llama4MoE(nn.Module): prefix: str = "", ): super().__init__() - 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.device_module = torch.get_device_module() @@ -221,8 +219,8 @@ class Llama4Attention(nn.Module): self.use_rope = (layer_id + 1) % 4 != 0 self.use_qk_norm = config.use_qk_norm and self.use_rope - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size self.total_num_heads = num_heads assert self.total_num_heads % attn_tp_size == 0 @@ -390,8 +388,8 @@ class Llama4DecoderLayer(nn.Module): rope_theta = config.rope_parameters["rope_theta"] rope_scaling = config.rope_parameters max_position_embeddings = config.max_position_embeddings - self.attn_tp_size = get_attention_tp_size() - self.attn_tp_rank = get_attention_tp_rank() + self.attn_tp_size = get_parallel().attn_tp_size + self.attn_tp_rank = get_parallel().attn_tp_rank self.self_attn = Llama4Attention( config=config, diff --git a/python/sglang/srt/models/longcat_flash.py b/python/sglang/srt/models/longcat_flash.py index 6536c46f0..29d3258c4 100644 --- a/python/sglang/srt/models/longcat_flash.py +++ b/python/sglang/srt/models/longcat_flash.py @@ -39,7 +39,6 @@ from torch import nn from sglang.srt.configs import LongcatFlashConfig from sglang.srt.distributed import ( - get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder @@ -48,8 +47,6 @@ from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.activation import SiluAndMul 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 @@ -89,6 +86,7 @@ from sglang.srt.model_loader.utils import ( ) from sglang.srt.model_loader.weight_utils import default_weight_loader 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, @@ -221,7 +219,7 @@ class LongcatFlashMoE(nn.Module): else: self.rounter_params_dtype = torch.bfloat16 - self.tp_size = get_tensor_model_parallel_world_size() + self.tp_size = get_parallel().tp_size if self.tp_size > config.n_routed_experts: raise ValueError( @@ -377,8 +375,8 @@ class LongcatFlashDecoderLayer(nn.Module): prefix=add_prefix("mlp", prefix), ) - self.attn_tp_size = get_attention_tp_size() - self.attn_tp_rank = get_attention_tp_rank() + self.attn_tp_size = get_parallel().attn_tp_size + self.attn_tp_rank = get_parallel().attn_tp_rank self.mlp_layer_scatter_modes = [ LayerScatterModes.init_new( @@ -607,7 +605,7 @@ class LongcatFlashForCausalLM(nn.Module): ] 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 = LongcatFlashModel( config, quant_config, prefix=add_prefix("model", prefix) diff --git a/python/sglang/srt/models/longcat_flash_nextn.py b/python/sglang/srt/models/longcat_flash_nextn.py index c5a630cf3..91fc84398 100644 --- a/python/sglang/srt/models/longcat_flash_nextn.py +++ b/python/sglang/srt/models/longcat_flash_nextn.py @@ -42,8 +42,6 @@ from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_r from sglang.srt.layers import deep_gemm_wrapper 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 @@ -70,6 +68,7 @@ from sglang.srt.model_loader.utils import should_deepgemm_weight_requant_ue8m0 from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA from sglang.srt.models.longcat_flash import LongcatFlashForCausalLM, LongcatFlashMLP +from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import ( BumpAllocator, add_prefix, @@ -154,8 +153,8 @@ class LongcatFlashDenseDecoderLayer(nn.Module): config.hidden_size, eps=config.rms_norm_eps ) - self.attn_tp_size = get_attention_tp_size() - self.attn_tp_rank = get_attention_tp_rank() + self.attn_tp_size = get_parallel().attn_tp_size + self.attn_tp_rank = get_parallel().attn_tp_rank self.layer_scatter_modes = LayerScatterModes.init_new( layer_id=self.layer_id, num_layers=config.num_hidden_layers, diff --git a/python/sglang/srt/models/mimo_mtp.py b/python/sglang/srt/models/mimo_mtp.py index 046fbd3c7..da6f2558d 100644 --- a/python/sglang/srt/models/mimo_mtp.py +++ b/python/sglang/srt/models/mimo_mtp.py @@ -8,7 +8,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.layernorm import RMSNorm from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.quantization.base_config import QuantizationConfig @@ -19,6 +18,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.models.qwen2 import Qwen2DecoderLayer +from sglang.srt.runtime_context import get_parallel class MiMoMultiTokenPredictorLayer(nn.Module): @@ -90,7 +90,7 @@ class MiMoMTP(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 self.model = MiMoMultiTokenPredictorLayer( diff --git a/python/sglang/srt/models/mimo_v2.py b/python/sglang/srt/models/mimo_v2.py index 33e35fe17..aa408bd90 100644 --- a/python/sglang/srt/models/mimo_v2.py +++ b/python/sglang/srt/models/mimo_v2.py @@ -22,9 +22,7 @@ from torch import nn from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo from sglang.srt.configs.model_config import get_mimo_v2_fused_qkv_expected_tp_size 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 @@ -38,8 +36,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_dp_attention_enabled, ) from sglang.srt.layers.layernorm import RMSNorm @@ -80,6 +76,7 @@ from sglang.srt.model_loader.weight_utils import ( ) from sglang.srt.models.mimo_audio import AudioEncoderMixin, MiMoAudioEncoderConfig from sglang.srt.models.mimo_vl import MiMoVisionTransformer, MiMoVLVisionConfig +from sglang.srt.runtime_context import get_parallel from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import ( LazyValue, @@ -107,8 +104,8 @@ def load_mimo_v2_qkv_proj_weight( f"expected sharded {tuple(param.shape)}" ) - tp_size = get_attention_tp_size() - tp_rank = get_attention_tp_rank() + tp_size = get_parallel().attn_tp_size + tp_rank = get_parallel().attn_tp_rank if expected_fused_tp_size is not None and tp_size != expected_fused_tp_size: raise ValueError( f"MiMoV2 fused qkv_proj checkpoint is TP={expected_fused_tp_size}-" @@ -229,7 +226,7 @@ class MiMoV2MoE(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.config = config self.layer_id = layer_id @@ -288,7 +285,7 @@ class MiMoV2MoE(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 @@ -469,8 +466,8 @@ class MiMoV2Attention(nn.Module): super().__init__() self.hidden_size = hidden_size - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size self.total_num_heads = num_heads assert self.total_num_heads % attn_tp_size == 0 @@ -973,8 +970,8 @@ class MiMoV2Model(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: - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size for layer_idx, scaling_factor in kv_cache_scales_loader( quantization_param_path, attn_tp_rank, @@ -1434,7 +1431,7 @@ class MiMoV2ForCausalLM(nn.Module, AudioEncoderMixin): if name in params_dict.keys(): param = params_dict[name] if "attention_sink_bias" in name: - start = get_attention_tp_rank() * param.numel() + start = get_parallel().attn_tp_rank * param.numel() param.data.copy_( loaded_weight[start : start + param.numel()] ) diff --git a/python/sglang/srt/models/mimo_v2_nextn.py b/python/sglang/srt/models/mimo_v2_nextn.py index d55a4a0f1..efeaddd2f 100644 --- a/python/sglang/srt/models/mimo_v2_nextn.py +++ b/python/sglang/srt/models/mimo_v2_nextn.py @@ -20,7 +20,6 @@ from torch import nn from transformers import PretrainedConfig from sglang.srt.configs.model_config import get_mimo_v2_fused_qkv_expected_tp_size -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.communicator import ( LayerCommunicator, @@ -28,7 +27,6 @@ from sglang.srt.layers.communicator import ( enable_moe_dense_fully_dp, ) from sglang.srt.layers.dp_attention import ( - get_attention_tp_rank, is_dp_attention_enabled, ) from sglang.srt.layers.layernorm import RMSNorm @@ -46,6 +44,7 @@ from sglang.srt.models.mimo_v2 import ( MiMoV2MLP, load_mimo_v2_qkv_proj_weight, ) +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 @@ -250,7 +249,7 @@ class MiMoV2MTP(MiMoV2ForCausalLM): ) -> 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 = MiMoV2ModelNextN( @@ -352,7 +351,7 @@ class MiMoV2MTP(MiMoV2ForCausalLM): if name in params_dict.keys(): param = params_dict[name] if "attention_sink_bias" in name: - start = get_attention_tp_rank() * param.numel() + start = get_parallel().attn_tp_rank * param.numel() param.data.copy_(loaded_weight[start : start + param.numel()]) else: weight_loader = getattr( diff --git a/python/sglang/srt/models/mindspore.py b/python/sglang/srt/models/mindspore.py index c560be72a..25056adf9 100644 --- a/python/sglang/srt/models/mindspore.py +++ b/python/sglang/srt/models/mindspore.py @@ -7,10 +7,6 @@ from typing import Any, Iterable, List, Optional, Tuple import torch -from sglang.srt.distributed import ( - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.model_executor.forward_batch_info import ForwardBatch @@ -19,6 +15,7 @@ from sglang.srt.model_executor.forward_context import ( get_token_to_kv_pool, ) from sglang.srt.models.registry import import_model_classes +from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import is_npu _is_npu = is_npu() @@ -185,10 +182,10 @@ class MindSporeForCausalLM(torch.nn.Module): logger.info( "MindSporeForCausalLM tp size %d tp rank %d", - get_tensor_model_parallel_world_size(), - get_tensor_model_parallel_rank(), + get_parallel().tp_size, + get_parallel().tp_rank, ) - if get_tensor_model_parallel_world_size() not in (1, 2, 4, 8): + if get_parallel().tp_size not in (1, 2, 4, 8): # MatMulAllReduce only support tp size in (1, 2, 4, 8) ms.set_context(graph_kernel_flags="--disable_pass=MatMulAllReduce") diff --git a/python/sglang/srt/models/minicpm.py b/python/sglang/srt/models/minicpm.py index 06ee8445c..63efbf1f2 100644 --- a/python/sglang/srt/models/minicpm.py +++ b/python/sglang/srt/models/minicpm.py @@ -19,7 +19,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 ( @@ -37,6 +36,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 @@ -94,7 +94,7 @@ class MiniCPMAttention(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 diff --git a/python/sglang/srt/models/minicpm3.py b/python/sglang/srt/models/minicpm3.py index ea24d6e65..fbcfab5fd 100644 --- a/python/sglang/srt/models/minicpm3.py +++ b/python/sglang/srt/models/minicpm3.py @@ -20,7 +20,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 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, is_cuda from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -147,7 +147,7 @@ class MiniCPM3AttentionMLA(nn.Module): self.q_lora_rank = q_lora_rank self.kv_lora_rank = kv_lora_rank self.num_heads = num_heads - tp_size = get_tensor_model_parallel_world_size() + tp_size = get_parallel().tp_size assert num_heads % tp_size == 0 self.num_local_heads = num_heads // tp_size self.scaling = self.qk_head_dim**-0.5 diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index a0a4f1e17..c7dd1f4b0 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -33,9 +33,7 @@ from sglang.jit_kernel.all_reduce import ( from sglang.kernel_api_logging import debug_kernel_api 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, tensor_model_parallel_all_reduce, ) from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder @@ -48,8 +46,6 @@ from sglang.srt.layers.communicator import ( 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.layernorm import RMSNorm @@ -85,6 +81,7 @@ from sglang.srt.model_loader.weight_utils import ( maybe_remap_kv_scale_name, narrow_padded_param_and_loaded_weight, ) +from sglang.srt.runtime_context import get_parallel from sglang.srt.server_args import get_global_server_args # get_bool_env_var is defined in sglang.srt.utils.common, not sglang.srt.distributed. @@ -289,8 +286,8 @@ class MiniMaxM2RMSNormTP(nn.Module): def __init__(self, hidden_size: int, num_heads: int, eps: float = 1e-6) -> None: super().__init__() - self.attn_tp_size = get_attention_tp_size() - self.attn_tp_rank = get_attention_tp_rank() + self.attn_tp_size = get_parallel().attn_tp_size + self.attn_tp_rank = get_parallel().attn_tp_rank # Align with QKVParallelLinear pattern if self.attn_tp_size >= num_heads: @@ -499,7 +496,7 @@ class MiniMaxM2MoE(nn.Module): prefix: str = "", ): super().__init__() - self.tp_size = get_tensor_model_parallel_world_size() + self.tp_size = get_parallel().tp_size if self.tp_size > config.num_local_experts: raise ValueError( f"Tensor parallel size {self.tp_size} is greater than " @@ -546,7 +543,7 @@ class MiniMaxM2MoE(nn.Module): self.layer_id = layer_id 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.top_k = config.num_experts_per_tok @staticmethod @@ -734,8 +731,8 @@ class MiniMaxM2Attention(nn.Module): self.hidden_size = config.hidden_size # Use attention TP rank/size for dp-attention support - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size # Get dimensions from config self.total_num_heads = config.num_attention_heads diff --git a/python/sglang/srt/models/mixtral.py b/python/sglang/srt/models/mixtral.py index db258623f..e797b7daf 100644 --- a/python/sglang/srt/models/mixtral.py +++ b/python/sglang/srt/models/mixtral.py @@ -27,7 +27,6 @@ from transformers import MixtralConfig from sglang.srt.distributed import ( get_pp_group, - get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) from sglang.srt.layers.layernorm import RMSNorm @@ -49,6 +48,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ( ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors 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, make_layers logger = logging.getLogger(__name__) @@ -76,7 +76,7 @@ class MixtralMoE(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 = hidden_size # Gate always runs at half / full precision for now. @@ -132,7 +132,7 @@ class MixtralAttention(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 diff --git a/python/sglang/srt/models/mixtral_quant.py b/python/sglang/srt/models/mixtral_quant.py index 61cff2e48..0535d0353 100644 --- a/python/sglang/srt/models/mixtral_quant.py +++ b/python/sglang/srt/models/mixtral_quant.py @@ -27,8 +27,6 @@ from torch import nn from transformers import MixtralConfig 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.layernorm import RMSNorm @@ -47,6 +45,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 @@ -107,8 +106,8 @@ class MixtralMoE(nn.Module): ): super().__init__() self.config = config - self.rank = get_tensor_model_parallel_rank() - self.tp_size = get_tensor_model_parallel_world_size() + self.rank = get_parallel().tp_rank + self.tp_size = get_parallel().tp_size self.num_total_experts = config.num_local_experts self.top_k = config.num_experts_per_tok if self.tp_size > self.num_total_experts: @@ -185,7 +184,7 @@ class MixtralAttention(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 diff --git a/python/sglang/srt/models/mllama.py b/python/sglang/srt/models/mllama.py index 9fca0b414..7b3d39a8d 100644 --- a/python/sglang/srt/models/mllama.py +++ b/python/sglang/srt/models/mllama.py @@ -21,7 +21,6 @@ from transformers.models.mllama.modeling_mllama import ( ) import sglang.srt.distributed.parallel_state as ps -from sglang.srt.distributed import get_tensor_model_parallel_world_size from sglang.srt.layers.activation import get_act_fn from sglang.srt.layers.attention.vision import VisionAttention from sglang.srt.layers.layernorm import RMSNorm @@ -43,6 +42,7 @@ from sglang.srt.managers.schedule_batch import MultimodalInputs 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.llama import LlamaDecoderLayer, LlamaMLP +from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import add_prefix @@ -491,7 +491,7 @@ class MllamaTextCrossAttention(nn.Module): ): super().__init__() self.config = config - self.model_parallel_size = get_tensor_model_parallel_world_size() + self.model_parallel_size = get_parallel().tp_size self.num_heads = self.config.num_attention_heads self.num_local_heads = self.num_heads // self.model_parallel_size self.num_key_value_heads = self.config.num_key_value_heads diff --git a/python/sglang/srt/models/moss_vl.py b/python/sglang/srt/models/moss_vl.py index fcab2fe10..8ea3be829 100644 --- a/python/sglang/srt/models/moss_vl.py +++ b/python/sglang/srt/models/moss_vl.py @@ -15,14 +15,10 @@ from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( Qwen2_5_VisionRotaryEmbedding, ) -from sglang.srt.distributed import ( - get_tensor_model_parallel_world_size, -) from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.attention.vision import VisionAttention from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes from sglang.srt.layers.conv import Conv3dLayer -from sglang.srt.layers.dp_attention import get_attention_tp_rank, get_attention_tp_size from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( ColumnParallelLinear, @@ -47,6 +43,7 @@ from sglang.srt.managers.schedule_batch import MultimodalInputs 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.server_args import get_global_server_args from sglang.srt.utils import add_prefix @@ -481,7 +478,7 @@ class MossVLTextCrossAttention(nn.Module): ): super().__init__() self.config = config - self.model_parallel_size = get_tensor_model_parallel_world_size() + self.model_parallel_size = get_parallel().tp_size self.num_heads = config.num_attention_heads self.num_local_heads = self.num_heads // self.model_parallel_size self.num_key_value_heads = config.num_key_value_heads @@ -753,10 +750,10 @@ class MossVLSelfAttention(nn.Module): ): super().__init__() self.hidden_size = config.hidden_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 - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size assert self.total_num_heads % attn_tp_size == 0 self.num_heads = self.total_num_heads // attn_tp_size diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index 841f3d14a..20b68e18d 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/python/sglang/srt/models/nemotron_h.py @@ -29,7 +29,6 @@ from sglang.srt.configs.nemotron_h import ATTENTION, MAMBA, MLP, MOE from sglang.srt.distributed import ( get_moe_ep_group, get_pp_group, - get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) from sglang.srt.layers.activation import ReLU2 @@ -40,8 +39,6 @@ from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( from sglang.srt.layers.attention.mamba.mamba import MambaMixer2 from sglang.srt.layers.dp_attention import ( attn_tp_all_reduce, - get_attention_tp_rank, - get_attention_tp_size, is_dp_attention_enabled, ) from sglang.srt.layers.layernorm import RMSNorm @@ -90,6 +87,7 @@ from sglang.srt.models.nemotron_h_utils import ( pad_to_original_num_tokens, ) 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 ( add_prefix, @@ -159,7 +157,7 @@ class NemotronHMoE(nn.Module): ) -> None: 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.device_module = torch.get_device_module() @@ -538,8 +536,8 @@ class NemotronHAttention(nn.Module): ) -> None: super().__init__() self.hidden_size = config.hidden_size - tp_rank = get_attention_tp_rank() - tp_size = get_attention_tp_size() + tp_rank = get_parallel().attn_tp_rank + 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 diff --git a/python/sglang/srt/models/nemotron_h_mtp.py b/python/sglang/srt/models/nemotron_h_mtp.py index 9120201e0..257e5cad3 100644 --- a/python/sglang/srt/models/nemotron_h_mtp.py +++ b/python/sglang/srt/models/nemotron_h_mtp.py @@ -22,8 +22,6 @@ from sglang.srt.distributed import get_pp_group 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.layernorm import RMSNorm @@ -41,6 +39,7 @@ from sglang.srt.models.nemotron_h import ( NemotronHMoEDecoderLayer, ) from sglang.srt.models.nemotron_h_utils import is_attn_layer +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 @@ -74,8 +73,8 @@ class NemotronHMTPAttentionDecoderLayer(NemotronHAttentionDecoderLayer): output_size=config.hidden_size, bias=False, gather_output=not _dp_attn, - tp_rank=get_attention_tp_rank() if _dp_attn else None, - tp_size=get_attention_tp_size() if _dp_attn else None, + tp_rank=get_parallel().attn_tp_rank if _dp_attn else None, + tp_size=get_parallel().attn_tp_size if _dp_attn else None, params_dtype=( config.dtype if hasattr(config, "dtype") else torch.bfloat16 ), @@ -159,8 +158,8 @@ class NemotronHMTPMoEDecoderLayer(NemotronHMoEDecoderLayer): output_size=config.hidden_size, bias=False, gather_output=not _dp_attn, - tp_rank=get_attention_tp_rank() if _dp_attn else None, - tp_size=get_attention_tp_size() if _dp_attn else None, + tp_rank=get_parallel().attn_tp_rank if _dp_attn else None, + tp_size=get_parallel().attn_tp_size if _dp_attn else None, params_dtype=( config.dtype if hasattr(config, "dtype") else torch.bfloat16 ), diff --git a/python/sglang/srt/models/olmo.py b/python/sglang/srt/models/olmo.py index 5f51e229a..15106fd76 100644 --- a/python/sglang/srt/models/olmo.py +++ b/python/sglang/srt/models/olmo.py @@ -25,7 +25,6 @@ import torch from torch import nn from transformers import OlmoConfig -from sglang.srt.distributed import get_tensor_model_parallel_world_size from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.linear import ( MergedColumnParallelLinear, @@ -42,6 +41,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, make_layers @@ -62,7 +62,7 @@ class OlmoAttention(nn.Module): super().__init__() self.config = config self.hidden_size = config.hidden_size - tensor_model_parallel_world_size = get_tensor_model_parallel_world_size() + tensor_model_parallel_world_size = get_parallel().tp_size self.total_num_heads = config.num_attention_heads assert self.hidden_size % self.total_num_heads == 0 diff --git a/python/sglang/srt/models/olmo2.py b/python/sglang/srt/models/olmo2.py index 91dfd12cd..761a08330 100644 --- a/python/sglang/srt/models/olmo2.py +++ b/python/sglang/srt/models/olmo2.py @@ -26,8 +26,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, split_tensor_along_last_dim, tensor_model_parallel_all_gather, ) @@ -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_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, is_cuda, make_layers _is_cuda = is_cuda() @@ -78,7 +77,7 @@ class Olmo2Attention(nn.Module): super().__init__() self.config = config self.hidden_size = config.hidden_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.hidden_size % self.total_num_heads == 0 @@ -113,7 +112,7 @@ class Olmo2Attention(nn.Module): quant_config=quant_config, prefix=add_prefix("qkv_proj", prefix), ) - self.tp_rank = get_tensor_model_parallel_rank() + self.tp_rank = get_parallel().tp_rank self.alt_stream = alt_stream self.k_norm = RMSNorm( diff --git a/python/sglang/srt/models/olmoe.py b/python/sglang/srt/models/olmoe.py index 46be4fe88..df7774b8a 100644 --- a/python/sglang/srt/models/olmoe.py +++ b/python/sglang/srt/models/olmoe.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.layernorm import RMSNorm from sglang.srt.layers.linear import ( QKVParallelLinear, @@ -44,6 +43,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, make_layers, print_warning_once @@ -122,7 +122,7 @@ class OlmoeAttention(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 diff --git a/python/sglang/srt/models/opt.py b/python/sglang/srt/models/opt.py index 7db9250af..0a6cc37f3 100644 --- a/python/sglang/srt/models/opt.py +++ b/python/sglang/srt/models/opt.py @@ -24,8 +24,6 @@ from transformers import OPTConfig from sglang.srt.distributed import ( get_pp_group, - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, ) from sglang.srt.layers.linear import ( ColumnParallelLinear, @@ -47,6 +45,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.utils import get_exception_traceback @@ -98,7 +97,7 @@ class OPTAttention(nn.Module): ) -> None: super().__init__() self.embed_dim = embed_dim - tensor_model_parallel_world_size = get_tensor_model_parallel_world_size() + tensor_model_parallel_world_size = get_parallel().tp_size total_num_heads = num_heads assert num_heads % tensor_model_parallel_world_size == 0 self.num_heads = total_num_heads // tensor_model_parallel_world_size @@ -364,8 +363,8 @@ class OPTModel(nn.Module): ) def load_kv_cache_scales(self, quantization_param_path: str) -> None: - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() + tp_size = get_parallel().tp_size + tp_rank = get_parallel().tp_rank for layer_idx, scaling_factor in kv_cache_scales_loader( quantization_param_path, tp_rank, diff --git a/python/sglang/srt/models/orion.py b/python/sglang/srt/models/orion.py index 1061c50b2..7eb2bc7cb 100644 --- a/python/sglang/srt/models/orion.py +++ b/python/sglang/srt/models/orion.py @@ -15,7 +15,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.distributed.parallel_state import get_pp_group from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.linear import ( @@ -34,6 +33,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ( ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors 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, make_layers from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -91,7 +91,7 @@ class OrionAttention(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 diff --git a/python/sglang/srt/models/persimmon.py b/python/sglang/srt/models/persimmon.py index 5d2585c63..8f5a2e6c5 100644 --- a/python/sglang/srt/models/persimmon.py +++ b/python/sglang/srt/models/persimmon.py @@ -5,7 +5,7 @@ import torch from torch import nn from transformers import PersimmonConfig -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 get_act_fn from sglang.srt.layers.linear import ( ColumnParallelLinear, @@ -23,6 +23,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, make_layers @@ -58,7 +59,7 @@ class PersimmonAttention(nn.Module): ): super().__init__() self.config = config - tensor_parallel_world_size = get_tensor_model_parallel_world_size() + tensor_parallel_world_size = get_parallel().tp_size self.hidden_size = config.hidden_size self.total_num_heads = config.num_attention_heads diff --git a/python/sglang/srt/models/phi.py b/python/sglang/srt/models/phi.py index c1a091f58..96867e0ca 100644 --- a/python/sglang/srt/models/phi.py +++ b/python/sglang/srt/models/phi.py @@ -7,7 +7,7 @@ import torch from torch import nn from transformers import PhiConfig -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 get_act_fn from sglang.srt.layers.linear import ( ColumnParallelLinear, @@ -24,6 +24,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, make_layers @@ -41,7 +42,7 @@ class PhiAttention(nn.Module): self.hidden_size = config.hidden_size self.head_size = self.hidden_size // self.total_num_heads - tensor_model_parallel_world_size = get_tensor_model_parallel_world_size() + tensor_model_parallel_world_size = get_parallel().tp_size assert self.total_num_heads % tensor_model_parallel_world_size == 0 self.num_heads = self.total_num_heads // tensor_model_parallel_world_size diff --git a/python/sglang/srt/models/phi3_small.py b/python/sglang/srt/models/phi3_small.py index cf049c43e..fe63b51a8 100644 --- a/python/sglang/srt/models/phi3_small.py +++ b/python/sglang/srt/models/phi3_small.py @@ -6,7 +6,7 @@ from torch import nn from transformers import Phi3Config from transformers.configuration_utils 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.linear import ( MergedColumnParallelLinear, QKVParallelLinear, @@ -25,6 +25,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, make_layers @@ -114,7 +115,7 @@ class Phi3SmallSelfAttention(nn.Module): self.num_heads = config.num_attention_heads self.head_dim = self.hidden_size // self.num_heads - self.tp_size = get_tensor_model_parallel_world_size() + self.tp_size = get_parallel().tp_size # Number of total Key Value Heads before tensor parallel self.num_key_value_heads = config.num_key_value_heads self.num_q_per_kv = self.num_heads // self.num_key_value_heads diff --git a/python/sglang/srt/models/phimoe.py b/python/sglang/srt/models/phimoe.py index a359483de..65246ede4 100644 --- a/python/sglang/srt/models/phimoe.py +++ b/python/sglang/srt/models/phimoe.py @@ -4,8 +4,6 @@ import torch from torch import nn from transformers.configuration_utils import PretrainedConfig -from sglang.srt.distributed import get_tensor_model_parallel_world_size -from sglang.srt.layers.dp_attention import get_attention_tp_rank, get_attention_tp_size from sglang.srt.layers.linear import ( QKVParallelLinear, ReplicatedLinear, @@ -28,6 +26,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, make_layers @@ -190,7 +189,7 @@ class PhiMoE(nn.Module): ): super().__init__() self.hidden_size = hidden_size - self.tp_size = get_tensor_model_parallel_world_size() + self.tp_size = get_parallel().tp_size # Gate always runs at half / full precision for now. self.gate = ReplicatedLinear( @@ -248,8 +247,8 @@ class PhiMoEAttention(nn.Module): super().__init__() self.hidden_size = hidden_size - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size self.total_num_heads = num_heads assert self.total_num_heads % attn_tp_size == 0 diff --git a/python/sglang/srt/models/qwen.py b/python/sglang/srt/models/qwen.py index 2c7545bdc..66fdf1b09 100644 --- a/python/sglang/srt/models/qwen.py +++ b/python/sglang/srt/models/qwen.py @@ -23,7 +23,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 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 from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -99,7 +99,7 @@ class QWenAttention(nn.Module): ): super().__init__() self.hidden_size = hidden_size - tensor_model_parallel_world_size = get_tensor_model_parallel_world_size() + tensor_model_parallel_world_size = get_parallel().tp_size self.total_num_heads = num_heads assert self.total_num_heads % tensor_model_parallel_world_size == 0 self.num_heads = self.total_num_heads // tensor_model_parallel_world_size diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py index b7ddc6536..744e9b1ce 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/python/sglang/srt/models/qwen2.py @@ -25,8 +25,6 @@ from torch import nn from sglang.srt.distributed import ( get_pp_group, get_pp_indices, - 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.server_args import get_global_server_args from sglang.srt.utils import add_prefix, make_layers from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -123,7 +122,7 @@ class Qwen2Attention(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 @@ -402,8 +401,8 @@ class Qwen2Model(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_rank = get_tensor_model_parallel_rank() + tp_size = get_parallel().tp_size + tp_rank = get_parallel().tp_rank for layer_idx, scaling_factor in kv_cache_scales_loader( quantization_param_path, tp_rank, diff --git a/python/sglang/srt/models/qwen2_5_vl.py b/python/sglang/srt/models/qwen2_5_vl.py index 58c4ea4d8..ccf13bd5a 100644 --- a/python/sglang/srt/models/qwen2_5_vl.py +++ b/python/sglang/srt/models/qwen2_5_vl.py @@ -42,10 +42,6 @@ from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( Qwen2_5_VisionRotaryEmbedding, ) -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.environ import envs from sglang.srt.layers.activation import SiluAndMul @@ -76,6 +72,7 @@ from sglang.srt.models.qwen2 import Qwen2Model from sglang.srt.models.utils import RotaryPosMixin, WeightsMapper, permute_inv from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner +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_npu @@ -96,10 +93,8 @@ class Qwen2_5_VLMLP(nn.Module): use_data_parallel: bool = False, ): super().__init__() - self.tp_size = ( - 1 if use_data_parallel else get_tensor_model_parallel_world_size() - ) - self.tp_rank = 0 if use_data_parallel else get_tensor_model_parallel_rank() + self.tp_size = 1 if use_data_parallel else get_parallel().tp_size + self.tp_rank = 0 if use_data_parallel else get_parallel().tp_rank self.gate_up_proj = MergedColumnParallelLinear( input_size=in_features, output_sizes=[hidden_features] * 2, # [gate_proj, up_proj] @@ -225,8 +220,8 @@ class Qwen2_5_VisionPatchMerger(nn.Module): super().__init__() self.hidden_size = context_dim * (spatial_merge_size**2) self.ln_q = RMSNorm(context_dim, eps=1e-6) - tp_size = 1 if use_data_parallel else get_tensor_model_parallel_world_size() - tp_rank = 0 if use_data_parallel else get_tensor_model_parallel_rank() + tp_size = 1 if use_data_parallel else get_parallel().tp_size + tp_rank = 0 if use_data_parallel else get_parallel().tp_rank self.mlp = nn.ModuleList( [ ColumnParallelLinear( @@ -327,9 +322,7 @@ class Qwen2_5_VisionTransformer(nn.Module, RotaryPosMixin): ) # Resource prepared for vit cuda graph - self.tp_size = ( - 1 if use_data_parallel else get_tensor_model_parallel_world_size() - ) + self.tp_size = 1 if use_data_parallel else get_parallel().tp_size self.max_context_len = max_context_len self.enable_cg = _is_cuda and envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get() diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 324e8678d..e82c039c5 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -29,16 +29,10 @@ from transformers import PretrainedConfig from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo from sglang.srt.distributed import ( - get_moe_data_parallel_world_size, - get_moe_expert_parallel_world_size, get_pp_group, get_pp_indices, - get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) -from sglang.srt.distributed.parallel_state import ( - get_attn_context_model_parallel_world_size, -) from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo @@ -49,8 +43,6 @@ from sglang.srt.layers.communicator import ( ScatterMode, ) from sglang.srt.layers.dp_attention import ( - get_attention_tp_rank, - get_attention_tp_size, is_dp_attention_enabled, ) from sglang.srt.layers.elementwise import fused_gate_sigmoid_mul_add @@ -96,6 +88,7 @@ from sglang.srt.model_executor.cuda_graph_config 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 ( add_prefix, @@ -237,7 +230,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module): enable_cuda_shared_expert_fusion: bool = False, ): super().__init__() - self.tp_size = get_tensor_model_parallel_world_size() + self.tp_size = get_parallel().tp_size self.layer_id = layer_id self.alt_stream = alt_stream if self.tp_size > config.num_experts: @@ -339,7 +332,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module): if get_moe_a2a_backend().is_deepep(): # 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.num_experts + get_global_server_args().ep_num_redundant_experts ) @@ -373,7 +366,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module): # post-experts all_reduce sums it ep_size times. Pre-scale the per-token # routing weight by 1/ep_size to cancel this, mirroring DeepSeek-V2's # fused_shared_experts_scaling_factor pattern. - moe_ep_size = get_moe_expert_parallel_world_size() + moe_ep_size = get_parallel().moe_ep_size if moe_ep_size > 1 and not is_deepep_class_backend(): w = w / float(moe_ep_size) return w @@ -611,8 +604,8 @@ class Qwen2MoeAttention(nn.Module): super().__init__() self.hidden_size = hidden_size - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size self.total_num_heads = num_heads assert self.total_num_heads % attn_tp_size == 0 @@ -725,8 +718,8 @@ class Qwen2MoeDecoderLayer(nn.Module): self.layer_id = layer_id - self.attn_tp_size = get_attention_tp_size() - self.attn_tp_rank = get_attention_tp_rank() + self.attn_tp_size = get_parallel().attn_tp_size + self.attn_tp_rank = get_parallel().attn_tp_rank # Qwen2MoE all layers are sparse and have no nextn now self.is_layer_sparse = True @@ -827,8 +820,8 @@ class Qwen2MoeModel(nn.Module): self.vocab_size = config.vocab_size self.pp_group = get_pp_group() - self.moe_dp_size = get_moe_data_parallel_world_size() - self.attn_cp_size = get_attn_context_model_parallel_world_size() + self.moe_dp_size = get_parallel().moe_dp_size + self.attn_cp_size = get_parallel().attn_cp_size if self.pp_group.is_first_rank: self.embed_tokens = VocabParallelEmbedding( diff --git a/python/sglang/srt/models/qwen3.py b/python/sglang/srt/models/qwen3.py index 14fb21dfc..ba6652915 100644 --- a/python/sglang/srt/models/qwen3.py +++ b/python/sglang/srt/models/qwen3.py @@ -7,11 +7,8 @@ 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.communicator import LayerCommunicator, LayerScatterModes -from sglang.srt.layers.dp_attention import get_attention_tp_rank, get_attention_tp_size from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import QKVParallelLinear, RowParallelLinear from sglang.srt.layers.logits_processor import LogitsProcessor @@ -36,6 +33,7 @@ from sglang.srt.model_loader.weight_utils import ( from sglang.srt.models.qwen2 import Qwen2MLP as Qwen3MLP from sglang.srt.models.qwen2 import Qwen2Model 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, get_bool_env_var, is_cuda, is_hip, is_npu @@ -84,10 +82,10 @@ class Qwen3Attention(nn.Module): super().__init__() self.hidden_size = hidden_size self.start_layer = start_layer - self.tp_size = get_tensor_model_parallel_world_size() + self.tp_size = get_parallel().tp_size self.total_num_heads = num_heads - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size assert self.total_num_heads % attn_tp_size == 0 self.num_heads = self.total_num_heads // attn_tp_size @@ -107,7 +105,7 @@ class Qwen3Attention(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 norm_kwargs = ( dict( diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 4273be43c..7157eb23e 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -43,8 +43,6 @@ from sglang.srt.layers.attention.fla.layernorm_gated import RMSNorm as RMSNormGa from sglang.srt.layers.attention.mamba.mamba import mamba_v2_sharded_weight_loader 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.elementwise import fused_sigmoid_mul @@ -93,6 +91,7 @@ from sglang.srt.models.utils import ( fused_qk_gemma_rmsnorm, fused_qk_gemma_rmsnorm_with_gate, ) +from sglang.srt.runtime_context import get_parallel from sglang.srt.server_args import get_global_server_args # Utils @@ -153,8 +152,8 @@ class Qwen3_5GatedDeltaNet(nn.Module): ) -> None: super().__init__() self.config = config - self.attn_tp_rank = get_attention_tp_rank() - self.attn_tp_size = get_attention_tp_size() + self.attn_tp_rank = get_parallel().attn_tp_rank + self.attn_tp_size = get_parallel().attn_tp_size self.hidden_size = config.hidden_size self.num_v_heads = ( config.linear_num_value_heads @@ -715,8 +714,8 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): super().__init__() self.config = config self.hidden_size = config.hidden_size - self.attn_tp_rank = get_attention_tp_rank() - self.attn_tp_size = get_attention_tp_size() + self.attn_tp_rank = get_parallel().attn_tp_rank + self.attn_tp_size = get_parallel().attn_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 diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index 93c44da2d..8bc637e4f 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -23,7 +23,7 @@ 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.environ import envs from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation @@ -34,6 +34,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead 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.qwen3_5 import Qwen3_5ForCausalLM +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 @@ -79,7 +80,7 @@ class Qwen3_5ForCausalLMMTP(nn.Module): quant_config = None 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() diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py index 5c6c76bf1..90ec529cd 100644 --- a/python/sglang/srt/models/qwen3_moe.py +++ b/python/sglang/srt/models/qwen3_moe.py @@ -26,13 +26,7 @@ from torch import nn from transformers import PretrainedConfig from sglang.srt.distributed import ( - get_attn_context_model_parallel_rank, - get_attn_context_model_parallel_world_size, - get_moe_data_parallel_world_size, - get_moe_expert_parallel_world_size, - get_moe_tensor_parallel_world_size, get_pp_group, - get_tensor_model_parallel_rank, moe_expert_parallel_all_reduce, moe_tensor_model_parallel_all_reduce, ) @@ -40,7 +34,6 @@ from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_r from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes -from sglang.srt.layers.dp_attention import get_attention_tp_rank, get_attention_tp_size from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( QKVParallelLinear, @@ -78,6 +71,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, @@ -239,8 +233,8 @@ class Qwen3MoeSparseMoeBlock(nn.Module): prefix: str = "", ): super().__init__() - self.tp_size = get_moe_tensor_parallel_world_size() - self.ep_size = get_moe_expert_parallel_world_size() + self.tp_size = get_parallel().moe_tp_size + self.ep_size = get_parallel().moe_ep_size self.layer_id = layer_id if self.tp_size > config.num_experts: raise ValueError( @@ -283,7 +277,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module): if get_moe_a2a_backend().is_deepep(): # 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.num_experts + get_global_server_args().ep_num_redundant_experts ) @@ -462,8 +456,8 @@ class Qwen3MoeAttention(nn.Module): self.hidden_size = hidden_size self.start_layer = start_layer - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size self.config = config self.total_num_heads = num_heads @@ -485,7 +479,7 @@ class Qwen3MoeAttention(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, @@ -763,8 +757,8 @@ class Qwen3MoeDecoderLayer(nn.Module): self.layer_id = layer_id - self.attn_tp_size = get_attention_tp_size() - self.attn_tp_rank = get_attention_tp_rank() + self.attn_tp_size = get_parallel().attn_tp_size + self.attn_tp_rank = get_parallel().attn_tp_rank # Qwen3MoE all layers are sparse and have no nextn now self.is_layer_sparse = True @@ -973,9 +967,9 @@ class Qwen3MoeForCausalLM(nn.Module): self.logits_processor = LogitsProcessor(config) self.capture_aux_hidden_states = False - self.attn_cp_size = get_attn_context_model_parallel_world_size() - self.attn_cp_rank = get_attn_context_model_parallel_rank() - self.moe_dp_size = get_moe_data_parallel_world_size() + self.attn_cp_size = get_parallel().attn_cp_size + self.attn_cp_rank = get_parallel().attn_cp_rank + self.moe_dp_size = get_parallel().moe_dp_size assert self.attn_cp_size % self.moe_dp_size == 0, ( f"attn_cp_size ({self.attn_cp_size}) must be divisible by " diff --git a/python/sglang/srt/models/qwen3_moe_mtp.py b/python/sglang/srt/models/qwen3_moe_mtp.py index e6f825eef..973d1adde 100644 --- a/python/sglang/srt/models/qwen3_moe_mtp.py +++ b/python/sglang/srt/models/qwen3_moe_mtp.py @@ -21,7 +21,7 @@ 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.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.logits_processor import LogitsProcessor @@ -30,6 +30,7 @@ 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.qwen3_moe import Qwen3MoeForCausalLM, Qwen3MoeModel +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 Qwen3MoeForCausalLMMTP(Qwen3MoeForCausalLM): 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() diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py index 8139bc8ac..f085fb43d 100644 --- a/python/sglang/srt/models/qwen3_next.py +++ b/python/sglang/srt/models/qwen3_next.py @@ -16,8 +16,6 @@ from sglang.srt.layers.attention.fla.layernorm_gated import RMSNorm as RMSNormGa from sglang.srt.layers.attention.mamba.mamba import mamba_v2_sharded_weight_loader 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 GemmaRMSNorm @@ -49,6 +47,7 @@ from sglang.srt.model_loader.weight_utils import ( sharded_weight_loader, ) from sglang.srt.models.qwen2_moe import Qwen2MoeMLP, Qwen2MoeSparseMoeBlock +from sglang.srt.runtime_context import get_parallel from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import ( LazyValue, @@ -93,8 +92,8 @@ class Qwen3GatedDeltaNet(nn.Module): ) -> None: super().__init__() self.config = config - self.attn_tp_rank = get_attention_tp_rank() - self.attn_tp_size = get_attention_tp_size() + self.attn_tp_rank = get_parallel().attn_tp_rank + self.attn_tp_size = get_parallel().attn_tp_size self.hidden_size = config.hidden_size self.num_v_heads = ( config.linear_num_value_heads @@ -611,8 +610,8 @@ class Qwen3HybridAttentionDecoderLayer(nn.Module): super().__init__() self.config = config self.hidden_size = config.hidden_size - self.attn_tp_rank = get_attention_tp_rank() - self.attn_tp_size = get_attention_tp_size() + self.attn_tp_rank = get_parallel().attn_tp_rank + self.attn_tp_size = get_parallel().attn_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 diff --git a/python/sglang/srt/models/qwen3_next_mtp.py b/python/sglang/srt/models/qwen3_next_mtp.py index f239339c6..69adf12a7 100644 --- a/python/sglang/srt/models/qwen3_next_mtp.py +++ b/python/sglang/srt/models/qwen3_next_mtp.py @@ -23,7 +23,7 @@ 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.environ import envs from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.layers.layernorm import GemmaRMSNorm @@ -32,6 +32,7 @@ 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.qwen3_next import Qwen3NextForCausalLM, Qwen3NextModel +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 @@ -50,7 +51,7 @@ class Qwen3NextForCausalLMMTP(Qwen3NextForCausalLM): # Deep-copy so MTP mutations below don't leak into the target's config. config = copy.deepcopy(config) 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 diff --git a/python/sglang/srt/models/qwen3_omni_moe.py b/python/sglang/srt/models/qwen3_omni_moe.py index fd488faae..18ae8475e 100644 --- a/python/sglang/srt/models/qwen3_omni_moe.py +++ b/python/sglang/srt/models/qwen3_omni_moe.py @@ -31,9 +31,6 @@ from sglang.srt.configs.qwen3_omni import ( Qwen3OmniMoeVisionEncoderConfig, ) from sglang.srt.configs.qwen3_vl import Qwen3VLMoeConfig -from sglang.srt.distributed import ( - get_tensor_model_parallel_world_size, -) from sglang.srt.layers.attention.vision import VisionAttention from sglang.srt.layers.linear import ( ColumnParallelLinear, @@ -50,6 +47,7 @@ from sglang.srt.models.qwen3_vl_moe import ( Qwen3VLMoeForConditionalGeneration, load_fused_expert_weights, ) +from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import add_prefix, is_cpu, is_npu, logger _is_cpu = is_cpu() @@ -104,7 +102,7 @@ class Qwen3OmniMoeAudioEncoderLayer(nn.Module): self.dropout = config.dropout self.activation_fn = ACT2FN[config.activation_function] self.activation_dropout = config.activation_dropout - tp_size = get_tensor_model_parallel_world_size() + tp_size = get_parallel().tp_size use_replicated = config.encoder_ffn_dim % tp_size != 0 fc1_cls = ReplicatedLinear if use_replicated else ColumnParallelLinear fc2_cls = ReplicatedLinear if use_replicated else RowParallelLinear diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py index 3ffe4dde7..76d6922ac 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -27,7 +27,6 @@ from einops import rearrange from transformers.activations import ACT2FN from sglang.srt.configs.qwen3_vl import Qwen3VLConfig, Qwen3VLVisionConfig -from sglang.srt.distributed import get_tensor_model_parallel_world_size from sglang.srt.distributed.parallel_state import get_pp_group from sglang.srt.environ import envs from sglang.srt.layers.attention.vision import ( @@ -38,8 +37,6 @@ from sglang.srt.layers.attention.vision import ( ) from sglang.srt.layers.conv import Conv3dLayer 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 ColumnParallelLinear, RowParallelLinear @@ -71,6 +68,7 @@ from sglang.srt.models.utils import ( ) from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner +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, @@ -110,8 +108,8 @@ class Qwen3_VisionMLP(nn.Module): use_data_parallel: bool = False, ): super().__init__() - self.tp_size = 1 if use_data_parallel else get_attention_tp_size() - self.tp_rank = 0 if use_data_parallel else get_attention_tp_rank() + self.tp_size = 1 if use_data_parallel else get_parallel().attn_tp_size + self.tp_rank = 0 if use_data_parallel else get_parallel().attn_tp_rank self.linear_fc1 = ColumnParallelLinear( in_features, hidden_features, @@ -270,8 +268,8 @@ class Qwen3VLMoeVisionPatchMerger(nn.Module): self.norm = norm_layer( self.hidden_size if use_postshuffle_norm else context_dim ) - self.tp_size = 1 if use_data_parallel else get_attention_tp_size() - self.tp_rank = 0 if use_data_parallel else get_attention_tp_rank() + self.tp_size = 1 if use_data_parallel else get_parallel().attn_tp_size + self.tp_rank = 0 if use_data_parallel else get_parallel().attn_tp_rank self.linear_fc1 = ColumnParallelLinear( self.hidden_size, self.padded_context_dim, @@ -423,9 +421,7 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin): ] ) - self.tp_size = ( - 1 if use_data_parallel else get_tensor_model_parallel_world_size() - ) + self.tp_size = 1 if use_data_parallel else get_parallel().tp_size self.graph_runners = graph_runners_dict[self.device.type](self) @property diff --git a/python/sglang/srt/models/sarvam_moe.py b/python/sglang/srt/models/sarvam_moe.py index ab93eb3d9..1e9a8c82f 100644 --- a/python/sglang/srt/models/sarvam_moe.py +++ b/python/sglang/srt/models/sarvam_moe.py @@ -14,7 +14,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.eplb.expert_distribution import get_global_expert_distribution_recorder @@ -27,8 +26,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_dp_attention_enabled, ) from sglang.srt.layers.layernorm import RMSNorm @@ -63,6 +60,7 @@ from sglang.srt.models.bailing_moe import BailingMoEForCausalLM from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mha import ( DeepseekMHAForwardMixin, ) +from sglang.srt.runtime_context import get_parallel from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import ( BumpAllocator, @@ -233,7 +231,7 @@ class SarvamMoESparseMoeBlock(nn.Module): super().__init__() self.config = config self.layer_id = layer_id - self.tp_size = get_tensor_model_parallel_world_size() + self.tp_size = get_parallel().tp_size self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 2.5) self.score_function = getattr(config, "score_function", "sigmoid") self.n_group = getattr(config, "n_group", None) @@ -449,8 +447,8 @@ class SarvamMoEMLAAttention(nn.Module): self.alt_stream = alt_stream self.quant_config = quant_config - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size self.qk_nope_head_dim = config.qk_nope_head_dim self.qk_rope_head_dim = config.qk_rope_head_dim @@ -1079,7 +1077,7 @@ class SarvamMoEMLADecoderLayer(nn.Module): config.hidden_size, eps=config.rms_norm_eps ) - self.attn_tp_size = get_attention_tp_size() + self.attn_tp_size = get_parallel().attn_tp_size self.layer_scatter_modes = LayerScatterModes.init_new( layer_id=layer_id, num_layers=config.num_hidden_layers, diff --git a/python/sglang/srt/models/sdar.py b/python/sglang/srt/models/sdar.py index 70ab59a48..ddd21950a 100644 --- a/python/sglang/srt/models/sdar.py +++ b/python/sglang/srt/models/sdar.py @@ -10,12 +10,10 @@ 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.activation import SiluAndMul 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 @@ -43,6 +41,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, make_layers @@ -99,9 +98,9 @@ class SDARAttention(nn.Module): self.layer_id = layer_id self.hidden_size = config.hidden_size self.total_num_heads = config.num_attention_heads - self.tp_size = get_tensor_model_parallel_world_size() - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() + self.tp_size = get_parallel().tp_size + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size assert self.total_num_heads % attn_tp_size == 0 self.num_heads = self.total_num_heads // attn_tp_size @@ -460,7 +459,7 @@ class SDARForCausalLM(nn.Module): ) if self.pp_group.is_last_rank: - tp_size = get_tensor_model_parallel_world_size() + tp_size = get_parallel().tp_size if ( self.pp_group.world_size == 1 and config.tie_word_embeddings diff --git a/python/sglang/srt/models/sdar_moe.py b/python/sglang/srt/models/sdar_moe.py index c09bfeb17..02a3ad6c5 100644 --- a/python/sglang/srt/models/sdar_moe.py +++ b/python/sglang/srt/models/sdar_moe.py @@ -11,9 +11,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 @@ -21,8 +19,6 @@ from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo 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 @@ -61,6 +57,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, add_prefix, is_cuda, make_layers @@ -85,7 +82,7 @@ class SDARMoeSparseMoeBlock(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 if self.tp_size > config.num_experts: raise ValueError( @@ -121,7 +118,7 @@ class SDARMoeSparseMoeBlock(nn.Module): # Deepep / FuseEP support 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 ) @@ -212,8 +209,8 @@ class SDARMoeAttention(nn.Module): self.hidden_size = config.hidden_size self.total_num_heads = config.num_attention_heads - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size assert self.total_num_heads % attn_tp_size == 0 self.num_heads = self.total_num_heads // attn_tp_size @@ -557,7 +554,7 @@ class SDARMoeForCausalLM(nn.Module): ) if self.pp_group.is_last_rank: - tp_size = get_tensor_model_parallel_world_size() + tp_size = get_parallel().tp_size if ( self.pp_group.world_size == 1 and getattr(config, "tie_word_embeddings", False) diff --git a/python/sglang/srt/models/solar.py b/python/sglang/srt/models/solar.py index b2b92c388..d66f22c7a 100644 --- a/python/sglang/srt/models/solar.py +++ b/python/sglang/srt/models/solar.py @@ -1,6 +1,37 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Iterable +from typing import Any, List, Optional, Tuple, Union + +import torch +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import get_pp_group +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput +from sglang.srt.layers.quantization import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.utils import PPMissingLayer +from sglang.srt.layers.vocab_parallel_embedding import ( + DEFAULT_VOCAB_PADDING_SIZE, + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.model_loader.weight_utils import ( + default_weight_loader, + kv_cache_scales_loader, +) + # Adapted from # https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py # Copyright 2023 The vLLM team. @@ -23,37 +54,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/solar.py -from collections.abc import Iterable -from typing import Any, List, Optional, Tuple, Union - -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.parallel_state import get_tensor_model_parallel_rank -from sglang.srt.layers.activation import SiluAndMul -from sglang.srt.layers.layernorm import RMSNorm -from sglang.srt.layers.linear import ( - MergedColumnParallelLinear, - QKVParallelLinear, - RowParallelLinear, -) -from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput -from sglang.srt.layers.quantization import QuantizationConfig -from sglang.srt.layers.radix_attention import RadixAttention -from sglang.srt.layers.rotary_embedding import get_rope -from sglang.srt.layers.utils import PPMissingLayer -from sglang.srt.layers.vocab_parallel_embedding import ( - DEFAULT_VOCAB_PADDING_SIZE, - ParallelLMHead, - VocabParallelEmbedding, -) -from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors -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 @@ -116,7 +117,7 @@ class SolarAttention(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 @@ -362,8 +363,8 @@ class SolarModel(nn.Module): return hidden_states def load_kv_cache_scales(self, quantization_param_path: str) -> None: - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() + tp_size = get_parallel().tp_size + tp_rank = get_parallel().tp_rank for layer_idx, scaling_factor in kv_cache_scales_loader( quantization_param_path, tp_rank, diff --git a/python/sglang/srt/models/stablelm.py b/python/sglang/srt/models/stablelm.py index b7cf47903..1ee4d53c4 100644 --- a/python/sglang/srt/models/stablelm.py +++ b/python/sglang/srt/models/stablelm.py @@ -28,7 +28,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 SiluAndMul from sglang.srt.layers.linear import ( MergedColumnParallelLinear, @@ -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, is_npu _is_npu = is_npu() @@ -95,7 +95,7 @@ class StablelmAttention(nn.Module): super().__init__() self.config = config 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 self.num_heads = self.total_num_heads // tp_size diff --git a/python/sglang/srt/models/starcoder2.py b/python/sglang/srt/models/starcoder2.py index e5cba190d..22250f63f 100644 --- a/python/sglang/srt/models/starcoder2.py +++ b/python/sglang/srt/models/starcoder2.py @@ -29,7 +29,7 @@ import torch from torch import nn from transformers import Starcoder2Config -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 get_act_fn from sglang.srt.layers.linear import ( ColumnParallelLinear, @@ -47,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, make_layers @@ -63,7 +64,7 @@ class Starcoder2Attention(nn.Module): self.config = config 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 diff --git a/python/sglang/srt/models/step3_vl.py b/python/sglang/srt/models/step3_vl.py index 1deb542f4..f659e926e 100644 --- a/python/sglang/srt/models/step3_vl.py +++ b/python/sglang/srt/models/step3_vl.py @@ -16,8 +16,6 @@ from sglang.srt.configs.step3_vl import ( Step3VLConfig, ) 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_location import ModelConfigForExpertLocation @@ -26,8 +24,6 @@ from sglang.srt.layers.attention.vision import VisionAttention from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes from sglang.srt.layers.conv import Conv2dLayer 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 @@ -60,6 +56,7 @@ from sglang.srt.managers.schedule_batch 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, log_info_on_rank0, make_layers from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -119,7 +116,7 @@ class Step3TextMoEMLP(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 if self.tp_size > config.moe_num_experts: raise ValueError( @@ -189,10 +186,10 @@ class Step3TextAttention(nn.Module): super().__init__() self.hidden_size = hidden_size - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size - self.all_tp_rank = get_tensor_model_parallel_rank() + self.all_tp_rank = get_parallel().tp_rank self.total_num_heads = num_heads self.attn_tp_rank = attn_tp_rank self.layer_id = layer_id @@ -541,8 +538,8 @@ class Step3VisionMLP(nn.Module): # Since this is a dense model, # the MLP component likewise adopts a DP-MLP approach modeled after DP Attention. # This choice may not represent the optimal solution and remains open to further deliberation. - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size self.fc1 = ColumnParallelLinear( dim, intermediate_size, diff --git a/python/sglang/srt/models/step3p5.py b/python/sglang/srt/models/step3p5.py index 843c2d862..cdc8944b9 100644 --- a/python/sglang/srt/models/step3p5.py +++ b/python/sglang/srt/models/step3p5.py @@ -5,10 +5,7 @@ import torch.nn.functional as F from torch import nn from sglang.srt.distributed import ( - get_moe_expert_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 @@ -17,8 +14,6 @@ from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo from sglang.srt.layers.activation import SiluAndMul 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 GemmaRMSNorm @@ -51,6 +46,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ( ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors 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, is_non_idle_and_non_empty, make_layers @@ -120,7 +116,7 @@ class Step3p5MoEMLP(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.need_fp32_gate = config.need_fp32_gate @@ -174,7 +170,7 @@ class Step3p5MoEMLP(nn.Module): if get_moe_a2a_backend().is_deepep(): # 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.moe_num_experts = ( config.moe_num_experts + get_global_server_args().ep_num_redundant_experts @@ -354,10 +350,10 @@ class Step3p5Attention(nn.Module): ) -> None: super().__init__() self.hidden_size = hidden_size - self.tp_size = get_tensor_model_parallel_world_size() + self.tp_size = get_parallel().tp_size self.total_num_heads = num_heads - attn_tp_rank = get_attention_tp_rank() - attn_tp_size = get_attention_tp_size() + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size assert self.total_num_heads % attn_tp_size == 0 self.num_heads = self.total_num_heads // attn_tp_size @@ -377,7 +373,7 @@ class Step3p5Attention(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.q_norm = GemmaRMSNorm(self.head_dim, eps=rms_norm_eps) self.k_norm = GemmaRMSNorm(self.head_dim, eps=rms_norm_eps) diff --git a/python/sglang/srt/models/step3p5_mtp.py b/python/sglang/srt/models/step3p5_mtp.py index 15e970383..070bbe1f6 100644 --- a/python/sglang/srt/models/step3p5_mtp.py +++ b/python/sglang/srt/models/step3p5_mtp.py @@ -6,7 +6,6 @@ import torch import torch.nn as nn from transformers import PretrainedConfig -from sglang.srt.distributed import get_tensor_model_parallel_world_size from sglang.srt.layers.layernorm import GemmaRMSNorm from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.quantization.base_config import QuantizationConfig @@ -17,6 +16,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.models.step3p5 import Step3p5DecoderLayer, Step3p5ForCausalLM +from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import add_prefix logger = logging.getLogger(__name__) @@ -148,7 +148,7 @@ class Step3p5MTP(Step3p5ForCausalLM): ) -> 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.draft_model_idx = draft_model_idx diff --git a/python/sglang/srt/models/torch_native_llama.py b/python/sglang/srt/models/torch_native_llama.py index 1685ee9fa..68bdbbf56 100644 --- a/python/sglang/srt/models/torch_native_llama.py +++ b/python/sglang/srt/models/torch_native_llama.py @@ -50,10 +50,6 @@ from torch import nn from torch.nn.parameter import Parameter from transformers import LlamaConfig -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.logits_processor import LogitsProcessor, LogitsProcessorOutput @@ -66,6 +62,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 tp_size: Optional[int] = None @@ -346,9 +343,9 @@ class LlamaModel(nn.Module): global tp_size, tp_rank if tp_size is None: - tp_size = get_tensor_model_parallel_world_size() + tp_size = get_parallel().tp_size if tp_rank is None: - tp_rank = get_tensor_model_parallel_rank() + tp_rank = get_parallel().tp_rank self.config = config self.padding_idx = config.pad_token_id diff --git a/python/sglang/srt/models/transformers.py b/python/sglang/srt/models/transformers.py index 79bbba865..dc8dd2dc8 100644 --- a/python/sglang/srt/models/transformers.py +++ b/python/sglang/srt/models/transformers.py @@ -34,10 +34,8 @@ from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS from sglang.srt.distributed import ( divide, - get_moe_expert_parallel_world_size, get_pp_group, get_pp_indices, - get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation @@ -67,6 +65,7 @@ from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInp from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.utils import AutoWeightsLoader, WeightsMapper +from sglang.srt.runtime_context import get_parallel from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import get_device from sglang.srt.utils.common import direct_register_custom_op @@ -642,7 +641,7 @@ class TransformersBase(nn.Module): # Pipeline parallel self.pipeline_parallel() # Module replacement (Linear → TP, RMSNorm → fused, MoE overridden by MoEMixin) - tp_size = get_tensor_model_parallel_world_size() + tp_size = get_parallel().tp_size self.recursive_replace() # Attention instances self.attention_instances = self._create_attention_instances(tp_size) @@ -755,7 +754,7 @@ class TransformersBase(nn.Module): # -- Recursive module replacement (Linear + RMSNorm) -------------------- def recursive_replace(self): - tp_size = get_tensor_model_parallel_world_size() + tp_size = get_parallel().tp_size tp_plan = self._normalize_tp_plan(self._get_model_tp_plan()) if not tp_plan and tp_size > 1: @@ -1232,7 +1231,7 @@ class MoEMixin: # EPLB / EP tracking num_redundant = get_global_server_args().ep_num_redundant_experts - ep_size = get_moe_expert_parallel_world_size() + ep_size = get_parallel().moe_ep_size self.mlp_moe_layers: list[nn.Module] = [] self.moe_layers: list[TransformersFusedMoE] = [] diff --git a/python/sglang/srt/models/whisper.py b/python/sglang/srt/models/whisper.py index 2c8f7aa43..d5135406c 100644 --- a/python/sglang/srt/models/whisper.py +++ b/python/sglang/srt/models/whisper.py @@ -6,7 +6,6 @@ from typing import Any, Iterable, Optional, Tuple import torch from transformers import WhisperConfig -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, @@ -20,6 +19,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead from sglang.srt.managers.schedule_batch import MultimodalInputs 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 class WhisperAttention(torch.nn.Module): @@ -41,7 +41,7 @@ class WhisperAttention(torch.nn.Module): self.is_cross_attention = is_cross_attention self.is_encoder = is_encoder - tp_size = get_tensor_model_parallel_world_size() + tp_size = get_parallel().tp_size assert ( num_heads % tp_size == 0 ), f"num_heads ({num_heads}) must be divisible by tp_size ({tp_size})" diff --git a/python/sglang/srt/models/xverse.py b/python/sglang/srt/models/xverse.py index 892ba8fb3..d4366c494 100644 --- a/python/sglang/srt/models/xverse.py +++ b/python/sglang/srt/models/xverse.py @@ -24,7 +24,6 @@ import torch from torch import nn from transformers import LlamaConfig -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 ( @@ -42,6 +41,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ( ) from sglang.srt.model_executor.model_runner 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 @@ -101,7 +101,7 @@ class XverseAttention(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 diff --git a/python/sglang/srt/models/xverse_moe.py b/python/sglang/srt/models/xverse_moe.py index 8418c83fc..943395e1c 100644 --- a/python/sglang/srt/models/xverse_moe.py +++ b/python/sglang/srt/models/xverse_moe.py @@ -20,8 +20,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.hardware_backend.npu.quantization.fused_moe_method_npu import ( @@ -48,6 +46,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 @@ -103,8 +102,8 @@ class XverseMoE(nn.Module): ): super().__init__() self.config = config - self.rank = get_tensor_model_parallel_rank() - self.tp_size = get_tensor_model_parallel_world_size() + self.rank = get_parallel().tp_rank + self.tp_size = get_parallel().tp_size self.n_routed_experts = config.num_experts self.top_k = config.moe_top_k if self.tp_size > self.n_routed_experts: @@ -211,7 +210,7 @@ class XverseAttention(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 diff --git a/python/sglang/srt/models/zaya.py b/python/sglang/srt/models/zaya.py index 5fb670f80..9fb63238c 100644 --- a/python/sglang/srt/models/zaya.py +++ b/python/sglang/srt/models/zaya.py @@ -50,11 +50,8 @@ from torch import nn from sglang.srt.configs.zaya import ZayaConfig 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.layers.dp_attention import get_attention_tp_size from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( ColumnParallelLinear, @@ -75,6 +72,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.forward_context import get_req_to_token_pool 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, make_layers, set_weight_attrs logger = logging.getLogger(__name__) @@ -202,9 +200,9 @@ class CCA(nn.Module): self.total_padding = self.padding0 + self.padding1 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 = int(tp_rank) self.tp_size = int(tp_size) @@ -833,8 +831,8 @@ class ZayaAttention(nn.Module): # divisible by tp_size; the KV-replicated GQA-TP variant (tp_size > # num_k_heads) is intentionally rejected with a clear error message # because both per-K-head paths assume each rank holds whole K heads. - self.tp_rank = get_tensor_model_parallel_rank() - self.tp_size = get_tensor_model_parallel_world_size() + self.tp_rank = get_parallel().tp_rank + self.tp_size = get_parallel().tp_size # The head split, the ``o_proj`` RowParallel all-reduce, and the # RadixAttention KV cache are all organized on the *global* TP group, # and ``ZayaConfig.mamba2_cache_params`` sizes the conv-state cache on @@ -843,7 +841,7 @@ class ZayaAttention(nn.Module): # ``use_dp_attention_reduce``), which this model does not wire up, so # require the two groups to coincide and fail fast instead of silently # mis-sizing the conv-state cache. - attn_tp_size = get_attention_tp_size() + attn_tp_size = get_parallel().attn_tp_size assert attn_tp_size == self.tp_size, ( f"ZAYA1 head-parallel attention requires the attention TP group " f"({attn_tp_size}) to equal the global TP group ({self.tp_size}); " @@ -1143,7 +1141,7 @@ class ZayaBlock(nn.Module): self.mlp_expansion = int(config.zaya_mlp_expansion) self.topk = int(getattr(config, "moe_router_topk", 1)) - self.tp_size = get_tensor_model_parallel_world_size() + self.tp_size = get_parallel().tp_size if self.tp_size > self.num_moe_experts: raise ValueError( f"Tensor parallel size {self.tp_size} is greater than the " diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py new file mode 100644 index 000000000..9cc28257a --- /dev/null +++ b/python/sglang/srt/runtime_context.py @@ -0,0 +1,226 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""A single structured accessor for process-static parallel-topology state. + +``get_parallel()`` returns a ``ParallelContext`` whose attributes — tp / pp / +moe / attn size and rank, plus the process-group handles — each delegate live to +the canonical getter in ``distributed.parallel_state`` / ``layers.dp_attention``. +Returned values are exactly what those getters return; this is a read-through +wrapper, not a cache. It gives call-sites one import and one naming scheme in +place of a dozen free functions, plus a test-only ``override()`` hook to force a +topology without monkeypatching the underlying getters. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import Any + + +# Imported lazily so this module has no import-time dependencies: any module can +# import get_parallel at module level without risking an import cycle. +def _ps(): + from sglang.srt.distributed import parallel_state + + return parallel_state + + +def _dp(): + from sglang.srt.layers import dp_attention + + return dp_attention + + +_PARALLEL_FIELDS = frozenset( + { + "world_size", + "world_rank", + "tp_size", + "tp_rank", + "pp_size", + "pp_rank", + "moe_ep_size", + "moe_ep_rank", + "moe_dp_size", + "moe_dp_rank", + "moe_tp_size", + "moe_tp_rank", + "attn_tp_size", + "attn_tp_rank", + "attn_cp_size", + "attn_cp_rank", + "attn_dp_size", + "attn_dp_rank", + "world_group", + "tp_group", + "pp_group", + "moe_ep_group", + "moe_dp_group", + "moe_tp_group", + "attn_tp_group", + "attn_cp_group", + } +) + + +class ParallelContext: + """Parallel-topology namespace; the only instance state is ``_overrides``.""" + + __slots__ = ("_overrides",) + + def __init__(self): + self._overrides = {} + + def _v(self, name, getter): + overrides = self._overrides + return overrides[name] if name in overrides else getter() + + @contextmanager + def override(self, **kwargs): + """Temporarily force parallel values, restoring on exit. Validates keys and + supports nesting.""" + unknown = set(kwargs) - _PARALLEL_FIELDS + if unknown: + raise ValueError(f"unknown parallel field(s): {sorted(unknown)}") + saved = dict(self._overrides) + self._overrides.update(kwargs) + try: + yield self + finally: + self._overrides = saved + + @property + def world_size(self) -> int: + return self._v("world_size", _ps().get_world_size) + + @property + def world_rank(self) -> int: + return self._v("world_rank", _ps().get_world_rank) + + @property + def tp_size(self) -> int: + return self._v("tp_size", _ps().get_tensor_model_parallel_world_size) + + @property + def tp_rank(self) -> int: + return self._v("tp_rank", _ps().get_tensor_model_parallel_rank) + + @property + def pp_size(self) -> int: + return self._v("pp_size", _ps().get_pipeline_model_parallel_world_size) + + @property + def pp_rank(self) -> int: + return self._v("pp_rank", _ps().get_pipeline_model_parallel_rank) + + @property + def moe_ep_size(self) -> int: + return self._v("moe_ep_size", _ps().get_moe_expert_parallel_world_size) + + @property + def moe_ep_rank(self) -> int: + return self._v("moe_ep_rank", _ps().get_moe_expert_parallel_rank) + + @property + def moe_dp_size(self) -> int: + return self._v("moe_dp_size", _ps().get_moe_data_parallel_world_size) + + @property + def moe_dp_rank(self) -> int: + return self._v("moe_dp_rank", _ps().get_moe_data_parallel_rank) + + @property + def moe_tp_size(self) -> int: + return self._v("moe_tp_size", _ps().get_moe_tensor_parallel_world_size) + + @property + def moe_tp_rank(self) -> int: + return self._v("moe_tp_rank", _ps().get_moe_tensor_parallel_rank) + + @property + def attn_tp_size(self) -> int: + return self._v("attn_tp_size", _ps().get_attn_tensor_model_parallel_world_size) + + @property + def attn_tp_rank(self) -> int: + return self._v("attn_tp_rank", _ps().get_attn_tensor_model_parallel_rank) + + @property + def attn_cp_size(self) -> int: + return self._v("attn_cp_size", _ps().get_attn_context_model_parallel_world_size) + + @property + def attn_cp_rank(self) -> int: + return self._v("attn_cp_rank", _ps().get_attn_context_model_parallel_rank) + + @property + def attn_dp_size(self) -> int: + return self._v("attn_dp_size", _dp().get_attention_dp_size) + + @property + def attn_dp_rank(self) -> int: + return self._v("attn_dp_rank", _dp().get_attention_dp_rank) + + @property + def world_group(self) -> Any: + return self._v("world_group", _ps().get_world_group) + + @property + def tp_group(self) -> Any: + return self._v("tp_group", _ps().get_tp_group) + + @property + def pp_group(self) -> Any: + return self._v("pp_group", _ps().get_pp_group) + + @property + def moe_ep_group(self) -> Any: + return self._v("moe_ep_group", _ps().get_moe_ep_group) + + @property + def moe_dp_group(self) -> Any: + return self._v("moe_dp_group", _ps().get_moe_dp_group) + + @property + def moe_tp_group(self) -> Any: + return self._v("moe_tp_group", _ps().get_moe_tp_group) + + @property + def attn_tp_group(self) -> Any: + return self._v("attn_tp_group", _ps().get_attn_tp_group) + + @property + def attn_cp_group(self) -> Any: + return self._v("attn_cp_group", _ps().get_attn_cp_group) + + +class RuntimeContext: + """Container for the structured runtime accessors; currently exposes ``parallel``.""" + + __slots__ = ("parallel",) + + def __init__(self, parallel: ParallelContext): + self.parallel = parallel + + +_PARALLEL = ParallelContext() +_CONTEXT = RuntimeContext(parallel=_PARALLEL) + + +def get_context() -> RuntimeContext: + return _CONTEXT + + +def get_parallel() -> ParallelContext: + return _PARALLEL diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py index 7d58293e2..8da926f89 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py @@ -7,7 +7,6 @@ import torch.nn.functional as F from torch import nn from sglang.srt.configs.model_config import AttentionArch -from sglang.srt.layers import dp_attention as _dp_attention from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, ReqToTokenPool @@ -19,6 +18,7 @@ from sglang.srt.model_executor.cuda_graph_config import ( from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.runtime_context import get_parallel from sglang.srt.server_args import set_global_server_args_for_scheduler from sglang.srt.speculative.spec_info import SpeculativeAlgorithm @@ -26,7 +26,8 @@ from ..mock_server_args import make_mock_server_args # Unit tests run without distributed initialization. Backends that size buffers by # attention tensor-parallel degree should see the single-rank default. -_dp_attention.get_attention_tp_size = lambda: 1 +_parallel_override = get_parallel().override(attn_tp_size=1) +_parallel_override.__enter__() DEFAULT_HEAD_DIM = 16 DEFAULT_HIDDEN_SIZE = 64 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py index 1eb20f02a..51230f641 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py @@ -6,7 +6,6 @@ import torch from torch import nn from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS -from sglang.srt.layers.attention.dsa import utils as _dsa_utils from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool, ReqToTokenPool from sglang.srt.model_executor.cuda_graph_config import ( @@ -17,6 +16,7 @@ from sglang.srt.model_executor.cuda_graph_config import ( from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.runtime_context import get_parallel from sglang.srt.server_args import set_global_server_args_for_scheduler from ..mock_server_args import make_mock_server_args @@ -37,8 +37,8 @@ from .dense_attention import ( # Unit tests run without distributed initialization. DSA context-parallel probes # should see the single-rank default. -_dsa_utils.get_attention_cp_size = lambda: 1 -_dsa_utils.get_attention_cp_rank = lambda: 0 +_parallel_override = get_parallel().override(attn_cp_size=1, attn_cp_rank=0) +_parallel_override.__enter__() DSA_PAGE_SIZE = 64 DSA_INDEX_HEAD_DIM = 128 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py index 6ed24703f..5a6ce04cd 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py @@ -19,7 +19,6 @@ from typing import Any import torch from torch import nn -from sglang.srt.layers import dp_attention as _dp_attention from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS from sglang.srt.layers.attention.dsv4.quant_k_cache import ( quant_to_nope_fp8_rope_bf16_pack_triton, @@ -34,15 +33,16 @@ from sglang.srt.model_executor.cuda_graph_config import ( ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_context import ForwardContext, forward_context +from sglang.srt.runtime_context import get_parallel from sglang.srt.server_args import set_global_server_args_for_scheduler from ..mock_server_args import make_mock_server_args # DSV4 backend pre-resolves attention TP at construction; pin to single-rank. -_dp_attention.get_attention_tp_size = lambda: 1 -_dp_attention.get_attention_tp_rank = lambda: 0 -_dp_attention.get_attention_cp_size = lambda: 1 -_dp_attention.get_attention_cp_rank = lambda: 0 +_parallel_override = get_parallel().override( + attn_tp_size=1, attn_tp_rank=0, attn_cp_size=1, attn_cp_rank=0 +) +_parallel_override.__enter__() # DSV4 hard-coded geometry. Do not change. DSV4_PAGE_SIZE = 256 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py index e5b9ed125..5a896abdb 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py @@ -18,6 +18,7 @@ from sglang.srt.model_executor.cuda_graph_config import ( from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.runtime_context import get_parallel from sglang.srt.server_args import set_global_server_args_for_scheduler from ..mock_server_args import make_mock_server_args @@ -78,7 +79,8 @@ DUAL_CHUNK_SPARSE_SUB_WINDOW_CONFIG = { # Unit tests run without distributed initialization. Sparse dual-chunk config # lookup should see the single-rank default. -_dual_chunk_backend.get_tensor_model_parallel_rank = lambda: 0 +_parallel_override = get_parallel().override(tp_rank=0) +_parallel_override.__enter__() @dataclass(frozen=True) diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py index 12a9c1270..e8ce45411 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py @@ -10,7 +10,6 @@ from sglang.srt.configs.mamba_utils import ( Mamba2StateShape, ) from sglang.srt.configs.model_config import AttentionArch -from sglang.srt.layers import dp_attention as _dp_attention from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( HybridLinearAttnBackend, @@ -30,10 +29,12 @@ from sglang.srt.model_executor.cuda_graph_config import ( from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.runtime_context import get_parallel from ..mock_server_args import make_mock_server_args -_dp_attention.get_attention_tp_size = lambda: 1 +_parallel_override = get_parallel().override(attn_tp_size=1) +_parallel_override.__enter__() DEFAULT_HEAD_K_DIM = 32 DEFAULT_HEAD_V_DIM = 32 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py index fed8cf940..ff6faa324 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py @@ -10,7 +10,6 @@ from sglang.srt.configs.mamba_utils import ( Mamba2StateDType, ) from sglang.srt.configs.model_config import AttentionArch -from sglang.srt.layers import dp_attention as _dp_attention from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( HybridLinearAttnBackend, @@ -30,10 +29,12 @@ from sglang.srt.model_executor.cuda_graph_config import ( from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.runtime_context import get_parallel from ..mock_server_args import make_mock_server_args -_dp_attention.get_attention_tp_size = lambda: 1 +_parallel_override = get_parallel().override(attn_tp_size=1) +_parallel_override.__enter__() DEFAULT_HEAD_K_DIM = 32 DEFAULT_HEAD_V_DIM = 32 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py index 66fd28037..23e729ace 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py @@ -10,7 +10,6 @@ from sglang.srt.configs.mamba_utils import ( Mamba2StateShape, ) from sglang.srt.configs.model_config import AttentionArch -from sglang.srt.layers import dp_attention as _dp_attention from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS from sglang.srt.layers.attention.linear.lightning_backend import ( LightningAttentionBackend, @@ -29,11 +28,12 @@ from sglang.srt.model_executor.cuda_graph_config import ( from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.runtime_context import get_parallel from ..mock_server_args import make_mock_server_args -_dp_attention.get_attention_tp_size = lambda: 1 -_dp_attention.get_attention_tp_rank = lambda: 0 +_parallel_override = get_parallel().override(attn_tp_size=1, attn_tp_rank=0) +_parallel_override.__enter__() # seg_la kernel constraints (see seg_la.py:683-694): # - decode (`seg_la_d_kernel`): K_SPLIT_DIM = 128, so head_dim must be >= 128 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py index 08695bbf3..11ff4e9c7 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py @@ -6,22 +6,13 @@ import torch.nn.functional as F from torch import nn # Patch TP world size / rank before importing modules that read them at __init__. -import sglang.srt.distributed as _distributed -import sglang.srt.layers.attention.mamba.mamba as _mamba_mod -import sglang.srt.layers.attention.mamba.mixer2_rms_norm_gated as _norm_mod import sglang.srt.layers.linear as _linear_mod -from sglang.srt.layers import dp_attention as _dp_attention +from sglang.srt.runtime_context import get_parallel -_distributed.get_tensor_model_parallel_world_size = lambda: 1 -_distributed.get_tensor_model_parallel_rank = lambda: 0 -_mamba_mod.get_tensor_model_parallel_world_size = lambda: 1 -_mamba_mod.get_tensor_model_parallel_rank = lambda: 0 -_norm_mod.get_tensor_model_parallel_world_size = lambda: 1 -_norm_mod.get_tensor_model_parallel_rank = lambda: 0 -_linear_mod.get_tensor_model_parallel_world_size = lambda: 1 -_linear_mod.get_tensor_model_parallel_rank = lambda: 0 -_dp_attention.get_attention_tp_size = lambda: 1 -_dp_attention.get_attention_tp_rank = lambda: 0 +_parallel_override = get_parallel().override( + tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0 +) +_parallel_override.__enter__() # RowParallelLinear.forward calls get_tp_group() to manage symmetric memory. # Provide a stub group with world_size=1 so use_symmetric_memory short-circuits. diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py index fa603f492..890be7273 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py @@ -7,7 +7,6 @@ import torch.nn.functional as F from torch import nn from sglang.srt.configs.model_config import AttentionArch -from sglang.srt.layers import dp_attention as _dp_attention from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool, ReqToTokenPool @@ -23,11 +22,13 @@ from sglang.srt.model_executor.forward_context import ( get_token_to_kv_pool, ) from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.runtime_context import get_parallel from sglang.srt.server_args import set_global_server_args_for_scheduler from ..mock_server_args import make_mock_server_args -_dp_attention.get_attention_tp_size = lambda: 1 +_parallel_override = get_parallel().override(attn_tp_size=1) +_parallel_override.__enter__() DEFAULT_HIDDEN_SIZE = 64 DEFAULT_KV_LORA_RANK = 32 diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py index e79205776..b12b33203 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py @@ -14,6 +14,7 @@ from sglang.srt.model_executor.forward_batch_info import ( ForwardBatch, ) from sglang.srt.model_executor.forward_context import ForwardContext, forward_context +from sglang.srt.runtime_context import get_parallel from sglang.srt.speculative.draft_utils import DraftBackendFactory from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import ( EAGLEDraftExtendCudaGraphRunner, @@ -560,10 +561,7 @@ def _capture_eagle_draft_extend_graph_runner( "sglang.srt.model_executor.runner.decode_cuda_graph_runner.get_available_gpu_memory", lambda *args, **kwargs: 0.0, ), - patch( - "sglang.srt.model_executor.runner.base_cuda_graph_runner.get_attention_cp_size", - lambda: 1, - ), + get_parallel().override(attn_cp_size=1), ): _reset_cuda_graph_test_buffers() return EAGLEDraftExtendCudaGraphRunner( diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py index 6c77f63cb..af2dca92a 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py @@ -16,6 +16,7 @@ from sglang.srt.model_executor.forward_batch_info import ( ) from sglang.srt.model_executor.input_buffers import _forward_input_buffer_pool from sglang.srt.model_executor.runner import set_global_graph_memory_pool +from sglang.srt.runtime_context import get_parallel from sglang.srt.server_args import set_global_server_args_for_scheduler from sglang.srt.speculative.draft_utils import DraftBackendFactory from sglang.srt.speculative.eagle_draft_cuda_graph_runner import ( @@ -450,10 +451,7 @@ def _capture_eagle_draft_graph_runner( "sglang.srt.model_executor.runner.decode_cuda_graph_runner.get_available_gpu_memory", lambda *args, **kwargs: 0.0, ), - patch( - "sglang.srt.model_executor.runner.base_cuda_graph_runner.get_attention_cp_size", - lambda: 1, - ), + get_parallel().override(attn_cp_size=1), ): _reset_cuda_graph_test_buffers() return EAGLEDraftCudaGraphRunner( @@ -479,10 +477,7 @@ def _capture_frozen_kv_mtp_graph_runner( "sglang.srt.model_executor.runner.decode_cuda_graph_runner.get_available_gpu_memory", lambda *args, **kwargs: 0.0, ), - patch( - "sglang.srt.model_executor.runner.base_cuda_graph_runner.get_attention_cp_size", - lambda: 1, - ), + get_parallel().override(attn_cp_size=1), ): _reset_cuda_graph_test_buffers() return FrozenKVMTPCudaGraphRunner(worker) diff --git a/test/manual/attention/test_trtllm_mla_backend.py b/test/manual/attention/test_trtllm_mla_backend.py index 5470b9053..7cd5abb4f 100755 --- a/test/manual/attention/test_trtllm_mla_backend.py +++ b/test/manual/attention/test_trtllm_mla_backend.py @@ -5,11 +5,10 @@ from types import SimpleNamespace import numpy as np import torch -from sglang.srt.layers import dp_attention as _dp_attn +from sglang.srt.runtime_context import get_parallel -# Patch DP-attention globals before importing backends -# TODO: change the interface of both trtllm_mla and flashinfer backends to take tp_size as an argument instead of patching -_dp_attn.get_attention_tp_size = lambda: 1 # TP size = 1 for unit test +_parallel_override = get_parallel().override(attn_tp_size=1) +_parallel_override.__enter__() from sglang.srt.configs.model_config import AttentionArch from sglang.srt.layers.attention.flashinfer_mla_backend import FlashInferMLAAttnBackend diff --git a/test/registered/kernels/test_dsa_indexer.py b/test/registered/kernels/test_dsa_indexer.py index e735030f3..c5cc81e47 100644 --- a/test/registered/kernels/test_dsa_indexer.py +++ b/test/registered/kernels/test_dsa_indexer.py @@ -5,11 +5,11 @@ from unittest.mock import MagicMock, patch import torch from sglang.srt.environ import envs -from sglang.srt.layers import dp_attention as _dp_attn +from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cuda_ci -# Patch DP-attention globals before importing backends -_dp_attn.get_attention_tp_size = lambda: 1 # TP size = 1 for unit test +_parallel_override = get_parallel().override(attn_tp_size=1) +_parallel_override.__enter__() from sglang.srt.configs.model_config import AttentionArch from sglang.srt.layers.attention.dsa.dsa_indexer import ( diff --git a/test/registered/layers/mamba/test_mamba2_mixer.py b/test/registered/layers/mamba/test_mamba2_mixer.py index d5a0c823d..4cd95c6dd 100644 --- a/test/registered/layers/mamba/test_mamba2_mixer.py +++ b/test/registered/layers/mamba/test_mamba2_mixer.py @@ -3,8 +3,6 @@ # Adapted from https://github.com/vllm-project/vllm/blob/2c58742dff8613a3bd7496f2008ce927e18d38d1/tests/kernels/mamba/test_mamba_mixer2.py -from unittest.mock import patch - import pytest import torch @@ -15,6 +13,7 @@ from sglang.srt.distributed.parallel_state import ( init_distributed_environment, initialize_model_parallel, ) +from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import get_device, get_device_count from sglang.test.ci.ci_register import register_cuda_ci @@ -109,10 +108,10 @@ def mixer2_gated_norm_tensor_parallel( gate_states = torch.randn(batch_size, seq_len, hidden_size) import sglang.srt.layers.attention.mamba.mixer2_rms_norm_gated as m2 - import sglang.srt.model_loader.weight_utils as wu - # Convenience: Avoid calling initialize_dp_attention - with patch.object(wu, "get_attention_tp_rank", return_value=local_rank): + # Force attn-TP rank through the context (the weight loader reads it via + # get_parallel().attn_tp_rank); avoids calling initialize_dp_attention. + with get_parallel().override(attn_tp_rank=local_rank): # create gated-norm with TP mixer = m2.Mixer2RMSNormGated( full_hidden_size=hidden_size, @@ -120,10 +119,8 @@ def mixer2_gated_norm_tensor_parallel( ) mixer.weight.weight_loader(mixer.weight, weight) - with ( - patch.object(m2, "get_tensor_model_parallel_world_size", return_value=1), - patch.object(m2, "get_tensor_model_parallel_rank", return_value=0), - ): + # m2 reads tp via get_parallel().tp_size/rank — force it through the context. + with get_parallel().override(tp_size=1, tp_rank=0): # create gated-norm without TP to compute reference mixer_single_gpu = m2.Mixer2RMSNormGated( full_hidden_size=hidden_size, diff --git a/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py b/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py index 179152e2f..a3cd88203 100644 --- a/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py +++ b/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py @@ -14,6 +14,7 @@ import torch import sglang.srt.batch_overlap.two_batch_overlap as tbo from sglang.srt.batch_overlap.two_batch_overlap import TboForwardBatchPreparer from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -37,7 +38,7 @@ def _make_target_verify_batch(bs: int) -> ForwardBatch: def _filter(batch: ForwardBatch, *, lo: int, hi: int) -> ForwardBatch: fake_args = SimpleNamespace(moe_dense_tp_size=None, attention_backend="fa3") - with patch.object(tbo, "get_attention_tp_size", lambda: 1), patch.object( + with get_parallel().override(attn_tp_size=1), patch.object( tbo, "get_global_server_args", lambda: fake_args ): return TboForwardBatchPreparer.filter_batch( diff --git a/test/registered/unit/eplb/test_deepep_waterfill_eplb.py b/test/registered/unit/eplb/test_deepep_waterfill_eplb.py index c193a4e0e..cb699eca4 100644 --- a/test/registered/unit/eplb/test_deepep_waterfill_eplb.py +++ b/test/registered/unit/eplb/test_deepep_waterfill_eplb.py @@ -14,6 +14,7 @@ from torch import nn from sglang.srt.layers.moe import topk as topk_module from sglang.srt.layers.moe.topk import TopKConfig from sglang.srt.models.deepseek_v2 import DeepseekV2MoE +from sglang.srt.runtime_context import get_parallel from sglang.test.test_utils import CustomTestCase @@ -74,10 +75,7 @@ class TestDeepEPWaterfillEPLB(CustomTestCase): patch.object(topk_module, "_is_cuda", True), patch.object(topk_module, "_use_aiter", False), patch.object(topk_module, "is_deepep_class_backend", return_value=True), - patch.object( - topk_module, "get_moe_expert_parallel_world_size", return_value=8 - ), - patch.object(topk_module, "get_moe_expert_parallel_rank", return_value=7), + get_parallel().override(moe_ep_size=8, moe_ep_rank=7), patch.object( topk_module, "_biased_grouped_topk_postprocess", diff --git a/test/registered/unit/layers/test_flashinfer_comm_fusion.py b/test/registered/unit/layers/test_flashinfer_comm_fusion.py index a0efc037e..87624be87 100644 --- a/test/registered/unit/layers/test_flashinfer_comm_fusion.py +++ b/test/registered/unit/layers/test_flashinfer_comm_fusion.py @@ -5,6 +5,7 @@ from unittest.mock import patch import torch from sglang.srt.layers import flashinfer_comm_fusion as fusion +from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=30, stage="base-c", runner_config="4-gpu-h100") @@ -149,11 +150,7 @@ class TestFlashInferCommFusion(unittest.TestCase): patch.object( fusion, "is_flashinfer_available", return_value=True ), - patch.object( - fusion, - "get_attn_tensor_model_parallel_world_size", - return_value=world_size, - ), + get_parallel().override(attn_tp_size=world_size), patch.object( fusion, "ensure_workspace_initialized", return_value=True ), diff --git a/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py b/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py index 66c65580f..13f2bd53c 100644 --- a/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py +++ b/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py @@ -904,12 +904,11 @@ class TestBuildDecodeRegistry(unittest.TestCase): ) def test_num_token_non_padded_gathered_dp_branch(self): - import unittest.mock as mock - from sglang.srt.model_executor import forward_batch_info as fbi from sglang.srt.model_executor.cuda_graph_buffer_registry import ( build_decode_registry, ) + from sglang.srt.runtime_context import get_parallel ntnp = torch.zeros(1, dtype=torch.int32) src = SimpleNamespace( @@ -926,9 +925,7 @@ class TestBuildDecodeRegistry(unittest.TestCase): ) # Gathered (DP) path: post_fill overwrites the FB copy with the local # count. Pin attn-TP (size=2, rank=0) so the result is deterministic. - with mock.patch.object( - fbi, "get_attention_tp_size", return_value=2 - ), mock.patch.object(fbi, "get_attention_tp_rank", return_value=0): + with get_parallel().override(attn_tp_size=2, attn_tp_rank=0): reg = build_decode_registry( device=torch.device("cpu"), max_bs=4, diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py new file mode 100644 index 000000000..5885e193e --- /dev/null +++ b/test/registered/unit/test_runtime_context.py @@ -0,0 +1,146 @@ +"""Unit tests for runtime_context: delegation, singletons, and override().""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +import unittest +from unittest.mock import patch + +from sglang.srt.runtime_context import ( + ParallelContext, + RuntimeContext, + get_context, + get_parallel, +) +from sglang.test.test_utils import CustomTestCase + +_PS = "sglang.srt.distributed.parallel_state" +_DP = "sglang.srt.layers.dp_attention" + +SIZE_RANK_DELEGATIONS = [ + ("world_size", f"{_PS}.get_world_size"), + ("world_rank", f"{_PS}.get_world_rank"), + ("tp_size", f"{_PS}.get_tensor_model_parallel_world_size"), + ("tp_rank", f"{_PS}.get_tensor_model_parallel_rank"), + ("pp_size", f"{_PS}.get_pipeline_model_parallel_world_size"), + ("pp_rank", f"{_PS}.get_pipeline_model_parallel_rank"), + ("moe_ep_size", f"{_PS}.get_moe_expert_parallel_world_size"), + ("moe_ep_rank", f"{_PS}.get_moe_expert_parallel_rank"), + ("moe_dp_size", f"{_PS}.get_moe_data_parallel_world_size"), + ("moe_dp_rank", f"{_PS}.get_moe_data_parallel_rank"), + ("moe_tp_size", f"{_PS}.get_moe_tensor_parallel_world_size"), + ("moe_tp_rank", f"{_PS}.get_moe_tensor_parallel_rank"), + ("attn_tp_size", f"{_PS}.get_attn_tensor_model_parallel_world_size"), + ("attn_tp_rank", f"{_PS}.get_attn_tensor_model_parallel_rank"), + ("attn_cp_size", f"{_PS}.get_attn_context_model_parallel_world_size"), + ("attn_cp_rank", f"{_PS}.get_attn_context_model_parallel_rank"), + ("attn_dp_size", f"{_DP}.get_attention_dp_size"), + ("attn_dp_rank", f"{_DP}.get_attention_dp_rank"), +] + +GROUP_DELEGATIONS = [ + ("world_group", f"{_PS}.get_world_group"), + ("tp_group", f"{_PS}.get_tp_group"), + ("pp_group", f"{_PS}.get_pp_group"), + ("moe_ep_group", f"{_PS}.get_moe_ep_group"), + ("moe_dp_group", f"{_PS}.get_moe_dp_group"), + ("moe_tp_group", f"{_PS}.get_moe_tp_group"), + ("attn_tp_group", f"{_PS}.get_attn_tp_group"), + ("attn_cp_group", f"{_PS}.get_attn_cp_group"), +] + + +class TestRuntimeContextSingletons(CustomTestCase): + def test_singletons(self): + self.assertIs(get_parallel(), get_parallel()) + self.assertIsInstance(get_parallel(), ParallelContext) + self.assertIsInstance(get_context(), RuntimeContext) + self.assertIs(get_context().parallel, get_parallel()) + + +class _IsolatedOverrides(CustomTestCase): + """Give each test a clean override map, restoring afterward only the overrides + installed outside it (e.g. by another test file sharing the process).""" + + def setUp(self): + super().setUp() + p = get_parallel() + self._saved_overrides = dict(p._overrides) + p._overrides.clear() + + def tearDown(self): + p = get_parallel() + p._overrides.clear() + p._overrides.update(self._saved_overrides) + super().tearDown() + + +class TestParallelDelegation(_IsolatedOverrides): + def test_size_rank_delegate_to_canonical_getters(self): + # Patch each getter to a distinct sentinel: a miswired attribute would read + # a different (unpatched) getter and fail. + for i, (attr, target) in enumerate(SIZE_RANK_DELEGATIONS): + sentinel = 1000 + i + with patch(target, return_value=sentinel): + self.assertEqual( + getattr(get_parallel(), attr), + sentinel, + msg=f"{attr} must delegate to {target}", + ) + + def test_groups_delegate_to_canonical_getters(self): + for attr, target in GROUP_DELEGATIONS: + sentinel = object() + with patch(target, return_value=sentinel): + self.assertIs( + getattr(get_parallel(), attr), + sentinel, + msg=f"{attr} must delegate to {target}", + ) + + def test_wrapper_holds_no_resolved_state(self): + # __slots__: no __dict__; the only instance state is the override hook. + self.assertFalse(hasattr(get_parallel(), "__dict__")) + # tp_group IS exposed: live delegation handles PD-multiplexing / the tp patch. + self.assertTrue(hasattr(ParallelContext, "tp_group")) + # local_attn_dp is intentionally not part of the wrapper surface. + self.assertFalse(hasattr(ParallelContext, "local_attn_dp_size")) + + +class TestParallelOverride(_IsolatedOverrides): + def test_override_takes_precedence(self): + p = get_parallel() + with p.override(tp_size=99, tp_rank=3, attn_dp_size=8): + self.assertEqual(p.tp_size, 99) + self.assertEqual(p.tp_rank, 3) + self.assertEqual(p.attn_dp_size, 8) + # same singleton: a fresh get_parallel() sees the override too + self.assertEqual(get_parallel().tp_size, 99) + self.assertEqual(p._overrides, {}) + + def test_override_can_force_groups(self): + sentinel = object() + with get_parallel().override(tp_group=sentinel): + self.assertIs(get_parallel().tp_group, sentinel) + + def test_override_nests_and_restores(self): + p = get_parallel() + with p.override(tp_size=2): + self.assertEqual(p.tp_size, 2) + with p.override(tp_size=4, pp_size=2): + self.assertEqual(p.tp_size, 4) + self.assertEqual(p.pp_size, 2) + self.assertEqual(p.tp_size, 2) + self.assertNotIn("pp_size", p._overrides) + + def test_override_unknown_key_raises_and_does_not_mutate(self): + p = get_parallel() + with self.assertRaises(ValueError): + with p.override(tp_sizee=1): # typo + pass + self.assertEqual(p._overrides, {}) + + +if __name__ == "__main__": + unittest.main()