diff --git a/python/sglang/srt/compilation/backend.py b/python/sglang/srt/compilation/backend.py index 9c926d4a4..08b871364 100644 --- a/python/sglang/srt/compilation/backend.py +++ b/python/sglang/srt/compilation/backend.py @@ -268,9 +268,6 @@ def split_graph( return split_gm, outputs -# we share the global graph pool among all the backends -global_graph_pool = None - compilation_start_time = 0.0 diff --git a/python/sglang/srt/eplb/expert_distribution.py b/python/sglang/srt/eplb/expert_distribution.py index 84d5d10a8..16fafcd11 100644 --- a/python/sglang/srt/eplb/expert_distribution.py +++ b/python/sglang/srt/eplb/expert_distribution.py @@ -284,18 +284,20 @@ class _ExpertDistributionRecorderReal(ExpertDistributionRecorder): return self._recording -_global_expert_distribution_recorder: Optional[ExpertDistributionRecorder] = ( - _ExpertDistributionRecorderNoop() -) - - def get_global_expert_distribution_recorder(): - return _global_expert_distribution_recorder + from sglang.srt.runtime_context import get_resources + + resources = get_resources() + if resources.expert_distribution_recorder is None: + # Call sites expect a recorder unconditionally; default to the noop. + resources.expert_distribution_recorder = _ExpertDistributionRecorderNoop() + return resources.expert_distribution_recorder def set_global_expert_distribution_recorder(value): - global _global_expert_distribution_recorder - _global_expert_distribution_recorder = value + from sglang.srt.runtime_context import get_resources + + get_resources().expert_distribution_recorder = value # --------------------------------------- SinglePassGatherer ----------------------------------------- diff --git a/python/sglang/srt/eplb/expert_location.py b/python/sglang/srt/eplb/expert_location.py index b5c464cbb..d45e6cb09 100644 --- a/python/sglang/srt/eplb/expert_location.py +++ b/python/sglang/srt/eplb/expert_location.py @@ -305,17 +305,18 @@ class ExpertLocationMetadata: ] -_global_expert_location_metadata: Optional[ExpertLocationMetadata] = None - - def get_global_expert_location_metadata(): - return _global_expert_location_metadata + from sglang.srt.runtime_context import get_resources + + return get_resources().expert_location_metadata def set_global_expert_location_metadata(value): - global _global_expert_location_metadata - assert _global_expert_location_metadata is None - _global_expert_location_metadata = value + from sglang.srt.runtime_context import get_resources + + resources = get_resources() + assert resources.expert_location_metadata is None + resources.expert_location_metadata = value def broadcast_global_expert_location_metadata( diff --git a/python/sglang/srt/eplb/lplb_solver.py b/python/sglang/srt/eplb/lplb_solver.py index 0186a370b..7da86d5d5 100644 --- a/python/sglang/srt/eplb/lplb_solver.py +++ b/python/sglang/srt/eplb/lplb_solver.py @@ -26,7 +26,6 @@ import torch logger = logging.getLogger(__name__) # Global per-layer LPLB solvers -_global_lplb_solvers: dict[int, LPLBSolver] = {} # LP dispatch requires every EP rank to call solver.solve() on every forward @@ -59,15 +58,21 @@ def assert_lplb_supported_model(architecture: str) -> None: def get_global_lplb_solver(layer_id: int) -> Optional[LPLBSolver]: - return _global_lplb_solvers.get(layer_id) + from sglang.srt.runtime_context import get_resources + + return get_resources().lplb_solvers.get(layer_id) def set_global_lplb_solver(layer_id: int, solver: LPLBSolver): - _global_lplb_solvers[layer_id] = solver + from sglang.srt.runtime_context import get_resources + + get_resources().lplb_solvers[layer_id] = solver def clear_global_lplb_solvers(): - _global_lplb_solvers.clear() + from sglang.srt.runtime_context import get_resources + + get_resources().lplb_solvers.clear() class LPLBSolver: diff --git a/python/sglang/srt/hardware_backend/musa/attention/flashattention_backend.py b/python/sglang/srt/hardware_backend/musa/attention/flashattention_backend.py index 92825453b..ab62b0145 100644 --- a/python/sglang/srt/hardware_backend/musa/attention/flashattention_backend.py +++ b/python/sglang/srt/hardware_backend/musa/attention/flashattention_backend.py @@ -32,7 +32,6 @@ if TYPE_CHECKING: # Global workspace buffer for MLA _MATE_MLA_WORKSPACE_SIZE_BYTES = 128 * 1024 * 1024 -_MATE_MLA_WORKSPACE_BUFFER: torch.Tensor | None = None # Cache for non-MLA scheduler metadata by prefix _MATE_NO_MLA_SCHEDULER_METADATA_DICT: dict = {} @@ -54,7 +53,7 @@ def _compute_scheduler_metadata( num_splits: int, ) -> Tuple[torch.Tensor, bool] | torch.Tensor: """Compute scheduler metadata based on backend's current state.""" - global _MATE_MLA_WORKSPACE_BUFFER, _MATE_NO_MLA_SCHEDULER_METADATA_DICT + global _MATE_NO_MLA_SCHEDULER_METADATA_DICT layer = backend._current_layer current_layer_id = layer.layer_id @@ -84,11 +83,15 @@ def _compute_scheduler_metadata( should_update = True if backend.use_mla: - if _MATE_MLA_WORKSPACE_BUFFER is None: - _MATE_MLA_WORKSPACE_BUFFER = torch.empty( + from sglang.srt.runtime_context import get_buffer + + workspace = get_buffer( + "musa_mate_mla_workspace", + lambda: torch.empty( _MATE_MLA_WORKSPACE_SIZE_BYTES, device=backend.device, dtype=torch.uint8 - ) - return (_MATE_MLA_WORKSPACE_BUFFER, not should_update) + ), + ) + return (workspace, not should_update) else: with _MATE_NO_MLA_SCHEDULER_METADATA_LOCK: if ( diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index c3dea8ab3..12fd629cc 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -123,9 +123,6 @@ class ForwardMetadata: swa_out_cache_loc: Optional[torch.Tensor] = None -global_workspace_buffer = None - - _AITER_PARTITION_SIZE_ROCM = 256 diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 1adf96cac..c4055f890 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -56,6 +56,7 @@ from sglang.srt.layers.utils.cp_utils import ( cp_split_and_rebuild_position, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.runtime_context import get_buffer from sglang.srt.utils import ( get_bool_env_var, is_cuda, @@ -136,7 +137,6 @@ def _to_2d_context_lens(seqlens_32: torch.Tensor, batch_size: int) -> torch.Tens # Reuse this workspace buffer across all DSA backend instances -global_workspace_buffer = None @dataclass(frozen=True) @@ -443,14 +443,14 @@ class DeepseekSparseAttnBackend( # Allocate global workspace buffer for TRT-LLM kernels (ragged attention on SM100/B200, or trtllm decode) if self.device_sm_major >= 10 or self.dsa_decode_impl == "trtllm": - global global_workspace_buffer - if global_workspace_buffer is None: - global_workspace_buffer = torch.empty( + self.workspace_buffer = get_buffer( + "dsa_trtllm_workspace", + lambda: torch.empty( envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.get(), dtype=torch.uint8, device=model_runner.device, - ) - self.workspace_buffer = global_workspace_buffer + ), + ) else: self.workspace_buffer = None diff --git a/python/sglang/srt/layers/attention/flash_mla_sm120.py b/python/sglang/srt/layers/attention/flash_mla_sm120.py index e6a6e5c96..12c74f105 100644 --- a/python/sglang/srt/layers/attention/flash_mla_sm120.py +++ b/python/sglang/srt/layers/attention/flash_mla_sm120.py @@ -284,9 +284,6 @@ _BYTES_PER_DST_PAGE = ( _BYTES_PER_DST_PAGE_PADDED = math.ceil(_BYTES_PER_DST_PAGE / 576) * 576 # 37440 -# Pre-allocated buffer for page-split output per device (lazily sized). -_split_buf = {} # device -> tensor - @triton.jit def _page_split_kernel( @@ -344,8 +341,13 @@ def _split_kv_pages_to_64(kv_u8: torch.Tensor, src_pbs: int) -> torch.Tensor: ratio = src_pbs // _PBS_DST num_dst_pages = N * ratio + from sglang.srt.runtime_context import get_resources + + # Pre-allocated grow-only buffer for page-split output per device. dev = kv_u8.device - buf = _split_buf.get(dev) + buffers = get_resources().buffers + key = f"flash_mla_sm120_split:{dev}" + buf = buffers.get(key) if buf is None or buf.shape[0] < num_dst_pages: buf = torch.empty( num_dst_pages, @@ -353,7 +355,7 @@ def _split_kv_pages_to_64(kv_u8: torch.Tensor, src_pbs: int) -> torch.Tensor: dtype=torch.uint8, device=dev, ) - _split_buf[dev] = buf + buffers[key] = buf out = buf[:num_dst_pages] # Get raw 2D view of source diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 9869cdd87..d06e8cda4 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -38,6 +38,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMo from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( is_in_tc_piecewise_cuda_graph, ) +from sglang.srt.runtime_context import get_buffer from sglang.srt.speculative.spec_info import SpecInput, SpecInputType from sglang.srt.speculative.spec_utils import ( draft_kv_indices_buffer_width, @@ -168,7 +169,6 @@ class PrefillMetadata: # Reuse this workspace buffer across all flashinfer wrappers -global_workspace_buffer = None # Safety margin on the computed split-kv worst case for the dedicated # full-CG prefill workspace (absorbs allocator alignment and minor @@ -383,15 +383,15 @@ class FlashInferAttnBackend(AttentionBackend): self.use_paged = envs.SGLANG_FLASHINFER_USE_PAGED.get() # Allocate buffers - global global_workspace_buffer - if global_workspace_buffer is None: - # different from flashinfer zero_init_global_workspace_buffer - global_workspace_size = envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.get() - global_workspace_buffer = torch.empty( - global_workspace_size, + # different from flashinfer zero_init_global_workspace_buffer + global_workspace_buffer = get_buffer( + "flashinfer_workspace", + lambda: torch.empty( + envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.get(), dtype=torch.uint8, device=model_runner.device, - ) + ), + ) if init_new_workspace: self.workspace_buffer = torch.empty( envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.get(), diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py index 47e3c46b9..3c172abc3 100644 --- a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py @@ -34,6 +34,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMo from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( is_in_tc_piecewise_cuda_graph, ) +from sglang.srt.runtime_context import get_buffer from sglang.srt.server_args import get_global_server_args from sglang.srt.speculative.spec_info import SpecInput from sglang.srt.speculative.spec_utils import ( @@ -80,7 +81,6 @@ class PrefillMetadata: # Reuse this workspace buffer across all flashinfer wrappers -global_workspace_buffer = None class FlashInferMhaChunkKVRunner: @@ -233,15 +233,15 @@ class FlashInferMLAAttnBackend(AttentionBackend): self.page_size = model_runner.page_size # Allocate buffers - global global_workspace_buffer - if global_workspace_buffer is None: - # different from flashinfer zero_init_global_workspace_buffer - global_workspace_buffer = torch.empty( + # different from flashinfer zero_init_global_workspace_buffer + self.workspace_buffer = get_buffer( + "flashinfer_mla_workspace", + lambda: torch.empty( envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.get(), dtype=torch.uint8, device=model_runner.device, - ) - self.workspace_buffer = global_workspace_buffer + ), + ) max_bs = model_runner.req_to_token_pool.size if kv_indptr_buf is None: diff --git a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py b/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py index 5a7279c3e..75b4ae06a 100644 --- a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py +++ b/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py @@ -62,12 +62,12 @@ logger = logging.getLogger(__name__) # MAX_Q_LEN=8 covers EAGLE3 num_draft_tokens=4 plus headroom. _TOKENSPEED_MAX_Q_LEN = 8 -_g_tokenspeed_workspace: dict[torch.device, torch.Tensor] = {} - def _get_tokenspeed_workspace( device: torch.device, num_heads: int, kv_lora_rank: int ) -> torch.Tensor: + from sglang.srt.runtime_context import get_resources + needed = ( tokenspeed_mla.get_num_sm(device) * num_heads @@ -75,12 +75,12 @@ def _get_tokenspeed_workspace( * (kv_lora_rank + 1) * 4 ) - existing = _g_tokenspeed_workspace.get(device) + buffers = get_resources().buffers + key = f"tokenspeed_mla_workspace:{device}" + existing = buffers.get(key) if existing is None or existing.numel() < needed: - _g_tokenspeed_workspace[device] = torch.empty( - needed, dtype=torch.int8, device=device - ) - return _g_tokenspeed_workspace[device] + buffers[key] = torch.empty(needed, dtype=torch.int8, device=device) + return buffers[key] # TODO(Qiaolin-Yu): Merge this attention backend into trtllm_mla_backend.py diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 50ec976cf..e19e59141 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -31,6 +31,7 @@ from sglang.srt.layers.attention.utils import canonicalize_stride from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.runtime_context import get_buffer from sglang.srt.utils import is_flashinfer_available from sglang.srt.utils.common import is_sm90_supported, is_sm120_supported @@ -50,7 +51,6 @@ if TYPE_CHECKING: DEFAULT_WORKSPACE_SIZE_MB = 512 # Reuse this workspace buffer across all TRTLLM MHA wrappers -global_zero_init_workspace_buffer = None @dataclass @@ -116,14 +116,14 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): # Workspace allocation self.workspace_size = workspace_size_bytes # Allocate buffers - global global_zero_init_workspace_buffer - if global_zero_init_workspace_buffer is None: - global_zero_init_workspace_buffer = torch.zeros( + self.workspace_buffer = get_buffer( + "trtllm_mha_zero_workspace", + lambda: torch.zeros( self.workspace_size, dtype=torch.uint8, device=model_runner.device, - ) - self.workspace_buffer = global_zero_init_workspace_buffer + ), + ) # CUDA graph state self.decode_cuda_graph_metadata = {} diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index cffc95da1..1d628f10e 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -38,7 +38,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMo from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( is_in_tc_piecewise_cuda_graph, ) -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_buffer, get_parallel from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import is_flashinfer_available, is_float4_e2m1fn_x2 @@ -93,7 +93,6 @@ def _quantize_fp8_qkv(q, k, v, layer): return q, k, v, k_scale, v_scale -global_zero_init_workspace_buffer = None # cute-dsl needs its own workspace: it overwrites the buffer with split-KV # partials, which corrupts the trtllm-gen multiCtasKv counters that rely on the # zero-init buffer (they share it under attention-backend=cutedsl_mla, where @@ -182,14 +181,14 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): ) self.workspace_buffer = global_cute_dsl_workspace_buffer else: - global global_zero_init_workspace_buffer - if global_zero_init_workspace_buffer is None: - global_zero_init_workspace_buffer = torch.zeros( + self.workspace_buffer = get_buffer( + "trtllm_mla_zero_workspace", + lambda: torch.zeros( self.workspace_size, dtype=torch.int8, device=model_runner.device, - ) - self.workspace_buffer = global_zero_init_workspace_buffer + ), + ) # CUDA graph state self.decode_cuda_graph_metadata = {} diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index ccac078c0..77b7992a1 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/python/sglang/srt/layers/dp_attention.py @@ -702,14 +702,10 @@ def dp_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor): # stream -> their collectives serialize in-order (no concurrent-collective # deadlock on the RCCL communicator), each overlapping the other's compute. # --------------------------------------------------------------------------- -_DP_TBO_COMM_STREAM: Optional[torch.cuda.Stream] = None - - def get_dp_tbo_comm_stream() -> torch.cuda.Stream: - global _DP_TBO_COMM_STREAM - if _DP_TBO_COMM_STREAM is None: - _DP_TBO_COMM_STREAM = torch.cuda.Stream() - return _DP_TBO_COMM_STREAM + from sglang.srt.runtime_context import get_stream + + return get_stream("dp_tbo_comm") # Persistent reusable CUDA events for non-EP DP TBO, keyed by (kind, subbatch). @@ -718,14 +714,14 @@ def get_dp_tbo_comm_stream() -> torch.cuda.Stream: # pool is exhausted after a few hundred forwards -> HSA_STATUS_ERROR_OUT_OF_RESOURCES # ("...create internal OS-specific events"). Reuse one event per (kind, subbatch) # and just re-record it (mirrors the mori CommStreamPool event reuse). -_TBO_EVENT_POOL: dict = {} - - def _tbo_event(key) -> torch.cuda.Event: - ev = _TBO_EVENT_POOL.get(key) + from sglang.srt.runtime_context import get_resources + + pool = get_resources().tbo_event_pool + ev = pool.get(key) if ev is None: ev = torch.cuda.Event() - _TBO_EVENT_POOL[key] = ev + pool[key] = ev return ev diff --git a/python/sglang/srt/lora/lora_moe_runner_marlin.py b/python/sglang/srt/lora/lora_moe_runner_marlin.py index 89b451fe6..3f7540f1f 100644 --- a/python/sglang/srt/lora/lora_moe_runner_marlin.py +++ b/python/sglang/srt/lora/lora_moe_runner_marlin.py @@ -6,7 +6,7 @@ LoRA deltas are injected via hooks. from __future__ import annotations -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING import torch @@ -38,9 +38,6 @@ if _is_cuda: from sglang.srt.layers.quantization.marlin_utils import marlin_make_workspace -_MARLIN_WORKSPACE: Optional[torch.Tensor] = None - - class MarlinLoraRunnerCore: """ MoE runner using Marlin kernels for base projections, with hooks for LoRA. @@ -64,7 +61,6 @@ class MarlinLoraRunnerCore: runner_config: MoeRunnerConfig, hooks=None, ) -> StandardCombineInput: - global _MARLIN_WORKSPACE from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput assert hooks is not None, "hooks must be provided for MarlinLoraRunnerCore" @@ -95,14 +91,13 @@ class MarlinLoraRunnerCore: topk_ids, block_size_m, E ) - if ( - _MARLIN_WORKSPACE is None - or _MARLIN_WORKSPACE.device != hidden_states.device - ): - _MARLIN_WORKSPACE = marlin_make_workspace( - hidden_states.device, max_blocks_per_sm=4 - ) - workspace = _MARLIN_WORKSPACE + from sglang.srt.runtime_context import get_resources + + buffers = get_resources().buffers + workspace = buffers.get("marlin_lora_workspace") + if workspace is None or workspace.device != hidden_states.device: + workspace = marlin_make_workspace(hidden_states.device, max_blocks_per_sm=4) + buffers["marlin_lora_workspace"] = workspace scalar_type1 = get_scalar_type(num_bits, quant_info.w13_qzeros is not None) scalar_type2 = get_scalar_type(num_bits, quant_info.w2_qzeros is not None) diff --git a/python/sglang/srt/lora/trtllm_lora_temp/__init__.py b/python/sglang/srt/lora/trtllm_lora_temp/__init__.py index cedc793d0..0d163c2ac 100644 --- a/python/sglang/srt/lora/trtllm_lora_temp/__init__.py +++ b/python/sglang/srt/lora/trtllm_lora_temp/__init__.py @@ -35,9 +35,6 @@ def is_two_stream_active(x: torch.Tensor) -> bool: return x.shape[0] <= lora_envs.SGLANG_TWO_STREAM_MAX_TOKENS.get() -_LORA_SIDE_STREAM: Optional[torch.cuda.Stream] = None - - def get_lora_side_stream() -> torch.cuda.Stream: """Lazily allocate a single shared LoRA side stream. @@ -45,10 +42,9 @@ def get_lora_side_stream() -> torch.cuda.Stream: run sequentially, so one stream suffices and avoids extra graph-capture nodes from per-site streams. """ - global _LORA_SIDE_STREAM - if _LORA_SIDE_STREAM is None: - _LORA_SIDE_STREAM = torch.cuda.Stream() - return _LORA_SIDE_STREAM + from sglang.srt.runtime_context import get_stream + + return get_stream("lora_side") def init_lora_two_stream_resources(device: Optional[torch.device] = None) -> None: diff --git a/python/sglang/srt/model_executor/runner_utils/pool.py b/python/sglang/srt/model_executor/runner_utils/pool.py index 21fc49e65..f0fd4cca4 100644 --- a/python/sglang/srt/model_executor/runner_utils/pool.py +++ b/python/sglang/srt/model_executor/runner_utils/pool.py @@ -20,22 +20,21 @@ from __future__ import annotations from typing import Any, Optional -_global_graph_memory_pool: Optional[Any] = None +from sglang.srt.runtime_context import get_resources def get_global_graph_memory_pool() -> Optional[Any]: - return _global_graph_memory_pool + return get_resources().graph_memory_pool def set_global_graph_memory_pool(val: Any) -> None: - global _global_graph_memory_pool - _global_graph_memory_pool = val + get_resources().graph_memory_pool = val def get_or_create_global_graph_memory_pool(device_module: Any) -> Any: """Return the shared graph memory pool, creating it on first use so later backends reuse the same handle.""" - global _global_graph_memory_pool - if _global_graph_memory_pool is None: - _global_graph_memory_pool = device_module.graph_pool_handle() - return _global_graph_memory_pool + resources = get_resources() + if resources.graph_memory_pool is None: + resources.graph_memory_pool = device_module.graph_pool_handle() + return resources.graph_memory_pool diff --git a/python/sglang/srt/models/bailing_moe.py b/python/sglang/srt/models/bailing_moe.py index 019656306..4ac27a69c 100644 --- a/python/sglang/srt/models/bailing_moe.py +++ b/python/sglang/srt/models/bailing_moe.py @@ -77,7 +77,7 @@ from sglang.srt.models.utils import ( create_fused_set_kv_buffer_arg, enable_fused_set_kv_buffer, ) -from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import add_prefix, is_cuda, is_non_idle_and_non_empty, make_layers @@ -813,7 +813,7 @@ class BailingMoEForCausalLM(nn.Module): self.pp_group = get_pp_group() self.config = config self.quant_config = quant_config - alt_stream = torch.cuda.Stream() if _is_cuda else None + alt_stream = get_stream("alt") if _is_cuda else None self.model = BailingMoEModel( config, diff --git a/python/sglang/srt/models/bailing_moe_linear.py b/python/sglang/srt/models/bailing_moe_linear.py index 3a983f0ab..bf46a1eb4 100644 --- a/python/sglang/srt/models/bailing_moe_linear.py +++ b/python/sglang/srt/models/bailing_moe_linear.py @@ -58,7 +58,7 @@ from sglang.srt.model_executor.runner import get_is_capture_mode from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA, DeepseekV2MLP, _is_hip from sglang.srt.models.utils import WeightsMapper -from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import ( BumpAllocator, @@ -957,7 +957,7 @@ class BailingMoELinearModel(nn.Module): else: self.word_embeddings = PPMissingLayer() - self.alt_stream = torch.cuda.Stream() if _is_cuda else None + self.alt_stream = get_stream("alt") if _is_cuda else None def layer_fn(idx, prefix): layer_idx = idx diff --git a/python/sglang/srt/models/exaone_moe.py b/python/sglang/srt/models/exaone_moe.py index 582568a10..32195bea3 100755 --- a/python/sglang/srt/models/exaone_moe.py +++ b/python/sglang/srt/models/exaone_moe.py @@ -62,7 +62,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ( from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_executor.runner import get_is_capture_mode from sglang.srt.model_loader.weight_utils import default_weight_loader -from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import LazyValue, add_prefix, is_cuda, make_layers @@ -637,7 +637,7 @@ class ExaoneMoEForCausalLM(nn.Module): self.pp_group = get_pp_group() self.config = config self.quant_config = quant_config - alt_stream = torch.cuda.Stream() if _is_cuda else None + alt_stream = get_stream("alt") if _is_cuda else None self.model = ExaoneMoEModel( config, quant_config=quant_config, diff --git a/python/sglang/srt/models/falcon_h1.py b/python/sglang/srt/models/falcon_h1.py index 3bd029a92..d215c21f7 100644 --- a/python/sglang/srt/models/falcon_h1.py +++ b/python/sglang/srt/models/falcon_h1.py @@ -33,7 +33,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ( from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_context import get_attn_backend from sglang.srt.model_loader.weight_utils import default_weight_loader -from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.utils import add_prefix, is_cuda, make_layers logger = logging.getLogger(__name__) @@ -387,7 +387,7 @@ class FalconH1Model(nn.Module): super().__init__() self.config = config - alt_stream = torch.cuda.Stream() if _is_cuda else None + alt_stream = get_stream("alt") if _is_cuda else None self.embedding_multiplier = config.embedding_multiplier self.embed_tokens = VocabParallelEmbedding( diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py index 97f20d3e3..1994e1249 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/python/sglang/srt/models/glm4_moe.py @@ -82,7 +82,7 @@ from sglang.srt.model_executor.runner import get_is_capture_mode from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM from sglang.srt.models.utils import apply_qk_norm -from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import ( add_prefix, @@ -1060,7 +1060,7 @@ class Glm4MoeModel(nn.Module): else: self.embed_tokens = PPMissingLayer() - self.alt_stream = torch.cuda.Stream() if _is_cuda else None + self.alt_stream = get_stream("alt") if _is_cuda else None pp_start_layer, _ = get_pp_indices( config.num_hidden_layers, self.pp_group.rank_in_group, diff --git a/python/sglang/srt/models/glm4_moe_lite.py b/python/sglang/srt/models/glm4_moe_lite.py index b0176a250..ecc1df2e2 100644 --- a/python/sglang/srt/models/glm4_moe_lite.py +++ b/python/sglang/srt/models/glm4_moe_lite.py @@ -74,7 +74,7 @@ from sglang.srt.models.deepseek_common.deepseek_weight_loader import ( ) from sglang.srt.models.deepseek_common.utils import _is_cuda, _use_aiter from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA -from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import ( BumpAllocator, @@ -789,7 +789,7 @@ class Glm4MoeLiteModel(nn.Module): else: self.embed_tokens = PPMissingLayer() - self.alt_stream = torch.cuda.Stream() if _is_cuda else None + self.alt_stream = get_stream("alt") if _is_cuda else None self.layers, self.start_layer, self.end_layer = make_layers( config.num_hidden_layers, lambda idx, prefix: Glm4MoeLiteDecoderLayer( diff --git a/python/sglang/srt/models/grok.py b/python/sglang/srt/models/grok.py index ebfef7714..ae2f156ef 100644 --- a/python/sglang/srt/models/grok.py +++ b/python/sglang/srt/models/grok.py @@ -58,7 +58,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.runner import get_is_capture_mode from sglang.srt.model_loader.loader import DefaultModelLoader from sglang.srt.model_loader.weight_utils import default_weight_loader -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_parallel, get_stream from sglang.srt.utils import add_prefix, is_npu _is_npu = is_npu() @@ -646,7 +646,7 @@ class Grok1Model(nn.Module): prefix=add_prefix("embed_tokens", prefix), ) - self.alt_stream = torch.cuda.Stream() + self.alt_stream = get_stream("alt") self.layers = nn.ModuleList( [ Grok1DecoderLayer( diff --git a/python/sglang/srt/models/hunyuan_v3.py b/python/sglang/srt/models/hunyuan_v3.py index a10e1e9b5..8fdeb1a32 100644 --- a/python/sglang/srt/models/hunyuan_v3.py +++ b/python/sglang/srt/models/hunyuan_v3.py @@ -44,7 +44,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ( from sglang.srt.managers.schedule_batch import ForwardBatch from sglang.srt.model_executor.runner import get_is_capture_mode from sglang.srt.model_loader.weight_utils import default_weight_loader -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_parallel, get_stream from sglang.srt.utils import is_cuda from sglang.srt.utils.hf_transformers_utils import get_rope_config @@ -423,7 +423,7 @@ class HYV3Model(nn.Module): prefix=f"{prefix}.embed_tokens", ) - self.alt_stream = torch.cuda.Stream() if is_cuda() else None + self.alt_stream = get_stream("alt") if is_cuda() else None self.layers = nn.ModuleList( [ diff --git a/python/sglang/srt/models/hunyuan_v3_nextn.py b/python/sglang/srt/models/hunyuan_v3_nextn.py index 6ed338428..c1b159982 100644 --- a/python/sglang/srt/models/hunyuan_v3_nextn.py +++ b/python/sglang/srt/models/hunyuan_v3_nextn.py @@ -32,6 +32,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ( from sglang.srt.managers.schedule_batch import ForwardBatch from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.hunyuan_v3 import HYV3DecoderLayer +from sglang.srt.runtime_context import get_stream from sglang.srt.utils import is_cuda logger = logging.getLogger(__name__) @@ -58,7 +59,7 @@ class HYV3ModelNextN(nn.Module): self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) - self.alt_stream = torch.cuda.Stream() if is_cuda() else None + self.alt_stream = get_stream("alt") if is_cuda() else None # Force MoE for the MTP layer: first_k_dense_replace=1 would make # layer_id=0 pick a dense MLP instead of MoE, so override it. diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index 004ba0c49..77c9e2384 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/python/sglang/srt/models/kimi_linear.py @@ -47,7 +47,7 @@ from sglang.srt.model_loader.weight_utils import ( from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA as KimiMLAAttention from sglang.srt.models.llama import LlamaMLP as KimiMLP from sglang.srt.models.transformers import maybe_prefix -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_parallel, get_stream from sglang.srt.utils import make_layers from sglang.srt.utils.common import BumpAllocator, add_prefix, set_weight_attrs @@ -527,7 +527,7 @@ class KimiLinearModel(nn.Module): else: self.embed_tokens = PPMissingLayer() - self.alt_stream = torch.cuda.Stream() + self.alt_stream = get_stream("alt") self.layers, self.start_layer, self.end_layer = make_layers( config.num_hidden_layers, diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index e77017606..81ecc3290 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -76,7 +76,7 @@ from sglang.srt.models.utils import ( create_fused_set_kv_buffer_arg, enable_fused_set_kv_buffer, ) -from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import ( add_prefix, @@ -778,7 +778,7 @@ class LLaDA2MoeModelLM(nn.Module): self.pp_group = get_pp_group() self.config = config self.quant_config = quant_config - alt_stream = torch.cuda.Stream() if _is_cuda else None + alt_stream = get_stream("alt") if _is_cuda else None self.model = LLaDA2MoeModel( config, diff --git a/python/sglang/srt/models/longcat_flash.py b/python/sglang/srt/models/longcat_flash.py index 69215d076..5221c8388 100644 --- a/python/sglang/srt/models/longcat_flash.py +++ b/python/sglang/srt/models/longcat_flash.py @@ -86,7 +86,7 @@ from sglang.srt.model_loader.utils import ( ) from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA -from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.utils import ( BumpAllocator, add_prefix, @@ -538,7 +538,7 @@ class LongcatFlashModel(nn.Module): use_attn_tp_group=is_dp_attention_enabled(), ) - self.alt_stream = torch.cuda.Stream() + self.alt_stream = get_stream("alt") self.layers = nn.ModuleList( [ LongcatFlashDecoderLayer( diff --git a/python/sglang/srt/models/longcat_flash_nextn.py b/python/sglang/srt/models/longcat_flash_nextn.py index 91fc84398..234f70a79 100644 --- a/python/sglang/srt/models/longcat_flash_nextn.py +++ b/python/sglang/srt/models/longcat_flash_nextn.py @@ -68,7 +68,7 @@ from sglang.srt.model_loader.utils import should_deepgemm_weight_requant_ue8m0 from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA from sglang.srt.models.longcat_flash import LongcatFlashForCausalLM, LongcatFlashMLP -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_parallel, get_stream from sglang.srt.utils import ( BumpAllocator, add_prefix, @@ -207,7 +207,7 @@ class LongcatFlashModelNextN(nn.Module): ) -> None: super().__init__() self.vocab_size = config.vocab_size - self.alt_stream = torch.cuda.Stream() + self.alt_stream = get_stream("alt") self.embed_tokens = VocabParallelEmbedding( config.vocab_size, diff --git a/python/sglang/srt/models/olmo2.py b/python/sglang/srt/models/olmo2.py index 761a08330..a743cd75a 100644 --- a/python/sglang/srt/models/olmo2.py +++ b/python/sglang/srt/models/olmo2.py @@ -47,7 +47,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ( from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.runner import get_is_capture_mode from sglang.srt.model_loader.weight_utils import default_weight_loader -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_parallel, get_stream from sglang.srt.utils import add_prefix, is_cuda, make_layers _is_cuda = is_cuda() @@ -332,7 +332,7 @@ class Olmo2Model(nn.Module): super().__init__() self.config = config if alt_stream is None and _is_cuda: - alt_stream = torch.cuda.Stream() + alt_stream = get_stream("alt") self.alt_stream = alt_stream self.embed_tokens = VocabParallelEmbedding( diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 123277bdd..fc18c3408 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -114,6 +114,7 @@ if is_npu(): ) from sglang.srt.environ import envs +from sglang.srt.runtime_context import get_stream from sglang.srt.utils.hf_transformers_utils import get_rope_config _SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get() @@ -991,7 +992,7 @@ class Qwen2MoeForCausalLM(nn.Module): self.pp_group = get_pp_group() self.config = config self.quant_config = quant_config - alt_stream = torch.cuda.Stream() if _is_cuda else None + alt_stream = get_stream("alt") if _is_cuda else None self.model = Qwen2MoeModel( config, quant_config, diff --git a/python/sglang/srt/models/qwen3.py b/python/sglang/srt/models/qwen3.py index 05fb5372a..8543bd4b0 100644 --- a/python/sglang/srt/models/qwen3.py +++ b/python/sglang/srt/models/qwen3.py @@ -33,7 +33,7 @@ from sglang.srt.model_loader.weight_utils import ( from sglang.srt.models.qwen2 import Qwen2MLP as Qwen3MLP from sglang.srt.models.qwen2 import Qwen2Model from sglang.srt.models.utils import apply_qk_norm -from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import add_prefix, get_bool_env_var, is_cuda, is_hip, is_npu @@ -440,7 +440,7 @@ class Qwen3Model(Qwen2Model): quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: - alt_stream = torch.cuda.Stream() if _is_cuda else None + alt_stream = get_stream("alt") if _is_cuda else None super().__init__( config=config, quant_config=quant_config, diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 283820957..7db4a819d 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -91,7 +91,7 @@ from sglang.srt.models.utils import ( fused_qk_gemma_rmsnorm, fused_qk_gemma_rmsnorm_with_gate, ) -from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream # Utils from sglang.srt.utils import ( @@ -1206,7 +1206,7 @@ class Qwen3_5ForCausalLM(nn.Module): if _is_hip: self._maybe_autodisable_shared_experts_fusion(config, quant_config) - alt_stream = torch.cuda.Stream() if _is_cuda or _hip_use_alt_stream else None + alt_stream = get_stream("alt") if _is_cuda or _hip_use_alt_stream else None # Embedding layer if self.pp_group.is_first_rank: diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py index 26c7f4964..9702e11dd 100644 --- a/python/sglang/srt/models/qwen3_moe.py +++ b/python/sglang/srt/models/qwen3_moe.py @@ -72,7 +72,7 @@ from sglang.srt.models.utils import ( create_fused_set_kv_buffer_arg, enable_fused_set_kv_buffer, ) -from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import ( LazyValue, @@ -916,7 +916,7 @@ class Qwen3MoeModel(Qwen2MoeModel): prefix: str = "", decoder_layer_type=Qwen3MoeDecoderLayer, ) -> None: - alt_stream = torch.cuda.Stream() if _is_cuda else None + alt_stream = get_stream("alt") if _is_cuda else None super().__init__( config=config, quant_config=quant_config, diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py index 233462d3f..d65ec3e24 100644 --- a/python/sglang/srt/models/qwen3_next.py +++ b/python/sglang/srt/models/qwen3_next.py @@ -47,7 +47,7 @@ from sglang.srt.model_loader.weight_utils import ( sharded_weight_loader, ) from sglang.srt.models.qwen2_moe import Qwen2MoeMLP, Qwen2MoeSparseMoeBlock -from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.utils import ( LazyValue, add_prefix, @@ -889,7 +889,7 @@ class Qwen3NextModel(nn.Module): super().__init__() self.config = config - alt_stream = torch.cuda.Stream() if _is_cuda else None + alt_stream = get_stream("alt") if _is_cuda else None self.embed_tokens = VocabParallelEmbedding( config.vocab_size, diff --git a/python/sglang/srt/models/sarvam_moe.py b/python/sglang/srt/models/sarvam_moe.py index 2eba09d64..afa7565c0 100644 --- a/python/sglang/srt/models/sarvam_moe.py +++ b/python/sglang/srt/models/sarvam_moe.py @@ -60,7 +60,7 @@ from sglang.srt.models.bailing_moe import BailingMoEForCausalLM from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mha import ( DeepseekMHAForwardMixin, ) -from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import ( BumpAllocator, @@ -1152,7 +1152,7 @@ class SarvamMLAModel(nn.Module): self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.pp_group = get_pp_group() - self.alt_stream = torch.cuda.Stream() if _is_cuda else None + self.alt_stream = get_stream("alt") if _is_cuda else None if self.pp_group.is_first_rank: self.embed_tokens = VocabParallelEmbedding( diff --git a/python/sglang/srt/models/sdar.py b/python/sglang/srt/models/sdar.py index ed52aec69..68bc12ce7 100644 --- a/python/sglang/srt/models/sdar.py +++ b/python/sglang/srt/models/sdar.py @@ -41,7 +41,7 @@ from sglang.srt.models.utils import ( create_fused_set_kv_buffer_arg, enable_fused_set_kv_buffer, ) -from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import add_prefix, is_cuda, make_layers @@ -449,7 +449,7 @@ class SDARForCausalLM(nn.Module): self.config = config self.quant_config = quant_config - alt_stream = torch.cuda.Stream() if _is_cuda else None + alt_stream = get_stream("alt") if _is_cuda else None self.model = SDARModel( config, diff --git a/python/sglang/srt/models/sdar_moe.py b/python/sglang/srt/models/sdar_moe.py index 858991d67..42c569d3f 100644 --- a/python/sglang/srt/models/sdar_moe.py +++ b/python/sglang/srt/models/sdar_moe.py @@ -57,7 +57,7 @@ from sglang.srt.models.utils import ( create_fused_set_kv_buffer_arg, enable_fused_set_kv_buffer, ) -from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import LazyValue, add_prefix, is_cuda, make_layers @@ -544,7 +544,7 @@ class SDARMoeForCausalLM(nn.Module): self.pp_group = get_pp_group() self.config = config self.quant_config = quant_config - alt_stream = torch.cuda.Stream() if _is_cuda else None + alt_stream = get_stream("alt") if _is_cuda else None self.model = SDARMoeModel( config, diff --git a/python/sglang/srt/models/step3p5.py b/python/sglang/srt/models/step3p5.py index 643f7f0c0..a55ce02f5 100644 --- a/python/sglang/srt/models/step3p5.py +++ b/python/sglang/srt/models/step3p5.py @@ -46,7 +46,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ( ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import default_weight_loader -from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.runtime_context import get_parallel, get_server_args, get_stream from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import add_prefix, is_cuda, is_non_idle_and_non_empty, make_layers @@ -670,7 +670,7 @@ class Step3p5Model(nn.Module): self.vocab_size = config.vocab_size self.pp_group = get_pp_group() - alt_stream = torch.cuda.Stream() if _is_cuda else None + alt_stream = get_stream("alt") if _is_cuda else None if self.pp_group.is_first_rank: self.embed_tokens = VocabParallelEmbedding( diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index 16dd824c8..cc8a112a2 100644 --- a/python/sglang/srt/runtime_context.py +++ b/python/sglang/srt/runtime_context.py @@ -321,16 +321,73 @@ class Flags(_FlagGroupBase): dp: DpFlags = dataclasses.field(default_factory=DpFlags) +@dataclasses.dataclass +class Resources(_FlagGroupBase): + """Process-level resource handles: named slots with one reset lifecycle, + scoped test injection via ``override()``, and the creation/publish + semantics kept in the owning modules' accessors (which are thin shims + over these slots).""" + + # CUDA graph memory pool shared across the prefill and decode graph + # backends (created lazily by model_executor.runner_utils.pool). + graph_memory_pool: Any = None + # EPLB: per-process recorder and the publish-once location metadata + # (owning accessors live in sglang.srt.eplb). + expert_distribution_recorder: Any = None + expert_location_metadata: Any = None + # LPLB: layer_id -> solver. + lplb_solvers: dict = dataclasses.field(default_factory=dict) + # Named side streams (see RuntimeContext.get_stream): name -> stream. + streams: dict = dataclasses.field(default_factory=dict) + # Named persistent buffers (see RuntimeContext.get_buffer): name -> tensor. + # Accessors with bespoke semantics (grow-only, per-device keys) manage + # their entries directly. + buffers: dict = dataclasses.field(default_factory=dict) + # Persistent reusable CUDA events for non-EP DP TBO, keyed by + # (kind, subbatch) — see dp_attention._tbo_event for why reuse matters. + tbo_event_pool: dict = dataclasses.field(default_factory=dict) + + class RuntimeContext: """Container for the structured runtime accessors; exposes ``parallel``, - ``server_args``, and ``flags``.""" + ``server_args``, ``flags``, and ``resources``.""" - __slots__ = ("parallel", "_server_args", "flags") + __slots__ = ("parallel", "_server_args", "flags", "resources") def __init__(self, parallel: ParallelContext): self.parallel = parallel self._server_args: ServerArgs | None = None self.flags = Flags() + self.resources = Resources() + + def get_stream(self, name: str) -> Any: + """Named process-level CUDA side stream: get-or-create, shared by + name (the keyed-lazy pattern of the persistent buffers). Creation is + a driver call that must stay outside cuda-graph capture — call sites + lease their stream at init/warmup time.""" + stream = self.resources.streams.get(name) + if stream is None: + import torch + + stream = torch.cuda.Stream() + self.resources.streams[name] = stream + return stream + + def set_stream(self, name: str, stream: Any) -> Any: + """Install (or replace) the named stream — explicit injection for + tests and backends that bring their own stream.""" + self.resources.streams[name] = stream + return stream + + def get_buffer(self, name: str, factory: Any) -> Any: + """Named process-level persistent buffer: get-or-create via + ``factory()``, shared by name (the keyed-lazy pattern of the + persistent buffers / named streams).""" + buf = self.resources.buffers.get(name) + if buf is None: + buf = factory() + self.resources.buffers[name] = buf + return buf @property def server_args(self) -> ServerArgs: @@ -378,11 +435,28 @@ def get_flags() -> Flags: return _CONTEXT.flags +def get_resources() -> Resources: + return _CONTEXT.resources + + +def get_stream(name: str) -> Any: + return _CONTEXT.get_stream(name) + + +def set_stream(name: str, stream: Any) -> Any: + return _CONTEXT.set_stream(name, stream) + + +def get_buffer(name: str, factory: Any) -> Any: + return _CONTEXT.get_buffer(name, factory) + + def reset_context() -> None: """Clear the context-owned store (unit-test teardown): drop the published - ``server_args`` and install a fresh ``Flags``. + ``server_args`` and install fresh ``Flags`` and ``Resources``. Wrapper subsystems (``parallel``) hold no state and are unaffected. """ _CONTEXT._server_args = None _CONTEXT.flags = Flags() + _CONTEXT.resources = Resources() diff --git a/python/sglang/srt/utils/offloader.py b/python/sglang/srt/utils/offloader.py index a2e1df8ac..66adcca99 100644 --- a/python/sglang/srt/utils/offloader.py +++ b/python/sglang/srt/utils/offloader.py @@ -12,6 +12,7 @@ from sglang.srt.distributed.naive_distributed import ( set_naive_distributed, ) from sglang.srt.layers.parameter import ModelWeightParameter +from sglang.srt.runtime_context import get_stream from sglang.srt.server_args import ServerArgs from sglang.srt.utils import MultiprocessingSerializer, is_pin_memory_available from sglang.srt.utils.host_shared_memory import ( @@ -198,7 +199,10 @@ class OffloaderV2(BaseOffloader): ): assert len(self.offloaders) == 0, "should only call wrap_modules once" - alt_stream = torch.cuda.Stream() + # The offloader's async prefetch/offload copies run on their own + # stream — sharing the models' "alt" overlap stream would serialize + # unrelated copy and compute work. + alt_stream = get_stream("offload") all_modules = [] offload_submodules = [] diff --git a/test/registered/moe/test_hash_topk.py b/test/registered/moe/test_hash_topk.py index 2ef89cfb1..93452911f 100644 --- a/test/registered/moe/test_hash_topk.py +++ b/test/registered/moe/test_hash_topk.py @@ -33,11 +33,9 @@ def test_hash_topk_remaps_per_rank_fused_shared_slots(monkeypatch): def on_select_experts(self, *, topk_ids): recorded["topk_ids"] = topk_ids.clone() - monkeypatch.setattr( - hash_topk_module, - "get_global_expert_distribution_recorder", - lambda: FakeRecorder(), - ) + from sglang.srt.runtime_context import get_resources + + monkeypatch.setattr(get_resources(), "expert_distribution_recorder", FakeRecorder()) topk = HashTopK( topk=3, diff --git a/test/registered/unit/test_module_state_ratchet.py b/test/registered/unit/test_module_state_ratchet.py index fb0c93572..126330114 100644 --- a/test/registered/unit/test_module_state_ratchet.py +++ b/test/registered/unit/test_module_state_ratchet.py @@ -36,8 +36,6 @@ _PINNED_GLOBALS = { "_ATTN_DP_SIZE", "_LOCAL_ATTN_DP_SIZE", "_LOCAL_ATTN_DP_RANK", - # Comm stream resource (resources vertical scope). - "_DP_TBO_COMM_STREAM", } ), } diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index a5e152689..b903b6a6e 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -367,6 +367,102 @@ class TestDpFlagsGroup(_IsolatedServerArgs): self.assertFalse(is_dp_attention_enabled()) +class TestResources(_IsolatedServerArgs): + """ctx.resources: named slots for process-level resource handles with one + reset lifecycle; owning accessors keep their creation/publish semantics.""" + + def test_graph_pool_lazy_create_and_reuse(self): + from types import SimpleNamespace + + from sglang.srt.model_executor.runner_utils.pool import ( + get_global_graph_memory_pool, + get_or_create_global_graph_memory_pool, + ) + + reset_context() + self.assertIsNone(get_global_graph_memory_pool()) + dev = SimpleNamespace(graph_pool_handle=lambda: object()) + handle = get_or_create_global_graph_memory_pool(dev) + self.assertIs(get_or_create_global_graph_memory_pool(dev), handle) + + def test_expert_recorder_noop_default_and_injection(self): + from sglang.srt.eplb.expert_distribution import ( + get_global_expert_distribution_recorder, + ) + from sglang.srt.runtime_context import get_resources + + reset_context() + self.assertEqual( + type(get_global_expert_distribution_recorder()).__name__, + "_ExpertDistributionRecorderNoop", + ) + with get_resources().override(expert_distribution_recorder="mock"): + self.assertEqual(get_global_expert_distribution_recorder(), "mock") + + def test_expert_location_metadata_publish_once_until_reset(self): + from sglang.srt.eplb.expert_location import ( + get_global_expert_location_metadata, + set_global_expert_location_metadata, + ) + + reset_context() + self.assertIsNone(get_global_expert_location_metadata()) + set_global_expert_location_metadata("meta") + with self.assertRaises(AssertionError): + set_global_expert_location_metadata("again") + reset_context() + self.assertIsNone(get_global_expert_location_metadata()) + + +class TestNamedStreams(_IsolatedServerArgs): + """ctx.get_stream(name): keyed get-or-create (the persistent-buffer + pattern); set_stream installs explicitly.""" + + def test_get_or_create_shares_by_name(self): + from unittest.mock import patch + + reset_context() + created = [] + + class _FakeStream: + def __init__(self): + created.append(self) + + with patch("torch.cuda.Stream", _FakeStream): + a = get_context().get_stream("alt") + b = get_context().get_stream("alt") + c = get_context().get_stream("other") + self.assertIs(a, b) + self.assertIsNot(a, c) + self.assertEqual(len(created), 2) + + def test_get_buffer_keyed_lazy(self): + reset_context() + created = [] + + def factory(): + created.append(object()) + return created[-1] + + a = get_context().get_buffer("ws", factory) + b = get_context().get_buffer("ws", factory) + self.assertIs(a, b) + self.assertEqual(len(created), 1) + self.assertIsNot(get_context().get_buffer("other", factory), a) + + def test_set_stream_installs_explicitly(self): + reset_context() + sentinel = object() + get_context().set_stream("alt", sentinel) + self.assertIs(get_context().get_stream("alt"), sentinel) + + def test_reset_clears_the_registry(self): + reset_context() + get_context().set_stream("alt", object()) + reset_context() + self.assertEqual(get_context().resources.streams, {}) + + class TestPublishLifecycle(_IsolatedServerArgs): """Publish installs the resolved server_args and seeds the capture tier."""