140 lines
5.1 KiB
Python
140 lines
5.1 KiB
Python
# Copyright 2023-2024 SGLang Team
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
from typing import Callable, Union
|
|
|
|
import torch
|
|
from packaging import version
|
|
from torch.multiprocessing import reductions
|
|
|
|
from sglang.srt.utils.common import is_npu
|
|
|
|
_is_npu = is_npu()
|
|
|
|
if _is_npu:
|
|
from torch_npu.multiprocessing import reductions as npu_reductions
|
|
|
|
def _rebuild_npu_tensor_modified(*args):
|
|
args = _modify_tuple(args, _REDUCE_TENSOR_ARG_DEVICE_INDEX, npu_verl_to_sglang)
|
|
return npu_reductions._rebuild_npu_tensor_original(*args)
|
|
|
|
def npu_verl_to_sglang(device: int):
|
|
assert (
|
|
SGLANG_TP_RANK is not None
|
|
), "SGLANG_TP_RANK is not registered. Please call register_sgl_tp_rank() first."
|
|
return SGLANG_TP_RANK
|
|
|
|
|
|
SGLANG_TP_RANK = None
|
|
|
|
|
|
def monkey_patch_torch_reductions():
|
|
"""Monkey patching before Torch https://github.com/pytorch/pytorch/pull/149248 is fixed"""
|
|
|
|
if not _is_npu:
|
|
if hasattr(reductions, "_reduce_tensor_original"):
|
|
return
|
|
reductions._reduce_tensor_original = reductions.reduce_tensor
|
|
reductions._rebuild_cuda_tensor_original = reductions.rebuild_cuda_tensor
|
|
|
|
reductions.reduce_tensor = _reduce_tensor_modified
|
|
reductions.rebuild_cuda_tensor = _rebuild_cuda_tensor_modified
|
|
reductions.init_reductions()
|
|
else:
|
|
# FIXME: This is a temp patch for npu as HDK does not support device uuid for now
|
|
if hasattr(npu_reductions, "_rebuild_npu_tensor_original"):
|
|
return
|
|
|
|
npu_reductions._rebuild_npu_tensor_original = npu_reductions.rebuild_npu_tensor
|
|
npu_reductions.rebuild_npu_tensor = _rebuild_npu_tensor_modified
|
|
|
|
|
|
# The signature has not been changed for years, and we will not need this when the next version is released,
|
|
# so it looks safe to use a constant.
|
|
_REDUCE_TENSOR_ARG_DEVICE_INDEX = 6
|
|
|
|
|
|
def register_sgl_tp_rank(rank: int):
|
|
global SGLANG_TP_RANK
|
|
SGLANG_TP_RANK = rank
|
|
|
|
|
|
def _reduce_tensor_modified(*args, **kwargs):
|
|
output_fn, output_args = reductions._reduce_tensor_original(*args, **kwargs)
|
|
output_args = _modify_tuple(
|
|
output_args, _REDUCE_TENSOR_ARG_DEVICE_INDEX, _device_to_uuid
|
|
)
|
|
return output_fn, output_args
|
|
|
|
|
|
def _rebuild_cuda_tensor_modified(*args):
|
|
args = _modify_tuple(args, _REDUCE_TENSOR_ARG_DEVICE_INDEX, _device_from_maybe_uuid)
|
|
return reductions._rebuild_cuda_tensor_original(*args)
|
|
|
|
|
|
def _device_to_uuid(device: int) -> str:
|
|
return str(torch.cuda.get_device_properties(device).uuid)
|
|
|
|
|
|
def _device_from_maybe_uuid(device_maybe_uuid: Union[int, str]) -> int:
|
|
if isinstance(device_maybe_uuid, int):
|
|
return device_maybe_uuid
|
|
|
|
if isinstance(device_maybe_uuid, str):
|
|
for device in range(torch.cuda.device_count()):
|
|
if str(torch.cuda.get_device_properties(device).uuid) == device_maybe_uuid:
|
|
return device
|
|
raise Exception("Invalid device_uuid=" + device_maybe_uuid)
|
|
|
|
raise Exception(f"Unknown type: {device_maybe_uuid=}")
|
|
|
|
|
|
def _modify_tuple(t, index: int, modifier: Callable):
|
|
return *t[:index], modifier(t[index]), *t[index + 1 :]
|
|
|
|
|
|
def monkey_patch_torch_compile():
|
|
if version.parse(torch.__version__) < version.parse("2.8.0"):
|
|
# These things are cacheable by torch.compile. torch.compile just doesn't know it.
|
|
# This was fixed in PyTorch 2.8, but until then, we monkey patch.
|
|
import torch._higher_order_ops.auto_functionalize as af
|
|
|
|
af.auto_functionalized_v2._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
|