diff --git a/python/pyproject.toml b/python/pyproject.toml index 5a90da6d6..8d9519e48 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -18,6 +18,7 @@ dependencies = [ "IPython", "aiohttp", "anthropic>=0.20.0", + "av ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'arm64' and platform_machine == 'armv7l')", "blobfile==3.0.0", "build", "compressed-tensors", @@ -38,6 +39,7 @@ dependencies = [ "ninja", "numpy", "nvidia-cutlass-dsl==4.2.1", + "nvidia-ml-py", "openai-harmony==0.0.4", "openai==2.6.1", "orjson", @@ -50,7 +52,6 @@ dependencies = [ "py-spy", "pybase64", "pydantic", - "nvidia-ml-py", "python-multipart", "pyzmq>=25.1.2", "requests", @@ -63,9 +64,8 @@ dependencies = [ "timm==1.0.16", "torch_memory_saver==0.0.9", "torch==2.9.1", - "torchcodec==0.7.0 ; sys_platform != 'linux' or (sys_platform == 'linux' and platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')", # torchcodec does not exist in those systems. If not provided, transformer will use torchvision instead by default. - "av ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'arm64' and platform_machine == 'armv7l')", "torchaudio==2.9.1", + "torchcodec==0.7.0 ; sys_platform != 'linux' or (sys_platform == 'linux' and platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')", # torchcodec does not exist in those systems. If not provided, transformer will use torchvision instead by default. "torchvision", "torchao==0.9.0", "tqdm", @@ -73,6 +73,7 @@ dependencies = [ "uvicorn", "uvloop", "xgrammar==0.1.27", + "grpcio==1.75.1", # keep it align with compile_proto.py "grpcio-tools==1.75.1", # keep it align with compile_proto.py "grpcio-reflection==1.75.1", # required by srt/entrypoints/grpc_server.py @@ -82,17 +83,17 @@ dependencies = [ [project.optional-dependencies] checkpoint-engine = ["checkpoint-engine==0.1.2"] diffusion = [ - "diffusers==0.35.2", - "yunchang==0.6.3.post1", - "opencv-python==4.10.0.84", - "imageio==2.36.0", - "imageio-ffmpeg==0.5.1", - "PyYAML==6.0.1", - "moviepy>=2.0.0", - "cloudpickle", - "remote-pdb", - "st_attn ==0.0.7", - "vsa==0.0.4", + "PyYAML==6.0.1", + "cloudpickle", + "diffusers==0.35.2", + "imageio==2.36.0", + "imageio-ffmpeg==0.5.1", + "moviepy>=2.0.0", + "opencv-python==4.10.0.84", + "remote-pdb", + "st_attn ==0.0.7", + "vsa==0.0.4", + "yunchang==0.6.3.post1", ] [tool.uv.extra-build-dependencies] diff --git a/python/sglang/compile_deep_gemm.py b/python/sglang/compile_deep_gemm.py index 11370acde..2386a46c0 100644 --- a/python/sglang/compile_deep_gemm.py +++ b/python/sglang/compile_deep_gemm.py @@ -19,12 +19,12 @@ import requests from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST from sglang.srt.entrypoints.http_server import launch_server +from sglang.srt.entrypoints.warmup import warmup from sglang.srt.environ import envs from sglang.srt.managers.io_struct import GenerateReqInput from sglang.srt.managers.tokenizer_manager import TokenizerManager from sglang.srt.server_args import ServerArgs from sglang.srt.utils import kill_process_tree -from sglang.srt.warmup import warmup multiprocessing.set_start_method("spawn", force=True) diff --git a/python/sglang/srt/operations.py b/python/sglang/srt/batch_overlap/operations.py similarity index 100% rename from python/sglang/srt/operations.py rename to python/sglang/srt/batch_overlap/operations.py diff --git a/python/sglang/srt/batch_overlap/operations_strategy.py b/python/sglang/srt/batch_overlap/operations_strategy.py new file mode 100644 index 000000000..152e4874d --- /dev/null +++ b/python/sglang/srt/batch_overlap/operations_strategy.py @@ -0,0 +1,211 @@ +from dataclasses import dataclass +from typing import List, Optional + +import torch + +from sglang.srt.batch_overlap import operations +from sglang.srt.batch_overlap.operations import Operation +from sglang.srt.layers.moe.token_dispatcher import DeepEPConfig +from sglang.srt.model_executor.forward_batch_info import ForwardMode + + +@dataclass +class OperationsStrategy: + operations: List[Operation] + deep_gemm_num_sms: Optional[int] = None + tbo_delta_stages: Optional[int] = None + + @classmethod + def concat(cls, items: List["OperationsStrategy"]) -> "OperationsStrategy": + return OperationsStrategy( + operations=[x for item in items for x in item.operations], + deep_gemm_num_sms=_assert_all_same( + [item.deep_gemm_num_sms for item in items] + ), + tbo_delta_stages=_assert_all_same( + [item.tbo_delta_stages for item in items] + ), + ) + + @staticmethod + def init_new_tbo( + layers: torch.nn.ModuleList, + forward_mode: ForwardMode, + ) -> "OperationsStrategy": + layer_name = layers[0].__class__.__name__ + if layer_name == "DeepseekV2DecoderLayer": + return OperationsStrategy.concat( + [ + _compute_moe_deepseek_layer_operations_strategy_tbo( + layer, forward_mode + ) + for layer in layers + ] + ) + elif layer_name == "Qwen3MoeDecoderLayer": + return OperationsStrategy.concat( + [ + _compute_moe_qwen3_layer_operations_strategy_tbo( + layer, forward_mode + ) + for layer in layers + ] + ) + else: + raise NotImplementedError + + +def _assert_all_same(items: List): + assert all(item == items[0] for item in items) + return items[0] + + +# -------------------------------- Strategy for DeepSeek --------------------------------------- + + +# TODO can refactor to make it more fancy if we have more complex strategies +def _compute_moe_deepseek_layer_operations_strategy_tbo( + layer: torch.nn.Module, + forward_mode: ForwardMode, +) -> OperationsStrategy: + assert layer.is_layer_sparse, "dense layer TBO not yet implemented" + if forward_mode == ForwardMode.EXTEND: + return _compute_moe_deepseek_blog_prefill(layer) + elif ( + forward_mode == ForwardMode.DECODE or forward_mode == ForwardMode.TARGET_VERIFY + ): + return _compute_moe_deepseek_blog_decode(layer) + else: + raise NotImplementedError(f"Unsupported {forward_mode=}") + + +def _compute_moe_deepseek_blog_prefill(layer): + device_properties = torch.cuda.get_device_properties(device="cuda") + total_num_sms = device_properties.multi_processor_count + deep_gemm_num_sms = total_num_sms - DeepEPConfig.get_instance().num_sms + + return OperationsStrategy( + deep_gemm_num_sms=deep_gemm_num_sms, + tbo_delta_stages=0, + operations=[ + layer.op_comm_prepare_attn, + layer.self_attn.op_prepare, + layer.self_attn.op_core, + layer.op_comm_prepare_mlp, + layer.mlp.op_gate, + layer.mlp.op_select_experts, + layer.mlp.op_dispatch_a, + operations.YieldOperation(), + layer.mlp.op_dispatch_b, + layer.mlp.op_experts, + layer.mlp.op_combine_a, + operations.YieldOperation(), + layer.mlp.op_shared_experts, + layer.mlp.op_combine_b, + layer.mlp.op_output, + layer.op_comm_postprocess_layer, + ], + ) + + +def _compute_moe_deepseek_blog_decode(layer): + return OperationsStrategy( + deep_gemm_num_sms=None, + tbo_delta_stages=2, + operations=[ + layer.op_comm_prepare_attn, + layer.self_attn.op_prepare, + operations.YieldOperation(), + layer.self_attn.op_core, + layer.op_comm_prepare_mlp, + layer.mlp.op_gate, + layer.mlp.op_select_experts, + operations.YieldOperation(), + layer.mlp.op_dispatch_a, + layer.mlp.op_shared_experts, + operations.YieldOperation(), + layer.mlp.op_dispatch_b, + layer.mlp.op_experts, + layer.mlp.op_combine_a, + operations.YieldOperation(), + layer.mlp.op_combine_b, + operations.YieldOperation(), + layer.mlp.op_output, + layer.op_comm_postprocess_layer, + ], + ) + + +# -------------------------------- Strategy for Qwen3 --------------------------------------- + + +# TODO: unstable, current strategy is almost the same as DeepSeek, keep redundant code here for +# convenience to adjust strategy +def _compute_moe_qwen3_layer_operations_strategy_tbo( + layer: torch.nn.Module, + forward_mode: ForwardMode, +) -> OperationsStrategy: + assert layer.is_layer_sparse, "qwen3 moe only support sparse layers" + if forward_mode == ForwardMode.EXTEND: + return _compute_moe_qwen3_prefill(layer) + elif ( + forward_mode == ForwardMode.DECODE or forward_mode == ForwardMode.TARGET_VERIFY + ): + return _compute_moe_qwen3_decode(layer) + else: + raise NotImplementedError(f"Unsupported {forward_mode=}") + + +def _compute_moe_qwen3_prefill(layer): + device_properties = torch.cuda.get_device_properties(device="cuda") + total_num_sms = device_properties.multi_processor_count + deep_gemm_num_sms = total_num_sms - DeepEPConfig.get_instance().num_sms + + return OperationsStrategy( + deep_gemm_num_sms=deep_gemm_num_sms, + tbo_delta_stages=0, + operations=[ + layer.op_comm_prepare_attn, + layer.self_attn.op_prepare, + layer.self_attn.op_core, + layer.op_comm_prepare_mlp, + layer.mlp.op_gate, + layer.mlp.op_select_experts, + layer.mlp.op_dispatch_a, + operations.YieldOperation(), + layer.mlp.op_dispatch_b, + layer.mlp.op_experts, + layer.mlp.op_combine_a, + operations.YieldOperation(), + layer.mlp.op_combine_b, + layer.mlp.op_output, + layer.op_comm_postprocess_layer, + ], + ) + + +def _compute_moe_qwen3_decode(layer): + return OperationsStrategy( + deep_gemm_num_sms=None, + tbo_delta_stages=2, + operations=[ + layer.op_comm_prepare_attn, + layer.self_attn.op_prepare, + operations.YieldOperation(), + layer.self_attn.op_core, + layer.op_comm_prepare_mlp, + layer.mlp.op_gate, + layer.mlp.op_select_experts, + operations.YieldOperation(), + layer.mlp.op_dispatch_a, + operations.YieldOperation(), + layer.mlp.op_dispatch_b, + layer.mlp.op_experts, + layer.mlp.op_combine_a, + operations.YieldOperation(), + layer.mlp.op_combine_b, + layer.mlp.op_output, + layer.op_comm_postprocess_layer, + operations.YieldOperation(), + ], + ) diff --git a/python/sglang/srt/single_batch_overlap.py b/python/sglang/srt/batch_overlap/single_batch_overlap.py similarity index 100% rename from python/sglang/srt/single_batch_overlap.py rename to python/sglang/srt/batch_overlap/single_batch_overlap.py diff --git a/python/sglang/srt/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py similarity index 99% rename from python/sglang/srt/two_batch_overlap.py rename to python/sglang/srt/batch_overlap/two_batch_overlap.py index 97cd7dff4..becb54e44 100644 --- a/python/sglang/srt/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -8,6 +8,11 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Sequence import torch +from sglang.srt.batch_overlap.operations import ( + execute_operations, + execute_overlapped_operations, +) +from sglang.srt.batch_overlap.operations_strategy import OperationsStrategy from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.communicator import ( @@ -32,15 +37,13 @@ from sglang.srt.model_executor.forward_batch_info import ( ForwardMode, compute_position, ) -from sglang.srt.operations import execute_operations, execute_overlapped_operations -from sglang.srt.operations_strategy import OperationsStrategy from sglang.srt.server_args import get_global_server_args from sglang.srt.speculative.spec_info import SpecInput from sglang.srt.utils import BumpAllocator, empty_context, get_bool_env_var, is_hip if TYPE_CHECKING: + from sglang.srt.batch_overlap.single_batch_overlap import CombineOverlapArgs from sglang.srt.layers.moe.token_dispatcher import DispatchOutput - from sglang.srt.single_batch_overlap import CombineOverlapArgs from sglang.srt.speculative.eagle_info import EagleVerifyInput _is_hip = is_hip() diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index f43843575..cf905378d 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -72,6 +72,7 @@ from sglang.srt.entrypoints.openai.serving_tokenize import ( OpenAIServingDetokenize, OpenAIServingTokenize, ) +from sglang.srt.entrypoints.warmup import execute_warmups from sglang.srt.environ import envs from sglang.srt.function_call.function_call_parser import FunctionCallParser from sglang.srt.managers.io_struct import ( @@ -126,7 +127,6 @@ from sglang.srt.utils import ( kill_process_tree, set_uvicorn_logging_configs, ) -from sglang.srt.warmup import execute_warmups from sglang.utils import get_exception_traceback from sglang.version import __version__ diff --git a/python/sglang/srt/warmup.py b/python/sglang/srt/entrypoints/warmup.py similarity index 100% rename from python/sglang/srt/warmup.py rename to python/sglang/srt/entrypoints/warmup.py diff --git a/python/sglang/srt/layers/attention/tbo_backend.py b/python/sglang/srt/layers/attention/tbo_backend.py index bdecfb380..04e7b07c7 100644 --- a/python/sglang/srt/layers/attention/tbo_backend.py +++ b/python/sglang/srt/layers/attention/tbo_backend.py @@ -2,7 +2,7 @@ from typing import TYPE_CHECKING, Callable, List, Optional import torch -from sglang.srt import two_batch_overlap +from sglang.srt.batch_overlap import two_batch_overlap from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.speculative.spec_info import SpecInput diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 32c4d2fd7..cb8b41ba5 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -6,6 +6,8 @@ from typing import List, Optional, Tuple import torch +from sglang.srt.batch_overlap.single_batch_overlap import DownGemmOverlapArgs +from sglang.srt.batch_overlap.two_batch_overlap import MaybeTboDeepEPDispatcher from sglang.srt.distributed import ( get_moe_expert_parallel_rank, get_moe_expert_parallel_world_size, @@ -46,8 +48,6 @@ from sglang.srt.layers.quantization.modelopt_quant import ModelOptNvFp4FusedMoEM from sglang.srt.layers.quantization.unquant import UnquantizedFusedMoEMethod from sglang.srt.model_loader.weight_utils import narrow_padded_param_and_loaded_weight from sglang.srt.server_args import get_global_server_args -from sglang.srt.single_batch_overlap import DownGemmOverlapArgs -from sglang.srt.two_batch_overlap import MaybeTboDeepEPDispatcher from sglang.srt.utils import ( cpu_has_amx_support, get_bool_env_var, diff --git a/python/sglang/srt/layers/moe/token_dispatcher/base.py b/python/sglang/srt/layers/moe/token_dispatcher/base.py index a5e61c107..48ebf22cf 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/base.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/base.py @@ -19,6 +19,7 @@ from typing import ( import torch if TYPE_CHECKING: + from sglang.srt.batch_overlap.single_batch_overlap import CombineOverlapArgs from sglang.srt.layers.moe.token_dispatcher import ( DeepEPLLCombineInput, DeepEPLLDispatchOutput, @@ -28,7 +29,6 @@ if TYPE_CHECKING: StandardDispatchOutput, ) from sglang.srt.layers.moe.topk import TopKOutput - from sglang.srt.single_batch_overlap import CombineOverlapArgs # ------------------------------ Dispatcher Hook ------------------------------------- diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index 17ca310ab..c73f2ce8b 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -34,7 +34,7 @@ from sglang.srt.utils import ( _is_npu = is_npu() if TYPE_CHECKING: - from sglang.srt.single_batch_overlap import CombineOverlapArgs + from sglang.srt.batch_overlap.single_batch_overlap import CombineOverlapArgs try: from deep_ep import Buffer, Config diff --git a/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py b/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py index d5a109072..c565ab855 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py @@ -19,7 +19,7 @@ from sglang.srt.layers.moe.utils import DeepEPMode from sglang.srt.utils import get_int_env_var if TYPE_CHECKING: - from sglang.srt.single_batch_overlap import CombineOverlapArgs + from sglang.srt.batch_overlap.single_batch_overlap import CombineOverlapArgs from enum import Enum, auto diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index efa8847ff..7546dfe1a 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -30,6 +30,11 @@ from typing import ( import torch +try: + from triton_kernels.routing import GatherIndx, RoutingData, ScatterIndx, routing +except ImportError: + pass + from sglang.srt.custom_op import CustomOp from sglang.srt.distributed import get_tp_group from sglang.srt.distributed.device_communicators.pynccl_allocator import ( @@ -57,13 +62,8 @@ from sglang.srt.utils.patch_torch import register_fake_if_exists if TYPE_CHECKING: from sglang.srt.layers.quantization import QuantizationConfig -try: - from triton_kernels.routing import GatherIndx, RoutingData, ScatterIndx, routing -except ImportError: - pass + logger = logging.getLogger(__name__) - - _is_cuda = is_cuda() _is_hip = is_hip() _is_cpu = is_cpu() @@ -80,7 +80,12 @@ if _is_cuda: pass if _is_cuda or _is_hip: - from sgl_kernel import topk_sigmoid, topk_softmax + from sgl_kernel import topk_softmax + + try: + from sgl_kernel import topk_sigmoid + except ImportError: + pass if _use_aiter: try: from aiter import biased_grouped_topk as aiter_biased_grouped_topk diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 878f70619..aa39fcbcc 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -88,13 +88,11 @@ _is_cuda = is_cuda() _is_npu = is_npu() _is_cpu_amx_available = cpu_has_amx_support() _is_cpu = is_cpu() - _is_fp8_fnuz = is_fp8_fnuz() - -_use_hip_int4 = get_bool_env_var("SGLANG_INT4_WEIGHT") +_use_hip_int4 = get_bool_env_var("SGLANG_INT4_WEIGHT") and _is_hip _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip -if _is_hip and (_use_aiter or _use_hip_int4): +if _use_aiter or _use_hip_int4: from aiter import ActivationType, QuantType from aiter.fused_moe import fused_moe from aiter.ops.shuffle import shuffle_weight diff --git a/python/sglang/srt/layers/quantization/fp8_kernel.py b/python/sglang/srt/layers/quantization/fp8_kernel.py index 814c2650a..86c0982d2 100644 --- a/python/sglang/srt/layers/quantization/fp8_kernel.py +++ b/python/sglang/srt/layers/quantization/fp8_kernel.py @@ -45,10 +45,6 @@ _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip if _is_cuda: from sgl_kernel import sgl_per_tensor_quant_fp8, sgl_per_token_quant_fp8 - @torch.library.register_fake("sgl_kernel::sgl_per_tensor_quant_fp8") - def _sgl_per_tensor_quant_fp8(input, output_q, output_s, is_static): - return - # Temporary try: from sgl_kernel import sgl_per_token_group_quant_8bit @@ -1861,3 +1857,7 @@ if _is_cuda: @torch.library.register_fake("sgl_kernel::sgl_per_token_quant_fp8") def _(input, output_q, output_s): return + + @torch.library.register_fake("sgl_kernel::sgl_per_tensor_quant_fp8") + def _sgl_per_tensor_quant_fp8(input, output_q, output_s, is_static): + return diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 6b106379e..0da0aba82 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -53,12 +53,12 @@ from sglang.srt.utils.common import ( from sglang.srt.utils.patch_torch import register_fake_if_exists if TYPE_CHECKING: + from sglang.srt.batch_overlap.single_batch_overlap import DownGemmOverlapArgs from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.moe.token_dispatcher import ( CombineInput, StandardDispatchOutput, ) - from sglang.srt.single_batch_overlap import DownGemmOverlapArgs try: if is_sm120_supported(): diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index e8f3eb611..e9347e811 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -21,7 +21,6 @@ from sglang.srt.managers.schedule_batch import ( ) from sglang.srt.mem_cache.multimodal_cache import MultiModalStaticCache from sglang.srt.model_executor.forward_batch_info import ForwardBatch -from sglang.srt.model_executor.piecewise_cuda_graph_runner import use_original_ca_comm from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import flatten_nested_list, is_npu, print_warning_once from sglang.utils import logger @@ -660,6 +659,10 @@ def general_mm_embed_routine( Returns: Hidden states from language model forward pass """ + # Lazy import to allow some monkey patch of piecewise_cuda_graph_runner + from sglang.srt.model_executor.piecewise_cuda_graph_runner import ( + use_original_ca_comm, + ) tp_group = get_tp_group() diff --git a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py index 13662cb37..16be97703 100644 --- a/python/sglang/srt/managers/scheduler_dp_attn_mixin.py +++ b/python/sglang/srt/managers/scheduler_dp_attn_mixin.py @@ -5,8 +5,8 @@ from typing import TYPE_CHECKING, Callable import torch +from sglang.srt.batch_overlap.two_batch_overlap import TboDPAttentionPreparer from sglang.srt.managers.schedule_batch import ScheduleBatch -from sglang.srt.two_batch_overlap import TboDPAttentionPreparer from sglang.srt.utils.common import require_mlp_tp_gather if TYPE_CHECKING: diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index b775b4535..52909c37c 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -28,6 +28,7 @@ import torch import tqdm from torch.profiler import ProfilerActivity, profile +from sglang.srt.batch_overlap.two_batch_overlap import TboCudaGraphRunnerPlugin from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH from sglang.srt.custom_op import CustomOp from sglang.srt.distributed import get_tensor_model_parallel_rank @@ -60,7 +61,6 @@ from sglang.srt.model_executor.forward_batch_info import ( ) from sglang.srt.model_executor.input_buffers import GraphInputBuffers from sglang.srt.multiplex.pdmux_context import get_current_stream_idx, get_stream_groups -from sglang.srt.two_batch_overlap import TboCudaGraphRunnerPlugin from sglang.srt.utils import ( empty_context, get_available_gpu_memory, diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 5245ba136..8baddc56f 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -719,7 +719,7 @@ class ForwardBatch: ) def prepare_mlp_sync_batch(self, model_runner: ModelRunner): - from sglang.srt.two_batch_overlap import TboForwardBatchPreparer + from sglang.srt.batch_overlap.two_batch_overlap import TboForwardBatchPreparer assert self.global_num_tokens_cpu is not None assert self.global_num_tokens_for_logprob_cpu is not None diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index a5c1f4d4e..5739603b7 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @@ -24,6 +24,7 @@ from typing import TYPE_CHECKING, Union import torch import tqdm +from sglang.srt.batch_overlap.two_batch_overlap import TboCudaGraphRunnerPlugin from sglang.srt.compilation.compilation_config import CompilationConfig from sglang.srt.compilation.compile import install_torch_compiled, set_compiled from sglang.srt.compilation.piecewise_context_manager import ( @@ -51,7 +52,6 @@ from sglang.srt.model_executor.forward_batch_info import ( ForwardMode, PPProxyTensors, ) -from sglang.srt.two_batch_overlap import TboCudaGraphRunnerPlugin from sglang.srt.utils import get_available_gpu_memory, log_info_on_rank0 logger = logging.getLogger(__name__) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index c74eef2f1..f849b7f48 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -29,6 +29,8 @@ import tqdm from torch import nn from transformers import PretrainedConfig +from sglang.srt.batch_overlap.single_batch_overlap import SboFlags, compute_overlap_args +from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph from sglang.srt.configs.model_config import ( get_nsa_index_head_dim, @@ -134,9 +136,7 @@ from sglang.srt.model_loader.utils import ( ) from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.server_args import get_global_server_args -from sglang.srt.single_batch_overlap import SboFlags, compute_overlap_args from sglang.srt.speculative.spec_info import SpeculativeAlgorithm -from sglang.srt.two_batch_overlap import model_forward_maybe_tbo from sglang.srt.utils import ( BumpAllocator, LazyValue, diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py index 1fcbfb45a..d160c8876 100644 --- a/python/sglang/srt/models/glm4_moe.py +++ b/python/sglang/srt/models/glm4_moe.py @@ -22,6 +22,7 @@ import torch.nn.functional as F from torch import nn from transformers import PretrainedConfig +from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo from sglang.srt.distributed import ( get_moe_expert_parallel_world_size, get_pp_group, @@ -75,7 +76,6 @@ from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode 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.server_args import get_global_server_args -from sglang.srt.two_batch_overlap import model_forward_maybe_tbo from sglang.srt.utils import ( add_prefix, cpu_has_amx_support, diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index 9b97b75ea..4ab5adbba 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -22,6 +22,7 @@ import torch from torch import nn from transformers import PretrainedConfig +from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo from sglang.srt.distributed import ( get_moe_expert_parallel_world_size, get_pp_group, @@ -61,7 +62,6 @@ from sglang.srt.model_loader.weight_utils import ( maybe_remap_kv_scale_name, ) from sglang.srt.server_args import get_global_server_args -from sglang.srt.two_batch_overlap import model_forward_maybe_tbo from sglang.srt.utils import ( BumpAllocator, add_prefix, diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 32a002682..656968192 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -25,6 +25,7 @@ import torch.nn.functional as F from torch import nn from transformers import PretrainedConfig +from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo from sglang.srt.distributed import ( get_moe_expert_parallel_world_size, get_pp_group, @@ -70,7 +71,6 @@ from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode 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.server_args import get_global_server_args -from sglang.srt.two_batch_overlap import model_forward_maybe_tbo from sglang.srt.utils import add_prefix, is_cuda, make_layers logger = logging.getLogger(__name__) diff --git a/python/sglang/srt/operations_strategy.py b/python/sglang/srt/operations_strategy.py index cbed560e3..152e4874d 100644 --- a/python/sglang/srt/operations_strategy.py +++ b/python/sglang/srt/operations_strategy.py @@ -3,10 +3,10 @@ from typing import List, Optional import torch -from sglang.srt import operations +from sglang.srt.batch_overlap import operations +from sglang.srt.batch_overlap.operations import Operation from sglang.srt.layers.moe.token_dispatcher import DeepEPConfig from sglang.srt.model_executor.forward_batch_info import ForwardMode -from sglang.srt.operations import Operation @dataclass diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index bd6e5b332..55d11f1bf 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -137,18 +137,22 @@ builtins.FP8_E4M3_MAX = FP8_E4M3_MAX builtins.FP8_E4M3_MIN = FP8_E4M3_MIN +@lru_cache(maxsize=1) def is_cuda(): return torch.cuda.is_available() and torch.version.cuda +@lru_cache(maxsize=1) def is_cuda_alike(): return is_cuda() or is_hip() +@lru_cache(maxsize=1) def is_hpu() -> bool: return hasattr(torch, "hpu") and torch.hpu.is_available() +@lru_cache(maxsize=1) def is_xpu() -> bool: return hasattr(torch, "xpu") and torch.xpu.is_available() @@ -158,6 +162,7 @@ def is_npu() -> bool: return hasattr(torch, "npu") and torch.npu.is_available() +@lru_cache(maxsize=1) def is_host_cpu_x86() -> bool: machine = platform.machine().lower() return ( @@ -167,6 +172,7 @@ def is_host_cpu_x86() -> bool: ) +@lru_cache(maxsize=1) def is_cpu() -> bool: return os.getenv("SGLANG_USE_CPU_ENGINE", "0") == "1" and is_host_cpu_x86()