From 1dc48c2c3bd345e3caa72befac3d8ffeadc1c5ab Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Tue, 14 Jul 2026 16:00:51 +0800 Subject: [PATCH] Absorb capturer setup and extract the shared-mooncake gate (#31160) --- .../mooncake_transfer_engine.py | 54 ++++++++++++- .../sglang/srt/model_executor/model_runner.py | 78 ++----------------- .../sglang/srt/state_capturer/indexer_topk.py | 34 ++++++++ .../srt/state_capturer/routed_experts.py | 13 +++- .../test_mooncake_transfer_engine_init.py | 2 +- 5 files changed, 103 insertions(+), 78 deletions(-) diff --git a/python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py b/python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py index cf205ee27..ce08c6fbc 100644 --- a/python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py +++ b/python/sglang/srt/distributed/device_communicators/mooncake_transfer_engine.py @@ -1,15 +1,20 @@ +from __future__ import annotations + import json import logging import os -from typing import Dict, List, Optional, Union +from typing import TYPE_CHECKING, Dict, List, Optional, Union from sglang.srt.environ import envs -from sglang.srt.utils.network import NetworkAddress, get_free_port +from sglang.srt.utils.network import NetworkAddress, get_free_port, get_local_ip_auto + +if TYPE_CHECKING: + from sglang.srt.server_args import ServerArgs logger = logging.getLogger(__name__) # Module-level shared engine instance, set by init_mooncake_transfer_engine(). -_mooncake_transfer_engine: Optional["MooncakeTransferEngine"] = None +_mooncake_transfer_engine: Optional[MooncakeTransferEngine] = None def parse_ib_device_config( @@ -295,3 +300,46 @@ def init_mooncake_transfer_engine( def get_mooncake_transfer_engine() -> Optional[MooncakeTransferEngine]: """Return the shared MooncakeTransferEngine if initialized, else None.""" return _mooncake_transfer_engine + + +def maybe_init_shared_mooncake_transfer_engine( + *, server_args: ServerArgs, gpu_id: int +) -> None: + """ + 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 = ( + ( + server_args.disaggregation_mode != "null" + and server_args.disaggregation_transfer_backend == "mooncake" + ) + or ( + server_args.enable_hierarchical_cache + and server_args.hicache_storage_backend == "mooncake" + and envs.SGLANG_HICACHE_MOONCAKE_REUSE_TE.get() + ) + or ( + server_args.encoder_only + and server_args.encoder_transfer_backend == "mooncake" + ) + or ( + server_args.language_only + and server_args.encoder_transfer_backend == "mooncake" + ) + or ( + server_args.enable_elastic_expert_backup + and server_args.elastic_ep_backend is not None + ) + ) + + if use_mooncake_te: + init_mooncake_transfer_engine( + hostname=get_local_ip_auto(), + gpu_id=gpu_id, + ib_device=( + server_args.disaggregation_ib_device or server_args.mooncake_ib_device + ), + ) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 081691665..5e09d154a 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -30,7 +30,6 @@ from sglang.srt.configs.model_config import ( AttentionArch, ModelConfig, ModelImpl, - get_num_indexer_layers, ) from sglang.srt.configs.update_config import adjust_config_with_unaligned_cpu_tp from sglang.srt.debug_utils.dumper import dumper @@ -38,6 +37,9 @@ from sglang.srt.distributed import ( bootstrap, get_world_group, ) +from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import ( + maybe_init_shared_mooncake_transfer_engine, +) from sglang.srt.distributed.device_communicators.pynccl_allocator import ( prealloc_symmetric_memory_pool, ) @@ -195,7 +197,6 @@ from sglang.srt.utils import ( set_cuda_arch, slow_rank_detector, ) -from sglang.srt.utils.network import get_local_ip_auto from sglang.srt.utils.nvtx_pytorch_hooks import PytHooks from sglang.srt.utils.nvtx_utils import profile_range from sglang.srt.utils.offloader import ( @@ -751,18 +752,10 @@ class ModelRunner(ModelRunnerKVCacheMixin): # would overwrite the target's process-global one. return - 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_server_args().enable_return_routed_experts, + model=self.model, 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, @@ -770,26 +763,9 @@ class ModelRunner(ModelRunnerKVCacheMixin): ) def init_indexer_capturer(self): - enable = get_server_args().enable_return_indexer_topk - # Producer wiring is CUDA-only (Indexer.forward_cuda + MLA skip_topk - # path); other backends would create a capturer but never feed it. - if enable and self.device != "cuda": - logger.warning( - "indexer-topk capture is CUDA-only; %s backend not yet wired. " - "Disabling capturer.", - self.device, - ) - set_global_indexer_capturer(None) - return - - hf_text_config = self.model_config.hf_text_config - num_indexer_layers = get_num_indexer_layers(hf_text_config) - index_topk = getattr(hf_text_config, "index_topk", 0) set_global_indexer_capturer( create_indexer_capturer( - enable=enable, - num_indexer_layers=num_indexer_layers, - index_topk=index_topk, + model_config=self.model_config, num_tokens=self.max_total_num_tokens + self.page_size, max_running_requests=self.max_running_requests, device=self.device, @@ -857,50 +833,10 @@ class ModelRunner(ModelRunnerKVCacheMixin): self.pre_model_load_memory = result.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 - ) + maybe_init_shared_mooncake_transfer_engine( + server_args=self.server_args, gpu_id=self.gpu_id ) - 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) diff --git a/python/sglang/srt/state_capturer/indexer_topk.py b/python/sglang/srt/state_capturer/indexer_topk.py index e37b5f243..afa652cec 100644 --- a/python/sglang/srt/state_capturer/indexer_topk.py +++ b/python/sglang/srt/state_capturer/indexer_topk.py @@ -5,6 +5,7 @@ import numpy as np import pybase64 import torch +from sglang.srt.configs.model_config import ModelConfig, get_num_indexer_layers from sglang.srt.runtime_context import get_parallel from sglang.srt.state_capturer.base import BaseTopkCapturer @@ -82,6 +83,39 @@ def extract_indexer_topk_from_meta_info(data): def create_indexer_capturer( + *, + model_config: ModelConfig, + num_tokens: int, + max_running_requests: int, + device: str, +) -> Optional[IndexerTopkCapturer]: + from sglang.srt.runtime_context import get_server_args + + enable = get_server_args().enable_return_indexer_topk + # Producer wiring is CUDA-only (Indexer.forward_cuda + MLA skip_topk + # path); other backends would create a capturer but never feed it. + if enable and device != "cuda": + logger.warning( + "indexer-topk capture is CUDA-only; %s backend not yet wired. " + "Disabling capturer.", + device, + ) + return None + + hf_text_config = model_config.hf_text_config + num_indexer_layers = get_num_indexer_layers(hf_text_config) + index_topk = getattr(hf_text_config, "index_topk", 0) + return _create_indexer_capturer_raw( + enable=enable, + num_indexer_layers=num_indexer_layers, + index_topk=index_topk, + num_tokens=num_tokens, + max_running_requests=max_running_requests, + device=device, + ) + + +def _create_indexer_capturer_raw( enable: bool, num_indexer_layers: int, index_topk: int, diff --git a/python/sglang/srt/state_capturer/routed_experts.py b/python/sglang/srt/state_capturer/routed_experts.py index 4e69f5dd3..65f6cc052 100644 --- a/python/sglang/srt/state_capturer/routed_experts.py +++ b/python/sglang/srt/state_capturer/routed_experts.py @@ -28,15 +28,22 @@ class RoutedExpertsCapturer(BaseTopkCapturer): @staticmethod def create( - enable: bool, + *, + model: torch.nn.Module, model_config: ModelConfig, - num_fused_shared_experts: int, num_tokens: int, max_running_requests: int, device: str, ) -> Optional["RoutedExpertsCapturer"]: - if not enable: + server_args = get_server_args() + if not server_args.enable_return_routed_experts: return None + if not server_args.disable_shared_experts_fusion and hasattr( + model, "num_fused_shared_experts" + ): + num_fused_shared_experts = model.num_fused_shared_experts + else: + num_fused_shared_experts = 0 return RoutedExpertsCapturer( model_config, num_tokens=num_tokens, diff --git a/test/manual/kv_transfer/test_mooncake_transfer_engine_init.py b/test/manual/kv_transfer/test_mooncake_transfer_engine_init.py index 64a45e2b0..bcffda5a0 100755 --- a/test/manual/kv_transfer/test_mooncake_transfer_engine_init.py +++ b/test/manual/kv_transfer/test_mooncake_transfer_engine_init.py @@ -65,7 +65,7 @@ def test_mooncake_te_condition(server_args: ServerArgs) -> bool: side_effect=_fake_init_mooncake_transfer_engine, ), patch( - "sglang.srt.model_executor.model_runner.get_local_ip_auto", + "sglang.srt.distributed.device_communicators.mooncake_transfer_engine.get_local_ip_auto", return_value="127.0.0.1", ), ):