[Fix] Register custom ops only if they exist (#13321)
This commit is contained in:
@@ -20,23 +20,23 @@
|
|||||||
| [**Slides**](https://github.com/sgl-project/sgl-learning-materials?tab=readme-ov-file#slides) |
|
| [**Slides**](https://github.com/sgl-project/sgl-learning-materials?tab=readme-ov-file#slides) |
|
||||||
|
|
||||||
## News
|
## News
|
||||||
- [2025/11] 🔥 SGLang diffusion is now available ([blog](https://lmsys.org/blog/2025-11-07-sglang-diffusion/)).
|
- [2025/11] 🔥 SGLang Diffusion accelerates video and image generation ([blog](https://lmsys.org/blog/2025-11-07-sglang-diffusion/)).
|
||||||
- [2025/10] 🔥 SGLang now runs natively on TPU with the SGLang-Jax backend ([blog](https://lmsys.org/blog/2025-10-29-sglang-jax/)).
|
- [2025/10] 🔥 SGLang now runs natively on TPU with the SGLang-Jax backend ([blog](https://lmsys.org/blog/2025-10-29-sglang-jax/)).
|
||||||
- [2025/10] AMD AI Dev Day 2025 SGLang ([slide](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/sglang_amd_ai_devday_2025.pdf)), PyTorch Conference 2025 SGLang ([slide](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/sglang_pytorch_2025.pdf)).
|
- [2025/10] PyTorch Conference 2025 SGLang Talk ([slide](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/sglang_pytorch_2025.pdf)).
|
||||||
- [2025/10] SGLang x Nvidia SF Meetup on 10/2 ([recap](https://x.com/lmsysorg/status/1975339501934510231)).
|
|
||||||
- [2025/09] 🔥 Deploying DeepSeek on GB200 NVL72 with PD and Large Scale EP (Part II): 3.8x Prefill, 4.8x Decode Throughput ([blog](https://lmsys.org/blog/2025-09-25-gb200-part-2/)).
|
- [2025/09] 🔥 Deploying DeepSeek on GB200 NVL72 with PD and Large Scale EP (Part II): 3.8x Prefill, 4.8x Decode Throughput ([blog](https://lmsys.org/blog/2025-09-25-gb200-part-2/)).
|
||||||
- [2025/09] SGLang Day 0 Support for DeepSeek-V3.2 with Sparse Attention ([blog](https://lmsys.org/blog/2025-09-29-deepseek-V32/)).
|
- [2025/09] SGLang Day 0 Support for DeepSeek-V3.2 with Sparse Attention ([blog](https://lmsys.org/blog/2025-09-29-deepseek-V32/)).
|
||||||
- [2025/08] SGLang x AMD SF Meetup on 8/22: Hands-on GPU workshop, tech talks by AMD/xAI/SGLang, and networking ([Roadmap](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_sglang_roadmap.pdf), [Large-scale EP](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_sglang_ep.pdf), [Highlights](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_highlights.pdf), [AITER/MoRI](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_aiter_mori.pdf), [Wave](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_wave.pdf)).
|
- [2025/08] SGLang x AMD SF Meetup on 8/22: Hands-on GPU workshop, tech talks by AMD/xAI/SGLang, and networking ([Roadmap](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_sglang_roadmap.pdf), [Large-scale EP](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_sglang_ep.pdf), [Highlights](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_highlights.pdf), [AITER/MoRI](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_aiter_mori.pdf), [Wave](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/amd_meetup_wave.pdf)).
|
||||||
- [2025/08] SGLang provides day-0 support for OpenAI gpt-oss model ([instructions](https://github.com/sgl-project/sglang/issues/8833))
|
- [2025/08] SGLang provides day-0 support for OpenAI gpt-oss model ([instructions](https://github.com/sgl-project/sglang/issues/8833))
|
||||||
- [2025/05] Deploying DeepSeek with PD Disaggregation and Large-scale Expert Parallelism on 96 H100 GPUs ([blog](https://lmsys.org/blog/2025-05-05-large-scale-ep/)).
|
- [2025/05] Deploying DeepSeek with PD Disaggregation and Large-scale Expert Parallelism on 96 H100 GPUs ([blog](https://lmsys.org/blog/2025-05-05-large-scale-ep/)).
|
||||||
- [2025/03] SGLang Joins PyTorch Ecosystem: Efficient LLM Serving Engine ([PyTorch blog](https://pytorch.org/blog/sglang-joins-pytorch/))
|
|
||||||
|
|
||||||
<details>
|
<details>
|
||||||
<summary>More</summary>
|
<summary>More</summary>
|
||||||
|
|
||||||
|
- [2025/10] SGLang x Nvidia SF Meetup on 10/2 ([recap](https://x.com/lmsysorg/status/1975339501934510231)).
|
||||||
- [2025/06] SGLang, the high-performance serving infrastructure powering trillions of tokens daily, has been awarded the third batch of the Open Source AI Grant by a16z ([a16z blog](https://a16z.com/advancing-open-source-ai-through-benchmarks-and-bold-experimentation/)).
|
- [2025/06] SGLang, the high-performance serving infrastructure powering trillions of tokens daily, has been awarded the third batch of the Open Source AI Grant by a16z ([a16z blog](https://a16z.com/advancing-open-source-ai-through-benchmarks-and-bold-experimentation/)).
|
||||||
- [2025/06] Deploying DeepSeek on GB200 NVL72 with PD and Large Scale EP (Part I): 2.7x Higher Decoding Throughput ([blog](https://lmsys.org/blog/2025-06-16-gb200-part-1/)).
|
- [2025/06] Deploying DeepSeek on GB200 NVL72 with PD and Large Scale EP (Part I): 2.7x Higher Decoding Throughput ([blog](https://lmsys.org/blog/2025-06-16-gb200-part-1/)).
|
||||||
- [2025/03] Supercharge DeepSeek-R1 Inference on AMD Instinct MI300X ([AMD blog](https://rocm.blogs.amd.com/artificial-intelligence/DeepSeekR1-Part2/README.html))
|
- [2025/03] Supercharge DeepSeek-R1 Inference on AMD Instinct MI300X ([AMD blog](https://rocm.blogs.amd.com/artificial-intelligence/DeepSeekR1-Part2/README.html))
|
||||||
|
- [2025/03] SGLang Joins PyTorch Ecosystem: Efficient LLM Serving Engine ([PyTorch blog](https://pytorch.org/blog/sglang-joins-pytorch/))
|
||||||
- [2025/02] Unlock DeepSeek-R1 Inference Performance on AMD Instinct™ MI300X GPU ([AMD blog](https://rocm.blogs.amd.com/artificial-intelligence/DeepSeekR1_Perf/README.html))
|
- [2025/02] Unlock DeepSeek-R1 Inference Performance on AMD Instinct™ MI300X GPU ([AMD blog](https://rocm.blogs.amd.com/artificial-intelligence/DeepSeekR1_Perf/README.html))
|
||||||
- [2025/01] SGLang provides day one support for DeepSeek V3/R1 models on NVIDIA and AMD GPUs with DeepSeek-specific optimizations. ([instructions](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3), [AMD blog](https://www.amd.com/en/developer/resources/technical-articles/amd-instinct-gpus-power-deepseek-v3-revolutionizing-ai-development-with-sglang.html), [10+ other companies](https://x.com/lmsysorg/status/1887262321636221412))
|
- [2025/01] SGLang provides day one support for DeepSeek V3/R1 models on NVIDIA and AMD GPUs with DeepSeek-specific optimizations. ([instructions](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3), [AMD blog](https://www.amd.com/en/developer/resources/technical-articles/amd-instinct-gpus-power-deepseek-v3-revolutionizing-ai-development-with-sglang.html), [10+ other companies](https://x.com/lmsysorg/status/1887262321636221412))
|
||||||
- [2024/12] v0.4 Release: Zero-Overhead Batch Scheduler, Cache-Aware Load Balancer, Faster Structured Outputs ([blog](https://lmsys.org/blog/2024-12-04-sglang-v0-4/)).
|
- [2024/12] v0.4 Release: Zero-Overhead Batch Scheduler, Cache-Aware Load Balancer, Faster Structured Outputs ([blog](https://lmsys.org/blog/2024-12-04-sglang-v0-4/)).
|
||||||
|
|||||||
@@ -683,7 +683,6 @@ class Engine(EngineBase):
|
|||||||
|
|
||||||
def _set_envs_and_config(server_args: ServerArgs):
|
def _set_envs_and_config(server_args: ServerArgs):
|
||||||
# Set global environments
|
# Set global environments
|
||||||
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
|
|
||||||
if "NCCL_CUMEM_ENABLE" not in os.environ or server_args.enable_symm_mem:
|
if "NCCL_CUMEM_ENABLE" not in os.environ or server_args.enable_symm_mem:
|
||||||
os.environ["NCCL_CUMEM_ENABLE"] = str(int(server_args.enable_symm_mem))
|
os.environ["NCCL_CUMEM_ENABLE"] = str(int(server_args.enable_symm_mem))
|
||||||
if (
|
if (
|
||||||
@@ -756,10 +755,13 @@ def _set_envs_and_config(server_args: ServerArgs):
|
|||||||
|
|
||||||
|
|
||||||
def _init_tokenizer_manager(
|
def _init_tokenizer_manager(
|
||||||
server_args: ServerArgs, port_args: PortArgs
|
server_args: ServerArgs,
|
||||||
|
port_args: PortArgs,
|
||||||
|
TokenizerManagerClass: Optional[TokenizerManager] = None,
|
||||||
) -> TokenizerManager:
|
) -> TokenizerManager:
|
||||||
# Launch tokenizer process
|
# Launch tokenizer process
|
||||||
tokenizer_manager = TokenizerManager(server_args, port_args)
|
TokenizerManagerClass = TokenizerManagerClass or TokenizerManager
|
||||||
|
tokenizer_manager = TokenizerManagerClass(server_args, port_args)
|
||||||
|
|
||||||
# Initialize templates
|
# Initialize templates
|
||||||
template_manager = TemplateManager()
|
template_manager = TemplateManager()
|
||||||
|
|||||||
@@ -32,6 +32,7 @@ from sglang.srt.layers.quantization.marlin_utils import (
|
|||||||
from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
|
from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
|
||||||
from sglang.srt.layers.quantization.utils import get_scalar_types, replace_parameter
|
from sglang.srt.layers.quantization.utils import get_scalar_types, replace_parameter
|
||||||
from sglang.srt.layers.quantization.w8a8_int8 import npu_fused_experts
|
from sglang.srt.layers.quantization.w8a8_int8 import npu_fused_experts
|
||||||
|
from sglang.srt.utils.patch_torch import register_fake_if_exists
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
|
from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
|
||||||
@@ -959,7 +960,7 @@ class AWQMoEAscendMethod(AWQMoEMethod):
|
|||||||
# Register fake implementations for torch.compile support
|
# Register fake implementations for torch.compile support
|
||||||
if _is_cuda:
|
if _is_cuda:
|
||||||
|
|
||||||
@torch.library.register_fake("sgl_kernel::awq_dequantize")
|
@register_fake_if_exists("sgl_kernel::awq_dequantize")
|
||||||
def _(
|
def _(
|
||||||
qweight,
|
qweight,
|
||||||
scales,
|
scales,
|
||||||
@@ -971,7 +972,7 @@ if _is_cuda:
|
|||||||
out_shape = qweight.shape[:-1] + (qweight.shape[-1] * 32 // num_bits,)
|
out_shape = qweight.shape[:-1] + (qweight.shape[-1] * 32 // num_bits,)
|
||||||
return qweight.new_empty(out_shape, dtype=scales.dtype)
|
return qweight.new_empty(out_shape, dtype=scales.dtype)
|
||||||
|
|
||||||
@torch.library.register_fake("sgl_kernel::awq_marlin_repack")
|
@register_fake_if_exists("sgl_kernel::awq_marlin_repack")
|
||||||
def _(b_q_weight, size_k, size_n, num_bits):
|
def _(b_q_weight, size_k, size_n, num_bits):
|
||||||
return b_q_weight.new_empty(
|
return b_q_weight.new_empty(
|
||||||
(size_k // 16, size_n * (num_bits // 2)), dtype=b_q_weight.dtype
|
(size_k // 16, size_n * (num_bits // 2)), dtype=b_q_weight.dtype
|
||||||
|
|||||||
@@ -42,16 +42,16 @@ from sglang.srt.layers.quantization.utils import (
|
|||||||
replace_parameter,
|
replace_parameter,
|
||||||
unpack_cols,
|
unpack_cols,
|
||||||
)
|
)
|
||||||
|
from sglang.srt.utils import is_cuda
|
||||||
|
from sglang.srt.utils.patch_torch import register_fake_if_exists
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
|
from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
|
||||||
from sglang.srt.layers.moe.token_dispatcher import (
|
from sglang.srt.layers.moe.token_dispatcher import (
|
||||||
StandardDispatchOutput,
|
|
||||||
CombineInput,
|
CombineInput,
|
||||||
|
StandardDispatchOutput,
|
||||||
)
|
)
|
||||||
|
|
||||||
from sglang.srt.utils import is_cuda
|
|
||||||
|
|
||||||
_is_cuda = is_cuda()
|
_is_cuda = is_cuda()
|
||||||
|
|
||||||
if _is_cuda:
|
if _is_cuda:
|
||||||
@@ -1099,21 +1099,21 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
|
|||||||
# Register fake implementations for torch.compile support
|
# Register fake implementations for torch.compile support
|
||||||
if _is_cuda:
|
if _is_cuda:
|
||||||
|
|
||||||
@torch.library.register_fake("sgl_kernel::gptq_gemm")
|
@register_fake_if_exists("sgl_kernel::gptq_gemm")
|
||||||
def _(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, use_shuffle, bit):
|
def _(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, use_shuffle, bit):
|
||||||
return a.new_empty((a.shape[0], b_q_weight.shape[-1]), dtype=a.dtype)
|
return a.new_empty((a.shape[0], b_q_weight.shape[-1]), dtype=a.dtype)
|
||||||
|
|
||||||
@torch.library.register_fake("sgl_kernel::gptq_marlin_repack")
|
@register_fake_if_exists("sgl_kernel::gptq_marlin_repack")
|
||||||
def _(b_q_weight, perm, size_k, size_n, num_bits):
|
def _(b_q_weight, perm, size_k, size_n, num_bits):
|
||||||
return b_q_weight.new_empty(
|
return b_q_weight.new_empty(
|
||||||
(size_k // 16, size_n * (num_bits // 2)), dtype=b_q_weight.dtype
|
(size_k // 16, size_n * (num_bits // 2)), dtype=b_q_weight.dtype
|
||||||
)
|
)
|
||||||
|
|
||||||
@torch.library.register_fake("sgl_kernel::gptq_shuffle")
|
@register_fake_if_exists("sgl_kernel::gptq_shuffle")
|
||||||
def _(q_weight, q_perm, bit):
|
def _(q_weight, q_perm, bit):
|
||||||
return
|
return
|
||||||
|
|
||||||
@torch.library.register_fake("sgl_kernel::moe_wna16_marlin_gemm")
|
@register_fake_if_exists("sgl_kernel::moe_wna16_marlin_gemm")
|
||||||
def _(
|
def _(
|
||||||
a,
|
a,
|
||||||
c,
|
c,
|
||||||
|
|||||||
@@ -27,7 +27,6 @@ import zmq
|
|||||||
from sglang.srt.managers.io_struct import (
|
from sglang.srt.managers.io_struct import (
|
||||||
BatchEmbeddingOutput,
|
BatchEmbeddingOutput,
|
||||||
BatchMultimodalDecodeReq,
|
BatchMultimodalDecodeReq,
|
||||||
BatchMultimodalOutput,
|
|
||||||
BatchStrOutput,
|
BatchStrOutput,
|
||||||
BatchTokenIDOutput,
|
BatchTokenIDOutput,
|
||||||
FreezeGCReq,
|
FreezeGCReq,
|
||||||
@@ -284,22 +283,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
|
|||||||
)
|
)
|
||||||
|
|
||||||
def handle_multimodal_decode_req(self, recv_obj: BatchMultimodalDecodeReq):
|
def handle_multimodal_decode_req(self, recv_obj: BatchMultimodalDecodeReq):
|
||||||
outputs = self.tokenizer.detokenize(recv_obj)
|
raise NotImplementedError()
|
||||||
return BatchMultimodalOutput(
|
|
||||||
rids=recv_obj.rids,
|
|
||||||
http_worker_ipcs=recv_obj.http_worker_ipcs,
|
|
||||||
finished_reasons=recv_obj.finished_reasons,
|
|
||||||
outputs=outputs,
|
|
||||||
prompt_tokens=recv_obj.prompt_tokens,
|
|
||||||
completion_tokens=recv_obj.completion_tokens,
|
|
||||||
cached_tokens=recv_obj.cached_tokens,
|
|
||||||
placeholder_tokens_idx=None,
|
|
||||||
placeholder_tokens_val=None,
|
|
||||||
queue_time=recv_obj.queue_time,
|
|
||||||
forward_entry_time=recv_obj.forward_entry_time,
|
|
||||||
prefill_launch_delay=recv_obj.prefill_launch_delay,
|
|
||||||
prefill_launch_latency=recv_obj.prefill_launch_latency,
|
|
||||||
)
|
|
||||||
|
|
||||||
def handle_freeze_gc_req(self, recv_req: FreezeGCReq):
|
def handle_freeze_gc_req(self, recv_req: FreezeGCReq):
|
||||||
freeze_gc("Detokenizer Manager")
|
freeze_gc("Detokenizer Manager")
|
||||||
|
|||||||
@@ -17,7 +17,7 @@ import torch
|
|||||||
from packaging import version
|
from packaging import version
|
||||||
from torch.multiprocessing import reductions
|
from torch.multiprocessing import reductions
|
||||||
|
|
||||||
from sglang.srt.utils import is_npu
|
from sglang.srt.utils.common import is_npu
|
||||||
|
|
||||||
_is_npu = is_npu()
|
_is_npu = is_npu()
|
||||||
|
|
||||||
@@ -88,3 +88,29 @@ def monkey_patch_torch_compile():
|
|||||||
|
|
||||||
af.auto_functionalized_v2._cacheable = True
|
af.auto_functionalized_v2._cacheable = True
|
||||||
af.auto_functionalized._cacheable = True
|
af.auto_functionalized._cacheable = True
|
||||||
|
|
||||||
|
|
||||||
|
def register_fake_if_exists(op_name):
|
||||||
|
"""
|
||||||
|
Decorator factory to conditionally register a fake for a custom op if it exists.
|
||||||
|
Parses op_name (e.g., 'sgl_kernel::gptq_gemm'), checks if the op exists via hasattr
|
||||||
|
on the namespace attribute of torch.ops. Registers the fake if present; otherwise,
|
||||||
|
returns the function unchanged.
|
||||||
|
Args:
|
||||||
|
op_name (str): Full operator name (e.g., 'sgl_kernel::gptq_gemm').
|
||||||
|
Returns:
|
||||||
|
callable: Decorator for the fake function.
|
||||||
|
Example:
|
||||||
|
@register_fake_if_exists('sgl_kernel::gptq_gemm')
|
||||||
|
def fake_gptq_gemm(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, use_shuffle, bit):
|
||||||
|
return a.new_empty((a.shape[0], b_q_weight.shape[-1]), dtype=a.dtype)
|
||||||
|
"""
|
||||||
|
|
||||||
|
def decorator(func):
|
||||||
|
namespace, bare_op = op_name.split("::")
|
||||||
|
ops_namespace = getattr(torch.ops, namespace, None)
|
||||||
|
if ops_namespace and hasattr(ops_namespace, bare_op):
|
||||||
|
torch.library.register_fake(op_name, func)
|
||||||
|
return func
|
||||||
|
|
||||||
|
return decorator
|
||||||
|
|||||||
Reference in New Issue
Block a user