137 lines
3.9 KiB
Python
137 lines
3.9 KiB
Python
import functools
|
|
import logging
|
|
from enum import IntEnum
|
|
from typing import TYPE_CHECKING, Callable
|
|
|
|
import torch
|
|
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.utils import get_npu_memory_capacity, is_npu
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.server_args import ServerArgs
|
|
|
|
logger = logging.getLogger(__name__)
|
|
_is_npu = is_npu()
|
|
indexer_weight_stream = None
|
|
|
|
|
|
class NPUACLFormat(IntEnum):
|
|
ACL_FORMAT_UNDEFINED = -1
|
|
ACL_FORMAT_ND = 2
|
|
ACL_FORMAT_FRACTAL_NZ = 29
|
|
|
|
|
|
def _call_once(fn: Callable):
|
|
|
|
@functools.wraps(fn)
|
|
def wrapper(*args, **kwargs):
|
|
if getattr(fn, "_has_been_called", False):
|
|
logger.debug("Function {} has already been called.", fn.__name__)
|
|
return
|
|
|
|
fn._has_been_called = True
|
|
return fn(*args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
|
|
def set_default_server_args(args: "ServerArgs"):
|
|
"""
|
|
Set default server arguments for NPU backend.
|
|
"""
|
|
|
|
# NPU only works with "ascend" attention backend for now
|
|
args.attention_backend = "ascend"
|
|
args.prefill_attention_backend = "ascend"
|
|
args.decode_attention_backend = "ascend"
|
|
if args.page_size is None:
|
|
args.page_size = 128
|
|
|
|
# NPU memory settings
|
|
npu_mem = get_npu_memory_capacity()
|
|
if npu_mem <= 32 * 1024:
|
|
# Ascend 910B4,910B4_1
|
|
# (chunked_prefill_size 4k, cuda_graph_max_bs 16 if tp < 4 else 64)
|
|
if args.chunked_prefill_size is None:
|
|
args.chunked_prefill_size = 4 * 1024
|
|
if args.cuda_graph_max_bs is None:
|
|
if args.tp_size < 4:
|
|
args.cuda_graph_max_bs = 16
|
|
else:
|
|
args.cuda_graph_max_bs = 64
|
|
elif npu_mem <= 64 * 1024:
|
|
# Ascend 910B1,910B2,910B2C,910B3,910_9391,910_9392,910_9381,910_9382,910_9372,910_9362
|
|
# (chunked_prefill_size 8k, cuda_graph_max_bs 64 if tp < 4 else 256)
|
|
if args.chunked_prefill_size is None:
|
|
args.chunked_prefill_size = 8 * 1024
|
|
if args.cuda_graph_max_bs is None:
|
|
if args.tp_size < 4:
|
|
args.cuda_graph_max_bs = 64
|
|
else:
|
|
args.cuda_graph_max_bs = 256
|
|
|
|
# NPU does not support CustomAllReduce
|
|
args.disable_custom_all_reduce = True
|
|
|
|
# handles hierarchical cache configs
|
|
if args.enable_hierarchical_cache:
|
|
args.hicache_io_backend = "kernel_ascend"
|
|
if args.use_mla_backend():
|
|
args.hicache_mem_layout = "page_first_kv_split"
|
|
else:
|
|
args.hicache_mem_layout = "page_first_direct"
|
|
|
|
|
|
@_call_once
|
|
def init_npu_backend():
|
|
"""
|
|
Initialize NPU backend. This function should be called only once.
|
|
"""
|
|
|
|
assert _is_npu, "NPU backend initialization called on non-NPU device."
|
|
|
|
import sgl_kernel_npu # noqa: F401
|
|
import torch_npu
|
|
from torch_npu.contrib import transfer_to_npu # noqa: F401
|
|
|
|
# Re-mock torch.cuda.is_available cuz transfer_to_npu mocks it True
|
|
torch.cuda.is_available = lambda: False
|
|
|
|
torch_npu.npu.config.allow_internal_format = True
|
|
torch_npu.npu.set_compile_mode(jit_compile=False)
|
|
|
|
|
|
def npu_format_cast(
|
|
tensor: torch.Tensor,
|
|
acl_format: NPUACLFormat = NPUACLFormat.ACL_FORMAT_FRACTAL_NZ,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Cast a tensor to a specific NPU ACL format.
|
|
|
|
Args:
|
|
tensor (torch.Tensor): The input tensor.
|
|
acl_format (NPUACLFormat): The target NPU ACL format.
|
|
|
|
Returns:
|
|
torch.Tensor: The tensor cast to the specified NPU ACL format.
|
|
"""
|
|
|
|
if not _is_npu:
|
|
return tensor
|
|
|
|
if envs.SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT.get():
|
|
return tensor
|
|
|
|
if tensor.device == torch.device("cpu"):
|
|
return torch.ops.npu.npu_format_cast(tensor.npu(), acl_format.value).cpu()
|
|
else:
|
|
return torch.ops.npu.npu_format_cast(tensor, acl_format.value)
|
|
|
|
|
|
def get_indexer_weight_stream():
|
|
global indexer_weight_stream
|
|
if indexer_weight_stream is None:
|
|
indexer_weight_stream = torch.npu.Stream()
|
|
return indexer_weight_stream
|