chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity
Files
f8fca72635512095a9cf65eb8af1550767e107d2
sglang/python/sglang/srt/model_executor/model_runner.py
laoyao0822 f8fca72635 Expand prefill CP KV capacity by sharding persistent NSA KV 
Prefill CP previously replicated NSA/MLA persistent KV on every CP rank, so CP8 consumed eight copies of KV memory while exposing only one rank of logical cache capacity. This change splits logical KV locs from per-rank physical storage, shards MLA latent KV and NSA index K/scale by deterministic page ownership, and keeps existing NSA attention kernels working through a full-view runtime materialization layer.

Mooncake PD transfer now sends each prefill CP rank's owned physical pages with explicit logical page positions so non-CP decode can reconstruct full-layout KV. The implementation is guarded by an explicit server flag and startup checks, and the design documentation records the implemented scope, debug environment, and Phase 3 boundary.

Constraint: Phase 2 must preserve existing NSA attention/index kernels via runtime full-view materialization
Constraint: Decode side remains non-CP and receives full KV through Mooncake
Rejected: Shard-aware NSA attention in this change | belongs to Phase 3 because it requires distributed topk/softmax/output contracts
Rejected: Request-contiguous CP ownership | unstable under chunked prefill and tied to attention split mode
Confidence: medium
Scope-risk: broad
Directive: Do not enable round-robin CP shared KV without wiring runtime materialization/PD transfer contracts for that split mode
Directive: Keep SGLANG_DEBUG_CP_SHARED_KV disabled for perf measurements; it intentionally enables CUDA-syncing diagnostics
Tested: Remote py_compile for shared-KV touched Python files in g0034 container
Tested: Remote pytest selected cp_shared/shared_kv/nsa suite: 37 passed, 34 deselected
Not-tested: Full GLM5 multi-node throughput/regression run after final doc update
Not-tested: Phase 3 shard-aware runtime, round-robin CP mode, and non-Mooncake PD backends
2026-04-26 04:11:17 +08:00

3024 lines
119 KiB
Python
Raw Blame History

# 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.
# ==============================================================================
"""ModelRunner runs the forward passes of the models."""
from __future__ import annotations
import datetime
import gc
import inspect
import json
import logging
import os
import socket
import threading
import time
from collections import defaultdict
from dataclasses import dataclass
from typing import Callable, List, Optional, Tuple, Union
import torch
import torch.distributed as dist
from torch import nn
from sglang.jit_kernel.ngram_embedding import update_token_table
from sglang.srt.configs import (
BailingHybridConfig,
FalconH1Config,
GraniteMoeHybridConfig,
JetNemotronConfig,
JetVLMConfig,
KimiLinearConfig,
Lfm2Config,
Lfm2MoeConfig,
NemotronH_Nano_VL_V2_Config,
NemotronHConfig,
Qwen3_5Config,
Qwen3_5MoeConfig,
Qwen3NextConfig,
)
from sglang.srt.configs.device_config import DeviceConfig
from sglang.srt.configs.load_config import LoadConfig, LoadFormat
from sglang.srt.configs.model_config import AttentionArch, ModelConfig, ModelImpl
from sglang.srt.configs.update_config import adjust_config_with_unaligned_cpu_tp
from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS
from sglang.srt.debug_utils.dumper import dumper
from sglang.srt.debug_utils.tensor_dump_forward_hook import (
register_forward_hook_for_model,
)
from sglang.srt.distributed import (
get_default_distributed_backend,
get_pp_group,
get_tp_group,
get_world_group,
init_distributed_environment,
initialize_model_parallel,
set_custom_all_reduce,
set_mscclpp_all_reduce,
set_torch_symm_mem_all_reduce,
)
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.distributed.parallel_state import monkey_patch_vllm_parallel_state
from sglang.srt.elastic_ep.elastic_ep import ElasticEPStateManager
from sglang.srt.elastic_ep.expert_backup_client import ExpertBackupClient
from sglang.srt.environ import envs
from sglang.srt.eplb.eplb_manager import EPLBManager
from sglang.srt.eplb.expert_distribution import (
ExpertDistributionMetrics,
ExpertDistributionRecorder,
get_global_expert_distribution_recorder,
set_global_expert_distribution_recorder,
)
from sglang.srt.eplb.expert_location import (
ExpertLocationMetadata,
compute_initial_expert_location_metadata,
get_global_expert_location_metadata,
set_global_expert_location_metadata,
)
from sglang.srt.eplb.expert_location_updater import ExpertLocationUpdater
from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import NPUGraphRunner
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.attention.attention_registry import (
ATTENTION_BACKENDS,
attn_backend_wrapper,
)
from sglang.srt.layers.attention.nsa.utils import is_nsa_enable_prefill_cp
from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
get_attention_cp_rank,
get_attention_tp_group,
get_attention_tp_size,
initialize_dp_attention,
set_dp_buffer_len,
set_is_extend_in_batch,
)
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.moe.routed_experts_capturer import (
RoutedExpertsCapturer,
get_global_experts_capturer,
set_global_experts_capturer,
)
from sglang.srt.layers.pooler import EmbeddingPoolerOutput
from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype
from sglang.srt.layers.sampler import create_sampler
from sglang.srt.layers.torchao_utils import apply_torchao_config_to_model
from sglang.srt.lora.lora_manager import LoRAManager
from sglang.srt.lora.lora_registry import LoRARef
from sglang.srt.managers.schedule_batch import sanity_check_mm_pad_shift_value
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.model_executor.cpu_graph_runner import CPUGraphRunner
from sglang.srt.model_executor.cuda_graph_runner import (
CudaGraphRunner,
DecodeInputBuffers,
set_torch_compile_config,
)
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
ForwardBatch,
ForwardMode,
PPProxyTensors,
)
from sglang.srt.model_executor.hook_manager import register_forward_hooks
from sglang.srt.model_executor.model_runner_kv_cache_mixin import (
MemoryPoolConfig,
ModelRunnerKVCacheMixin,
)
from sglang.srt.model_executor.piecewise_cuda_graph_runner import (
PiecewiseCudaGraphRunner,
)
from sglang.srt.model_loader.loader import DefaultModelLoader, get_model_loader
from sglang.srt.model_loader.remote_instance_weight_loader_utils import (
RemoteInstanceWeightLoaderBackend,
register_memory_region,
trigger_init_weights_send_group_for_remote_instance_request,
)
from sglang.srt.model_loader.utils import set_default_torch_dtype
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
from sglang.srt.server_args import (
ServerArgs,
get_global_server_args,
set_global_server_args_for_scheduler,
)
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.utils import (
MultiprocessingSerializer,
cpu_has_amx_support,
dynamic_import,
empty_context,
enable_show_time_cost,
get_available_gpu_memory,
get_cpu_ids_by_node,
init_custom_process_group,
is_hip,
is_host_cpu_arm64,
is_npu,
log_info_on_rank0,
monkey_patch_p2p_access_check,
require_attn_tp_gather,
require_gathered_buffer,
require_mlp_tp_gather,
reserve_rope_cache_for_long_sequences,
set_cuda_arch,
slow_rank_detector,
)
from sglang.srt.utils.network import NetworkAddress, get_local_ip_auto
from sglang.srt.utils.nvtx_pytorch_hooks import PytHooks
from sglang.srt.utils.offloader import (
create_offloader_from_server_args,
get_offloader,
set_offloader,
)
from sglang.srt.utils.patch_torch import (
monkey_patch_torch_reductions,
register_sgl_tp_rank,
)
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
from sglang.srt.utils.weight_checker import WeightChecker
from sglang.srt.weight_sync.tensor_bucket import (
FlattenedTensorBucket,
FlattenedTensorMetadata,
)
_is_hip = is_hip()
_is_npu = is_npu()
_is_cpu_amx_available = cpu_has_amx_support()
_is_cpu_arm64 = is_host_cpu_arm64()
if _is_npu:
from sglang.srt.hardware_backend.npu.utils import init_npu_backend
init_npu_backend()
MLA_ATTENTION_BACKENDS = [
"aiter",
"flashinfer",
"fa3",
"fa4",
"triton",
"flashmla",
"cutlass_mla",
"trtllm_mla",
"ascend",
"nsa",
]
CHUNKED_PREFIX_CACHE_SUPPORTED_ATTENTION_BACKENDS = [
"flashinfer",
"fa3",
"fa4",
"flashmla",
"cutlass_mla",
"trtllm_mla",
]
TORCH_DTYPE_TO_KV_CACHE_STR = {
torch.float8_e4m3fn: "fp8_e4m3",
torch.float8_e4m3fnuz: "fp8_e4m3",
torch.float8_e5m2: "fp8_e5m2",
torch.bfloat16: "bf16",
}
def add_mla_attention_backend(backend_name):
if backend_name not in MLA_ATTENTION_BACKENDS:
MLA_ATTENTION_BACKENDS.append(backend_name)
logger.info(f"Added {backend_name} to MLA_ATTENTION_BACKENDS.")
def add_chunked_prefix_cache_attention_backend(backend_name):
if backend_name not in CHUNKED_PREFIX_CACHE_SUPPORTED_ATTENTION_BACKENDS:
CHUNKED_PREFIX_CACHE_SUPPORTED_ATTENTION_BACKENDS.append(backend_name)
logger.info(
f"Added {backend_name} to CHUNKED_PREFIX_CACHE_SUPPORTED_ATTENTION_BACKENDS."
)
# Detect stragger ranks in model loading
UNBALANCED_MODEL_LOADING_TIMEOUT_S = 480 # leave more time for post data processing
logger = logging.getLogger(__name__)
def resolve_language_model(model: nn.Module) -> nn.Module:
model_cls_name = model.__class__.__name__
if model_cls_name == "Qwen3OmniMoeForConditionalGeneration":
return model.thinker.model
return model.model
class RankZeroFilter(logging.Filter):
"""Filter that only allows INFO level logs from rank 0, but allows all other levels from any rank."""
def __init__(self, is_rank_zero):
super().__init__()
self.is_rank_zero = is_rank_zero
def filter(self, record):
if record.levelno == logging.INFO:
return self.is_rank_zero
return True
@dataclass
class ModelRunnerOutput:
logits_output: Union[LogitsProcessorOutput, PPProxyTensors]
can_run_graph: bool
expert_distribution_metrics: Optional[ExpertDistributionMetrics] = None
class ModelRunner(ModelRunnerKVCacheMixin):
"""ModelRunner runs the forward passes of the models."""
def __init__(
self,
model_config: ModelConfig,
mem_fraction_static: float,
gpu_id: int,
tp_rank: int,
tp_size: int,
moe_ep_rank: int,
moe_ep_size: int,
pp_rank: int,
pp_size: int,
nccl_port: int,
server_args: ServerArgs,
dp_rank: Optional[int] = None,
attn_cp_rank: Optional[int] = None,
moe_dp_rank: Optional[int] = None,
is_draft_worker: bool = False,
req_to_token_pool: Optional[ReqToTokenPool] = None,
token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator] = None,
memory_pool_config: Optional[MemoryPoolConfig] = None,
draft_model_idx: Optional[int] = None,
):
# Parse args
self.mem_fraction_static = mem_fraction_static
self.device = server_args.device
self.gpu_id = gpu_id
self.tp_rank = tp_rank
self.tp_size = tp_size
self.moe_ep_rank = moe_ep_rank
self.moe_ep_size = moe_ep_size
self.dp_size = server_args.dp_size if server_args.enable_dp_attention else 1
self.pp_rank = pp_rank
self.pp_size = pp_size
self.attn_cp_rank = attn_cp_rank
self.attn_cp_size = server_args.attn_cp_size
self.moe_dp_rank = moe_dp_rank
self.moe_dp_size = server_args.moe_dp_size
self.model_config = model_config
self.dist_port = nccl_port
self.server_args = server_args
self.is_draft_worker = is_draft_worker
self.memory_pool_config = memory_pool_config
self.is_generation = model_config.is_generation
self.is_multimodal = model_config.is_multimodal
self.is_multimodal_chunked_prefill_supported = (
model_config.is_multimodal_chunked_prefill_supported
)
self.spec_algorithm = SpeculativeAlgorithm.from_string(
server_args.speculative_algorithm
)
self.page_size = server_args.page_size
self.req_to_token_pool = req_to_token_pool
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
self.physical_max_total_num_tokens = None
self.uses_cp_shared_kv = server_args.enable_nsa_prefill_cp_shared_kv
self.cp_shared_kv_layout = None
self.is_hybrid_swa = model_config.is_hybrid_swa
self.is_hybrid_swa_compress = model_config.is_hybrid_swa_compress
self.use_mla_backend = self.model_config.attention_arch == AttentionArch.MLA
self.attention_chunk_size = model_config.attention_chunk_size
self.forward_pass_id = 0
self.init_new_workspace = False
self.draft_model_idx = draft_model_idx
self.enable_hisparse = server_args.enable_hisparse
self.remote_instance_transfer_engine = None
self.remote_instance_transfer_engine_session_id = ""
self.remote_instance_transfer_engine_weight_info = None
# auxiliary hidden capture mode. TODO: expose this to server args?
self.eagle_use_aux_hidden_state = False
if self.spec_algorithm.is_eagle3() and not self.is_draft_worker:
# load draft config
draft_model_config = ModelConfig.from_server_args(
server_args,
model_path=(server_args.speculative_draft_model_path),
model_revision=server_args.speculative_draft_model_revision,
is_draft_model=True,
)
self.eagle_use_aux_hidden_state = True
try:
# get the aux layer from draft model config
eagle_config = getattr(
draft_model_config.hf_config, "eagle_config", None
)
self.eagle_use_aux_hidden_state = eagle_config.get(
"use_aux_hidden_state", True
)
self.eagle_aux_hidden_state_layer_ids = eagle_config[
"eagle_aux_hidden_state_layer_ids"
]
except:
# if there is no aux layer, set to None
self.eagle_aux_hidden_state_layer_ids = None
# Apply the rank zero filter to logger
if server_args.show_time_cost:
enable_show_time_cost()
# Model-specific adjustment
self.model_specific_adjustment()
# Set the global server_args in the scheduler process
set_global_server_args_for_scheduler(server_args)
global_server_args = get_global_server_args()
# FIXME: hacky set `use_mla_backend`
global_server_args.use_mla_backend = self.use_mla_backend
# Init OpenMP threads binding for CPU
if self.device == "cpu":
self.init_threads_binding()
# Initialize MooncakeTransferEngine
self.init_shared_mooncake_transfer_engine()
# Get available memory before model loading
pre_model_load_memory = self.init_torch_distributed()
# Init forward stream for overlap schedule
self.forward_stream = torch.get_device_module(self.device).Stream()
# CPU offload
set_offloader(create_offloader_from_server_args(server_args, dp_rank=dp_rank))
self._weight_checker = WeightChecker(model_runner=self)
if envs.SGLANG_DETECT_SLOW_RANK.get():
slow_rank_detector.execute()
# Init mindspore running environment when model impl is "mindspore"
self.init_mindspore_runner()
# Update deep gemm configure
if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM:
deep_gemm_wrapper.update_deep_gemm_config(gpu_id, server_args)
# For hisparse (must be set before initialize() so CUDA graph capture can see it)
self.hisparse_coordinator = None
# Initialize the model runner
self.initialize(pre_model_load_memory)
self.check_quantized_moe_compatibility()
if self.is_multimodal:
sanity_check_mm_pad_shift_value(self.model_config.vocab_size)
# Temporary cached values
self.support_pp = (
"pp_proxy_tensors" in inspect.signature(self.model.forward).parameters
)
if self.pp_size > 1:
assert (
self.support_pp
), "Pipeline Parallel is not compatible with this model."
# For weight updates
self._model_update_group = {}
self._weights_send_group = {}
def init_mindspore_runner(self):
# Init the mindspore runner
# for now, there is only some communication initialization work
if self.server_args.model_impl.lower() == ModelImpl.MINDSPORE and _is_npu:
from sglang.srt.model_executor.mindspore_runner import init_ms_distributed
init_ms_distributed(
world_size=self.tp_size * self.pp_size,
rank=self.tp_size * self.pp_rank + self.tp_rank,
local_rank=self.gpu_id,
server_args=self.server_args,
port=self.dist_port,
)
def initialize(self, pre_model_load_memory: float):
server_args = self.server_args
self.memory_saver_adapter = TorchMemorySaverAdapter.create(
enable=self.server_args.enable_memory_saver
)
if self.server_args.remote_instance_weight_loader_use_transfer_engine():
self.remote_instance_init_transfer_engine()
if not self.is_draft_worker:
set_global_expert_location_metadata(
compute_initial_expert_location_metadata(
server_args=server_args,
model_config=self.model_config,
moe_ep_rank=self.moe_ep_rank,
)
)
if self.tp_rank == 0 and envs.SGLANG_LOG_EXPERT_LOCATION_METADATA.get():
logger.info(
f"Initial expert_location_metadata: {get_global_expert_location_metadata()}"
)
set_global_expert_distribution_recorder(
ExpertDistributionRecorder.init_new(
server_args,
get_global_expert_location_metadata(),
rank=self.tp_rank,
)
)
# Expert parallelism
self.eplb_manager = (
EPLBManager(self)
if self.server_args.enable_eplb and (not self.is_draft_worker)
else None
)
self.expert_location_updater = ExpertLocationUpdater()
(
ElasticEPStateManager.init(self.server_args)
if self.server_args.elastic_ep_backend
else None
)
# Load the model
self.sampler = create_sampler()
self.load_model()
# Load the expert backup client
self.expert_backup_client = (
ExpertBackupClient(self.server_args, self)
if (
self.server_args.enable_elastic_expert_backup
and self.server_args.elastic_ep_backend is not None
)
else None
)
if (
self.server_args.remote_instance_weight_loader_use_transfer_engine()
and self.remote_instance_transfer_engine is not None
and self.remote_instance_transfer_engine_weight_info is None
):
self.remote_instance_transfer_engine_weight_info = register_memory_region(
self.model, self.remote_instance_transfer_engine
)
# For MTP models like DeepSeek-V3 or GLM-4.5, the MTP layer(s) are used separately as draft
# models for speculative decoding. In those cases, `num_nextn_predict_layers` is used to
# determine the number of layers.
model_has_mtp_layers = self.model_config.num_nextn_predict_layers is not None
model_num_layers = (
self.model_config.num_nextn_predict_layers
if self.is_draft_worker and model_has_mtp_layers
else max(
self.model_config.num_hidden_layers,
self.model_config.num_attention_layers,
)
)
if self.model_config.hf_config.architectures[0] == "MiMoV2MTP":
model_num_layers = 1
elif self.model_config.hf_config.architectures[0] == "Step3p5MTP":
model_num_layers = 1
self.start_layer = getattr(self.model, "start_layer", 0)
self.end_layer = getattr(self.model, "end_layer", model_num_layers)
self.num_effective_layers = self.end_layer - self.start_layer
# For LoopCoder models, each loop has its own layer_id, so we need to multiply by loop_num
loop_num = getattr(self.model_config.hf_config, "loop_num", 1)
if loop_num > 1:
self.num_effective_layers = self.num_effective_layers * loop_num
assert (
(not model_has_mtp_layers)
or (self.spec_algorithm.is_none())
or (
(not self.spec_algorithm.is_none())
and (self.num_effective_layers == model_num_layers)
)
), "PP is not compatible with MTP models."
# Consider PP, so use start_layer and end_layer.
full_attention_layer_ids = [
layer_idx
for layer_idx in range(self.start_layer, self.end_layer + 1)
if hasattr(self.model_config, "full_attention_layer_ids")
and layer_idx in self.model_config.full_attention_layer_ids
]
swa_attention_layer_ids = [
layer_idx
for layer_idx in range(self.start_layer, self.end_layer + 1)
if hasattr(self.model_config, "swa_attention_layer_ids")
and layer_idx in self.model_config.swa_attention_layer_ids
]
# Update back to model_config.
self.model_config.swa_attention_layer_ids = swa_attention_layer_ids
self.model_config.full_attention_layer_ids = full_attention_layer_ids
# Apply torchao quantization
torchao_applied = getattr(self.model, "torchao_applied", False)
# In layered loading, torchao may have been applied
if not torchao_applied:
apply_torchao_config_to_model(
self.model, get_global_server_args().torchao_config
)
# Apply torch TP if the model supports it
supports_torch_tp = getattr(self.model, "supports_torch_tp", False)
if self.tp_size > 1 and supports_torch_tp:
self.apply_torch_tp()
# Init lora
if server_args.enable_lora:
self.init_lora_manager()
# Init Double Sparsity
if server_args.enable_double_sparsity:
if server_args.ds_heavy_channel_type is None:
raise ValueError(
"Please specify the heavy channel type for double sparsity optimization."
)
self.init_double_sparsity_channel_config(server_args.ds_heavy_channel_type)
# Enable batch invariant mode
if server_args.enable_deterministic_inference:
from sglang.srt.batch_invariant_ops import enable_batch_invariant_mode
enable_batch_invariant_mode()
# Deduce KV cache dtype
self.configure_kv_cache_dtype()
# Init memory pool and attention backends
self.init_memory_pool(pre_model_load_memory)
# Init ngram embedding token table
self.maybe_init_ngram_embedding()
# Init hisparse coordinator (must happen before CUDA graph capture)
if self.enable_hisparse:
from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator
from sglang.srt.mem_cache.sparsity import parse_hisparse_config
hisparse_cfg = parse_hisparse_config(self.server_args)
self.hisparse_coordinator = HiSparseCoordinator(
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
top_k=hisparse_cfg.top_k,
device_buffer_size=hisparse_cfg.device_buffer_size,
device=self.device,
tp_group=(
self.attention_tp_group.cpu_group
if self.server_args.enable_dp_attention
else self.tp_group.cpu_group
),
host_to_device_ratio=hisparse_cfg.host_to_device_ratio,
)
# Init routed experts capturer
self.init_routed_experts_capturer()
if self.device == "cuda" or self.device == "musa":
self.init_cublas()
self.init_attention_backend()
self.kernel_warmup()
self.init_device_graphs()
self.init_deep_gemm()
elif self.device in ["npu", "cpu"]:
self.init_attention_backend()
self.init_device_graphs()
else:
self.graph_runner = None
self.graph_mem_usage = 0
self.init_attention_backend()
if server_args.forward_hooks:
register_forward_hooks(self.model, server_args.forward_hooks)
if self.eagle_use_aux_hidden_state:
self.model.set_eagle3_layers_to_capture(
self.eagle_aux_hidden_state_layer_ids
)
# Initialize piecewise CUDA graph
self.init_piecewise_cuda_graphs()
self.prealloc_symmetric_memory_pool()
def init_routed_experts_capturer(self):
if not self.server_args.disable_shared_experts_fusion and hasattr(
self.model, "num_fused_shared_experts"
):
num_fused_shared_experts = self.model.num_fused_shared_experts
else:
num_fused_shared_experts = 0
set_global_experts_capturer(
RoutedExpertsCapturer.create(
enable=get_global_server_args().enable_return_routed_experts,
model_config=self.model_config,
num_fused_shared_experts=num_fused_shared_experts,
num_tokens=self.max_total_num_tokens + self.page_size,
max_running_requests=self.max_running_requests,
device=self.device,
)
)
def remote_instance_init_transfer_engine(self):
try:
from mooncake.engine import TransferEngine
except ImportError as e:
logger.warning(
"Please install mooncake for using remote instance transfer engine: pip install mooncake"
)
return
self.remote_instance_transfer_engine = TransferEngine()
local_ip = get_local_ip_auto()
self.remote_instance_transfer_engine.initialize(
local_ip, "P2PHANDSHAKE", "rdma", envs.MOONCAKE_DEVICE.get()
)
self.remote_instance_transfer_engine_session_id = NetworkAddress(
local_ip, self.remote_instance_transfer_engine.get_rpc_port()
).to_host_port_str()
def _publish_modelexpress_metadata(self):
"""Publish TransferEngine metadata to ModelExpress server (seed mode)."""
try:
from modelexpress import p2p_pb2
from modelexpress.client import MxClient
except ImportError as exc:
raise ImportError(
"ModelExpress support requires the 'modelexpress' package. "
"Install it with: pip install modelexpress"
) from exc
model_name = (
self.server_args.modelexpress_model_name or self.server_args.model_path
)
mx_url = self.server_args.modelexpress_url
session_id = self.remote_instance_transfer_engine_session_id
weight_info = self.remote_instance_transfer_engine_weight_info
if not session_id or weight_info is None:
logger.warning(
"ModelExpress source: skipping publish -- "
"TransferEngine not initialized or no weight info"
)
return
# Build tensor descriptors from weight_info dict
tensors = []
for name, (addr, numel, element_size) in weight_info.items():
tensors.append(
p2p_pb2.TensorDescriptor(
name=name,
addr=addr,
size=numel * element_size,
device_id=self.gpu_id,
)
)
worker = p2p_pb2.WorkerMetadata(
worker_rank=self.tp_rank,
transfer_engine_session_id=session_id,
tensors=tensors,
)
mx_client = MxClient(server_url=mx_url)
try:
logger.info(
"ModelExpress source: publishing metadata for model=%s, "
"tp_rank=%d, session=%s, %d tensors",
model_name,
self.tp_rank,
session_id,
len(tensors),
)
mx_client.publish_metadata(model_name, [worker])
mx_client.publish_ready(
model_name,
worker_id=self.tp_rank,
session_id=mx_client.session_id,
metadata_hash="",
)
logger.info(
"ModelExpress source: published ready for model=%s, tp_rank=%d",
model_name,
self.tp_rank,
)
finally:
mx_client.close()
def model_specific_adjustment(self):
server_args = self.server_args
if server_args.enable_double_sparsity:
logger.info(
"Double sparsity optimization is turned on. Use triton backend without CUDA graph."
)
server_args.attention_backend = "triton"
server_args.disable_cuda_graph = True
if self.is_multimodal:
if not self.is_multimodal_chunked_prefill_supported:
server_args.chunked_prefill_size = -1
logger.info(
f"Automatically turn off --chunked-prefill-size as it is not supported for "
f"{self.model_config.hf_config.model_type}"
)
if (
not self.use_mla_backend
or server_args.attention_backend
not in CHUNKED_PREFIX_CACHE_SUPPORTED_ATTENTION_BACKENDS
):
server_args.disable_chunked_prefix_cache = True
if not server_args.disable_chunked_prefix_cache:
log_info_on_rank0(logger, "Chunked prefix cache is turned on.")
def check_quantized_moe_compatibility(self):
if (
quantization_config := getattr(
self.model_config.hf_config, "quantization_config", None
)
) is not None and (
weight_block_size := quantization_config.get("weight_block_size", None)
) is not None:
weight_block_size_n = weight_block_size[0]
if self.tp_size % self.moe_ep_size != 0:
raise ValueError(
f"tp_size {self.tp_size} must be divisible by ep_size {self.moe_ep_size}"
)
moe_tp_size = self.tp_size // self.moe_ep_size // self.moe_dp_size
moe_intermediate_size = getattr(
self.model_config.hf_text_config, "moe_intermediate_size", None
)
if moe_intermediate_size is None:
return
if moe_intermediate_size % moe_tp_size != 0:
raise ValueError(
f"moe_intermediate_size {moe_intermediate_size} must be divisible by moe_tp_size ({moe_tp_size}) which is tp_size ({self.tp_size}) divided by moe_ep_size ({self.moe_ep_size})."
)
if (moe_intermediate_size // moe_tp_size) % weight_block_size_n != 0:
raise ValueError(
f"For quantized MoE models, please make sure ({moe_intermediate_size=} / {moe_tp_size=}) % {weight_block_size_n=} == 0 "
f"where moe_tp_size is equal to tp_size ({self.tp_size}) divided by ep_size ({self.moe_ep_size}). "
f"You can fix this by setting arguments `--tp` and `--ep` correctly."
)
def init_torch_distributed(self):
tic = time.perf_counter()
logger.info("Init torch distributed begin.")
try:
torch.get_device_module(self.device).set_device(self.gpu_id)
except Exception:
logger.warning(
f"Context: {self.device=} {self.gpu_id=} {os.environ.get('CUDA_VISIBLE_DEVICES')=} {self.tp_rank=} {self.tp_size=}"
)
raise
backend = get_default_distributed_backend(self.device)
if self.device == "cuda" and self.server_args.elastic_ep_backend == "mooncake":
backend = "mooncake"
if self.server_args.mooncake_ib_device:
mooncake_ib_device = self.server_args.mooncake_ib_device.split(",")
try:
from mooncake import ep as mooncake_ep
mooncake_ep.set_device_filter(mooncake_ib_device)
except:
pass # A warning will be raised in `init_distributed_environment`
before_avail_memory = get_available_gpu_memory(self.device, self.gpu_id)
if not self.server_args.enable_p2p_check:
monkey_patch_p2p_access_check()
# Allow external orchestrators (e.g. trainpi) to override the distributed
# init method. When set to "env://", torch uses MASTER_ADDR/MASTER_PORT
# env-vars and an externally-created TCPStore, completely avoiding port
# conflicts with intra-host collocation.
dist_init_method_override = envs.SGLANG_DISTRIBUTED_INIT_METHOD_OVERRIDE.get()
if dist_init_method_override:
dist_init_method = dist_init_method_override
elif self.server_args.dist_init_addr:
na = NetworkAddress.parse(self.server_args.dist_init_addr)
dist_init_method = na.to_tcp()
else:
dist_init_method = NetworkAddress(
self.server_args.host or "127.0.0.1", self.dist_port
).to_tcp()
set_custom_all_reduce(not self.server_args.disable_custom_all_reduce)
set_mscclpp_all_reduce(self.server_args.enable_mscclpp)
set_torch_symm_mem_all_reduce(self.server_args.enable_torch_symm_mem)
if not self.is_draft_worker:
if self.device == "cpu":
if _is_cpu_amx_available or _is_cpu_arm64:
# Bind OpenMP threads to CPU cores
torch.ops.sgl_kernel.init_cpu_threads_env(self.local_omp_cpuid)
# Set local size to hint SGLang to use shared memory based AllReduce
os.environ["LOCAL_SIZE"] = str(self.tp_size)
torch.ops.sgl_kernel.initialize(self.tp_size, self.tp_rank)
@torch.library.register_fake("sgl_kernel::shm_allgather")
def _(data, dim):
return torch.cat([data] * self.tp_size, dim=dim)
else:
logger.warning(
"init_cpu_threads_env and shared memory based AllReduce is disabled, only intel amx backend and arm64 are supported"
)
# Only initialize the distributed environment on the target model worker.
init_distributed_environment(
backend=backend,
world_size=self.tp_size * self.pp_size,
rank=self.tp_size * self.pp_rank + self.tp_rank,
local_rank=self.gpu_id,
distributed_init_method=dist_init_method,
timeout=self.server_args.dist_timeout,
moe_a2a_backend=self.server_args.moe_a2a_backend,
)
initialize_model_parallel(
tensor_model_parallel_size=self.tp_size,
attention_data_parallel_size=self.dp_size,
pipeline_model_parallel_size=self.pp_size,
expert_model_parallel_size=self.moe_ep_size,
attention_context_model_parallel_size=self.attn_cp_size,
moe_data_model_parallel_size=self.moe_dp_size,
duplicate_tp_group=self.server_args.enable_pdmux,
)
initialize_dp_attention(
server_args=self.server_args,
model_config=self.model_config,
)
if self.uses_cp_shared_kv:
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
self.cp_shared_kv_layout = CpSharedKVLayout(
page_size=self.page_size,
cp_size=self.server_args.attn_cp_size,
cp_rank=get_attention_cp_rank(),
)
if is_npu():
register_sgl_tp_rank(self.gpu_id)
# Pre-warm NCCL/RCCL to eliminate cold-start latency in first request
# Controlled by --pre-warm-nccl flag (default: enabled on AMD GPUs)
if self.server_args.pre_warm_nccl and (
self.tp_size > 1 or self.pp_size > 1 or self.moe_ep_size > 1
):
warmup_start = time.perf_counter()
tp_group_handle = get_tp_group().device_group
# Single warmup all_reduce to initialize NCCL/RCCL communicator
warmup_tensor = torch.zeros(1, device=torch.cuda.current_device())
dist.all_reduce(warmup_tensor, group=tp_group_handle)
torch.cuda.synchronize()
warmup_elapsed = time.perf_counter() - warmup_start
logger.info(
f"NCCL/RCCL warmup completed in {warmup_elapsed:.3f}s "
f"(tp_size={self.tp_size}, pp_size={self.pp_size}, ep_size={self.moe_ep_size})"
)
pre_model_load_memory = get_available_gpu_memory(
self.device,
self.gpu_id,
distributed=get_world_group().world_size > 1,
cpu_group=get_world_group().cpu_group,
)
self.tp_group = get_tp_group()
self.pp_group = get_pp_group()
self.attention_tp_group = get_attention_tp_group()
# Check memory for tensor parallelism
local_gpu_memory = get_available_gpu_memory(self.device, self.gpu_id)
if self.tp_size > 1 and not self.is_draft_worker:
if pre_model_load_memory < local_gpu_memory * 0.9:
msg = "The memory capacity is unbalanced. Some GPUs may be occupied by other processes. "
msg += f"{pre_model_load_memory=}, {local_gpu_memory=}, {local_gpu_memory * 0.9=}"
if envs.SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK.get():
raise RuntimeError(msg)
else:
logger.warning(msg)
logger.info(
f"Init torch distributed ends. elapsed={time.perf_counter() - tic:.2f} s, "
f"mem usage={(before_avail_memory - local_gpu_memory):.2f} GB"
)
return pre_model_load_memory
def init_shared_mooncake_transfer_engine(self):
"""
Need MooncakeTransferEngine when:
1) PD disaggregation uses mooncake for KV transfer (prefill/decode)
2) HiCache uses mooncake storage backend
3) Encoder disaggregation uses mooncake
"""
use_mooncake_te = (
(
self.server_args.disaggregation_mode != "null"
and self.server_args.disaggregation_transfer_backend == "mooncake"
)
or (
self.server_args.enable_hierarchical_cache
and self.server_args.hicache_storage_backend == "mooncake"
and envs.SGLANG_HICACHE_MOONCAKE_REUSE_TE.get()
)
or (
self.server_args.encoder_only
and self.server_args.encoder_transfer_backend == "mooncake"
)
or (
self.server_args.language_only
and self.server_args.encoder_transfer_backend == "mooncake"
)
or (
self.server_args.enable_elastic_expert_backup
and self.server_args.elastic_ep_backend is not None
)
)
if use_mooncake_te:
from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
init_mooncake_transfer_engine,
)
init_mooncake_transfer_engine(
hostname=get_local_ip_auto(),
gpu_id=self.gpu_id,
ib_device=(
self.server_args.disaggregation_ib_device
or self.server_args.mooncake_ib_device
),
)
def load_model(self):
tic_total = time.perf_counter()
before_avail_memory = get_available_gpu_memory(self.device, self.gpu_id)
logger.info(
f"Load weight begin. avail mem={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB"
)
# This can reduce thread conflicts and speed up weight loading.
if self.device != "cpu":
torch.set_num_threads(1)
if self.device == "cuda":
if torch.cuda.get_device_capability()[0] < 8:
logger.info(
"Compute capability below sm80. Use float16 due to lack of bfloat16 support."
)
self.server_args.dtype = "float16"
self.model_config.dtype = torch.float16
if torch.cuda.get_device_capability()[1] < 5:
raise RuntimeError("SGLang only supports sm75 and above.")
set_cuda_arch()
# Prepare the model config
from sglang.srt.configs.modelopt_config import ModelOptConfig
modelopt_config = ModelOptConfig(
quant=self.server_args.modelopt_quant,
checkpoint_restore_path=self.server_args.modelopt_checkpoint_restore_path,
checkpoint_save_path=self.server_args.modelopt_checkpoint_save_path,
export_path=self.server_args.modelopt_export_path,
quantize_and_serve=self.server_args.quantize_and_serve,
)
self.load_config = LoadConfig(
load_format=self.server_args.load_format,
download_dir=self.server_args.download_dir,
model_loader_extra_config=self.server_args.model_loader_extra_config,
tp_rank=self.tp_rank,
remote_instance_weight_loader_seed_instance_ip=self.server_args.remote_instance_weight_loader_seed_instance_ip,
remote_instance_weight_loader_seed_instance_service_port=self.server_args.remote_instance_weight_loader_seed_instance_service_port,
remote_instance_weight_loader_send_weights_group_ports=self.server_args.remote_instance_weight_loader_send_weights_group_ports,
remote_instance_weight_loader_backend=self.server_args.remote_instance_weight_loader_backend,
remote_instance_weight_loader_transfer_engine=self.remote_instance_transfer_engine,
modelexpress_url=self.server_args.modelexpress_url,
modelexpress_model_name=self.server_args.modelexpress_model_name
or self.server_args.model_path,
modelopt_config=modelopt_config,
rl_quant_profile=self.server_args.rl_quant_profile,
draft_model_idx=self.draft_model_idx,
)
if self.device == "cpu":
self.model_config = adjust_config_with_unaligned_cpu_tp(
self.model_config, self.load_config, self.tp_size
)
if (
self.server_args.load_format == LoadFormat.REMOTE_INSTANCE
and self.server_args.remote_instance_weight_loader_backend
== RemoteInstanceWeightLoaderBackend.NCCL
):
if self.tp_rank == 0:
instance_ip = NetworkAddress.resolve_host(socket.gethostname())
t = threading.Thread(
target=trigger_init_weights_send_group_for_remote_instance_request,
args=(
self.server_args.remote_instance_weight_loader_seed_instance_ip,
self.server_args.remote_instance_weight_loader_seed_instance_service_port,
self.server_args.remote_instance_weight_loader_send_weights_group_ports,
instance_ip,
),
)
t.start()
# Load the model
# Remove monkey_patch when linear.py quant remove dependencies with vllm
monkey_patch_vllm_parallel_state()
enable_cpu_backup = self.server_args.enable_weights_cpu_backup or (
self.is_draft_worker and self.server_args.enable_draft_weights_cpu_backup
)
with self.memory_saver_adapter.region(
GPU_MEMORY_TYPE_WEIGHTS,
enable_cpu_backup=enable_cpu_backup,
):
self.loader = get_model_loader(
load_config=self.load_config,
model_config=self.model_config,
)
self.model = self.loader.load_model(
model_config=self.model_config,
device_config=DeviceConfig(self.device, self.gpu_id),
)
if hasattr(self.loader, "remote_instance_transfer_engine_weight_info"):
self.remote_instance_transfer_engine_weight_info = (
self.loader.remote_instance_transfer_engine_weight_info
)
monkey_patch_vllm_parallel_state(reverse=True)
# Publish metadata to ModelExpress if running as seed source
if self.server_args.modelexpress_source:
# Seed loads via DefaultModelLoader (load_format=auto), which doesn't
# call register_memory_region(). Do it here so weight_info is populated.
if (
self.remote_instance_transfer_engine_weight_info is None
and self.remote_instance_transfer_engine is not None
):
from sglang.srt.model_loader.remote_instance_weight_loader_utils import (
register_memory_region,
)
self.remote_instance_transfer_engine_weight_info = (
register_memory_region(
self.model, self.remote_instance_transfer_engine
)
)
self._publish_modelexpress_metadata()
get_offloader().post_init()
# Register model for layerwise NVTX profiling if enabled
if self.server_args.enable_layerwise_nvtx_marker:
self.pyt_hooks = PytHooks()
self.pyt_hooks.register_hooks(self.model, module_prefix="model")
if self.server_args.kv_cache_dtype == "fp8_e4m3":
if self.server_args.quantization_param_path is not None:
if callable(getattr(self.model, "load_kv_cache_scales", None)):
self.model.load_kv_cache_scales(
self.server_args.quantization_param_path
)
logger.info(
"Loaded KV cache scaling factors from %s",
self.server_args.quantization_param_path,
)
else:
raise RuntimeError(
"Using FP8 KV cache and scaling factors provided but "
"model %s does not support loading scaling factors.",
self.model.__class__,
)
else:
logger.warning(
"Using FP8 KV cache but no scaling factors "
"provided. Defaulting to scaling factors of 1.0. "
"This may lead to less accurate results!"
)
# Parse other args
self.sliding_window_size = None
if hasattr(self.model, "get_attention_sliding_window_size"):
self.sliding_window_size = self.model.get_attention_sliding_window_size()
elif (
self.model_config.is_hybrid_swa
and self.model_config.sliding_window_size is not None
):
# sliding window field in model config may have different meaning for different kinds of models (e.g., dllm), here we only consider the sliding window in SWA model
self.sliding_window_size = self.model_config.sliding_window_size
elif self.model_config.attention_chunk_size is not None:
self.sliding_window_size = self.model_config.attention_chunk_size
logger.info(
f"Setting sliding_window_size to be attention_chunk_size: {self.sliding_window_size}"
)
self.dtype = self.model_config.dtype
after_avail_memory = get_available_gpu_memory(self.device, self.gpu_id)
self.weight_load_mem_usage = before_avail_memory - after_avail_memory
# Get quantization config from ModelConfig
# This handles both config.json (standard) and hf_quant_config.json (ModelOpt)
quant_str = self.model_config.get_quantization_config_log_str()
logger.info(
f"Load weight end. "
f"elapsed={time.perf_counter() - tic_total:.2f} s, "
f"type={type(self.model).__name__}, "
f"{quant_str + ', ' if quant_str else ''}"
f"avail mem={after_avail_memory:.2f} GB, "
f"mem usage={self.weight_load_mem_usage:.2f} GB."
)
if self.server_args.debug_tensor_dump_output_folder is not None:
register_forward_hook_for_model(
self.model,
self.server_args.debug_tensor_dump_output_folder,
self.server_args.debug_tensor_dump_layers,
self.tp_size,
self.tp_rank,
self.pp_rank,
)
if dumper.may_enable:
dumper.apply_source_patches()
dumper.register_non_intrusive_dumper(self.model)
# Pre-expand RoPE cache before CUDA Graph capture
reserve_rope_cache_for_long_sequences(
self.model,
self.server_args,
self.model_config,
logger,
)
if self.server_args.elastic_ep_backend == "mooncake":
# Mooncake does not support `monitored_barrier`
dist.barrier(group=get_tp_group().cpu_group)
else:
# Handle the case where some ranks do not finish loading.
try:
dist.monitored_barrier(
group=get_tp_group().cpu_group,
timeout=datetime.timedelta(
seconds=UNBALANCED_MODEL_LOADING_TIMEOUT_S
),
wait_all_ranks=True,
)
except RuntimeError:
raise ValueError(
f"TP rank {self.tp_rank} could finish the model loading, but there are other ranks that didn't finish loading. It is likely due to unexpected failures (e.g., OOM) or a slow node."
) from None
def update_expert_location(
self,
new_expert_location_metadata: ExpertLocationMetadata,
update_layer_ids: List[int],
):
p2p_missing_logical_experts = self.expert_location_updater.update(
self.model.routed_experts_weights_of_layer,
new_expert_location_metadata,
update_layer_ids=update_layer_ids,
nnodes=self.server_args.nnodes,
rank=self.tp_rank,
)
if len(p2p_missing_logical_experts) > 0:
# Load the missing expert weights from disk
if callable(getattr(self.model, "generate_weight_name_filter", None)):
# Filter and load only missing expert weights
weight_name_filter = self.model.generate_weight_name_filter(
p2p_missing_logical_experts
)
else:
# Do a full reload from disk/DRAM
logger.info(
"[Elastic EP] Model does not implement generate_weight_name_filter. "
"Performing full weight reload."
)
weight_name_filter = None
if (
self.expert_backup_client is not None
and self.expert_backup_client.use_backup
):
# Load the missing weights from the DRAM backup
self.expert_backup_client.update_weights(weight_name_filter)
else:
# Load the missing weights from disk
self.update_weights_from_disk(
get_global_server_args().model_path,
get_global_server_args().load_format,
weight_name_filter=weight_name_filter,
)
def update_weights_from_disk(
self,
model_path: str,
load_format: str,
weight_name_filter: Optional[Callable[[str], bool]] = None,
recapture_cuda_graph: bool = False,
) -> tuple[bool, str]:
"""Update engine weights in-place from the disk."""
logger.info(
f"Update engine weights online from disk begin. "
f"avail mem={get_available_gpu_memory(self.device, self.gpu_id, empty_cache=False):.2f} GB"
)
target_device = torch.device(self.device)
self.model_config.model_path = model_path
load_config = LoadConfig(load_format=load_format)
# Only support DefaultModelLoader for now
loader = get_model_loader(load_config, self.model_config)
if not isinstance(loader, DefaultModelLoader):
message = f"Failed to get model loader: {loader}."
return False, message
def get_weight_iter(config):
iter = loader._get_weights_iterator(
DefaultModelLoader.Source.init_new(config, self.model)
)
if weight_name_filter is not None:
iter = (
(name, weight) for name, weight in iter if weight_name_filter(name)
)
return iter
def model_load_weights(model, iter):
loader.load_weights_and_postprocess(model, iter, target_device)
return model
with set_default_torch_dtype(self.model_config.dtype):
try:
iter = get_weight_iter(self.model_config)
except Exception as e:
message = f"Failed to get weights iterator: {e}."
return False, message
try:
model = model_load_weights(self.model, iter)
except Exception as e:
message = (
f"Failed to update weights: {e}.\nRolling back to original weights."
)
del iter
gc.collect()
iter = get_weight_iter(self.model_config)
self.model = model_load_weights(self.model, iter)
return False, message
self.model = model
self.server_args.model_path = model_path
self.server_args.load_format = load_format
self.load_config = load_config
if recapture_cuda_graph and (self.device == "cuda" or self.device == "musa"):
self.init_device_graphs()
logger.info("Update weights end.")
return True, "Succeeded to update model weights."
def init_weights_send_group_for_remote_instance(
self,
master_address,
ports,
group_rank,
world_size,
group_name,
backend="nccl",
):
assert (
torch.distributed.is_initialized()
), "Default torch process group must be initialized"
assert group_name != "", "Group name cannot be empty"
ports_list = ports.split(",")
assert (
len(ports_list) == self.tp_size
), f"Expected {self.tp_size} ports, but got {len(ports_list)} ports."
group_port = ports_list[self.tp_rank]
group_name = f"{group_name}_{group_port}_{self.tp_rank}"
logger.info(
f"init custom process group: tp_rank={self.tp_rank}, gpu_id={self.gpu_id}, master_address={master_address}, master_port={group_port}, "
f"group_rank={group_rank}, world_size={world_size}, group_name={group_name}, backend={backend}"
)
torch.cuda.empty_cache()
success = False
message = ""
try:
na = NetworkAddress(master_address, group_port)
self._weights_send_group[group_name] = init_custom_process_group(
backend=backend,
init_method=na.to_tcp(),
world_size=world_size,
rank=group_rank,
group_name=group_name,
device_id=torch.device("cuda", self.gpu_id),
)
dist.barrier(group=self._weights_send_group[group_name])
success = True
message = f"Succeeded to init group through {na.to_host_port_str()} group."
except Exception as e:
message = f"Failed to init group: {e}."
logger.error(message)
torch.cuda.empty_cache()
return success, message
def send_weights_to_remote_instance(
self,
master_address,
ports,
group_name,
):
assert (
torch.distributed.is_initialized()
), "Default torch process group must be initialized"
assert group_name != "", "Group name cannot be empty"
ports_list = ports.split(",")
assert (
len(ports_list) == self.tp_size
), f"Expected {self.tp_size} ports, but got {len(ports_list)} ports."
group_port = ports_list[self.tp_rank]
group_name = f"{group_name}_{group_port}_{self.tp_rank}"
if self._weights_send_group[group_name] is not None:
send_group = self._weights_send_group[group_name]
else:
message = f"Group {group_name} not in _weights_send_group list. Please call `init_weights_send_group_for_remote_instance` first."
logger.error(message)
return False, message
torch.cuda.empty_cache()
success = False
na = NetworkAddress(master_address, group_port)
message = ""
try:
for _, weights in self.model.named_parameters():
torch.distributed.broadcast(
weights,
src=0,
group=send_group,
)
success = True
message = f"Succeeded to send weights through {na.to_host_port_str()} {group_name}."
except Exception as e:
message = f"Failed to send weights: {e}."
logger.error(message)
# destroy the process group after sending weights
del self._weights_send_group[group_name]
torch.distributed.distributed_c10d.destroy_process_group(send_group)
torch.cuda.empty_cache()
return success, message
def init_weights_update_group(
self,
master_address,
master_port,
rank_offset,
world_size,
group_name,
backend="nccl",
):
"""Initialize the Torch process group for model parameter updates.
`_model_update_group` is used in the RLHF workflow, where rank
0 is the actor model in the training engine, and the other ranks are
the inference engine, which is used for rollout.
In the RLHF workflow, the training engine updates the model
weights/parameters online, and broadcasts them to the inference
engine through the `_model_update_group` process group.
"""
assert (
torch.distributed.is_initialized()
), "Default torch process group must be initialized"
assert group_name != "", "Group name cannot be empty"
rank = rank_offset + self.tp_rank
logger.info(
f"init custom process group: master_address={master_address}, master_port={master_port}, "
f"rank_offset={rank_offset}, rank={rank}, world_size={world_size}, group_name={group_name}, backend={backend}"
)
try:
na = NetworkAddress(master_address, master_port)
self._model_update_group[group_name] = init_custom_process_group(
backend=backend,
init_method=na.to_tcp(),
world_size=world_size,
rank=rank,
group_name=group_name,
)
return True, "Succeeded to initialize custom process group."
except Exception as e:
message = f"Failed to initialize custom process group: {e}."
logger.error(message)
return False, message
def destroy_weights_update_group(self, group_name):
try:
if group_name in self._model_update_group:
pg = self._model_update_group.pop(group_name)
torch.distributed.destroy_process_group(pg)
return True, "Succeeded to destroy custom process group."
else:
return False, "The group to be destroyed does not exist."
except Exception as e:
message = f"Failed to destroy custom process group: {e}."
logger.error(message)
return False, message
def update_weights_from_distributed(
self,
names,
dtypes,
shapes,
group_name,
load_format: Optional[str] = None,
):
"""
Update specific parameter in the model weights online
through `_model_update_group` process group.
Args:
name: the name of the parameter to be updated.
dtype: the data type of the parameter to be updated.
shape: the shape of the parameter to be updated.
"""
assert group_name in self._model_update_group, (
f"Group {group_name} not in {list(self._model_update_group.keys())}. "
"Please call `init_weights_update_group` first."
)
if load_format == "flattened_bucket":
return self._update_bucketed_weights_from_distributed(
names, dtypes, shapes, group_name
)
try:
weights = []
handles = []
for name, dtype, shape in zip(names, dtypes, shapes):
target_dtype = (
dtype if isinstance(dtype, torch.dtype) else getattr(torch, dtype)
)
weight = torch.empty(shape, dtype=target_dtype, device=self.device)
handles.append(
torch.distributed.broadcast(
weight,
src=0,
group=self._model_update_group[group_name],
async_op=True,
)
)
weights.append((name, weight))
for handle in handles:
handle.wait()
self.model.load_weights(weights)
return True, "Succeeded to update parameter online."
except Exception as e:
error_msg = (
f"Failed to update parameter online: {e}. "
f"The full weights of the ModelRunner are partially updated. "
f"Please discard the whole weights."
)
logger.error(error_msg)
return False, error_msg
def _update_bucketed_weights_from_distributed(
self, names, dtypes, shapes, group_name
):
try:
named_tensors = []
for name, dtype, shape in zip(names, dtypes, shapes):
target_dtype = (
dtype if isinstance(dtype, torch.dtype) else getattr(torch, dtype)
)
named_tensors.append(
(name, torch.empty(shape, dtype=target_dtype, device=self.device))
)
bucket = FlattenedTensorBucket(named_tensors=named_tensors)
flattened_tensor = bucket.get_flattened_tensor()
torch.distributed.broadcast(
flattened_tensor,
src=0,
group=self._model_update_group[group_name],
)
reconstructed_tensors = bucket.reconstruct_tensors()
self.model.load_weights(reconstructed_tensors)
return True, f"Succeeded to update parameter online."
except Exception as e:
error_msg = (
f"Failed to update parameter online: {e}. "
f"The full weights of the ModelRunner are partially updated. "
f"Please discard the whole weights."
)
logger.error(error_msg)
return False, error_msg
def update_weights_from_tensor(
self,
named_tensors: List[Tuple[str, Union[torch.Tensor, "LocalSerializedTensor"]]],
load_format: Optional[str] = None,
):
monkey_patch_torch_reductions()
if load_format == "flattened_bucket":
# Handle flattened bucket format
return self._update_weights_from_flattened_bucket(
flattened_tensor_bucket_dict=named_tensors
)
# We need to get device after patch otherwise the device would be wrong
self.device_module = torch.get_device_module(self.device)
infered_device = self.device_module.current_device()
named_tensors = [
(name, _unwrap_tensor(tensor, tp_rank=self.tp_rank, device=infered_device))
for name, tensor in named_tensors
]
if load_format == "direct":
_model_load_weights_direct(self.model, named_tensors)
elif load_format in self.server_args.custom_weight_loader:
custom_loader = dynamic_import(load_format)
custom_loader(self.model, named_tensors)
elif load_format is None:
self.model.load_weights(named_tensors)
else:
raise NotImplementedError(f"Unknown load_format={load_format}")
return True, "Success"
def _update_weights_from_flattened_bucket(
self,
flattened_tensor_bucket_dict,
):
"""Handle flattened bucket format for weight updates"""
flattened_tensor = flattened_tensor_bucket_dict["flattened_tensor"]
metadata = flattened_tensor_bucket_dict["metadata"]
# Convert metadata dict to our format
converted_metadata = []
for meta in metadata:
converted_meta = FlattenedTensorMetadata(
name=meta.name,
shape=meta.shape,
dtype=meta.dtype,
start_idx=meta.start_idx,
end_idx=meta.end_idx,
numel=meta.numel,
)
converted_metadata.append(converted_meta)
# Create bucket and reconstruct tensors
bucket = FlattenedTensorBucket(
flattened_tensor=flattened_tensor, metadata=converted_metadata
)
reconstructed_tensors = bucket.reconstruct_tensors()
# Load the reconstructed tensors using the standard method
self.model.load_weights(reconstructed_tensors)
return True, "Success"
def get_weights_by_name(
self, name: str, truncate_size: int = 100
) -> Optional[torch.Tensor]:
"""Get the weights of the parameter by its name. Similar to `get_parameter` in Hugging Face.
Only used for unit test with an unoptimized performance.
For optimized performance, please use torch.save and torch.load.
"""
# TODO: (chenyang) Add support for Qwen models.
try:
return self.model.get_weights_by_name(
name, truncate_size, tp_size=self.tp_size
)
except Exception as e:
logger.error(f"Error when getting parameter {name}: {e}")
return None
def init_lora_manager(self):
self.lora_manager = LoRAManager(
base_model=self.model,
base_hf_config=self.model_config.hf_config,
max_loras_per_batch=self.server_args.max_loras_per_batch,
load_config=self.load_config,
dtype=self.dtype,
server_args=self.server_args,
lora_backend=self.server_args.lora_backend,
tp_size=self.tp_size,
tp_rank=self.tp_rank,
max_lora_rank=self.server_args.max_lora_rank,
target_modules=self.server_args.lora_target_modules,
lora_paths=self.server_args.lora_paths,
)
def load_lora_adapter(self, lora_ref: LoRARef):
"""Load a new lora adapter from disk or huggingface."""
logger.info(
f"LoRA adapter loading starts: {lora_ref}. "
f"avail mem={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB"
)
result = self.lora_manager.load_lora_adapter(lora_ref)
logger.info(
f"LoRA adapter loading completes: {lora_ref}. "
f"avail mem={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB"
)
return result
def load_lora_adapter_from_tensors(
self, lora_ref: LoRARef, tensors, config_dict, added_tokens_config=None
):
logger.info(f"LoRA adapter loading from tensors starts: {lora_ref}.")
result = self.lora_manager.load_lora_adapter_from_tensors(
lora_ref, tensors, config_dict, added_tokens_config
)
logger.info(f"LoRA adapter loading from tensors completes: {lora_ref}.")
return result
def unload_lora_adapter(self, lora_ref: LoRARef):
"""Unload a lora adapter that was previously loaded during initialization or dynamic loading."""
logger.info(
f"LoRA adapter unloading starts: {lora_ref}. "
f"avail mem={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB"
)
result = self.lora_manager.unload_lora_adapter(lora_ref)
logger.info(
f"LoRA adapter unloading completes: {lora_ref}. "
f"avail mem={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB"
)
return result
@property
def qwen3_next_config(self):
config = self.model_config.hf_config
if isinstance(config, Qwen3NextConfig):
return config
return None
@property
def hybrid_lightning_config(self):
config = self.model_config.hf_config
if isinstance(config, BailingHybridConfig):
return config
return None
@property
def hybrid_gdn_config(self):
config = self.model_config.hf_config.get_text_config()
if isinstance(
config,
Qwen3NextConfig
| Qwen3_5Config
| Qwen3_5MoeConfig
| JetNemotronConfig
| JetVLMConfig,
):
return config
return None
@property
def mamba2_config(self):
config = self.model_config.hf_config
if isinstance(config, NemotronHConfig) and self.is_draft_worker:
# NemotronH MTP draft models have no Mamba layers (pattern like "*E")
# so they shouldn't use HybridLinearAttnBackend
pattern = getattr(config, "mtp_hybrid_override_pattern", None)
if pattern is not None and "M" not in pattern:
return None
if isinstance(
config, FalconH1Config | NemotronHConfig | Lfm2Config | Lfm2MoeConfig
):
return config
if isinstance(config, NemotronH_Nano_VL_V2_Config):
return config.llm_config
if isinstance(config, GraniteMoeHybridConfig):
has_mamba = any(
layer_type == "mamba"
for layer_type in getattr(config, "layer_types", [])
)
if not has_mamba:
return None
else:
return config
return None
@property
def max_token_pool_size(self):
"""Return the max token pool size considering hybrid swa settings."""
if self.is_hybrid_swa:
return min(self.swa_max_total_num_tokens, self.max_total_num_tokens)
else:
return self.max_total_num_tokens
@property
def kimi_linear_config(self):
config = self.model_config.hf_config
if isinstance(config, KimiLinearConfig):
return config
return None
@property
def mambaish_config(self):
return (
self.mamba2_config
or self.hybrid_gdn_config
or self.kimi_linear_config
or self.hybrid_lightning_config
)
def configure_kv_cache_dtype(self):
if self.server_args.kv_cache_dtype == "auto":
quant_config = getattr(self.model, "quant_config", None)
kv_cache_quant_algo = getattr(quant_config, "kv_cache_quant_algo", None)
if (
isinstance(kv_cache_quant_algo, str)
and kv_cache_quant_algo.upper() == "FP8"
):
if _is_hip:
self.kv_cache_dtype = fp8_dtype
self.server_args.kv_cache_dtype = TORCH_DTYPE_TO_KV_CACHE_STR[
self.kv_cache_dtype
]
else:
self.kv_cache_dtype = torch.float8_e4m3fn
self.server_args.kv_cache_dtype = TORCH_DTYPE_TO_KV_CACHE_STR[
self.kv_cache_dtype
]
else:
self.kv_cache_dtype = self.dtype
elif self.server_args.kv_cache_dtype == "fp8_e5m2":
if _is_hip: # Using natively supported format
self.kv_cache_dtype = fp8_dtype
else:
self.kv_cache_dtype = torch.float8_e5m2
elif self.server_args.kv_cache_dtype == "fp8_e4m3":
if _is_hip: # Using natively supported format
self.kv_cache_dtype = fp8_dtype
else:
self.kv_cache_dtype = torch.float8_e4m3fn
elif self.server_args.kv_cache_dtype in ("bf16", "bfloat16"):
self.kv_cache_dtype = torch.bfloat16
elif self.server_args.kv_cache_dtype == "fp4_e2m1":
if hasattr(torch, "float4_e2m1fn_x2"):
self.kv_cache_dtype = torch.float4_e2m1fn_x2
logger.warning(f"FP4 (E2M1) KV Cache might lead to a accuracy drop!")
else:
logger.warning(
f"--kv-cache-dtype falls back to 'auto' because this torch version does not support torch.float4_e2m1fn_x2"
)
self.kv_cache_dtype = self.dtype
else:
raise ValueError(
f"Unsupported kv_cache_dtype: {self.server_args.kv_cache_dtype}."
)
log_info_on_rank0(logger, f"Using KV cache dtype: {self.kv_cache_dtype}")
def init_cublas(self):
"""We need to run a small matmul to init cublas. Otherwise, it will raise some errors later."""
dtype = torch.float16
device = "cuda"
a = torch.ones((16, 16), dtype=dtype, device=device)
b = torch.ones((16, 16), dtype=dtype, device=device)
c = a @ b
return c
def init_attention_backend(self):
"""Init attention kernel backend."""
if self.server_args.enable_pdmux:
self.attn_backend = self._get_attention_backend(init_new_workspace=True)
self.decode_attn_backend_group = []
for _ in range(self.server_args.sm_group_num):
self.decode_attn_backend_group.append(self._get_attention_backend())
self.decode_attn_backend = self.decode_attn_backend_group[0]
elif self.server_args.enable_two_batch_overlap and not self.is_draft_worker:
self.attn_backend = TboAttnBackend.init_new(self._get_attention_backend)
else:
self.attn_backend = self._get_attention_backend()
def _get_attention_backend(self, init_new_workspace: bool = False):
"""Init attention kernel backend."""
draft_attn_backend = self.server_args.speculative_draft_attention_backend
if self.is_draft_worker and draft_attn_backend:
logger.warning(
f"Overriding draft attention backend to {draft_attn_backend}."
)
return self._get_attention_backend_from_str(
draft_attn_backend,
init_new_workspace=init_new_workspace,
)
(
self.prefill_attention_backend_str,
self.decode_attention_backend_str,
) = self.server_args.get_attention_backends()
if self.decode_attention_backend_str != self.prefill_attention_backend_str:
from sglang.srt.layers.attention.hybrid_attn_backend import (
HybridAttnBackend,
)
attn_backend = HybridAttnBackend(
self,
decode_backend=self._get_attention_backend_from_str(
self.decode_attention_backend_str,
init_new_workspace=init_new_workspace,
),
prefill_backend=self._get_attention_backend_from_str(
self.prefill_attention_backend_str,
init_new_workspace=init_new_workspace,
),
)
logger.info(
f"Using hybrid attention backend for decode and prefill: "
f"decode_backend={self.decode_attention_backend_str}, "
f"prefill_backend={self.prefill_attention_backend_str}."
)
logger.warning(
"Warning: Attention backend specified by --attention-backend or default backend might be overridden."
"The feature of hybrid attention backend is experimental and unstable. Please raise an issue if you encounter any problem."
)
else:
attn_backend = self._get_attention_backend_from_str(
self.server_args.attention_backend,
init_new_workspace=init_new_workspace,
)
(
get_global_server_args().prefill_attention_backend,
get_global_server_args().decode_attention_backend,
) = (self.prefill_attention_backend_str, self.decode_attention_backend_str)
return attn_backend
def _get_attention_backend_from_str(
self, backend_str: str, init_new_workspace: bool = False
):
if backend_str not in ATTENTION_BACKENDS:
raise ValueError(f"Invalid attention backend: {backend_str}")
self.init_new_workspace = init_new_workspace
full_attention_backend = ATTENTION_BACKENDS[backend_str](self)
return attn_backend_wrapper(self, full_attention_backend)
def init_double_sparsity_channel_config(self, selected_channel):
selected_channel = "." + selected_channel + "_proj"
self.sorted_channels = []
# load channel config
with open(self.server_args.ds_channel_config_path, "r") as f:
channel_config = json.load(f)
for i in range(self.start_layer, self.end_layer):
key = "model.layers." + str(i) + ".self_attn" + selected_channel
self.sorted_channels.append(
torch.tensor(channel_config[key])[
:, : self.server_args.ds_heavy_channel_num
]
.contiguous()
.cuda()
)
def kernel_warmup(self):
"""
Warmup and tune kernels before cuda graph capture.
Currently only doing FlashInfer autotune.
"""
if self.device != "cuda":
return
if self._should_run_flashinfer_autotune():
self._flashinfer_autotune()
def _should_run_flashinfer_autotune(self) -> bool:
"""Check if flashinfer autotune should be run."""
if self.server_args.disable_flashinfer_autotune:
return False
backend_str = self.server_args.moe_runner_backend
# TODO smor- support other cases for flashinfer autotune, such as, mamba backend
if backend_str not in [
"flashinfer_trtllm",
"flashinfer_mxfp4",
# TODO: flashinfer_cutlass will cause some flashinfer compilation errors. To be fixed.
# "flashinfer_cutlass",
]:
return False
major, _ = torch.cuda.get_device_capability()
if major < 9:
return False
if (
self.spec_algorithm.is_eagle()
or self.spec_algorithm.is_standalone()
or self.spec_algorithm.is_ngram()
):
return not self.is_draft_worker
return True
def _flashinfer_autotune(self):
"""Run flashinfer autotune."""
from flashinfer.autotuner import autotune
logger.info("Running FlashInfer autotune...")
# Run warmup on the non-default stream to avoid NCCL 2.29+ cudaMemcpyBatchAsync
# calls on default stream (unsupported by CUDA) when --enable-symm-mem is used.
self.forward_stream.wait_stream(torch.cuda.current_stream())
with torch.get_device_module(self.device).stream(self.forward_stream):
with torch.inference_mode(), autotune():
self._dummy_run(
batch_size=self.req_to_token_pool.size, run_ctx=autotune()
)
torch.cuda.current_stream().wait_stream(self.forward_stream)
logger.info("FlashInfer autotune completed.")
def _dummy_run(self, batch_size: int, run_ctx=None):
"""Run a dummy forward pass for warmup/profiling."""
if self.is_generation:
capture_forward_mode = ForwardMode.DECODE
else:
capture_forward_mode = ForwardMode.EXTEND
capture_hidden_mode = CaptureHiddenMode.NULL
num_tokens_per_bs = 1
if (
self.spec_algorithm.is_eagle()
or self.spec_algorithm.is_standalone()
or self.spec_algorithm.is_ngram()
):
if self.is_draft_worker:
raise RuntimeError("This should not happen")
else:
capture_forward_mode = ForwardMode.TARGET_VERIFY
num_tokens_per_bs = self.server_args.speculative_num_draft_tokens
if self.server_args.enable_return_hidden_states:
capture_hidden_mode = CaptureHiddenMode.FULL
num_tokens = batch_size * num_tokens_per_bs
if require_gathered_buffer(self.server_args):
attn_tp_size = get_attention_tp_size()
if attn_tp_size > 1 and num_tokens % attn_tp_size != 0:
num_tokens = num_tokens // attn_tp_size * attn_tp_size
batch_size = num_tokens // num_tokens_per_bs
seq_len_fill_value = self.attn_backend.get_cuda_graph_seq_len_fill_value()
if self.server_args.enable_torch_compile:
set_torch_compile_config()
if self.eagle_use_aux_hidden_state:
self.model.set_eagle3_layers_to_capture()
require_mlp_tp_gather_ = require_mlp_tp_gather(self.server_args)
if require_gathered_buffer(self.server_args):
assert require_mlp_tp_gather_ or require_attn_tp_gather(self.server_args)
buffers: DecodeInputBuffers = DecodeInputBuffers.create(
device=self.device,
max_bs=batch_size,
max_num_token=num_tokens,
hidden_size=self.model_config.hidden_size,
vocab_size=self.model_config.vocab_size,
dtype=self.model_config.dtype,
dp_size=self.server_args.dp_size,
pp_size=self.server_args.pp_size,
is_encoder_decoder=self.model_config.is_encoder_decoder,
require_mlp_tp_gather=require_mlp_tp_gather_,
seq_len_fill_value=seq_len_fill_value,
encoder_len_fill_value=0,
num_tokens_per_bs=num_tokens_per_bs,
cache_loc_dtype=torch.int64,
enable_mamba_track=False,
)
buffers.num_token_non_padded[...] = num_tokens
# For extend mode
if not self.is_generation:
extend_prefix_lens_cpu = [0] * batch_size
extend_seq_lens_cpu = [seq_len_fill_value] * batch_size
extend_num_tokens = num_tokens
extend_seq_lens = torch.full(
(batch_size,), seq_len_fill_value, dtype=torch.int32, device=self.device
)
extend_prefix_lens = torch.zeros(
(batch_size,), dtype=torch.int32, device=self.device
)
extend_start_loc = torch.arange(
0, num_tokens, num_tokens_per_bs, dtype=torch.int32, device=self.device
)
else:
extend_prefix_lens_cpu = None
extend_seq_lens_cpu = None
extend_num_tokens = None
extend_seq_lens = None
extend_prefix_lens = None
extend_start_loc = None
if self.server_args.pp_size > 1:
pp_proxy_tensors = PPProxyTensors(
{k: v[:num_tokens] for k, v in buffers.pp_proxy_tensors.items()}
)
if require_mlp_tp_gather_:
buffers.global_num_tokens_gpu.copy_(
torch.tensor(
[num_tokens] * self.server_args.dp_size,
dtype=torch.int32,
device=self.device,
)
)
buffers.global_num_tokens_for_logprob_gpu.copy_(
torch.tensor(
[num_tokens] * self.server_args.dp_size,
dtype=torch.int32,
device=self.device,
)
)
global_dp_buffer_len = num_tokens * self.server_args.dp_size
elif require_attn_tp_gather(self.server_args):
buffers.global_num_tokens_gpu.copy_(
torch.tensor(
[num_tokens],
dtype=torch.int32,
device=self.device,
)
)
buffers.global_num_tokens_for_logprob_gpu.copy_(
torch.tensor(
[num_tokens],
dtype=torch.int32,
device=self.device,
)
)
global_dp_buffer_len = num_tokens
else:
global_dp_buffer_len = None
def get_spec_info():
spec_info = None
if self.spec_algorithm.is_eagle() or self.spec_algorithm.is_standalone():
from sglang.srt.speculative.eagle_info import EagleVerifyInput
if self.is_draft_worker:
raise RuntimeError("This should not happen.")
else:
spec_info = EagleVerifyInput(
draft_token=None,
custom_mask=buffers.custom_mask,
positions=None,
retrive_index=None,
retrive_next_token=None,
retrive_next_sibling=None,
retrive_cum_len=None,
spec_steps=self.server_args.speculative_num_steps,
topk=self.server_args.speculative_eagle_topk,
draft_token_num=self.server_args.speculative_num_draft_tokens,
capture_hidden_mode=CaptureHiddenMode.FULL,
seq_lens_sum=None,
seq_lens_cpu=None,
)
elif self.spec_algorithm.is_ngram():
from sglang.srt.speculative.ngram_info import NgramVerifyInput
spec_info = NgramVerifyInput(
draft_token=None,
tree_mask=buffers.custom_mask,
positions=None,
retrive_index=None,
retrive_next_token=None,
retrive_next_sibling=None,
draft_token_num=num_tokens_per_bs,
)
spec_info.capture_hidden_mode = CaptureHiddenMode.NULL
return spec_info
spec_info = get_spec_info()
if capture_hidden_mode != CaptureHiddenMode.FULL:
capture_hidden_mode = (
spec_info.capture_hidden_mode if spec_info else CaptureHiddenMode.NULL
)
if self.server_args.enable_lora:
lora_ids = [None] * batch_size
else:
lora_ids = None
forward_batch = ForwardBatch(
forward_mode=capture_forward_mode,
batch_size=batch_size,
input_ids=buffers.input_ids,
req_pool_indices=buffers.req_pool_indices,
seq_lens=buffers.seq_lens,
seq_lens_cpu=buffers.seq_lens_cpu,
next_token_logits_buffer=buffers.next_token_logits_buffer,
orig_seq_lens=buffers.seq_lens,
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool=self.token_to_kv_pool,
attn_backend=self.attn_backend,
out_cache_loc=buffers.out_cache_loc,
seq_lens_sum=buffers.seq_lens.sum().item(),
encoder_lens=buffers.encoder_lens,
return_logprob=False,
positions=buffers.positions,
extend_num_tokens=extend_num_tokens,
extend_seq_lens=extend_seq_lens,
extend_prefix_lens=extend_prefix_lens,
extend_start_loc=extend_start_loc,
extend_prefix_lens_cpu=extend_prefix_lens_cpu,
extend_seq_lens_cpu=extend_seq_lens_cpu,
global_num_tokens_gpu=buffers.global_num_tokens_gpu,
global_num_tokens_for_logprob_gpu=buffers.global_num_tokens_for_logprob_gpu,
dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
global_dp_buffer_len=global_dp_buffer_len,
mrope_positions=buffers.mrope_positions,
spec_algorithm=self.spec_algorithm,
spec_info=spec_info,
capture_hidden_mode=capture_hidden_mode,
num_token_non_padded=buffers.num_token_non_padded,
global_forward_mode=capture_forward_mode,
lora_ids=lora_ids,
)
forward_batch.uses_cp_shared_kv = self.uses_cp_shared_kv
forward_batch.cp_shared_kv_layout = self.cp_shared_kv_layout
if lora_ids is not None:
self.lora_manager.prepare_lora_batch(forward_batch)
self.attn_backend.init_forward_metadata(forward_batch)
def run_once():
forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
set_dp_buffer_len(
global_dp_buffer_len,
num_tokens,
forward_batch.dp_padding_mode.is_max_len(),
)
set_is_extend_in_batch(False)
kwargs = {}
if (
self.server_args.pp_size > 1
and "pp_proxy_tensors"
in inspect.signature(self.model.forward).parameters
):
kwargs["pp_proxy_tensors"] = PPProxyTensors(
{k: v.clone() for k, v in pp_proxy_tensors.tensors.items()}
)
if not self.is_generation:
kwargs["get_embedding"] = True
logits_output_or_pp_proxy_tensors = self.model.forward(
buffers.input_ids,
forward_batch.positions,
forward_batch,
**kwargs,
)
return logits_output_or_pp_proxy_tensors
torch.get_device_module(self.device).synchronize()
self.tp_group.barrier()
with torch.inference_mode(), run_ctx or empty_context():
run_once()
def maybe_init_ngram_embedding(self):
self.use_ngram_embedding = self.model_config.use_ngram_embedding
if self.use_ngram_embedding:
from sglang.srt.layers.n_gram_embedding import NgramEmbedding
self.token_table = torch.empty(
self.req_to_token_pool.size,
self.model_config.context_len,
dtype=torch.int32,
device=self.device,
)
chunked_prefill_size = self.server_args.chunked_prefill_size
assert (
chunked_prefill_size is not None and chunked_prefill_size > 0
), "Ngram embedding requires chunked prefill to be enabled (chunked_prefill_size > 0)"
for module in self.model.modules():
if isinstance(module, NgramEmbedding):
module.init_buffers(
self.max_running_requests, chunked_prefill_size, self.device
)
def maybe_update_ngram_token_table(
self,
next_token_ids: torch.Tensor,
forward_batch: "ForwardBatch",
):
"""Update the ngram embedding token table after sampling."""
ngram_embedding_info = forward_batch.ngram_embedding_info
if ngram_embedding_info is None:
return
ngram_embedding_info.out_column_starts[: forward_batch.batch_size] = (
forward_batch.seq_lens
)
ngram_embedding_info.out_req_lens[: forward_batch.batch_size] = 1
update_token_table(
ne_token_table=ngram_embedding_info.token_table,
tokens=next_token_ids.to(torch.int32),
row_indices=forward_batch.req_pool_indices,
column_starts=ngram_embedding_info.out_column_starts,
req_lens=torch.ones_like(ngram_embedding_info.out_column_starts),
ignore_tokens=None,
)
def init_device_graphs(self):
"""Capture device graphs."""
self.graph_runner = None
self.graph_mem_usage = 0
if not self.is_generation:
# TODO: Currently, cuda graph only captures decode steps, which only exists for generation models
return
if self.server_args.model_impl.lower() == ModelImpl.MINDSPORE:
return
if self.device != "cpu" and self.server_args.disable_cuda_graph:
return
if self.device == "cpu" and not self.server_args.enable_torch_compile:
return
tic = time.perf_counter()
before_mem = get_available_gpu_memory(self.device, self.gpu_id)
graph_backend = defaultdict(
lambda: "cuda graph",
{
"cpu": "cpu graph",
"npu": "npu graph",
},
)
logger.info(
f"Capture {graph_backend[self.device]} begin. This can take up to several minutes. avail mem={before_mem:.2f} GB"
)
graph_runners = defaultdict(
lambda: CudaGraphRunner,
{
"cpu": CPUGraphRunner,
"npu": NPUGraphRunner,
},
)
self.graph_runner = graph_runners[self.device](self)
after_mem = get_available_gpu_memory(self.device, self.gpu_id)
self.graph_mem_usage = before_mem - after_mem
logger.info(
f"Capture {graph_backend[self.device]} end. Time elapsed: {time.perf_counter() - tic:.2f} s. "
f"mem usage={self.graph_mem_usage:.2f} GB. avail mem={after_mem:.2f} GB."
)
def init_piecewise_cuda_graphs(self):
"""Initialize piecewise CUDA graph runner."""
self.piecewise_cuda_graph_runner = None
if self.server_args.disable_piecewise_cuda_graph:
logger.info(
"Disable piecewise CUDA graph because --disable-piecewise-cuda-graph is set"
)
return
# Disable piecewise CUDA graph for non-language models
if not hasattr(self.model, "model"):
logger.warning(
"Disable piecewise CUDA graph because the model is not a language model"
)
return
# Disable piecewise CUDA graph for non capture size
if not self.server_args.piecewise_cuda_graph_tokens:
logger.warning(
"Disable piecewise CUDA graph because the capture size is not set"
)
return
# Collect attention layers and moe layers from the model
self.model.model = resolve_language_model(self.model)
language_model = getattr(self.model, "language_model", self.model)
self.attention_layers = []
self.moe_layers = []
self.moe_fusions = []
for layer in language_model.model.layers:
attn_layer = None
if hasattr(layer, "self_attn"):
if hasattr(layer.self_attn, "attn"):
attn_layer = layer.self_attn.attn
elif hasattr(layer.self_attn, "attn_mqa"):
# For DeepSeek model
attn_layer = layer.self_attn.attn_mqa
# For hybrid model
elif hasattr(layer, "attn"):
attn_layer = layer.attn
elif hasattr(layer, "linear_attn"):
if hasattr(layer.linear_attn, "attn"):
attn_layer = layer.linear_attn.attn
else:
attn_layer = layer.linear_attn
# For InternVL model
elif hasattr(layer, "attention"):
if hasattr(layer.attention, "attn"):
attn_layer = layer.attention.attn
# For NemotronH and similar hybrid models using 'mixer' attribute
elif hasattr(layer, "mixer"):
if hasattr(layer.mixer, "attn"):
attn_layer = layer.mixer.attn
elif hasattr(layer, "_forward_mamba"):
# Mamba layer with split op support - store the layer itself
attn_layer = layer
if attn_layer is not None:
self.attention_layers.append(attn_layer)
moe_block = None
moe_fusion = None
if hasattr(layer, "mlp") and hasattr(layer.mlp, "experts"):
moe_block = layer.mlp.experts
moe_fusion = layer.mlp
if hasattr(layer, "block_sparse_moe") and hasattr(
layer.block_sparse_moe, "experts"
):
moe_block = layer.block_sparse_moe.experts
moe_fusion = layer.block_sparse_moe
if hasattr(layer, "moe") and hasattr(layer.moe, "experts"):
moe_block = layer.moe.experts
moe_fusion = layer.moe
# For NemotronH MoE layers using 'mixer' attribute
if hasattr(layer, "mixer") and hasattr(layer.mixer, "experts"):
moe_block = layer.mixer.experts
moe_fusion = layer.mixer
self.moe_layers.append(moe_block)
self.moe_fusions.append(moe_fusion)
if len(self.attention_layers) < self.model_config.num_hidden_layers:
# TODO(yuwei): support Non-Standard GQA
log_info_on_rank0(
logger,
"Disable piecewise CUDA graph because some layers do not apply Standard GQA",
)
return
tic = time.perf_counter()
before_mem = get_available_gpu_memory(self.device, self.gpu_id)
logger.info(
f"Capture piecewise CUDA graph begin. avail mem={before_mem:.2f} GB"
)
self.piecewise_cuda_graph_runner = PiecewiseCudaGraphRunner(self)
after_mem = get_available_gpu_memory(self.device, self.gpu_id)
mem_usage = before_mem - after_mem
logger.info(
f"Capture piecewise CUDA graph end. Time elapsed: {time.perf_counter() - tic:.2f} s. "
f"mem usage={mem_usage:.2f} GB. avail mem={after_mem:.2f} GB."
)
def init_deep_gemm(self):
logger.info("[DeepGEMM Debug] Entering init_deep_gemm")
if not deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM:
return
try:
from sglang.srt.layers.deep_gemm_wrapper.compile_utils import (
DeepGemmKernelType,
_maybe_compile_deep_gemm_one_type_all,
)
import deep_gemm
except ImportError:
return
# Find DeepSeek V3 MoE layers
model = self.model
if hasattr(model, "model"):
model = model.model
if not hasattr(model, "layers"):
return
# We only need to compile once for each unique shape.
# Since _maybe_compile_deep_gemm_one_type_all handles caching, we can just call it for every MoE layer.
for layer in model.layers:
if hasattr(layer, "mlp") and hasattr(layer.mlp, "experts"):
experts = layer.mlp.experts
w13 = None
w2 = None
if hasattr(experts, "w13_weight_fp8") and hasattr(
experts, "w2_weight_fp8"
):
w13 = experts.w13_weight_fp8[0]
w2 = experts.w2_weight_fp8[0]
elif hasattr(experts, "w13_weight") and hasattr(experts, "w2_weight"):
w13 = experts.w13_weight
w2 = experts.w2_weight
if w13 is not None and w2 is not None:
# w13: [num_groups, n, k]
num_groups, n, k = w13.shape
_maybe_compile_deep_gemm_one_type_all(
DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_MASKED,
n,
k,
num_groups,
)
# w2: [num_groups, n, k]
num_groups, n, k = w2.shape
_maybe_compile_deep_gemm_one_type_all(
DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_MASKED,
n,
k,
num_groups,
)
# Also compile Dense GEMM kernels for other layers (e.g. Attention) if they use FP8 and DeepGEMM
for module in self.model.modules():
if hasattr(module, "weight") and module.weight is not None:
# Check if it is FP8
is_fp8 = module.weight.dtype in [
torch.float8_e4m3fn,
torch.float8_e4m3fnuz,
]
if is_fp8 and module.weight.dim() == 2:
n, k = module.weight.shape
_maybe_compile_deep_gemm_one_type_all(
DeepGemmKernelType.GEMM_NT_F8F8BF16,
n,
k,
1, # num_groups=1 for standard GEMM
)
if module.weight.dim() == 2 and module.weight.dtype == torch.bfloat16:
n, k = module.weight.shape
_maybe_compile_deep_gemm_one_type_all(
DeepGemmKernelType.GEMM_NT_BF16BF16F32,
n,
k,
1,
)
# deep_gemm.preload_kernels()
def init_threads_binding(self):
omp_cpuids = os.environ.get("SGLANG_CPU_OMP_THREADS_BIND", "all")
cpu_ids_by_node = get_cpu_ids_by_node()
n_numa_node = len(cpu_ids_by_node)
if omp_cpuids == "all":
assert self.tp_size <= n_numa_node, (
f"SGLANG_CPU_OMP_THREADS_BIND is not set, in this case, "
f"tp_size {self.tp_size} should be smaller than or equal to number of numa node on the machine {n_numa_node}. "
f"If you need tp_size to be larger than number of numa node, please set the CPU cores for each tp rank via SGLANG_CPU_OMP_THREADS_BIND explicitly. "
f"For example, on a machine with 2 numa nodes, where core 0-31 are on numa node 0 and core 32-63 are on numa node 1, "
f"it is suggested to use -tp 2 and bind tp rank 0 to core 0-31 and tp rank 1 to core 32-63. "
f"This is the default behavior if SGLANG_CPU_OMP_THREADS_BIND is not set and it is the same as setting SGLANG_CPU_OMP_THREADS_BIND=0-31|32-63. "
f"If you do need tp_size to be larger than the number of numa nodes, you could set SGLANG_CPU_OMP_THREADS_BIND explicitly for example SGLANG_CPU_OMP_THREADS_BIND=0-15|16-31|32-47|48-63 and run with -tp 4. "
f"If you don't want each tp rank to use all the cores on one numa node, you could set for example SGLANG_CPU_OMP_THREADS_BIND=0-15|32-47 and run with -tp 2."
)
if self.tp_size < n_numa_node:
logger.warning(
f"Detected the current machine has {n_numa_node} numa nodes available, but tp_size is set to {self.tp_size}, so only {self.tp_size} numa nodes are used."
)
self.local_omp_cpuid = cpu_ids_by_node[self.tp_rank]
else:
threads_bind_list = omp_cpuids.split("|")
assert self.tp_size == len(threads_bind_list), (
f"SGLANG_CPU_OMP_THREADS_BIND setting must be aligned with TP size parameter ({self.tp_size}). "
f"Please double check your settings."
)
self.local_omp_cpuid = threads_bind_list[self.tp_rank]
if self.tp_size > n_numa_node:
logger.warning(
f"TP size ({self.tp_size})is larger than numa node number ({n_numa_node}), "
f"in this case the available memory amount of each rank cannot be determined in prior. "
f"Please set proper `--max-total-tokens` to avoid the out-of-memory error."
)
def apply_torch_tp(self):
logger.info(f"Enabling torch tensor parallelism on {self.tp_size} devices.")
from sglang.srt.layers.model_parallel import tensor_parallel
device_mesh = torch.distributed.init_device_mesh(self.device, (self.tp_size,))
tensor_parallel(self.model, device_mesh)
def update_decode_attn_backend(self, stream_idx: int):
self.decode_attn_backend = self.decode_attn_backend_group[stream_idx]
def forward_decode(
self,
forward_batch: ForwardBatch,
skip_attn_backend_init: bool = False,
pp_proxy_tensors=None,
) -> Union[LogitsProcessorOutput, PPProxyTensors]:
if not skip_attn_backend_init:
if self.server_args.enable_pdmux:
self.decode_attn_backend.init_forward_metadata(forward_batch)
forward_batch.attn_backend = self.decode_attn_backend
else:
self.attn_backend.init_forward_metadata(forward_batch)
# FIXME: add pp_proxy_tensors arg to all models
kwargs = {}
if self.support_pp:
kwargs["pp_proxy_tensors"] = pp_proxy_tensors
return self.model.forward(
forward_batch.input_ids,
forward_batch.positions,
forward_batch,
**kwargs,
)
def forward_extend(
self,
forward_batch: ForwardBatch,
skip_attn_backend_init: bool = False,
pp_proxy_tensors=None,
) -> Tuple[
Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput], bool
]:
kwargs = {}
if self.support_pp:
kwargs["pp_proxy_tensors"] = pp_proxy_tensors
if forward_batch.input_embeds is not None:
kwargs["input_embeds"] = forward_batch.input_embeds.bfloat16()
if not self.is_generation:
kwargs["get_embedding"] = True
can_run_graph = (
self.piecewise_cuda_graph_runner is not None
and self.piecewise_cuda_graph_runner.can_run(forward_batch)
)
if can_run_graph:
return (
self.piecewise_cuda_graph_runner.replay(forward_batch, **kwargs),
can_run_graph,
)
if not skip_attn_backend_init:
self.attn_backend.init_forward_metadata(forward_batch)
return (
self.model.forward(
forward_batch.input_ids,
forward_batch.positions,
forward_batch,
**kwargs,
),
can_run_graph,
)
def forward_idle(
self, forward_batch: ForwardBatch, pp_proxy_tensors=None
) -> Union[LogitsProcessorOutput, PPProxyTensors]:
# In DP Attention, IDLE batches are padded (batch_size > 0) for MLP sync.
# in this case, we need to reinit the forward metadata, otherwise the stale
# metadata causes batch_size mismatch in attention kernel(e.g. NSA Indexer).
if forward_batch.batch_size > 0:
self.attn_backend.init_forward_metadata(forward_batch)
kwargs = {}
if self.support_pp:
kwargs["pp_proxy_tensors"] = pp_proxy_tensors
return self.model.forward(
forward_batch.input_ids,
forward_batch.positions,
forward_batch,
**kwargs,
)
def forward_split_prefill(
self,
forward_batch: ForwardBatch,
reinit_attn_backend: bool = False,
forward_count: int = 1,
) -> LogitsProcessorOutput:
if forward_batch.split_index == 0 or reinit_attn_backend:
self.attn_backend.init_forward_metadata(forward_batch)
next_split_index = min(
forward_batch.split_index + forward_count,
self.model_config.num_hidden_layers,
)
ret = self.model.forward_split_prefill(
forward_batch.input_ids,
forward_batch.positions,
forward_batch,
(forward_batch.split_index, next_split_index),
)
forward_batch.split_index = next_split_index
return ret
def forward(
self,
forward_batch: ForwardBatch,
skip_attn_backend_init: bool = False,
pp_proxy_tensors: Optional[PPProxyTensors] = None,
reinit_attn_backend: bool = False,
split_forward_count: int = 1,
) -> ModelRunnerOutput:
self.forward_pass_id += 1
with get_global_expert_distribution_recorder().with_forward_pass(
self.forward_pass_id,
forward_batch,
) as recorder_outputs:
output = self._forward_raw(
forward_batch,
skip_attn_backend_init,
pp_proxy_tensors,
reinit_attn_backend,
split_forward_count,
)
elastic_ep_state = ElasticEPStateManager.instance()
if (
elastic_ep_state is not None
and not elastic_ep_state.is_active_equal_last()
):
elastic_ep_state.snapshot_active_to_last()
elastic_ep_state.sync_active_to_cpu()
logging.info("EPLB due to rank faults")
gen = self.eplb_manager.rebalance()
while True:
try:
next(gen)
except StopIteration:
break
output = self._forward_raw(
forward_batch,
skip_attn_backend_init,
pp_proxy_tensors,
reinit_attn_backend,
split_forward_count,
)
output.expert_distribution_metrics = recorder_outputs.get("metrics")
# Copy cached routing experts' buffers back to CPU cache
get_global_experts_capturer().on_forward_end(
forward_batch=forward_batch,
can_run_graph=output.can_run_graph,
cuda_graph_batch=getattr(self.graph_runner, "bs", None),
)
if self.eplb_manager is not None:
self.eplb_manager.on_forward_pass_end()
if dumper.may_enable:
dumper.step()
return output
def _forward_raw(
self,
forward_batch: ForwardBatch,
skip_attn_backend_init: bool,
pp_proxy_tensors: Optional[PPProxyTensors],
reinit_attn_backend: bool = False,
split_forward_count: int = 1,
) -> ModelRunnerOutput:
mode_check = (
forward_batch.forward_mode.is_cpu_graph
if self.device == "cpu"
else forward_batch.forward_mode.is_cuda_graph
)
can_run_graph = bool(
mode_check()
and self.graph_runner
and self.graph_runner.can_run(forward_batch)
)
if can_run_graph:
ret = self.graph_runner.replay(
forward_batch,
skip_attn_backend_init=skip_attn_backend_init,
pp_proxy_tensors=pp_proxy_tensors,
)
return ModelRunnerOutput(logits_output=ret, can_run_graph=can_run_graph)
# For MLP sync
if forward_batch.global_num_tokens_cpu is not None:
forward_batch.prepare_mlp_sync_batch(self)
else:
forward_batch.prepare_attn_tp_scatter_input(self)
# Normalize num_token_non_padded to be local to this attention TP rank if needed.
if (
forward_batch.num_token_non_padded is not None
and forward_batch.global_num_tokens_gpu is not None
and require_gathered_buffer(self.server_args)
and not is_nsa_enable_prefill_cp()
):
forward_batch.adjust_num_token_non_padded_for_attn_tp(
server_args=self.server_args,
)
# Use precomputed SWA cache location
if forward_batch.out_cache_loc_swa is not None:
self.token_to_kv_pool.set_swa_loc(forward_batch.out_cache_loc_swa)
forward_batch.hisparse_coordinator = self.hisparse_coordinator
if self.hisparse_coordinator is not None:
self.hisparse_coordinator.num_real_reqs.fill_(forward_batch.batch_size)
if forward_batch.forward_mode.is_decode():
ret = self.forward_decode(
forward_batch,
skip_attn_backend_init=skip_attn_backend_init,
pp_proxy_tensors=pp_proxy_tensors,
)
elif forward_batch.forward_mode.is_split_prefill():
ret = self.forward_split_prefill(
forward_batch,
reinit_attn_backend=reinit_attn_backend,
forward_count=split_forward_count,
)
elif forward_batch.forward_mode.is_extend(include_draft_extend_v2=True):
ret, can_run_graph = self.forward_extend(
forward_batch,
skip_attn_backend_init=skip_attn_backend_init,
pp_proxy_tensors=pp_proxy_tensors,
)
elif forward_batch.forward_mode.is_idle():
ret = self.forward_idle(forward_batch, pp_proxy_tensors=pp_proxy_tensors)
else:
raise ValueError(f"Invalid forward mode: {forward_batch.forward_mode}")
if (
forward_batch.global_num_tokens_cpu is not None
and self.pp_group.is_last_rank
):
forward_batch.post_forward_mlp_sync_batch(ret)
return ModelRunnerOutput(logits_output=ret, can_run_graph=can_run_graph)
def _preprocess_logits(
self, logits_output: LogitsProcessorOutput, sampling_info: SamplingBatchInfo
):
# NOTE: In overlap mode, the function update_regex_vocab_mask (in sample)
# was executed after we processed last batch's results.
# Calculate logits bias and apply it to next_token_logits.
sampling_info.update_regex_vocab_mask()
sampling_info.apply_logits_bias(logits_output.next_token_logits)
# Release the vocab_mask GPU tensor immediately after it has been applied
# to the logits. In overlap scheduling, the sampling_info (and its
# vocab_mask) can be kept alive by the delay_sample_func closure and
# batch_record_buf until the next iteration, causing a steady VRAM leak
# when structured output (grammar) is used.
sampling_info.vocab_mask = None
def sample(
self,
logits_output: LogitsProcessorOutput,
forward_batch: ForwardBatch,
) -> torch.Tensor:
"""Sample and compute logprobs and update logits_output.
Args:
logits_output: The logits output from the model forward
forward_batch: The forward batch that generates logits_output
Returns:
A list of next_token_ids
"""
# For duplex models with multiple output streams.
if isinstance(logits_output, tuple):
return torch.stack(
[self.sample(values, forward_batch) for values in logits_output],
axis=-1,
)
self._preprocess_logits(logits_output, forward_batch.sampling_info)
# Sample the next tokens
next_token_ids = self.sampler(
logits_output,
forward_batch.sampling_info,
forward_batch.return_logprob,
forward_batch.top_logprobs_nums,
forward_batch.token_ids_logprobs,
# For prefill, we only use the position of the last token.
(
forward_batch.positions
if forward_batch.forward_mode.is_decode()
else forward_batch.seq_lens - 1
),
)
self.maybe_update_ngram_token_table(next_token_ids, forward_batch)
return next_token_ids
def compute_logprobs_only(
self,
logits_output: LogitsProcessorOutput,
forward_batch: ForwardBatch,
) -> None:
"""
Compute token_ids_logprobs without performing sampling.
Optimized path for prefill-only requests that need token_ids_logprobs but don't
require next token generation. Skips expensive sampling operations
while still providing requested probability information.
Args:
logits_output: The logits output from the model forward
forward_batch: The forward batch that generates logits_output
"""
if not forward_batch.token_ids_logprobs:
return
# Preprocess logits (same as in sample method)
self._preprocess_logits(logits_output, forward_batch.sampling_info)
# Delegate to sampler for logprob-only computation
# This populates logits_output with requested token probabilities
self.sampler.compute_logprobs_only(
logits_output,
forward_batch.sampling_info,
forward_batch.return_logprob,
forward_batch.top_logprobs_nums,
forward_batch.token_ids_logprobs,
)
@property
def model_is_mrope(self) -> bool:
"""Detect if the model has "mrope" rope_scaling type.
mrope requires keep "rope_deltas" between prompt and decoding phases."""
rope_scaling = getattr(
self.model_config.hf_text_config, "rope_parameters", None
) or getattr(self.model_config.hf_text_config, "rope_scaling", {})
if rope_scaling is None:
return False
is_mrope_enabled = "mrope_section" in rope_scaling
return is_mrope_enabled
def save_remote_model(self, url: str):
from sglang.srt.model_loader.loader import RemoteModelLoader
logger.info(f"Saving model to {url}")
RemoteModelLoader.save_model(self.model, self.model_config.model_path, url)
def save_sharded_model(
self, path: str, pattern: Optional[str] = None, max_size: Optional[int] = None
):
from sglang.srt.model_loader.loader import ShardedStateLoader
logger.info(
f"Save sharded model to {path} with pattern {pattern} and max_size {max_size}"
)
ShardedStateLoader.save_model(self.model, path, pattern, max_size)
def check_weights(self, action: str):
self._weight_checker.handle(action=action)
def update_weights_from_ipc(self, recv_req):
"""Update weights from IPC for checkpoint-engine integration."""
try:
from sglang.srt.checkpoint_engine.checkpoint_engine_worker import (
SGLangCheckpointEngineWorkerExtensionImpl,
)
# Create a worker extension that integrates with SGLang's model
worker = SGLangCheckpointEngineWorkerExtensionImpl(self)
worker.update_weights_from_ipc(recv_req.zmq_handles)
return True, "IPC weight update completed successfully"
except ImportError as e:
return False, f"IPC weight update failed: ImportError {e}"
except Exception as e:
logger.error(f"IPC weight update failed: {e}")
return False, str(e)
def prealloc_symmetric_memory_pool(self):
# PyTorch mempools never de-fragment memory in OOM scenarios, so we need to pre-allocate a large chunk of memory to limit fragmentation.
if (
self.is_draft_worker
or not self.server_args.enable_symm_mem
or envs.SGLANG_SYMM_MEM_PREALLOC_GB_SIZE.get() <= 0
):
return
# Memory allocation is tied to a cuda stream, use the forward stream
with torch.get_device_module(self.device).stream(self.forward_stream):
logger.info(
f"Pre-allocating symmetric memory pool with {envs.SGLANG_SYMM_MEM_PREALLOC_GB_SIZE.get()} GiB"
)
with use_symmetric_memory(get_tp_group()):
torch.empty(
(envs.SGLANG_SYMM_MEM_PREALLOC_GB_SIZE.get() * 1024 * 1024 * 1024,),
dtype=torch.uint8,
device=self.device,
)
def _model_load_weights_direct(model, named_tensors: List[Tuple[str, torch.Tensor]]):
params_dict = dict(model.named_parameters())
for name, tensor in named_tensors:
default_weight_loader(params_dict[name], tensor)
def _unwrap_tensor(tensor, tp_rank, device):
if isinstance(tensor, LocalSerializedTensor):
tensor = tensor.get(tp_rank)
return tensor.to(device)
@dataclass
class LocalSerializedTensor:
"""torch.Tensor that gets serialized by MultiprocessingSerializer (which only serializes a pointer and not the data).
The i-th element in the list corresponds to i-th rank's GPU."""
values: List[bytes]
def get(self, rank: int):
return MultiprocessingSerializer.deserialize(self.values[rank])
Reference in New Issue View Git Blame Copy Permalink