diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index d1fa31cf3..26e567ee8 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -14,15 +14,13 @@ /python/sglang/srt/eplb @fzyzcjy @ch-wan /python/sglang/srt/function_call @CatherineSue @JustinTong0323 /python/sglang/srt/grpc @CatherineSue @slin1237 +/python/sglang/srt/hardware_backend/npu @ping1jing2 @iforgetmyname /python/sglang/srt/layers @merrymercy @Ying1123 @Fridge003 @ispobock @HaiShaw @ch-wan @BBuf @Edwardf0t1 /python/sglang/srt/layers/quantization @ch-wan @BBuf @Edwardf0t1 @FlamingoPg @AniZpZ -/python/sglang/srt/layers/attention/ascend_backend.py @ping1jing2 @iforgetmyname /python/sglang/srt/lora @Ying1123 @Fridge003 @lifuhuang /python/sglang/srt/managers @merrymercy @Ying1123 @hnyls2002 @xiezhq-hermann @zhyncs /python/sglang/srt/mem_cache @merrymercy @Ying1123 @hnyls2002 @xiezhq-hermann -/python/sglang/srt/mem_cache/allocator_ascend.py @ping1jing2 @iforgetmyname /python/sglang/srt/model_executor @merrymercy @Ying1123 @hnyls2002 @Fridge003 @ispobock -/python/sglang/srt/model_executor/npu_graph_runner.py @ping1jing2 @iforgetmyname /python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @hebiao064 /python/sglang/srt/multimodal @mickqian @JustinTong0323 @yhyang201 @yuan-luo /python/sglang/srt/speculative @Ying1123 @merrymercy @hnyls2002 diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 65234d7b0..947191e3b 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -232,6 +232,9 @@ class Envs: SGLANG_ROCM_FUSED_DECODE_MLA = EnvBool(False) SGLANG_ROCM_DISABLE_LINEARQUANT = EnvBool(False) + # NPU + SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT = EnvBool(False) + # Quantization SGLANG_INT4_WEIGHT = EnvBool(False) SGLANG_CPU_QUANTIZATION = EnvBool(False) diff --git a/python/sglang/srt/eplb/expert_distribution.py b/python/sglang/srt/eplb/expert_distribution.py index bfec3802e..d40439409 100644 --- a/python/sglang/srt/eplb/expert_distribution.py +++ b/python/sglang/srt/eplb/expert_distribution.py @@ -31,9 +31,7 @@ from sglang.srt.environ import envs from sglang.srt.metrics.collector import ExpertDispatchCollector from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.server_args import ServerArgs -from sglang.srt.utils import Withable, get_int_env_var, is_npu - -_is_npu = is_npu() +from sglang.srt.utils import Withable, get_int_env_var if TYPE_CHECKING: from sglang.srt.eplb.expert_location import ExpertLocationMetadata @@ -465,10 +463,6 @@ def _list_sum(a: List, b: List) -> List: class _LayerBasedGpuSinglePassGatherer(_SinglePassGatherer): def __init__(self, *args, enable_global_physical_experts: bool, **kwargs): super().__init__(*args, **kwargs) - if not _is_npu: - device = "cuda" - else: - device = "npu" self._enable_global_physical_experts = enable_global_physical_experts self._data = torch.zeros( ( @@ -480,7 +474,7 @@ class _LayerBasedGpuSinglePassGatherer(_SinglePassGatherer): ), ), dtype=torch.int, - device=device, + device="cuda", ) def reset(self): diff --git a/python/sglang/srt/mem_cache/allocator_ascend.py b/python/sglang/srt/hardware_backend/npu/allocator_npu.py similarity index 96% rename from python/sglang/srt/mem_cache/allocator_ascend.py rename to python/sglang/srt/hardware_backend/npu/allocator_npu.py index 74ac8983d..dd7221287 100644 --- a/python/sglang/srt/mem_cache/allocator_ascend.py +++ b/python/sglang/srt/hardware_backend/npu/allocator_npu.py @@ -1,17 +1,15 @@ -from __future__ import annotations - from typing import TYPE_CHECKING import torch -if TYPE_CHECKING: - from sglang.srt.mem_cache.memory_pool import KVCache - from sglang.srt.mem_cache.allocator import PagedTokenToKVPoolAllocator from sglang.srt.utils import get_num_new_pages, next_power_of_2 +if TYPE_CHECKING: + from sglang.srt.mem_cache.memory_pool import KVCache -def alloc_extend_kernel_ascend( + +def _alloc_extend_naive( prefix_lens, seq_lens, last_loc, @@ -65,14 +63,14 @@ def alloc_extend_kernel_ascend( ).view(-1) -class AscendPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator): +class NPUPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator): def __init__( self, size: int, page_size: int, dtype: torch.dtype, device: str, - kvcache: KVCache, + kvcache: "KVCache", need_sort: bool, ): super().__init__(size, page_size, dtype, device, kvcache, need_sort) @@ -130,7 +128,7 @@ class AscendPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator): dtype=torch.int32, device=self.device, ) - alloc_extend_kernel_ascend( + _alloc_extend_naive( prefix_lens, seq_lens, last_loc, diff --git a/python/sglang/srt/layers/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py similarity index 99% rename from python/sglang/srt/layers/attention/ascend_backend.py rename to python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index d2b18c638..5536cf80c 100644 --- a/python/sglang/srt/layers/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -7,8 +7,10 @@ import torch import torch_npu from sglang.srt.configs.model_config import AttentionArch +from sglang.srt.hardware_backend.npu.attention.mla_preprocess import ( + is_mla_preprocess_enabled, +) from sglang.srt.layers.attention.base_attn_backend import AttentionBackend -from sglang.srt.layers.attention.npu_ops.mla_preprocess import is_mla_preprocess_enabled from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBackend from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.layers.radix_attention import AttentionType diff --git a/python/sglang/srt/layers/attention/npu_ops/mla_preprocess.py b/python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py similarity index 92% rename from python/sglang/srt/layers/attention/npu_ops/mla_preprocess.py rename to python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py index 76f802bd2..d1780133e 100644 --- a/python/sglang/srt/layers/attention/npu_ops/mla_preprocess.py +++ b/python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py @@ -1,23 +1,19 @@ +from functools import lru_cache +from typing import TYPE_CHECKING, Optional + import torch import torch.nn.functional as F -from sglang.srt.utils import get_bool_env_var, is_npu +from sglang.srt.hardware_backend.npu.utils import npu_format_cast +from sglang.srt.utils import get_bool_env_var -_is_npu = is_npu() -_ENABLE_MLA_PREPROCESS_FLAG = get_bool_env_var("SGLANG_NPU_USE_MLAPO") -_NPU_FORMAT_NZ = 29 +if TYPE_CHECKING: + from sglang.srt.layers.quantization.base_config import QuantizationConfig +@lru_cache(maxsize=1) def is_mla_preprocess_enabled() -> bool: - return _is_npu and _ENABLE_MLA_PREPROCESS_FLAG - - -if is_mla_preprocess_enabled(): - import sgl_kernel_npu # noqa: F401 - import torch_npu - - torch.npu.config.allow_internal_format = True - torch.npu.set_compile_mode(jit_compile=False) + return get_bool_env_var("SGLANG_NPU_USE_MLAPO") def round_up(val: int, align: int) -> int: @@ -66,6 +62,7 @@ class NPUFusedMLAPreprocess(torch.nn.Module): num_local_heads, qk_nope_head_dim, qk_rope_head_dim, + quant_config: Optional["QuantizationConfig"] = None, ): super().__init__() self.qkv_a_proj = fused_qkv_a_proj_with_mqa @@ -75,6 +72,7 @@ class NPUFusedMLAPreprocess(torch.nn.Module): self.w_kc = w_kc.contiguous() self.rotary_emb = rotary_emb self.layer_id = layer_id + self.quant_config = quant_config self.has_preprocess_weights = False self.dtype = None @@ -124,9 +122,7 @@ class NPUFusedMLAPreprocess(torch.nn.Module): .unsqueeze(0) .contiguous() ) - self.qkv_a_proj_weight_nz = torch_npu.npu_format_cast( - fused_qkv_a_proj_with_mqa_weight_nz, _NPU_FORMAT_NZ - ) + self.qkv_a_proj_weight_nz = npu_format_cast(fused_qkv_a_proj_with_mqa_weight_nz) # matmul_0 deq_scale [2112] fused_qkv_a_proj_with_mqa_deq_scale_q = self.qkv_a_proj.deq_scale.data[ @@ -198,9 +194,7 @@ class NPUFusedMLAPreprocess(torch.nn.Module): q_b_proj_weight_nz = ( transdata(q_b_proj_weight, block_size=(16, 32)).unsqueeze(0).contiguous() ) - self.q_b_proj_weight_nz = torch_npu.npu_format_cast( - q_b_proj_weight_nz, _NPU_FORMAT_NZ - ) + self.q_b_proj_weight_nz = npu_format_cast(q_b_proj_weight_nz) # matmul_1 deq_scale [num_head * 192] q_b_proj_deq_scale = self.q_b_proj.deq_scale.data.clone() @@ -280,7 +274,7 @@ class NPUFusedMLAPreprocess(torch.nn.Module): q_pe = q_pe.view(-1, self.num_local_heads, 1, self.qk_rope_head_dim) cos = cos.view(-1, 1, 1, self.qk_rope_head_dim) sin = sin.view(-1, 1, 1, self.qk_rope_head_dim) - q_pe = torch_npu.npu_interleave_rope(q_pe, cos, sin) # (B,N,S,D) + q_pe = torch.ops.npu.npu_interleave_rope(q_pe, cos, sin) # (B,N,S,D) q_pe = q_pe.view(cos.shape[0], self.num_local_heads, self.qk_rope_head_dim) latent_cache = latent_cache.view( @@ -300,7 +294,7 @@ class NPUFusedMLAPreprocess(torch.nn.Module): 1, forward_batch.attn_backend.qk_rope_head_dim, ) - k_rope, k_nope, _, _ = torch_npu.npu_kv_rmsnorm_rope_cache( + k_rope, k_nope, _, _ = torch.ops.npu.npu_kv_rmsnorm_rope_cache( latent_cache, self.kv_a_layernorm.weight, cos, @@ -378,10 +372,11 @@ class NPUFusedMLAPreprocess(torch.nn.Module): ) def forward(self, positions, hidden_states, forward_batch, zero_allocator): + assert self.quant_config and self.quant_config.get_name() == "modelslim" + # route by `qkv_a_proj` quant type as MTP layers can be unquantized _is_w8a8 = ( - hasattr(self.qkv_a_proj.quant_method, "quantization_config") - and self.qkv_a_proj.quant_method.quantization_config.get_name() - == "w8a8_int8" + hasattr(self.qkv_a_proj.quant_method, "quant_config") + and self.qkv_a_proj.quant_method.quant_config.get_name() == "modelslim" ) if _is_w8a8: return self.forward_mlapo( diff --git a/python/sglang/srt/hardware_backend/npu/cmo.py b/python/sglang/srt/hardware_backend/npu/cmo.py new file mode 100644 index 000000000..40f3b4f16 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/cmo.py @@ -0,0 +1,54 @@ +import torch + +cmo_stream = None + + +def get_cmo_stream(): + """ + Cache Management Operation(CMO). + Launch a new stream to prefetch the weight of matmul when running other + AIV or communication kernels, aiming to overlap the memory access time. + """ + global cmo_stream + return cmo_stream + + +def set_cmo_stream(stream): + global cmo_stream + cmo_stream = stream + + +def prepare_weight_cache(handle, cache, PREFETCH_MAX_SIZE=1000000000): + """ + PREFETCH_MAX_SIZE: maximum size (bytes) for each prefetch operation. + This affects the time spent in prefetch: + time ≈ PREFETCH_MAX_SIZE / system_bandwidth + """ + import torch_npu + + stream = get_cmo_stream() + if stream is None: + stream = torch.npu.Stream() + set_cmo_stream(stream) + stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(stream): + if isinstance(cache, list): + for weight in cache: + torch_npu.npu_prefetch( + weight, + handle, + PREFETCH_MAX_SIZE, + ) + else: + torch_npu.npu_prefetch( + cache, + handle, + PREFETCH_MAX_SIZE, + ) + + +def wait_cmo_stream(): + stream = get_cmo_stream() + if stream is not None: + cur_stream = torch.npu.current_stream() + cur_stream.wait_stream(stream) diff --git a/python/sglang/srt/speculative/eagle_draft_extend_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py similarity index 100% rename from python/sglang/srt/speculative/eagle_draft_extend_npu_graph_runner.py rename to python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py diff --git a/python/sglang/srt/speculative/eagle_draft_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py similarity index 100% rename from python/sglang/srt/speculative/eagle_draft_npu_graph_runner.py rename to python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py diff --git a/python/sglang/srt/model_executor/npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py similarity index 100% rename from python/sglang/srt/model_executor/npu_graph_runner.py rename to python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py diff --git a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py new file mode 100644 index 000000000..943476786 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py @@ -0,0 +1,300 @@ +from typing import TYPE_CHECKING, Optional + +import torch +import torch_npu + +from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE +from sglang.srt.mem_cache.memory_pool import ( + MHATokenToKVPool, + MLATokenToKVPool, + get_tensor_size_bytes, +) + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + + +class NPUMHATokenToKVPool(MHATokenToKVPool): + + def _create_buffers(self): + with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): + # [size, head_num, head_dim] for each layer + # The padded slot 0 is used for writing dummy outputs from padded tokens. + # Continuous memory improves the efficiency of Ascend`s transmission backend, + # while other backends remain unchanged. + self.kv_buffer = torch.zeros( + ( + 2, + self.layer_num, + self.size // self.page_size + 1, + self.page_size, + self.head_num, + self.head_dim, + ), + dtype=self.store_dtype, + device=self.device, + ) + self.k_buffer = self.kv_buffer[0] + self.v_buffer = self.kv_buffer[1] + + # for disagg + def get_contiguous_buf_infos(self): + # layer_num x [seq_len, head_num, head_dim] + # layer_num x [page_num, page_size, head_num, head_dim] + kv_data_ptrs = [ + self.get_key_buffer(i).data_ptr() + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] + [ + self.get_value_buffer(i).data_ptr() + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] + kv_data_lens = [ + self.get_key_buffer(i).nbytes + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] + [ + self.get_value_buffer(i).nbytes + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] + kv_item_lens = [ + self.get_key_buffer(i)[0].nbytes + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] + [ + self.get_value_buffer(i)[0].nbytes + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] + return kv_data_ptrs, kv_data_lens, kv_item_lens + + def set_kv_buffer( + self, + layer: "RadixAttention", + loc: torch.Tensor, + cache_k: torch.Tensor, + cache_v: torch.Tensor, + k_scale: Optional[float] = None, + v_scale: Optional[float] = None, + layer_id_override: Optional[int] = None, + ): + if layer_id_override is not None: + layer_id = layer_id_override + else: + layer_id = layer.layer_id + if cache_k.dtype != self.dtype: + if k_scale is not None: + cache_k.div_(k_scale) + if v_scale is not None: + cache_v.div_(v_scale) + cache_k = cache_k.to(self.dtype) + cache_v = cache_v.to(self.dtype) + + if self.store_dtype != self.dtype: + cache_k = cache_k.view(self.store_dtype) + cache_v = cache_v.view(self.store_dtype) + + torch_npu._npu_reshape_and_cache( + key=cache_k, + value=cache_v, + key_cache=self.k_buffer[layer_id - self.start_layer].view( + -1, self.page_size, self.head_num, self.head_dim + ), + value_cache=self.v_buffer[layer_id - self.start_layer].view( + -1, self.page_size, self.head_num, self.head_dim + ), + slot_indices=loc, + ) + + +class NPUMLATokenToKVPool(MLATokenToKVPool): + + def __init__( + self, + size: int, + page_size: int, + dtype: torch.dtype, + kv_lora_rank: int, + qk_rope_head_dim: int, + index_head_dim: Optional[int], + layer_num: int, + device: str, + enable_memory_saver: bool, + start_layer: Optional[int] = None, + end_layer: Optional[int] = None, + ): + super(MLATokenToKVPool, self).__init__( + size, + page_size, + dtype, + layer_num, + device, + enable_memory_saver, + start_layer, + end_layer, + ) + + self.kv_lora_rank = kv_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.index_head_dim = index_head_dim + + self.custom_mem_pool = None + + with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): + # The padded slot 0 is used for writing dummy outputs from padded tokens. + self.k_buffer = torch.zeros( + ( + layer_num, + self.size // self.page_size + 1, + self.page_size, + 1, + self.kv_lora_rank, + ), + dtype=self.store_dtype, + device=self.device, + ) + self.v_buffer = torch.zeros( + ( + layer_num, + self.size // self.page_size + 1, + self.page_size, + 1, + self.qk_rope_head_dim, + ), + dtype=self.store_dtype, + device=self.device, + ) + if self.index_head_dim is not None: + self.index_k_buffer = torch.zeros( + ( + layer_num, + self.size // self.page_size + 1, + self.page_size, + 1, + self.index_head_dim, + ), + dtype=self.store_dtype, + device=self.device, + ) + + self._finalize_allocation_log(size) + + def get_kv_size_bytes(self): + assert hasattr(self, "k_buffer") + assert hasattr(self, "v_buffer") + kv_size_bytes = 0 + for k_cache in self.k_buffer: + kv_size_bytes += get_tensor_size_bytes(k_cache) + for v_cache in self.v_buffer: + kv_size_bytes += get_tensor_size_bytes(v_cache) + if self.index_head_dim is not None: + assert hasattr(self, "index_k_buffer") + for index_k_cache in self.index_k_buffer: + kv_size_bytes += get_tensor_size_bytes(index_k_cache) + return kv_size_bytes + + def get_kv_buffer(self, layer_id: int): + if self.layer_transfer_counter is not None: + self.layer_transfer_counter.wait_until(layer_id - self.start_layer) + return ( + self.k_buffer[layer_id - self.start_layer], + self.v_buffer[layer_id - self.start_layer], + ) + + def get_key_buffer(self, layer_id: int): + if self.layer_transfer_counter is not None: + self.layer_transfer_counter.wait_until(layer_id - self.start_layer) + + if self.store_dtype != self.dtype: + return self.k_buffer[layer_id - self.start_layer].view(self.dtype) + return self.k_buffer[layer_id - self.start_layer] + + def get_value_buffer(self, layer_id: int): + if self.layer_transfer_counter is not None: + self.layer_transfer_counter.wait_until(layer_id - self.start_layer) + + if self.store_dtype != self.dtype: + return self.v_buffer[layer_id - self.start_layer].view(self.dtype) + return self.v_buffer[layer_id - self.start_layer] + + def get_index_k_buffer(self, layer_id: int): + if self.layer_transfer_counter is not None: + self.layer_transfer_counter.wait_until(layer_id - self.start_layer) + + if self.store_dtype != self.dtype: + return self.index_k_buffer[layer_id - self.start_layer].view(self.dtype) + return self.index_k_buffer[layer_id - self.start_layer] + + # for disagg + def get_contiguous_buf_infos(self): + # MLA has only one kv_buffer, so only the information of this buffer needs to be returned. + kv_data_ptrs = [self.k_buffer[i].data_ptr() for i in range(self.layer_num)] + [ + self.v_buffer[i].data_ptr() for i in range(self.layer_num) + ] + kv_data_lens = [self.k_buffer[i].nbytes for i in range(self.layer_num)] + [ + self.v_buffer[i].nbytes for i in range(self.layer_num) + ] + kv_item_lens = [self.k_buffer[i][0].nbytes for i in range(self.layer_num)] + [ + self.v_buffer[i][0].nbytes for i in range(self.layer_num) + ] + if self.index_head_dim is not None: + kv_data_ptrs += [ + self.index_k_buffer[i].data_ptr() for i in range(self.layer_num) + ] + kv_data_lens += [ + self.index_k_buffer[i].nbytes for i in range(self.layer_num) + ] + kv_item_lens += [ + self.index_k_buffer[i][0].nbytes for i in range(self.layer_num) + ] + return kv_data_ptrs, kv_data_lens, kv_item_lens + + def set_kv_buffer( + self, + layer: "RadixAttention", + loc: torch.Tensor, + cache_k: torch.Tensor, + cache_v: torch.Tensor, + ): + layer_id = layer.layer_id + if cache_k.dtype != self.dtype: + cache_k = cache_k.to(self.dtype) + cache_v = cache_v.to(self.dtype) + + if self.store_dtype != self.dtype: + cache_k = cache_k.view(self.store_dtype) + cache_v = cache_v.view(self.store_dtype) + + if cache_v is None: + cache_k, cache_v = cache_k.split( + [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + + torch_npu.npu_scatter_nd_update_( + self.k_buffer[layer_id - self.start_layer].view(-1, 1, self.kv_lora_rank), + loc.view(-1, 1), + cache_k.view(-1, 1, self.kv_lora_rank), + ) + torch_npu.npu_scatter_nd_update_( + self.v_buffer[layer_id - self.start_layer].view( + -1, 1, self.qk_rope_head_dim + ), + loc.view(-1, 1), + cache_v.view(-1, 1, self.qk_rope_head_dim), + ) + + def set_index_k_buffer( + self, + layer_id: int, + loc: torch.Tensor, + index_k: torch.Tensor, + ): + if index_k.dtype != self.dtype: + index_k = index_k.to(self.dtype) + + if self.store_dtype != self.dtype: + index_k = index_k.view(self.store_dtype) + + torch_npu.npu_scatter_nd_update_( + self.index_k_buffer[layer_id - self.start_layer].view( + -1, 1, self.index_head_dim + ), + loc.view(-1, 1), + index_k.view(-1, 1, self.index_head_dim), + ) diff --git a/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py b/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py new file mode 100644 index 000000000..470617585 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py @@ -0,0 +1,380 @@ +from typing import TYPE_CHECKING + +import torch + +from sglang.srt.hardware_backend.npu.attention.mla_preprocess import ( + NPUFusedMLAPreprocess, + is_mla_preprocess_enabled, +) +from sglang.srt.layers.attention.nsa.utils import ( + cp_split_and_rebuild_position, + enable_prefill_cp, +) +from sglang.srt.layers.communicator import get_attn_tp_context + +if TYPE_CHECKING: + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA + from sglang.srt.utils import BumpAllocator + + +# region MHA +def forward_mha_prepare_npu( + m: "DeepseekV2AttentionMLA", + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: "ForwardBatch", + zero_allocator: "BumpAllocator", +): + if m.q_lora_rank is not None: + q, latent_cache = ( + get_attn_tp_context() + .fetch_qkv_latent() + .split( + [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], + dim=-1, + ) + ) + + # NSA Indexer: cache quantized keys, auto-skip topk for sequences <= nsa_index_topk + + if m.use_nsa: + q_lora = m.q_a_layernorm(q) + q = m.q_b_proj(q_lora)[0].view(-1, m.num_local_heads, m.qk_head_dim) + _ = m.indexer( + x=hidden_states, + q_lora=q_lora, + positions=positions, + forward_batch=forward_batch, + layer_id=m.layer_id, + return_indices=False, + ) + + else: + q = m.q_a_layernorm(q) + q = m.q_b_proj(q)[0].view(-1, m.num_local_heads, m.qk_head_dim) + + else: + q = m.q_proj(hidden_states)[0].view(-1, m.num_local_heads, m.qk_head_dim) + latent_cache = m.kv_a_proj_with_mqa(hidden_states)[0] + + _, q_pe = q.split([m.qk_nope_head_dim, m.qk_rope_head_dim], dim=-1) + kv_a, _ = latent_cache.split([m.kv_lora_rank, m.qk_rope_head_dim], dim=-1) + latent_cache = latent_cache.unsqueeze(1) + + kv_a = m.kv_a_layernorm(kv_a) + kv = m.kv_b_proj(kv_a)[0] + + k_pe = latent_cache[:, :, m.kv_lora_rank :] + if m.rotary_emb is not None: + q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe) + q[..., m.qk_nope_head_dim :] = q_pe + + m._set_mla_kv_buffer(latent_cache, kv_a, k_pe, forward_batch) + if forward_batch.mha_one_shot and sum(forward_batch.extend_prefix_lens_cpu) != 0: + if m.use_nsa and m.kv_cache_dtype == "fp8_e4m3": + # FP8 path: dequantize NSA-specific FP8 format to BF16 + kv_a, k_pe = m._get_mla_kv_buffer_from_fp8(forward_batch) + else: + # BF16/FP16 path: directly fetch from cache + kv_a, k_pe = m._get_mla_kv_buffer( + forward_batch.fetch_mha_one_shot_kv_indices(), + q.dtype, + forward_batch, + ) + kv = m.kv_b_proj(kv_a)[0] + kv = kv.view(-1, m.num_local_heads, m.qk_nope_head_dim + m.v_head_dim) + k_nope = kv[..., : m.qk_nope_head_dim] + v = kv[..., m.qk_nope_head_dim :] + + k = m._concat_and_cast_mha_k(k_nope, k_pe, forward_batch) + return q, k, v, forward_batch + + +def forward_mha_core_npu( + m: "DeepseekV2AttentionMLA", + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + forward_batch: "ForwardBatch", +) -> torch.Tensor: + attn_output = m.attn_mha(q, k, v, forward_batch, save_kv_cache=False) + attn_output = attn_output.reshape(-1, m.num_local_heads * m.v_head_dim) + output, _ = m.o_proj(attn_output) + return output + + +# endregion + + +# region MLA +def forward_mla_prepare_npu( + m: "DeepseekV2AttentionMLA", + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: "ForwardBatch", + zero_allocator: "BumpAllocator", +): + if is_mla_preprocess_enabled(): + if not hasattr(m, "mla_preprocess"): + m.mla_preprocess = NPUFusedMLAPreprocess( + m.fused_qkv_a_proj_with_mqa, + m.q_a_layernorm, + m.kv_a_layernorm, + m.q_b_proj, + m.w_kc, + m.rotary_emb, + m.layer_id, + m.num_local_heads, + m.qk_nope_head_dim, + m.qk_rope_head_dim, + m.quant_config, + ) + ( + q_pe, + k_pe, + q_nope_out, + k_nope, + forward_batch, + zero_allocator, + positions, + ) = m.mla_preprocess.forward( + positions, hidden_states, forward_batch, zero_allocator + ) + topk_indices = None + else: + q_lora = None + if m.q_lora_rank is not None: + q, latent_cache = ( + get_attn_tp_context() + .fetch_qkv_latent() + .split( + [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], + dim=-1, + ) + ) + k_nope = latent_cache[..., : m.kv_lora_rank] + + q = m.q_a_layernorm(q) + k_nope = m.kv_a_layernorm(k_nope) + + # q_lora needed by indexer + if m.use_nsa: + q_lora = q + + k_nope = k_nope.unsqueeze(1) + q = m.q_b_proj(q)[0].view(-1, m.num_local_heads, m.qk_head_dim) + else: + q = m.q_proj(hidden_states)[0].view(-1, m.num_local_heads, m.qk_head_dim) + latent_cache = m.kv_a_proj_with_mqa(hidden_states)[0] + k_nope = latent_cache[..., : m.kv_lora_rank] + k_nope = m.kv_a_layernorm(k_nope).unsqueeze(1) + + q_nope, q_pe = q.split([m.qk_nope_head_dim, m.qk_rope_head_dim], dim=-1) + k_pe = latent_cache[..., m.kv_lora_rank :].unsqueeze(1) + + q_nope_out = torch.bmm(q_nope.transpose(0, 1), m.w_kc) + + q_nope_out = q_nope_out.transpose(0, 1) + + if enable_prefill_cp(forward_batch, m.nsa_enable_prefill_cp): + positions = cp_split_and_rebuild_position(forward_batch, positions) + + q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe) + + if enable_prefill_cp(forward_batch, m.nsa_enable_prefill_cp): + # support allgather+rerrange + k_nope, k_pe = m.rebuild_cp_kv_cache( + latent_cache, forward_batch, k_nope, k_pe + ) + topk_indices = None + if q_lora is not None: + topk_indices = m.indexer( + x=hidden_states, + q_lora=q_lora, + positions=positions, + forward_batch=forward_batch, + layer_id=m.layer_id, + ) + + return ( + q_pe, + k_pe, + q_nope_out, + k_nope, + forward_batch, + zero_allocator, + positions, + topk_indices, + ) + + +def forward_mla_core_npu( + m: "DeepseekV2AttentionMLA", + q_pe: torch.Tensor, + k_pe: torch.Tensor, + q_nope_out: torch.Tensor, + k_nope: torch.Tensor, + forward_batch: "ForwardBatch", + zero_allocator: "BumpAllocator", + positions: torch.Tensor, + topk_indices: torch.Tensor, +) -> torch.Tensor: + attn_output = m.attn_mqa( + q_nope_out, + k_nope, + k_nope, + forward_batch, + q_rope=q_pe, + k_rope=k_pe, + **(dict(topk_indices=topk_indices) if topk_indices is not None else {}), + ) + + attn_output = attn_output.view(-1, m.num_local_heads, m.kv_lora_rank) + + attn_bmm_output = torch.empty( + (attn_output.shape[0], m.num_local_heads * m.v_head_dim), + dtype=attn_output.dtype, + device=attn_output.device, + ) + torch.bmm( + attn_output.transpose(0, 1), + m.w_vc, + out=attn_bmm_output.view(-1, m.num_local_heads, m.v_head_dim).transpose(0, 1), + ) + output, _ = m.o_proj(attn_bmm_output) + + return output + + +# endregion + + +# region DSA +def forward_dsa_prepare_npu( + m: "DeepseekV2AttentionMLA", + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: "ForwardBatch", + zero_allocator: "BumpAllocator", +): + if is_mla_preprocess_enabled() and forward_batch.forward_mode.is_decode(): + if not hasattr(m, "mla_preprocess"): + m.mla_preprocess = NPUFusedMLAPreprocess( + m.fused_qkv_a_proj_with_mqa, + m.q_a_layernorm, + m.kv_a_layernorm, + m.q_b_proj, + m.w_kc, + m.rotary_emb, + m.layer_id, + m.num_local_heads, + m.qk_nope_head_dim, + m.qk_rope_head_dim, + m.quant_config, + ) + ( + q_pe, + k_pe, + q_nope_out, + k_nope, + forward_batch, + zero_allocator, + positions, + ) = m.mla_preprocess.forward( + positions, hidden_states, forward_batch, zero_allocator + ) + + fused_qkv_a_proj_out = m.fused_qkv_a_proj_with_mqa(hidden_states)[0] + q, _ = fused_qkv_a_proj_out.split( + [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1 + ) + q_lora = m.q_a_layernorm(q) + else: + fused_qkv_a_proj_out = m.fused_qkv_a_proj_with_mqa(hidden_states)[0] + q, latent_cache = fused_qkv_a_proj_out.split( + [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1 + ) + k_nope = latent_cache[..., : m.kv_lora_rank] + + q = m.q_a_layernorm(q) + k_nope = m.kv_a_layernorm(k_nope) + + q_lora = q.clone() # required for topk_indices + k_nope = k_nope.unsqueeze(1) + q = m.q_b_proj(q)[0].view(-1, m.num_local_heads, m.qk_head_dim) + + q_nope, q_pe = q.split([m.qk_nope_head_dim, m.qk_rope_head_dim], dim=-1) + k_pe = latent_cache[..., m.kv_lora_rank :].unsqueeze(1) + + q_nope_out = torch.bmm(q_nope.transpose(0, 1), m.w_kc) + + q_nope_out = q_nope_out.transpose(0, 1) + + q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe) + + topk_indices = m.indexer( + hidden_states, q_lora, positions, forward_batch, m.layer_id + ) + + return ( + q_pe, + k_pe, + q_nope_out, + k_nope, + topk_indices, + forward_batch, + zero_allocator, + positions, + ) + + +def forward_dsa_core_npu( + m: "DeepseekV2AttentionMLA", + q_pe: torch.Tensor, + k_pe: torch.Tensor, + q_nope_out: torch.Tensor, + k_nope: torch.Tensor, + topk_indices: torch.Tensor, + forward_batch: "ForwardBatch", + zero_allocator: "BumpAllocator", + positions: torch.Tensor, +) -> torch.Tensor: + attn_output = m.attn_mqa( + q_nope_out.contiguous(), + k_nope.contiguous(), + k_nope.contiguous(), + forward_batch, + save_kv_cache=True, # False if forward_batch.forward_mode.is_extend() else True, + q_rope=q_pe.contiguous(), + k_rope=k_pe.contiguous(), + topk_indices=topk_indices, + ) + attn_output = attn_output.view(-1, m.num_local_heads, m.kv_lora_rank) + + attn_bmm_output = torch.empty( + (attn_output.shape[0], m.num_local_heads, m.v_head_dim), + dtype=attn_output.dtype, + device=attn_output.device, + ) + + if not forward_batch.forward_mode.is_decode(): + attn_output = attn_output.transpose(0, 1) + torch.bmm( + attn_output, + m.w_vc, + out=attn_bmm_output.view(-1, m.num_local_heads, m.v_head_dim).transpose( + 0, 1 + ), + ) + else: + attn_output = attn_output.contiguous() + torch.ops.npu.batch_matmul_transpose(attn_output, m.w_vc, attn_bmm_output) + + attn_bmm_output = attn_bmm_output.reshape(-1, m.num_local_heads * m.v_head_dim) + + output, _ = m.o_proj(attn_bmm_output) + return output + + +# endregion diff --git a/python/sglang/srt/hardware_backend/npu/moe/topk.py b/python/sglang/srt/hardware_backend/npu/moe/topk.py new file mode 100644 index 000000000..e00c33a26 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/moe/topk.py @@ -0,0 +1,79 @@ +from typing import TYPE_CHECKING, Optional + +import torch +from sgl_kernel_npu.norm.l1_norm import l1_norm + +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.eplb.expert_location_dispatch import topk_ids_logical_to_physical +from sglang.srt.layers.moe.topk import StandardTopKOutput, select_experts + +if TYPE_CHECKING: + from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo + from sglang.srt.layers.moe.topk import TopKConfig, TopKOutput + + +def fused_topk_npu( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + topk_config: "TopKConfig", + num_token_non_padded: Optional[torch.Tensor] = None, + expert_location_dispatch_info: Optional["ExpertLocationDispatchInfo"] = None, +) -> "TopKOutput": + + use_grouped_topk = topk_config.use_grouped_topk + renormalize = topk_config.renormalize + correction_bias = topk_config.correction_bias + + if not use_grouped_topk: + topk_weights, topk_ids, _ = torch.ops.npu.npu_moe_gating_top_k_softmax( + router_logits, + k=topk_config.top_k, + ) + + if renormalize: + topk_weights = l1_norm( + topk_weights + if topk_config.num_fused_shared_experts == 0 + else topk_weights[:, :-1] + ) + topk_weights = topk_weights.to(torch.float32) + + elif use_grouped_topk and correction_bias is not None: + routed_scaling_factor = topk_config.routed_scaling_factor or 1 + + topk_weights, topk_ids, _ = torch.ops.npu.npu_moe_gating_top_k( + router_logits.to(torch.float32), + k=topk_config.top_k, + bias=correction_bias.to(torch.float32), + k_group=topk_config.topk_group, + group_count=topk_config.num_expert_group, + group_select_mode=1, + renorm=0, + norm_type=1, + routed_scaling_factor=routed_scaling_factor, + eps=float(1e-20), + ) + + if renormalize: + topk_weights_sum = ( + topk_weights.sum(dim=-1, keepdim=True) + if topk_config.num_fused_shared_experts == 0 + else topk_weights[:, :-1].sum(dim=-1, keepdim=True) + ) + topk_weights = topk_weights / topk_weights_sum + + else: + topk_config.torch_native = True + return select_experts( + hidden_states=hidden_states, + router_logits=router_logits, + topk_config=topk_config, + num_token_non_padded=num_token_non_padded, + expert_location_dispatch_info=expert_location_dispatch_info, + ) + + if expert_location_dispatch_info is not None: + topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info) + get_global_expert_distribution_recorder().on_select_experts(topk_ids=topk_ids) + + return StandardTopKOutput(topk_weights, topk_ids, router_logits) diff --git a/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py b/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py new file mode 100644 index 000000000..938314b0f --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py @@ -0,0 +1,916 @@ +from typing import TYPE_CHECKING + +import numpy as np +import torch + +from sglang.srt.hardware_backend.npu.utils import npu_format_cast +from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase +from sglang.srt.utils import set_weight_attrs + +if TYPE_CHECKING: + from sglang.srt.layers.moe import MoeRunnerConfig + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) + + +def npu_fused_experts( + hidden_states: torch.Tensor, + w13: torch.Tensor, + w13_scale: torch.Tensor, + w2: torch.Tensor, + w2_scale: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + top_k: int, + **kwargs, +): + w13_offset = kwargs.get("w13_offset", None) + w2_offset = kwargs.get("w2_offset", None) + use_wna16 = kwargs.get("use_wna16", False) + + original_shape = hidden_states.shape + original_dtype = hidden_states.dtype + scale_dtype = original_dtype if original_dtype == torch.bfloat16 else torch.float32 + if len(original_shape) == 3: + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + num_tokens = hidden_states.shape[0] + num_experts = w13.shape[0] + row_idx_len = num_tokens * top_k + row_idx = ( + torch.arange(0, row_idx_len, dtype=torch.int32, device=topk_weights.device) + .view(top_k, -1) + .permute(1, 0) + .contiguous() + ) + hidden_states, expanded_row_idx, expanded_expert_idx = ( + torch.ops.npu.npu_moe_init_routing( + hidden_states, row_idx=row_idx, expert_idx=topk_ids, active_num=num_tokens + ) + ) + expert_tokens = torch.ops.npu.npu_moe_compute_expert_tokens( + expanded_expert_idx, num_experts + ) + expert_tokens = expert_tokens.to(torch.int64) + # gmm1: gate_up_proj + if not use_wna16: + hidden_states, pertoken_scale = torch.ops.npu.npu_dynamic_quant(hidden_states) + scale_args13 = { + "scale": [w13_scale.to(scale_dtype)], + "per_token_scale": [pertoken_scale], + } + else: + scale_args13 = { + "antiquant_scale": [w13_scale], + "antiquant_offset": [w13_offset], + } + + hidden_states = torch.ops.npu.npu_grouped_matmul( + x=[hidden_states], + weight=[w13], + **scale_args13, + split_item=2, + group_list_type=0, + group_type=0, + group_list=expert_tokens, + output_dtype=original_dtype, + )[0] + # act_fn: swiglu + hidden_states = torch.ops.npu.npu_swiglu(hidden_states) + if not use_wna16: + hidden_states, pertoken_scale = torch.ops.npu.npu_dynamic_quant(hidden_states) + + scale_args2 = { + "scale": [w2_scale.to(scale_dtype)], + "per_token_scale": [pertoken_scale], + } + else: + scale_args2 = {"antiquant_scale": [w2_scale], "antiquant_offset": [w2_offset]} + # gmm2: down_proj + hidden_states = torch.ops.npu.npu_grouped_matmul( + x=[hidden_states], + weight=[w2], + **scale_args2, + split_item=2, + group_list_type=0, + group_type=0, + group_list=expert_tokens, + output_dtype=original_dtype, + )[0] + + final_hidden_states = torch.ops.npu.npu_moe_finalize_routing( + hidden_states, + skip1=None, + skip2=None, + bias=None, + scales=topk_weights, + expanded_src_to_dst_row=expanded_row_idx, + export_for_source_row=topk_ids, + ) + if len(original_shape) == 3: + final_hidden_states = final_hidden_states.view(original_shape) + return final_hidden_states + + +def npu_fused_moe_without_routing_weights_bf16( + layer, hidden_states, group_list_type, group_list, output_dtype +): + # gmm1: gate_up_proj + hidden_states = torch.ops.npu.npu_grouped_matmul( + x=[hidden_states], + weight=[layer.w13_weight.permute(0, 2, 1)], + split_item=2, + group_list_type=group_list_type, + group_type=0, + group_list=group_list, + output_dtype=output_dtype, + )[0] + hidden_states = torch.ops.npu.npu_swiglu(hidden_states) + # gmm2: down_proj + hidden_states = torch.ops.npu.npu_grouped_matmul( + x=[hidden_states], + weight=[layer.w2_weight.permute(0, 2, 1)], + split_item=2, + group_list_type=group_list_type, + group_type=0, + group_list=group_list, + output_dtype=output_dtype, + )[0] + return hidden_states + + +class NPUW8A8Int8DynamicMoEMethod(FusedMoEMethodBase): + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + self.num_experts = num_experts + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} + ) + + # weight + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size, + dtype=torch.int8, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + dtype=torch.int8, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + # scale + w13_weight_scale = torch.nn.Parameter( + torch.empty( + num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + w2_weight_scale = torch.nn.Parameter( + torch.empty(num_experts, hidden_size, 1, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + # offset + w13_weight_offset = torch.nn.Parameter( + torch.empty( + num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_offset", w13_weight_offset) + set_weight_attrs(w13_weight_offset, extra_weight_attrs) + w2_weight_offset = torch.nn.Parameter( + torch.empty(num_experts, hidden_size, 1, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w2_weight_offset", w2_weight_offset) + set_weight_attrs(w2_weight_offset, extra_weight_attrs) + + def release_weight_cache(self, weight: torch.Tensor): + # .contiguous() introduces additional memory overhead and needs to be released using resize_(0) + origin_weight = weight.data.transpose(1, 2) + new_weight = origin_weight.contiguous() + origin_weight.untyped_storage().resize_(0) + return new_weight + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + weight_data = self.release_weight_cache(layer.w13_weight.data) + layer.w13_weight = torch.nn.Parameter(weight_data, requires_grad=False) + + weight_data = self.release_weight_cache(layer.w2_weight.data) + layer.w2_weight = torch.nn.Parameter(weight_data, requires_grad=False) + + layer.w13_weight_scale = torch.nn.Parameter( + layer.w13_weight_scale.data.squeeze(-1).contiguous().to(torch.float32), + requires_grad=False, + ) + layer.w2_weight_scale = torch.nn.Parameter( + layer.w2_weight_scale.data.squeeze(-1).contiguous(), requires_grad=False + ) + layer.w13_weight_offset = torch.nn.Parameter( + layer.w13_weight_offset.data.squeeze(-1).contiguous(), requires_grad=False + ) + layer.w2_weight_offset = torch.nn.Parameter( + layer.w2_weight_offset.data.squeeze(-1).contiguous(), requires_grad=False + ) + + layer.w13_weight.data = npu_format_cast(layer.w13_weight.data) + layer.w2_weight.data = npu_format_cast(layer.w2_weight.data) + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig" + ): + self.moe_runner_config = moe_runner_config + + def apply( + self, + layer, + dispatch_output: "StandardDispatchOutput", + ) -> "CombineInput": + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + topk_weights, topk_ids, _ = topk_output + topk_ids = topk_ids.to(torch.int32) + topk_weights = topk_weights.to(x.dtype) + output = npu_fused_experts( + hidden_states=x, + w13=layer.w13_weight, + w13_scale=layer.w13_weight_scale, + w2=layer.w2_weight, + w2_scale=layer.w2_weight_scale, + topk_weights=topk_weights, + topk_ids=topk_ids, + top_k=topk_ids.shape[1], + ) + return StandardCombineInput(hidden_states=output) + + def apply_without_routing_weights( + self, + layer, + hidden_states, + hidden_states_scale, + group_list_type, + group_list, + output_dtype, + ): + # gmm1: gate_up_proj + hidden_states = torch.ops.npu.npu_grouped_matmul( + x=[hidden_states], + weight=[layer.w13_weight], + split_item=2, + group_list_type=group_list_type, + group_type=0, + group_list=group_list, + output_dtype=torch.int32, + )[0] + + # act_fn: swiglu + hidden_states, swiglu_out_scale = torch.ops.npu.npu_dequant_swiglu_quant( + x=hidden_states, + weight_scale=layer.w13_weight_scale, + activation_scale=hidden_states_scale, + bias=None, + quant_scale=None, + quant_offset=None, + group_index=group_list, + activate_left=True, + quant_mode=1, + ) + + # gmm2: down_proj + hidden_states = torch.ops.npu.npu_grouped_matmul( + x=[hidden_states], + weight=[layer.w2_weight], + scale=[layer.w2_weight_scale.to(output_dtype)], + per_token_scale=[swiglu_out_scale], + split_item=2, + group_list_type=group_list_type, + group_type=0, + group_list=group_list, + output_dtype=output_dtype, + )[0] + return hidden_states + + +class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase): + + def __init__(self) -> None: + self.group_size = 256 + self.tp_size = 1 + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + self.num_experts = num_experts + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} + ) + + # >> weight + w13_output_size = intermediate_size_per_partition + w2_output_size = hidden_size // 2 + w13_weight = torch.nn.Parameter( + torch.empty(num_experts, w13_output_size, hidden_size, dtype=torch.int8), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + w2_output_size, + intermediate_size_per_partition, + dtype=torch.int8, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # >> scale + w13_weight_scale = torch.nn.Parameter( + torch.empty( + num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + + w2_weight_scale = torch.nn.Parameter( + torch.empty(num_experts, hidden_size, 1, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + # >> offset + w13_weight_offset = torch.nn.Parameter( + torch.empty( + num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_offset", w13_weight_offset) + set_weight_attrs(w13_weight_offset, extra_weight_attrs) + + w2_weight_offset = torch.nn.Parameter( + torch.empty(num_experts, hidden_size, 1, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w2_weight_offset", w2_weight_offset) + set_weight_attrs(w2_weight_offset, extra_weight_attrs) + + # >>> special param for w4a8 + w13_weight_scale_second = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // self.group_size, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale_second", w13_weight_scale_second) + set_weight_attrs(w13_weight_scale_second, extra_weight_attrs) + w13_weight_offset_second = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // self.group_size, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_offset_second", w13_weight_offset_second) + set_weight_attrs(w13_weight_offset_second, extra_weight_attrs) + + w2_weight_scale_second = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition // self.group_size, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale_second", w2_weight_scale_second) + set_weight_attrs(w2_weight_scale_second, extra_weight_attrs) + + w2_weight_offset_second = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition // self.group_size, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_offset_second", w2_weight_offset_second) + set_weight_attrs(w2_weight_offset_second, extra_weight_attrs) + + w13_scale_bias = torch.nn.Parameter( + torch.empty( + num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 + ), + requires_grad=False, + ) + layer.register_parameter("w13_scale_bias", w13_scale_bias) + set_weight_attrs(w13_scale_bias, extra_weight_attrs) + + w2_scale_bias = torch.nn.Parameter( + torch.empty( + num_experts, hidden_size, 16 // self.tp_size, dtype=torch.float32 + ), + requires_grad=False, + ) + layer.register_parameter("w2_scale_bias", w2_scale_bias) + set_weight_attrs(w2_scale_bias, extra_weight_attrs) + + def process_scale(self, weight: torch.Tensor, scale, per_group_scale): + scale = scale.transpose(1, 2).contiguous() + per_group_scale = per_group_scale.transpose(1, 2).contiguous() + group_num, k, n = weight.shape + # the weight of the new version is reduced by half by pack n, so it needs to be restored + n = n * 2 + per_group_scale = per_group_scale.reshape(group_num, -1, n) + group_num, quantgroup_num, n = per_group_scale.shape + bias = None + + scale_fp32 = (scale * per_group_scale).to(torch.float16).to(torch.float32) + scale_fp32_np = scale_fp32.cpu().numpy() + scale_fp32_np.dtype = np.uint32 + sscale_uint64 = np.zeros((group_num, quantgroup_num, n * 2), dtype=np.uint32) + + sscale_uint64[..., ::2] = scale_fp32_np + + sscale_uint64_buffer = np.frombuffer( + sscale_uint64.tobytes(), dtype=np.int64 + ).copy() + sscale_uint64_tensor = torch.from_numpy(sscale_uint64_buffer).reshape( + group_num, quantgroup_num, n + ) + sscale_uint64_tensor = sscale_uint64_tensor.npu() + return sscale_uint64_tensor, bias + + def update_bias(self, layer, w13_bias, w2_bias): + layer.w13_scale_bias.data = ( + layer.w13_scale_bias.data.transpose(1, 2).contiguous().sum(axis=1) + ) + layer.w2_scale_bias.data = ( + layer.w2_scale_bias.data.transpose(1, 2).contiguous().sum(axis=1) + ) + + def pack_to_int32(self, weight: torch.Tensor): + # pack 4 int8(int4*2) to int32, because in pytorch, we need to use int32 to represent int4 + assert ( + weight.shape[-1] % 4 == 0 + ), "the last dim of weight needs to be divided by 4" + return weight.view(torch.int32).contiguous() + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.w13_weight = torch.nn.Parameter( + layer.w13_weight.data.transpose(1, 2).contiguous(), requires_grad=False + ) + layer.w2_weight = torch.nn.Parameter( + layer.w2_weight.data.transpose(1, 2).contiguous(), requires_grad=False + ) + + w13_weight_scale_second = ( + layer.w13_weight_scale_second.data + if hasattr(layer, "w13_weight_scale_second") + else None + ) + w2_weight_scale_second = ( + layer.w2_weight_scale_second.data + if hasattr(layer, "w2_weight_scale_second") + else None + ) + layer.w13_weight_scale.data, w13_bias = self.process_scale( + layer.w13_weight, layer.w13_weight_scale.data, w13_weight_scale_second + ) + layer.w2_weight_scale.data, w2_bias = self.process_scale( + layer.w2_weight, layer.w2_weight_scale.data, w2_weight_scale_second + ) + if hasattr(layer, "w13_weight_scale_second"): + # scale_second is no longer used, release this part of the memory + del layer.w13_weight_scale_second + del layer.w2_weight_scale_second + del layer.w13_weight_offset_second + del layer.w2_weight_offset_second + + self.update_bias(layer, w13_bias, w2_bias) + + layer.w13_weight.data = npu_format_cast(layer.w13_weight.data) + layer.w2_weight.data = npu_format_cast(layer.w2_weight.data) + layer.w13_weight.data = self.pack_to_int32(layer.w13_weight.data) + layer.w2_weight.data = self.pack_to_int32(layer.w2_weight.data) + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig" + ): + self.moe_runner_config = moe_runner_config + + def apply( + self, + layer, + dispatch_output: "StandardDispatchOutput", + ) -> "CombineInput": + # FIXME W4A8 only support with deepep + raise NotImplementedError( + f"W4A8 only support with deepep for now, please enable --moe-a2a-backend deepep" + ) + + def apply_without_routing_weights( + self, + layer, + hidden_states, + hidden_states_scale, + group_list_type, + group_list, + output_dtype, + ): + hidden_states = torch.ops.npu.npu_grouped_matmul( + x=[hidden_states], + weight=[self.w13_weight], + scale=[self.w13_weight_scale], + bias=[self.w13_scale_bias], + per_token_scale=[hidden_states_scale], + group_list=group_list, + split_item=2, + group_type=0, + group_list_type=group_list_type, + output_dtype=output_dtype, + )[0] + + # act_fn: swiglu + hidden_states = torch.ops.npu.npu_swiglu(hidden_states) + hidden_states, swiglu_out_scale = torch.ops.npu.npu_dynamic_quant(hidden_states) + + hidden_states = torch.ops.npu.npu_grouped_matmul( + x=[hidden_states], + weight=[self.w2_weight], + scale=[self.w2_weight_scale], + bias=[self.w2_scale_bias], + per_token_scale=[swiglu_out_scale], + group_list=group_list, + split_item=2, + group_type=0, + group_list_type=group_list_type, + output_dtype=output_dtype, + )[0] + + return hidden_states + + +class NPUW4A16Int4DynamicMoEMethod(FusedMoEMethodBase): + + def __init__(self, quantization_config) -> None: + self.pack_factor = 8 # weight dtype is int4, but use int32 to create + target = ( + "MoEGMM" if "MoEGMM" in quantization_config.target_scheme_map else "Linear" + ) + if target in quantization_config.target_scheme_map: + self.group_size = quantization_config.target_scheme_map[target][ + "weights" + ].group_size + else: + self.group_size = 128 + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + self.num_experts = num_experts + if ( + extra_weight_attrs.get( + "intermediate_size_full", intermediate_size_per_partition + ) + // intermediate_size_per_partition + > 1 + ): + quant_method = FusedMoeWeightScaleSupported.GROUP.value + else: + quant_method = FusedMoeWeightScaleSupported.CHANNEL.value + extra_weight_attrs.update({"quant_method": quant_method}) + # weight + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // self.pack_factor, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition // self.pack_factor, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # scale + weight_scale_dtype = torch.bfloat16 + w13_weight_scale = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // self.group_size, + dtype=weight_scale_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + w2_weight_scale = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition // self.group_size, + dtype=weight_scale_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + # offset + w13_weight_offset = torch.nn.Parameter( + torch.zeros( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // self.group_size, + dtype=weight_scale_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_offset", w13_weight_offset) + set_weight_attrs(w13_weight_offset, extra_weight_attrs) + + w2_weight_offset = torch.nn.Parameter( + torch.zeros( + num_experts, + hidden_size, + intermediate_size_per_partition // self.group_size, + dtype=weight_scale_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_offset", w2_weight_offset) + set_weight_attrs(w2_weight_offset, extra_weight_attrs) + + def pack_to_int32(self, weight: torch.Tensor): + assert weight.dim() == 3 + if weight.dtype == torch.int32: + # pack 8 int4 to int32, we use a int32 to represent a int4 + assert ( + weight.shape[-1] % 8 == 0 + ), "the last dim of weight needs to be divided by 8" + new_weight = torch.ops.npu.npu_convert_weight_to_int4pack( + weight.flatten(0, 1) + ) + new_weight = new_weight.view(weight.shape[0], weight.shape[1], -1) + elif weight.dtype == torch.int8: + # pack 4 int8(int4*2) to int32, because in pytorch, we need to use int32 to represent int4 + assert ( + weight.shape[-1] % 4 == 0 + ), "the last dim of weight needs to be divided by 4" + new_weight = weight.view(torch.int32).contiguous() + else: + raise ValueError(f"{weight.dtype=} is not supported !") + return new_weight + + def unpack_from_int32( + self, + value: torch.Tensor, + num_bits: int, + shape: torch.Size = None, + packed_dim=1, + ) -> torch.Tensor: + """ + Unpacks a tensor of packed int32 weights into individual int8s, maintaining the + original bit range. + + Return tensors in int8 + + :param value: tensor to unpack + :param num_bits: number of bits to unpack each data point into + :param shape: shape to unpack into, used to remove padding + :returns: unpacked int8 tensor + """ + if value.dtype is not torch.int32: + raise ValueError( + f"Expected {torch.int32} but got {value.dtype}, Aborting unpack." + ) + + if num_bits > 8: + raise ValueError("Unpacking is only supported for less than 8 bits") + + pack_factor = 32 // num_bits + + # unpack + mask = (1 << num_bits) - 1 + + if packed_dim == 1: + unpacked = torch.zeros( + (value.shape[0], value.shape[1] * pack_factor), + device=value.device, + dtype=torch.int32, + ) + for i in range(pack_factor): + unpacked[:, i::pack_factor] = (value >> (num_bits * i)) & mask + + # remove padding + if shape is not None: + original_row_size = int(shape[1]) + unpacked = unpacked[:, :original_row_size] + else: + unpacked = torch.zeros( + (value.shape[0] * pack_factor, value.shape[1]), + device=value.device, + dtype=torch.int32, + ) + for i in range(pack_factor): + unpacked[i::pack_factor, :] = (value >> (num_bits * i)) & mask + + # remove padding + original_row_size = int(shape[0]) + unpacked = unpacked[:original_row_size, :] + + # bits are packed in unsigned format, reformat to signed + # update the value range from unsigned to signed + offset = pow(2, num_bits) // 2 + unpacked = (unpacked - offset).to(torch.int8) + + return unpacked + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + w13_weight_scale = layer.w13_weight_scale.data.transpose(-1, -2).contiguous() + w2_weight_scale = layer.w2_weight_scale.data.transpose(-1, -2).contiguous() + layer.w13_weight_scale = torch.nn.Parameter( + w13_weight_scale, requires_grad=False + ) + layer.w2_weight_scale = torch.nn.Parameter(w2_weight_scale, requires_grad=False) + + layer.w13_weight_offset = torch.nn.Parameter( + layer.w13_weight_offset.data.transpose(-1, -2).contiguous(), + requires_grad=False, + ) + layer.w2_weight_offset = torch.nn.Parameter( + layer.w2_weight_offset.data.transpose(-1, -2).contiguous(), + requires_grad=False, + ) + + # w = [n, k // 8] --> [k, n // 8] + # w13_weight = layer.w13_weight.data.transpose(1, 2).contiguous() + # w2_weight = layer.w2_weight.data.transpose(1, 2).contiguous() + unpacked_w13_weight = ( + self.unpack_from_int32(layer.w13_weight.data.flatten(0, 1), 4) + .view(layer.w13_weight.data.shape[0], layer.w13_weight.data.shape[1], -1) + .transpose(1, 2) + .contiguous() + .int() + ) + unpacked_w2_weight = ( + self.unpack_from_int32(layer.w2_weight.data.flatten(0, 1), 4) + .view(layer.w2_weight.data.shape[0], layer.w2_weight.data.shape[1], -1) + .transpose(1, 2) + .contiguous() + .int() + ) + + w13_weight = self.pack_to_int32(unpacked_w13_weight) + w2_weight = self.pack_to_int32(unpacked_w2_weight) + + layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False) + layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False) + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig" + ): + self.moe_runner_config = moe_runner_config + + def apply( + self, + layer, + dispatch_output: "StandardDispatchOutput", + ) -> "CombineInput": + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + topk_weights, topk_ids, _ = topk_output + topk_ids = topk_ids.to(torch.int32) + topk_weights = topk_weights.to(x.dtype) + output = npu_fused_experts( + hidden_states=x, + w13=layer.w13_weight, + w13_scale=layer.w13_weight_scale, + w13_offset=layer.w13_weight_offset, + w2=layer.w2_weight, + w2_scale=layer.w2_weight_scale, + w2_offset=layer.w2_weight_offset, + topk_weights=topk_weights, + topk_ids=topk_ids, + top_k=topk_ids.shape[1], + use_wna16=True, + ) + return StandardCombineInput(hidden_states=output) + + def apply_without_routing_weights( + self, + layer, + hidden_states, + hidden_states_scale, + group_list_type, + group_list, + output_dtype, + ): + if hidden_states_scale is None: + # gmm1: gate_up_proj + hidden_states = torch.ops.npu.npu_grouped_matmul( + x=[hidden_states], + weight=[layer.w13_weight], + antiquant_scale=[layer.w13_weight_scale], + antiquant_offset=[layer.w13_weight_offset], + split_item=2, + group_list_type=group_list_type, + group_type=0, + group_list=group_list, + output_dtype=output_dtype, + )[0] + + # act_fn: swiglu + hidden_states = torch.ops.npu.npu_swiglu(hidden_states) + + # gmm2: down_proj + out_hidden = torch.ops.npu.npu_grouped_matmul( + x=[hidden_states], + weight=[layer.w2_weight], + antiquant_scale=[layer.w2_weight_scale], + antiquant_offset=[layer.w2_weight_offset], + split_item=2, + group_list_type=group_list_type, + group_type=0, + group_list=group_list, + output_dtype=output_dtype, + )[0] + else: + raise ValueError( + "when weight is int4, hidden_states only supports non-quant dtype!" + ) + + return out_hidden diff --git a/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py b/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py new file mode 100644 index 000000000..46db893b3 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py @@ -0,0 +1,215 @@ +from typing import TYPE_CHECKING, List, Optional + +import torch + +from sglang.srt.hardware_backend.npu.utils import npu_format_cast +from sglang.srt.layers.parameter import ( + ChannelQuantScaleParameter, + ModelWeightParameter, + PerTensorScaleParameter, +) +from sglang.srt.layers.quantization.base_config import LinearMethodBase + +if TYPE_CHECKING: + from sglang.srt.layers.quantization.base_config import QuantizationConfig + + +class _NPULinearMethodBase(LinearMethodBase): + + def __init__( + self, + quant_config: Optional["QuantizationConfig"] = None, + ): + super().__init__() + self.quant_config = quant_config + + +class NPUW8A8Int8LinearMethod(_NPULinearMethodBase): + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + weight_loader = extra_weight_attrs.get("weight_loader") + output_size_per_partition = sum(output_partition_sizes) + + weight = ModelWeightParameter( + data=torch.empty( + (output_size_per_partition, input_size_per_partition), dtype=torch.int8 + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + weight_scale = ChannelQuantScaleParameter( + data=torch.empty((output_size_per_partition, 1), dtype=params_dtype), + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight_scale", weight_scale) + + weight_offset = ChannelQuantScaleParameter( + data=torch.empty((output_size_per_partition, 1), dtype=params_dtype), + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight_offset", weight_offset) + + input_scale = PerTensorScaleParameter( + data=torch.empty(1, dtype=params_dtype), + weight_loader=weight_loader, + ) + input_scale.ignore_warning = True + layer.register_parameter("input_scale", input_scale) + + input_offset = PerTensorScaleParameter( + data=torch.empty(1, dtype=params_dtype), + weight_loader=weight_loader, + ) + input_offset.ignore_warning = True + layer.register_parameter("input_offset", input_offset) + + quant_bias = ChannelQuantScaleParameter( + data=torch.empty(output_size_per_partition, dtype=torch.int32), + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("quant_bias", quant_bias) + + if params_dtype == torch.bfloat16: + deq_scale_dtype = torch.float32 + elif params_dtype == torch.float16: + deq_scale_dtype = torch.int64 + else: + raise ValueError(f"Unsupported params_dtype: {params_dtype}") + deq_scale = ChannelQuantScaleParameter( + data=torch.empty(output_size_per_partition, dtype=deq_scale_dtype), + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("deq_scale", deq_scale) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + from sglang.srt.layers.linear import RowParallelLinear + + original_dtype = x.dtype + if original_dtype != torch.int8: + x = torch.ops.npu.npu_quantize( + x, + layer.aclnn_input_scale_reciprocal, + layer.aclnn_input_offset, + torch.qint8, + -1, + False, + ) + # Only fuse bias add into GEMM for rank 0 (this ensures that + # bias will not get added more than once in Attention TP>1 case) + if isinstance(layer, RowParallelLinear) and layer.tp_rank > 0: + quant_bias = None + else: + quant_bias = layer.quant_bias + return torch.ops.npu.npu_quant_matmul( + x, + layer.weight, + layer.deq_scale, + bias=quant_bias, + output_dtype=original_dtype, + ) + + def process_weights_after_loading(self, layer: torch.nn.Module): + layer.weight.data = layer.weight.data.transpose(0, 1).contiguous() + layer.weight.data = npu_format_cast(layer.weight.data) + + layer.weight_scale.data = torch.flatten(layer.weight_scale.data) + layer.weight_offset.data = torch.flatten(layer.weight_offset.data) + + expanding_factor = layer.weight.data.shape[0] + layer.aclnn_input_scale = torch.nn.Parameter( + layer.input_scale.data.repeat(expanding_factor).to(device="npu"), + requires_grad=False, + ) + layer.aclnn_input_scale_reciprocal = 1 / torch.nn.Parameter( + layer.input_scale.data.repeat(expanding_factor).to(device="npu"), + requires_grad=False, + ) + layer.aclnn_input_offset = torch.nn.Parameter( + layer.input_offset.data.repeat(expanding_factor).to(device="npu"), + requires_grad=False, + ) + + +class NPUW8A8Int8DynamicLinearMethod(_NPULinearMethodBase): + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + weight_loader = extra_weight_attrs.get("weight_loader") + output_size_per_partition = sum(output_partition_sizes) + + weight = ModelWeightParameter( + data=torch.empty( + (output_size_per_partition, input_size_per_partition), dtype=torch.int8 + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + weight_scale = ChannelQuantScaleParameter( + data=torch.empty((output_size_per_partition, 1), dtype=params_dtype), + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight_scale", weight_scale) + + weight_offset = ChannelQuantScaleParameter( + data=torch.empty((output_size_per_partition, 1), dtype=params_dtype), + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight_offset", weight_offset) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + original_dtype = x.dtype + quant_out, dynamic_scale = torch.ops.npu.npu_dynamic_quant(x) + return torch.ops.npu.npu_quant_matmul( + quant_out, + layer.weight, + layer.weight_scale, + pertoken_scale=dynamic_scale, + bias=bias, + output_dtype=original_dtype, + ) + + def process_weights_after_loading(self, layer: torch.nn.Module): + layer.weight.data = layer.weight.data.transpose(0, 1).contiguous() + layer.weight.data = npu_format_cast(layer.weight.data) + + layer.weight_scale.data = layer.weight_scale.data.flatten() + layer.weight_offset.data = layer.weight_offset.data.flatten() diff --git a/python/sglang/srt/hardware_backend/npu/quantization/modelslim.py b/python/sglang/srt/hardware_backend/npu/quantization/modelslim.py new file mode 100644 index 000000000..aae786836 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/quantization/modelslim.py @@ -0,0 +1,241 @@ +from __future__ import annotations + +from types import MappingProxyType +from typing import Any, Dict, List, Mapping, Optional, Tuple, Union, cast + +import torch +from compressed_tensors.quantization import QuantizationStrategy + +from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( + NPUW4A8Int4DynamicMoEMethod, + NPUW4A16Int4DynamicMoEMethod, + NPUW8A8Int8DynamicMoEMethod, +) +from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import ( + NPUW8A8Int8DynamicLinearMethod, + NPUW8A8Int8LinearMethod, +) +from sglang.srt.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import ( + CompressedTensorsConfig, +) +from sglang.srt.layers.quantization.compressed_tensors.utils import should_ignore_layer +from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod +from sglang.srt.utils import apply_module_patch + + +# func refers to RMSNorm.__init__ +def npu_wrapper_rmsnorm_init(func): + def init(self, hidden_size: int, **extra_args) -> None: + func(self, hidden_size, **extra_args) + self.ignore_anti = True + # The Ascend w8a8_int8 quantization requires adding a bias in rmsnorm + self.bias = torch.nn.Parameter(torch.zeros(hidden_size), requires_grad=False) + + return init + + +# func refers to RMSNorm.forward_oot +def npu_wrapper_rmsnorm_forward(func): + def _rmsnorm_forward_oot( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + from sgl_kernel_npu.norm.add_rmsnorm_bias import add_rmsnorm_bias + + if not x.is_contiguous(): + x = x.contiguous() + if residual is not None: + out, residual_out = add_rmsnorm_bias( + x, + residual, + self.weight.data, + self.bias, + self.variance_epsilon, + ) + return out.to(x.dtype), residual_out + + out = torch.ops.npu.npu_rms_norm(x, self.weight.data, self.variance_epsilon)[0] + out = out + self.bias + return out.to(x.dtype) + + return _rmsnorm_forward_oot + + +class ModelSlimConfig(QuantizationConfig): + """ + Config class for ModelSlim Quantization, a NPU-specific quantization type. + """ + + def __init__(self, quant_config: Dict[str, Any] = {}): + super().__init__() + self.quant_description = quant_config + self.is_dynamic = quant_config.get("is_dynamic", False) + self.is_moe_w4_dynamic = False + ignore = cast(List[str], quant_config.get("ignore", [])) + self.ignore = ignore if ignore is not None else [] + packed_modules_mapping = quant_config.get("packed_modules_mapping", {}) + self.packed_modules_mapping = ( + packed_modules_mapping if packed_modules_mapping is not None else {} + ) + self.target_scheme_map = ( + CompressedTensorsConfig._quantization_scheme_map_from_config( + config=quant_config + ) + ) + target = "MoEGMM" if "MoEGMM" in self.target_scheme_map else "Linear" + target_scheme = self.target_scheme_map.get(target, None) + if target_scheme is None: + self.is_moe_w4_dynamic = False + else: + weight_quant = target_scheme.get("weights") + input_quant = target_scheme.get("input_activations") + self.is_moe_w4_dynamic = self.is_dynamic_token_w4(weight_quant, input_quant) + self.is_moe_input_quant = input_quant + + for name in self.quant_description.keys(): + if "norm.bias" in name: + apply_module_patch( + "sglang.srt.layers.layernorm.RMSNorm", + "__init__", + [npu_wrapper_rmsnorm_init], + ) + apply_module_patch( + "sglang.srt.layers.layernorm.RMSNorm", + "forward_npu", + [npu_wrapper_rmsnorm_forward], + ) + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.int8, torch.float16, torch.bfloat16] + + @classmethod + def get_min_capability(cls) -> int: + return 0 + + @classmethod + def get_name(self) -> str: + return "modelslim" + + @classmethod + def get_config_filenames(cls) -> List[str]: + filenames = ["quant_model_description.json"] + return filenames + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> ModelSlimConfig: + return cls(config) + + def get_quant_method( + self, + layer: torch.nn.Module, + prefix: str, + ) -> Optional[QuantizeMethodBase]: + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + if isinstance(layer, LinearBase): + if should_ignore_layer( + prefix, + ignore=self.ignore, + fused_mapping=self.packed_modules_mapping, + ): + return UnquantizedLinearMethod() + key = "model" + if "vision_model" in prefix: + key = "vision_model" + elif "visual" in prefix: + key = "visual" + packed_modules_mapping_subset = self.packed_modules_mapping.get(key, {}) + prefix_in_quant_config = prefix + proj_name = prefix.split(".")[-1] + if proj_name in packed_modules_mapping_subset: + prefix_in_quant_config = prefix.replace( + proj_name, packed_modules_mapping_subset[proj_name][0] + ) + self.is_dynamic = ( + self.quant_description[prefix_in_quant_config + ".weight"] + == "W8A8_DYNAMIC" + ) + if self.is_layer_skipped(prefix, packed_modules_mapping_subset): + return UnquantizedLinearMethod() + return ( + NPUW8A8Int8DynamicLinearMethod(self) + if self.is_dynamic + else NPUW8A8Int8LinearMethod(self) + ) + elif isinstance(layer, FusedMoE): + prefix_in_quant_config = prefix + ".0.down_proj.weight" + is_moe_w4a8_dynamic = ( + self.quant_description.get(prefix_in_quant_config, "STATIC") + == "W4A8_DYNAMIC" + ) + if ( + self.is_moe_w4_dynamic and self.is_moe_input_quant is not None + ) or is_moe_w4a8_dynamic: + return NPUW4A8Int4DynamicMoEMethod() + elif self.is_moe_w4_dynamic and self.is_moe_input_quant is None: + return NPUW4A16Int4DynamicMoEMethod(self) + else: + return NPUW8A8Int8DynamicMoEMethod() + return None + + def is_layer_skipped( + self, prefix: str, fused_mapping: Mapping[str, List[str]] = MappingProxyType({}) + ): + # adapted from vllm.model_executor.layers.quantization.utils.quant_utils.is_layer_skipped + proj_name = prefix.split(".")[-1] + if proj_name in fused_mapping: + shard_prefixes = [ + prefix.replace(proj_name, shard_proj_name) + for shard_proj_name in fused_mapping[proj_name] + ] + + is_skipped = None + for shard_prefix in shard_prefixes: + is_shard_skipped = ( + self.quant_description[shard_prefix + ".weight"] == "FLOAT" + ) + + if is_skipped is None: + is_skipped = is_shard_skipped + elif is_shard_skipped != is_skipped: + raise ValueError( + f"Detected some but not all shards of {prefix} " + "are quantized. All shards of fused layers " + "to have the same precision." + ) + else: + is_skipped = self.quant_description[prefix + ".weight"] == "FLOAT" + + assert is_skipped is not None + return is_skipped + + def get_scaled_act_names(self) -> List[str]: + return [] + + def is_dynamic_token_w4(self, weight_quant, input_quant) -> bool: + is_w4 = weight_quant.num_bits == 4 + weight_strategy = ( + weight_quant.strategy == QuantizationStrategy.TENSOR.value + or weight_quant.strategy == QuantizationStrategy.CHANNEL.value + or weight_quant.strategy == QuantizationStrategy.GROUP.value + ) + if input_quant is not None: + is_token = ( + weight_strategy + and input_quant.strategy == QuantizationStrategy.TOKEN.value + ) + is_dynamic = not weight_quant.dynamic and input_quant.dynamic + else: + is_token = weight_strategy + is_dynamic = not weight_quant.dynamic + + # Both symmetric and asymmetric input quantization supported. + # Only symmetric weight quantization supported. + return is_w4 and weight_quant.symmetric and is_token and is_dynamic diff --git a/python/sglang/srt/hardware_backend/npu/utils.py b/python/sglang/srt/hardware_backend/npu/utils.py new file mode 100644 index 000000000..97c126db5 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/utils.py @@ -0,0 +1,104 @@ +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 is_npu + +if TYPE_CHECKING: + from sglang.srt.server_args import ServerArgs + +logger = logging.getLogger(__name__) +_is_npu = is_npu() + + +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 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 + + import torch_npu + + return torch_npu.npu_format_cast(tensor, acl_format.value) diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 41b37f08d..246e2554f 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -70,7 +70,9 @@ def create_wave_backend(runner): @register_attention_backend("ascend") def create_ascend_backend(runner): - from sglang.srt.layers.attention.ascend_backend import AscendAttnBackend + from sglang.srt.hardware_backend.npu.attention.ascend_backend import ( + AscendAttnBackend, + ) return AscendAttnBackend(runner) diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 29d3ac8af..932f52aeb 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -56,9 +56,9 @@ from sglang.srt.utils import ( is_flashinfer_available, is_gfx95_supported, is_hip, + is_npu, is_sm90_supported, is_sm100_supported, - prepare_weight_cache, ) _is_cuda = is_cuda() @@ -67,11 +67,14 @@ _is_sm90_supported = _is_cuda and is_sm90_supported() _is_sm100_supported = _is_cuda and is_sm100_supported() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip() _is_gfx95_supported = is_gfx95_supported() +_is_npu = is_npu() if _use_aiter and _is_gfx95_supported: from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant from sglang.srt.layers.quantization.rocm_mxfp4_utils import fused_rms_mxfp4_quant +elif _is_npu: + from sglang.srt.hardware_backend.npu.cmo import prepare_weight_cache FUSE_ALLREDUCE_MAX_BATCH_SIZE = 2048 @@ -776,7 +779,7 @@ class CommunicateWithAllReduceAndLayerNormFn: ) else: hidden_states = tensor_model_parallel_all_reduce(hidden_states) - if context.cache is not None: + if _is_npu and context.cache is not None: _ = prepare_weight_cache(hidden_states, context.cache) hidden_states, residual = layernorm(hidden_states, residual) return hidden_states, residual diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index 8cb34ba79..2aba8e1e2 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -318,6 +318,9 @@ class DeepEPMoE(FusedMoE): assert self.quant_method is not None assert self.moe_runner_config.activation == "silu" + from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( + npu_fused_moe_without_routing_weights_bf16, + ) from sglang.srt.layers.moe.token_dispatcher import DispatchOutputChecker # NOTE: Ascend's Dispatch & Combine does not support FP16 @@ -388,33 +391,6 @@ class DeepEPMoE(FusedMoE): return hidden_states -def npu_fused_moe_without_routing_weights_bf16( - layer, hidden_states, group_list_type, group_list, output_dtype -): - # gmm1: gate_up_proj - hidden_states = torch_npu.npu_grouped_matmul( - x=[hidden_states], - weight=[layer.w13_weight.permute(0, 2, 1)], - split_item=2, - group_list_type=group_list_type, - group_type=0, - group_list=group_list, - output_dtype=output_dtype, - )[0] - hidden_states = torch_npu.npu_swiglu(hidden_states) - # gmm2: down_proj - hidden_states = torch_npu.npu_grouped_matmul( - x=[hidden_states], - weight=[layer.w2_weight.permute(0, 2, 1)], - split_item=2, - group_list_type=group_list_type, - group_type=0, - group_list=group_list, - output_dtype=output_dtype, - )[0] - return hidden_states - - class NpuFuseEPMoE(DeepEPMoE): def __init__( self, diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 7546dfe1a..a802647e8 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -91,9 +91,6 @@ if _use_aiter: from aiter import biased_grouped_topk as aiter_biased_grouped_topk except ImportError: raise ImportError("aiter is required when SGLANG_USE_AITER is set to True") -if _is_npu: - import torch_npu - from sgl_kernel_npu.norm.l1_norm import l1_norm # -------------------------------- TopKConfig --------------------------------------- @@ -344,76 +341,15 @@ class TopK(CustomOp): expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, ) -> TopKOutput: - use_grouped_topk = self.topk_config.use_grouped_topk - torch_native = self.topk_config.torch_native - renormalize = self.topk_config.renormalize + from sglang.srt.hardware_backend.npu.moe.topk import fused_topk_npu - if not use_grouped_topk and not torch_native: - topk_weights, topk_ids, _ = torch_npu.npu_moe_gating_top_k_softmax( - router_logits, - k=self.topk_config.top_k, - ) - - if renormalize: - topk_weights = l1_norm( - topk_weights - if self.topk_config.num_fused_shared_experts == 0 - else topk_weights[:, :-1] - ) - topk_weights = topk_weights.to(torch.float32) - - if expert_location_dispatch_info is not None: - topk_ids = topk_ids_logical_to_physical( - topk_ids, expert_location_dispatch_info - ) - get_global_expert_distribution_recorder().on_select_experts( - topk_ids=topk_ids - ) - - return StandardTopKOutput(topk_weights, topk_ids, _) - if use_grouped_topk and not torch_native and router_logits.shape[-1] == 256: - # NOTE: now npu_moe_gating_top_k can only support `group_count=256` pattern - routed_scaling_factor = self.topk_config.routed_scaling_factor or 1 - - topk_weights, topk_ids, _ = torch_npu.npu_moe_gating_top_k( - router_logits.to(torch.float32), - k=self.topk_config.top_k, - bias=self.topk_config.correction_bias.to(torch.float32), - k_group=self.topk_config.topk_group, - group_count=self.topk_config.num_expert_group, - group_select_mode=1, - renorm=0, - norm_type=1, - routed_scaling_factor=routed_scaling_factor, - eps=float(1e-20), - ) - - if renormalize: - topk_weights_sum = ( - topk_weights.sum(dim=-1, keepdim=True) - if self.topk_config.num_fused_shared_experts == 0 - else topk_weights[:, :-1].sum(dim=-1, keepdim=True) - ) - topk_weights = topk_weights / topk_weights_sum - - if expert_location_dispatch_info is not None: - topk_ids = topk_ids_logical_to_physical( - topk_ids, expert_location_dispatch_info - ) - get_global_expert_distribution_recorder().on_select_experts( - topk_ids=topk_ids - ) - - return StandardTopKOutput(topk_weights, topk_ids, _) - else: - self.topk_config.torch_native = True - return select_experts( - hidden_states=hidden_states, - router_logits=router_logits, - topk_config=self.topk_config, - num_token_non_padded=num_token_non_padded, - expert_location_dispatch_info=expert_location_dispatch_info, - ) + return fused_topk_npu( + hidden_states=hidden_states, + router_logits=router_logits, + topk_config=self.topk_config, + num_token_non_padded=num_token_non_padded, + expert_location_dispatch_info=expert_location_dispatch_info, + ) def empty_topk_output(self, device: torch.device) -> TopKOutput: topk = self.topk_config.top_k - self.topk_config.num_fused_shared_experts @@ -548,7 +484,7 @@ def fused_topk( # This is used by the Deepseek V2/V3/R1 series models -@torch.compile(dynamic=True, backend=get_compiler_backend()) +@torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu) def grouped_topk_gpu( hidden_states: torch.Tensor, gating_output: torch.Tensor, @@ -565,9 +501,6 @@ def grouped_topk_gpu( assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" scores = torch.softmax(gating_output, dim=-1) - # NPU compiler limitation - if _is_npu and scores.dtype == torch.bfloat16: - scores = scores.to(torch.float16) num_token = scores.shape[0] num_experts = scores.shape[1] group_scores = ( @@ -584,7 +517,6 @@ def grouped_topk_gpu( .reshape(num_token, -1) ) # [n, e] tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0) # [n, e] - # TODO: NPU can't support directly evaluating a comparison for now topk_weights, topk_ids = torch.topk( tmp_scores, k=topk, @@ -729,7 +661,6 @@ def biased_grouped_topk_impl( tmp_scores = scores_for_choice.masked_fill( ~score_mask.bool(), float("-inf") ) # [n, e] - # TODO: NPU can't support directly evaluating a comparison for now _, topk_ids = torch.topk( tmp_scores, k=topk, diff --git a/python/sglang/srt/layers/quantization/__init__.py b/python/sglang/srt/layers/quantization/__init__.py index 87d6eb03d..f4ec7d8c4 100644 --- a/python/sglang/srt/layers/quantization/__init__.py +++ b/python/sglang/srt/layers/quantization/__init__.py @@ -39,7 +39,7 @@ from sglang.srt.layers.quantization.quark.quark import QuarkConfig from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config from sglang.srt.layers.quantization.w8a8_fp8 import W8A8Fp8Config from sglang.srt.layers.quantization.w8a8_int8 import W8A8Int8Config -from sglang.srt.utils import is_cuda, is_hip, mxfp_supported +from sglang.srt.utils import is_cuda, is_hip, is_npu, mxfp_supported _is_mxfp_supported = mxfp_supported() @@ -78,6 +78,15 @@ if is_cuda() or (_is_mxfp_supported and is_hip()): } ) +if is_npu(): + from sglang.srt.hardware_backend.npu.quantization.modelslim import ModelSlimConfig + + BASE_QUANTIZATION_METHODS.update( + { + "modelslim": ModelSlimConfig, + } + ) + QUANTIZATION_METHODS = {**BASE_QUANTIZATION_METHODS} diff --git a/python/sglang/srt/layers/quantization/awq.py b/python/sglang/srt/layers/quantization/awq.py index 0aa7da0ea..ab3105ec8 100644 --- a/python/sglang/srt/layers/quantization/awq.py +++ b/python/sglang/srt/layers/quantization/awq.py @@ -7,6 +7,9 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional import torch +from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( + npu_fused_experts, +) from sglang.srt.layers.linear import LinearBase, set_weight_attrs from sglang.srt.layers.parameter import GroupQuantScaleParameter, PackedvLLMParameter from sglang.srt.layers.quantization.base_config import ( @@ -31,7 +34,6 @@ from sglang.srt.layers.quantization.marlin_utils import ( ) 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.w8a8_int8 import npu_fused_experts from sglang.srt.utils.patch_torch import register_fake_if_exists if TYPE_CHECKING: diff --git a/python/sglang/srt/layers/quantization/w8a8_int8.py b/python/sglang/srt/layers/quantization/w8a8_int8.py index afb7d4e6c..7c145e310 100644 --- a/python/sglang/srt/layers/quantization/w8a8_int8.py +++ b/python/sglang/srt/layers/quantization/w8a8_int8.py @@ -2,210 +2,48 @@ from __future__ import annotations import logging from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Tuple, Union, cast +from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, cast import torch -from compressed_tensors.quantization import QuantizationStrategy from torch.nn.parameter import Parameter from sglang.srt.distributed import get_tensor_model_parallel_world_size from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo -from sglang.srt.layers.parameter import ( - ChannelQuantScaleParameter, - ModelWeightParameter, - PerTensorScaleParameter, -) +from sglang.srt.layers.parameter import ChannelQuantScaleParameter, ModelWeightParameter from sglang.srt.layers.quantization.base_config import ( FusedMoEMethodBase, LinearMethodBase, QuantizationConfig, QuantizeMethodBase, ) -from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import ( - CompressedTensorsConfig, -) from sglang.srt.layers.quantization.compressed_tensors.utils import should_ignore_layer from sglang.srt.layers.quantization.int8_kernel import per_token_quant_int8 from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod from sglang.srt.utils import ( - apply_module_patch, cpu_has_amx_support, - get_bool_env_var, is_cpu, is_cuda, - is_npu, set_weight_attrs, use_intel_amx_backend, ) if TYPE_CHECKING: - from sglang.srt.layers.moe.token_dispatcher import ( - CombineInput, - StandardDispatchOutput, - ) + from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput _is_cuda = is_cuda() _is_cpu_amx_available = cpu_has_amx_support() _is_cpu = is_cpu() -_is_npu = is_npu() if _is_cuda: from sgl_kernel import int8_scaled_mm -elif _is_npu: - import torch_npu - - try: - from mindie_turbo import _ops as ops - from mindie_turbo.quantize.quant_utils import quant_per_tensor - except ImportError: - useMindIETurbo = False - else: - useMindIETurbo = True - - from sgl_kernel_npu.norm.add_rmsnorm_bias import add_rmsnorm_bias logger = logging.getLogger(__name__) -# func refers to RMSNorm.__init__ -def npu_wrapper_rmsnorm_init(func): - def init(self, hidden_size: int, **extra_args) -> None: - func(self, hidden_size, **extra_args) - self.ignore_anti = True - # The Ascend w8a8_int8 quantization requires adding a bias in rmsnorm - self.bias = torch.nn.Parameter(torch.zeros(hidden_size), requires_grad=False) - - return init - - -# func refers to RMSNorm.forward_oot -def npu_wrapper_rmsnorm_forward(func): - def _rmsnorm_forward_oot( - self, - x: torch.Tensor, - residual: Optional[torch.Tensor] = None, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: - if not x.is_contiguous(): - x = x.contiguous() - if residual is not None: - out, residual_out = add_rmsnorm_bias( - x, - residual, - self.weight.data, - self.bias, - self.variance_epsilon, - ) - return out.to(x.dtype), residual_out - - out = torch_npu.npu_rms_norm(x, self.weight.data, self.variance_epsilon)[0] - out = out + self.bias - return out.to(x.dtype) - - return _rmsnorm_forward_oot - - -def npu_fused_experts( - hidden_states: torch.Tensor, - w13: torch.Tensor, - w13_scale: torch.Tensor, - w2: torch.Tensor, - w2_scale: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - top_k: int, - **kwargs, -): - w13_offset = kwargs.get("w13_offset", None) - w2_offset = kwargs.get("w2_offset", None) - use_wna16 = kwargs.get("use_wna16", False) - - original_shape = hidden_states.shape - original_dtype = hidden_states.dtype - scale_dtype = original_dtype if original_dtype == torch.bfloat16 else torch.float32 - if len(original_shape) == 3: - hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - num_tokens = hidden_states.shape[0] - num_experts = w13.shape[0] - row_idx_len = num_tokens * top_k - row_idx = ( - torch.arange(0, row_idx_len, dtype=torch.int32, device=topk_weights.device) - .view(top_k, -1) - .permute(1, 0) - .contiguous() - ) - hidden_states, expanded_row_idx, expanded_expert_idx = ( - torch_npu.npu_moe_init_routing( - hidden_states, row_idx=row_idx, expert_idx=topk_ids, active_num=num_tokens - ) - ) - expert_tokens = torch_npu.npu_moe_compute_expert_tokens( - expanded_expert_idx, num_experts - ) - expert_tokens = expert_tokens.to(torch.int64) - # gmm1: gate_up_proj - if not use_wna16: - hidden_states, pertoken_scale = torch_npu.npu_dynamic_quant(hidden_states) - scale_args13 = { - "scale": [w13_scale.to(scale_dtype)], - "per_token_scale": [pertoken_scale], - } - else: - scale_args13 = { - "antiquant_scale": [w13_scale], - "antiquant_offset": [w13_offset], - } - - hidden_states = torch_npu.npu_grouped_matmul( - x=[hidden_states], - weight=[w13], - **scale_args13, - split_item=2, - group_list_type=0, - group_type=0, - group_list=expert_tokens, - output_dtype=original_dtype, - )[0] - # act_fn: swiglu - hidden_states = torch_npu.npu_swiglu(hidden_states) - if not use_wna16: - hidden_states, pertoken_scale = torch_npu.npu_dynamic_quant(hidden_states) - - scale_args2 = { - "scale": [w2_scale.to(scale_dtype)], - "per_token_scale": [pertoken_scale], - } - else: - scale_args2 = {"antiquant_scale": [w2_scale], "antiquant_offset": [w2_offset]} - # gmm2: down_proj - hidden_states = torch_npu.npu_grouped_matmul( - x=[hidden_states], - weight=[w2], - **scale_args2, - split_item=2, - group_list_type=0, - group_type=0, - group_list=expert_tokens, - output_dtype=original_dtype, - )[0] - - final_hidden_states = torch_npu.npu_moe_finalize_routing( - hidden_states, - skip1=None, - skip2=None, - bias=None, - scales=topk_weights, - expanded_src_to_dst_row=expanded_row_idx, - export_for_source_row=topk_ids, - ) - if len(original_shape) == 3: - final_hidden_states = final_hidden_states.view(original_shape) - return final_hidden_states - - class W8A8Int8Config(QuantizationConfig): - """Config class for W8A8 or W4A16 Quantization. + """Config class for W8A8 Quantization. - Weight: static, per-channel, symmetric - Activation: dynamic, per-token, symmetric @@ -215,59 +53,20 @@ class W8A8Int8Config(QuantizationConfig): super().__init__() self.quant_description = quant_config self.is_dynamic = quant_config.get("is_dynamic", False) - self.is_moe_w4_dynamic = False ignore = cast(List[str], quant_config.get("ignore", [])) self.ignore = ignore if ignore is not None else [] packed_modules_mapping = quant_config.get("packed_modules_mapping", {}) self.packed_modules_mapping = ( packed_modules_mapping if packed_modules_mapping is not None else {} ) - self.target_scheme_map = ( - CompressedTensorsConfig._quantization_scheme_map_from_config( - config=quant_config - ) - ) - target = "MoEGMM" if "MoEGMM" in self.target_scheme_map else "Linear" - target_scheme = self.target_scheme_map.get(target, None) - if target_scheme is None: - self.is_moe_w4_dynamic = False - else: - weight_quant = target_scheme.get("weights") - input_quant = target_scheme.get("input_activations") - self.is_moe_w4_dynamic = self.is_dynamic_token_w4(weight_quant, input_quant) - self.is_moe_input_quant = input_quant - - if _is_npu: - # Ascend w8a8_int8 quantization with bias, use wrappers to isolate the effects between models - for name in self.quant_description.keys(): - if "norm.bias" in name: - apply_module_patch( - "sglang.srt.layers.layernorm.RMSNorm", - "__init__", - [npu_wrapper_rmsnorm_init], - ) - apply_module_patch( - "sglang.srt.layers.layernorm.RMSNorm", - "forward_npu", - [npu_wrapper_rmsnorm_forward], - ) @classmethod def get_supported_act_dtypes(cls) -> List[torch.dtype]: - return ( - [torch.float16, torch.bfloat16] - if not _is_npu - else [torch.int8, torch.float16, torch.bfloat16] - ) + return [torch.float16, torch.bfloat16] @classmethod def get_min_capability(cls) -> int: - if _is_npu: - raise NotImplementedError( - 'NPU hardware does not support "get_min_capability" feature.' - ) - else: - return 75 + return 75 @classmethod def get_name(self) -> str: @@ -276,15 +75,13 @@ class W8A8Int8Config(QuantizationConfig): @classmethod def get_config_filenames(cls) -> List[str]: filenames = [] - if _is_npu: - filenames.append("quant_model_description.json") return filenames @classmethod def from_config(cls, config: Dict[str, Any]) -> W8A8Int8Config: return cls(config) - def _get_quant_method_npu( + def get_quant_method( self, layer: torch.nn.Module, prefix: str, @@ -292,73 +89,16 @@ class W8A8Int8Config(QuantizationConfig): from sglang.srt.layers.linear import LinearBase from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + if should_ignore_layer( + prefix, ignore=self.ignore, fused_mapping=self.packed_modules_mapping + ): + return UnquantizedLinearMethod() if isinstance(layer, LinearBase): - if should_ignore_layer( - prefix, - ignore=self.ignore, - fused_mapping=self.packed_modules_mapping, - ): - return UnquantizedLinearMethod() - key = "model" - if "vision_model" in prefix: - key = "vision_model" - elif "visual" in prefix: - key = "visual" - packed_modules_mapping_subset = self.packed_modules_mapping.get(key, {}) - prefix_in_quant_config = prefix - proj_name = prefix.split(".")[-1] - if proj_name in packed_modules_mapping_subset: - prefix_in_quant_config = prefix.replace( - proj_name, packed_modules_mapping_subset[proj_name][0] - ) - self.is_dynamic = ( - self.quant_description[prefix_in_quant_config + ".weight"] - == "W8A8_DYNAMIC" - ) - if self.is_layer_skipped(prefix, packed_modules_mapping_subset): - return UnquantizedLinearMethod() - return ( - NPU_W8A8DynamicLinearMethod(self) - if self.is_dynamic - else NPU_W8A8LinearMethod(self) - ) + return W8A8Int8LinearMethod(self) elif isinstance(layer, FusedMoE): - prefix_in_quant_config = prefix + ".0.down_proj.weight" - is_moe_w4a8_dynamic = ( - self.quant_description.get(prefix_in_quant_config, "STATIC") - == "W4A8_DYNAMIC" - ) - if ( - self.is_moe_w4_dynamic and self.is_moe_input_quant is not None - ) or is_moe_w4a8_dynamic: - raise ValueError("npu does not support W4A8 currently!") - elif self.is_moe_w4_dynamic and self.is_moe_input_quant is None: - return NPU_W4A16MoEMethod(self) - else: - return NPU_W8A8MoEMethod(self) + return W8A8Int8MoEMethod(self) return None - def get_quant_method( - self, - layer: torch.nn.Module, - prefix: str, - ) -> Optional[QuantizeMethodBase]: - if _is_npu: - return self._get_quant_method_npu(layer, prefix) - else: - from sglang.srt.layers.linear import LinearBase - from sglang.srt.layers.moe.fused_moe_triton import FusedMoE - - if should_ignore_layer( - prefix, ignore=self.ignore, fused_mapping=self.packed_modules_mapping - ): - return UnquantizedLinearMethod() - if isinstance(layer, LinearBase): - return W8A8Int8LinearMethod(self) - elif isinstance(layer, FusedMoE): - return W8A8Int8MoEMethod(self) - return None - def is_layer_skipped( self, prefix: str, fused_mapping: Mapping[str, List[str]] = MappingProxyType({}) ): @@ -393,27 +133,6 @@ class W8A8Int8Config(QuantizationConfig): def get_scaled_act_names(self) -> List[str]: return [] - def is_dynamic_token_w4(self, weight_quant, input_quant) -> bool: - is_w4 = weight_quant.num_bits == 4 - weight_strategy = ( - weight_quant.strategy == QuantizationStrategy.TENSOR.value - or weight_quant.strategy == QuantizationStrategy.CHANNEL.value - or weight_quant.strategy == QuantizationStrategy.GROUP.value - ) - if input_quant is not None: - is_token = ( - weight_strategy - and input_quant.strategy == QuantizationStrategy.TOKEN.value - ) - is_dynamic = not weight_quant.dynamic and input_quant.dynamic - else: - is_token = weight_strategy - is_dynamic = not weight_quant.dynamic - - # Both symmetric and asymmetric input quantization supported. - # Only symmetric weight quantization supported. - return is_w4 and weight_quant.symmetric and is_token and is_dynamic - class W8A8Int8LinearMethod(LinearMethodBase): @@ -640,872 +359,3 @@ class W8A8Int8MoEMethod(FusedMoEMethodBase): a2_scale=layer.w2_input_scale, ) return self.runner.run(dispatch_output, quant_info) - - -class NPU_W8A8LinearMethodImpl: - """Linear method for NPU W8A8.""" - - def __init__(self) -> None: - # aclnn quant matmul requires to transpose matrix B, set to true by default. - self.transpose_weight = True - - @staticmethod - def get_weight( - input_size: int, - output_size: int, - params_dtype: torch.dtype = torch.bfloat16, - ) -> Dict[str, Any]: - params_dict = {"weight": torch.empty(output_size, input_size, dtype=torch.int8)} - return params_dict - - @staticmethod - def get_pertensor_param(params_dtype: torch.dtype) -> Dict[str, Any]: - params_dict = {} - params_dict["input_scale"] = torch.empty(1, dtype=params_dtype) - params_dict["input_offset"] = torch.empty(1, dtype=params_dtype) - return params_dict - - @staticmethod - def get_perchannel_param( - output_size: int, - params_dtype: torch.dtype, - ) -> Dict[str, Any]: - params_dict = {} - params_dict["quant_bias"] = torch.empty(output_size, dtype=torch.int32) - if params_dtype == torch.bfloat16: - params_dict["deq_scale"] = torch.empty(output_size, dtype=torch.float32) - elif params_dtype == torch.float16: - params_dict["deq_scale"] = torch.empty(output_size, dtype=torch.int64) - params_dict["weight_scale"] = torch.empty(output_size, 1, dtype=params_dtype) - params_dict["weight_offset"] = torch.empty(output_size, 1, dtype=params_dtype) - return params_dict - - @staticmethod - def apply( - layer: torch.nn.Module, - x: torch.Tensor, - bias: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - # To prevent import loops - from sglang.srt.layers.linear import RowParallelLinear - - original_dtype = x.dtype - if original_dtype != torch.int8: - x = torch_npu.npu_quantize( - x, - layer.aclnn_input_scale_reciprocal, - layer.aclnn_input_offset, - torch.qint8, - -1, - False, - ) - # Only fuse bias add into GEMM for rank 0 (this ensures that - # bias will not get added more than once in Attention TP>1 case) - if isinstance(layer, RowParallelLinear) and layer.tp_rank > 0: - quant_bias = None - else: - quant_bias = layer.quant_bias - return torch_npu.npu_quant_matmul( - x, - layer.weight, - layer.deq_scale, - bias=quant_bias, - output_dtype=original_dtype, - ) - - def process_weights_after_loading(self, layer): - expanding_factor = layer.weight.data.shape[1] - layer.aclnn_input_scale = torch.nn.Parameter( - layer.input_scale.data.repeat(expanding_factor).to(device="npu"), - requires_grad=False, - ) - layer.aclnn_input_scale_reciprocal = 1 / torch.nn.Parameter( - layer.input_scale.data.repeat(expanding_factor).to(device="npu"), - requires_grad=False, - ) - layer.aclnn_input_offset = torch.nn.Parameter( - layer.input_offset.data.repeat(expanding_factor).to(device="npu"), - requires_grad=False, - ) - if self.transpose_weight: - layer.weight.data = layer.weight.data.transpose(0, 1).contiguous() - layer.weight_scale.data = torch.flatten(layer.weight_scale.data) - layer.weight_offset.data = torch.flatten(layer.weight_offset.data) - layer.weight.data = torch_npu.npu_format_cast(layer.weight.data, 29) - - -class NPU_W8A8LinearMethodMTImpl: - """Linear method for NPU W8A8.""" - - def __init__(self) -> None: - self.transpose_weight = True - - @staticmethod - def get_weight( - input_size: int, - output_size: int, - params_dtype: torch.dtype = torch.bfloat16, - ) -> Dict[str, Any]: - params_dict = {"weight": torch.empty(output_size, input_size, dtype=torch.int8)} - return params_dict - - @staticmethod - def get_pertensor_param(params_dtype: torch.dtype) -> Dict[str, Any]: - params_dict = {} - params_dict["input_scale"] = torch.empty(1, dtype=params_dtype) - params_dict["input_offset"] = torch.empty(1, dtype=torch.int8) - return params_dict - - @staticmethod - def get_perchannel_param( - output_size: int, - params_dtype: torch.dtype, - ) -> Dict[str, Any]: - params_dict = {} - params_dict["quant_bias"] = torch.empty(output_size, dtype=torch.int32) - if params_dtype == torch.bfloat16: - params_dict["deq_scale"] = torch.empty(output_size, dtype=torch.float32) - elif params_dtype == torch.float16: - params_dict["deq_scale"] = torch.empty(output_size, dtype=torch.int64) - params_dict["weight_scale"] = torch.empty(output_size, 1, dtype=params_dtype) - params_dict["weight_offset"] = torch.empty(output_size, 1, dtype=params_dtype) - return params_dict - - @staticmethod - def apply( - layer: torch.nn.Module, - x: torch.Tensor, - bias: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - # To prevent import loops - from sglang.srt.layers.linear import RowParallelLinear - - original_dtype = x.dtype - if original_dtype != torch.int8: - x = quant_per_tensor(x, layer.input_scale, layer.input_offset) - - # Only fuse bias add into GEMM for rank 0 (this ensures that - # bias will not get added more than once in Attention TP>1 case) - if isinstance(layer, RowParallelLinear) and layer.tp_rank > 0: - quant_bias = None - else: - quant_bias = layer.quant_bias - - return ops.quant_matmul( - x=x, weight=layer.weight, deq_scale=layer.deq_scale, deq_bias=quant_bias - ) - - def process_weights_after_loading(self, layer): - layer.aclnn_deq_scale = torch.nn.Parameter( - torch_npu.npu_trans_quant_param(layer.deq_scale.npu()).to(device="npu"), - requires_grad=False, - ) - - -class NPU_W8A8LinearMethod(LinearMethodBase): - """Linear method for NPU quantization. - - This class search for specific quantization - implementation supported on NPU hardware for linear methods. - - Args: - quant_config: The NPU quantization config. - """ - - def __init__(self, quantization_config: W8A8Int8Config) -> None: - self.quantization_config = quantization_config - self.quant_method = ( - NPU_W8A8LinearMethodMTImpl() - if useMindIETurbo - else NPU_W8A8LinearMethodImpl() - ) - - def create_weights( - self, - layer: torch.nn.Module, - input_size_per_partition: int, - output_partition_sizes: List[int], - input_size: int, - output_size: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ) -> None: - output_size_per_partition = sum(output_partition_sizes) - weight_loader = extra_weight_attrs.get("weight_loader") - - weight_dict = self.quant_method.get_weight( - input_size_per_partition, output_size_per_partition, params_dtype - ) - for weight_name, weight_param in weight_dict.items(): - param = torch.nn.Parameter(weight_param, requires_grad=False) - set_weight_attrs(param, {"input_dim": 1, "output_dim": 0}) - layer.register_parameter(weight_name, param) - set_weight_attrs(param, extra_weight_attrs) - - pertensor_dict = self.quant_method.get_pertensor_param(params_dtype) - for pertensor_name, pertensor_param in pertensor_dict.items(): - param = PerTensorScaleParameter( - data=pertensor_param, weight_loader=weight_loader - ) - # disable warning - param.ignore_warning = True - layer.register_parameter(pertensor_name, param) - - perchannel_dict = self.quant_method.get_perchannel_param( - output_size_per_partition, params_dtype - ) - for perchannel_name, perchannel_param in perchannel_dict.items(): - param = torch.nn.Parameter(perchannel_param, requires_grad=False) - set_weight_attrs(param, {"output_dim": 0}) - layer.register_parameter(perchannel_name, param) - set_weight_attrs(param, extra_weight_attrs) - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - if hasattr(self.quant_method, "process_weights_after_loading"): - self.quant_method.process_weights_after_loading(layer) - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - return self.quant_method.apply(layer, x, bias) - - -class NPU_W8A8DynamicLinearMethodImpl: - """Linear method for NPU W8A8_DYNAMIC.""" - - def __init__(self): - self.transpose_weight = True - - @staticmethod - def get_weight( - input_size: int, output_size: int, params_dtype: torch.dtype - ) -> Dict[str, Any]: - params_dict = {"weight": torch.empty(output_size, input_size, dtype=torch.int8)} - return params_dict - - @staticmethod - def get_pertensor_param(params_dtype: torch.dtype) -> Dict[str, Any]: - return {} - - @staticmethod - def get_perchannel_param( - output_size: int, - params_dtype: torch.dtype, - ) -> Dict[str, Any]: - params_dict = {} - params_dict["weight_scale"] = torch.empty(output_size, 1, dtype=params_dtype) - params_dict["weight_offset"] = torch.empty(output_size, 1, dtype=params_dtype) - return params_dict - - @staticmethod - def apply( - layer: torch.nn.Module, - x: torch.Tensor, - bias: Optional[torch.Tensor] = None, - tp_rank: Optional[int] = 0, - ) -> torch.Tensor: - original_dtype = x.dtype - quant_out, dynamic_scale = torch_npu.npu_dynamic_quant(x) - return torch_npu.npu_quant_matmul( - quant_out, - layer.weight, - layer.weight_scale, - pertoken_scale=dynamic_scale, - bias=bias, - output_dtype=original_dtype, - ) - - def process_weights_after_loading(self, layer): - if self.transpose_weight: - layer.weight.data = layer.weight.data.transpose(0, 1).contiguous() - layer.weight_scale.data = layer.weight_scale.data.flatten() - layer.weight_scale_fp32 = layer.weight_scale.data.to(torch.float32) - layer.weight_offset.data = layer.weight_offset.data.flatten() - layer.weight.data = torch_npu.npu_format_cast(layer.weight.data, 29) - - -class NPU_W8A8DynamicLinearMethod(LinearMethodBase): - """Linear method for NPU quantization. - - This class search for specific quantization - implementations supported on NPU hardware for linear methods. - - Args: - quant_config: The NPU quantization config. - """ - - def __init__(self, quantization_config: W8A8Int8Config) -> None: - self.quantization_config = quantization_config - self.quant_method = NPU_W8A8DynamicLinearMethodImpl() - - def create_weights( - self, - layer: torch.nn.Module, - input_size_per_partition: int, - output_partition_sizes: List[int], - input_size: int, - output_size: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ) -> None: - output_size_per_partition = sum(output_partition_sizes) - weight_loader = extra_weight_attrs.get("weight_loader") - - weight_dict = self.quant_method.get_weight( - input_size_per_partition, output_size_per_partition, params_dtype - ) - for weight_name, weight_param in weight_dict.items(): - param = torch.nn.Parameter(weight_param, requires_grad=False) - set_weight_attrs(param, {"input_dim": 1, "output_dim": 0}) - layer.register_parameter(weight_name, param) - set_weight_attrs(param, extra_weight_attrs) - - pertensor_dict = self.quant_method.get_pertensor_param(params_dtype) - for pertensor_name, pertensor_param in pertensor_dict.items(): - param = PerTensorScaleParameter( - data=pertensor_param, weight_loader=weight_loader - ) - # disable warning - param.ignore_warning = True - layer.register_parameter(pertensor_name, param) - - perchannel_dict = self.quant_method.get_perchannel_param( - output_size_per_partition, params_dtype - ) - for perchannel_name, perchannel_param in perchannel_dict.items(): - param = torch.nn.Parameter(perchannel_param, requires_grad=False) - set_weight_attrs(param, {"output_dim": 0}) - layer.register_parameter(perchannel_name, param) - set_weight_attrs(param, extra_weight_attrs) - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - if hasattr(self.quant_method, "process_weights_after_loading"): - self.quant_method.process_weights_after_loading(layer) - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - return self.quant_method.apply(layer, x, bias) - - -class NPU_W8A8MoEMethod(FusedMoEMethodBase): - """MoE method for NPU quantization. - - This class search for specific quantization - implementations supported on NPU hardware for moe methods. - - Args: - quant_config: The NPU quantization config. - """ - - def __init__(self, quantization_config: W8A8Int8Config) -> None: - self.quantization_config = quantization_config - self.quant_method = self - - def create_weights( - self, - layer: torch.nn.Module, - num_experts: int, - hidden_size: int, - intermediate_size_per_partition: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ) -> None: - from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported - - self.num_experts = num_experts - extra_weight_attrs.update( - {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} - ) - - # weight - w13_weight = torch.nn.Parameter( - torch.empty( - num_experts, - 2 * intermediate_size_per_partition, - hidden_size, - dtype=torch.int8, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight", w13_weight) - set_weight_attrs(w13_weight, extra_weight_attrs) - w2_weight = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - intermediate_size_per_partition, - dtype=torch.int8, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight", w2_weight) - set_weight_attrs(w2_weight, extra_weight_attrs) - # scale - w13_weight_scale = torch.nn.Parameter( - torch.empty( - num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_scale", w13_weight_scale) - set_weight_attrs(w13_weight_scale, extra_weight_attrs) - w2_weight_scale = torch.nn.Parameter( - torch.empty(num_experts, hidden_size, 1, dtype=torch.float32), - requires_grad=False, - ) - layer.register_parameter("w2_weight_scale", w2_weight_scale) - set_weight_attrs(w2_weight_scale, extra_weight_attrs) - # offset - w13_weight_offset = torch.nn.Parameter( - torch.empty( - num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_offset", w13_weight_offset) - set_weight_attrs(w13_weight_offset, extra_weight_attrs) - w2_weight_offset = torch.nn.Parameter( - torch.empty(num_experts, hidden_size, 1, dtype=torch.float32), - requires_grad=False, - ) - layer.register_parameter("w2_weight_offset", w2_weight_offset) - set_weight_attrs(w2_weight_offset, extra_weight_attrs) - - def release_weight_cache(self, weight: torch.Tensor): - # .contiguous() introduces additional memory overhead and needs to be released using resize_(0) - origin_weight = weight.data.transpose(1, 2) - new_weight = origin_weight.contiguous() - origin_weight.untyped_storage().resize_(0) - return new_weight - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - weight_data = self.release_weight_cache(layer.w13_weight.data) - layer.w13_weight = Parameter(weight_data, requires_grad=False) - - weight_data = self.release_weight_cache(layer.w2_weight.data) - layer.w2_weight = Parameter(weight_data, requires_grad=False) - - layer.w13_weight_scale = Parameter( - layer.w13_weight_scale.data.squeeze(-1).contiguous().to(torch.float32), - requires_grad=False, - ) - layer.w2_weight_scale = Parameter( - layer.w2_weight_scale.data.squeeze(-1).contiguous(), requires_grad=False - ) - layer.w13_weight_offset = Parameter( - layer.w13_weight_offset.data.squeeze(-1).contiguous(), requires_grad=False - ) - layer.w2_weight_offset = Parameter( - layer.w2_weight_offset.data.squeeze(-1).contiguous(), requires_grad=False - ) - - if get_bool_env_var("ENABLE_ASCEND_MOE_NZ"): - layer.w13_weight.data = torch_npu.npu_format_cast(layer.w13_weight.data, 29) - layer.w2_weight.data = torch_npu.npu_format_cast(layer.w2_weight.data, 29) - - def create_moe_runner( - self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig - ): - self.moe_runner_config = moe_runner_config - - def apply( - self, - layer, - dispatch_output: StandardDispatchOutput, - ) -> CombineInput: - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - - x = dispatch_output.hidden_states - topk_output = dispatch_output.topk_output - - topk_weights, topk_ids, _ = topk_output - topk_ids = topk_ids.to(torch.int32) - topk_weights = topk_weights.to(x.dtype) - output = npu_fused_experts( - hidden_states=x, - w13=layer.w13_weight, - w13_scale=layer.w13_weight_scale, - w2=layer.w2_weight, - w2_scale=layer.w2_weight_scale, - topk_weights=topk_weights, - topk_ids=topk_ids, - top_k=topk_ids.shape[1], - ) - return StandardCombineInput(hidden_states=output) - - def apply_without_routing_weights( - self, - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ): - # gmm1: gate_up_proj - hidden_states = torch_npu.npu_grouped_matmul( - x=[hidden_states], - weight=[layer.w13_weight], - split_item=2, - group_list_type=group_list_type, - group_type=0, - group_list=group_list, - output_dtype=torch.int32, - )[0] - - # act_fn: swiglu - hidden_states, swiglu_out_scale = torch_npu.npu_dequant_swiglu_quant( - x=hidden_states, - weight_scale=layer.w13_weight_scale, - activation_scale=hidden_states_scale, - bias=None, - quant_scale=None, - quant_offset=None, - group_index=group_list, - activate_left=True, - quant_mode=1, - ) - - # gmm2: down_proj - hidden_states = torch_npu.npu_grouped_matmul( - x=[hidden_states], - weight=[layer.w2_weight], - scale=[layer.w2_weight_scale.to(output_dtype)], - per_token_scale=[swiglu_out_scale], - split_item=2, - group_list_type=group_list_type, - group_type=0, - group_list=group_list, - output_dtype=output_dtype, - )[0] - return hidden_states - - -class NPU_W4A16MoEMethod(FusedMoEMethodBase): - """MoE method for NPU W4A16 quantization. - - This class search for specific quantization - implementations supported on NPU hardware for moe methods. - - Args: - quant_config: The NPU quantization config. - """ - - def __init__(self, quantization_config: W8A8Int8Config) -> None: - self.quantization_config = quantization_config - self.quant_method = self - self.pack_factor = 8 # weight dtype is int4, but use int32 to create - target = ( - "MoEGMM" if "MoEGMM" in quantization_config.target_scheme_map else "Linear" - ) - if target in quantization_config.target_scheme_map: - self.group_size = quantization_config.target_scheme_map[target][ - "weights" - ].group_size - else: - self.group_size = 128 - logger.warning_once("NPU_W4A16MoEMethod !!!") - - def create_weights( - self, - layer: torch.nn.Module, - num_experts: int, - hidden_size: int, - intermediate_size_per_partition: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ) -> None: - from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported - - self.num_experts = num_experts - if ( - extra_weight_attrs.get( - "intermediate_size_full", intermediate_size_per_partition - ) - // intermediate_size_per_partition - > 1 - ): - quant_method = FusedMoeWeightScaleSupported.GROUP.value - else: - quant_method = FusedMoeWeightScaleSupported.CHANNEL.value - extra_weight_attrs.update({"quant_method": quant_method}) - # weight - w13_weight = torch.nn.Parameter( - torch.empty( - num_experts, - 2 * intermediate_size_per_partition, - hidden_size // self.pack_factor, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight", w13_weight) - set_weight_attrs(w13_weight, extra_weight_attrs) - w2_weight = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - intermediate_size_per_partition // self.pack_factor, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight", w2_weight) - set_weight_attrs(w2_weight, extra_weight_attrs) - - # scale - weight_scale_dtype = torch.bfloat16 - w13_weight_scale = torch.nn.Parameter( - torch.empty( - num_experts, - 2 * intermediate_size_per_partition, - hidden_size // self.group_size, - dtype=weight_scale_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_scale", w13_weight_scale) - set_weight_attrs(w13_weight_scale, extra_weight_attrs) - w2_weight_scale = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - intermediate_size_per_partition // self.group_size, - dtype=weight_scale_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_scale", w2_weight_scale) - set_weight_attrs(w2_weight_scale, extra_weight_attrs) - - # offset - w13_weight_offset = torch.nn.Parameter( - torch.zeros( - num_experts, - 2 * intermediate_size_per_partition, - hidden_size // self.group_size, - dtype=weight_scale_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_offset", w13_weight_offset) - set_weight_attrs(w13_weight_offset, extra_weight_attrs) - - w2_weight_offset = torch.nn.Parameter( - torch.zeros( - num_experts, - hidden_size, - intermediate_size_per_partition // self.group_size, - dtype=weight_scale_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_offset", w2_weight_offset) - set_weight_attrs(w2_weight_offset, extra_weight_attrs) - - def pack_to_int32(self, weight: torch.Tensor): - assert weight.dim() == 3 - if weight.dtype == torch.int32: - # pack 8 int4 to int32, we use a int32 to represent a int4 - assert ( - weight.shape[-1] % 8 == 0 - ), "the last dim of weight needs to be divided by 8" - new_weight = torch_npu.npu_convert_weight_to_int4pack(weight.flatten(0, 1)) - new_weight = new_weight.view(weight.shape[0], weight.shape[1], -1) - elif weight.dtype == torch.int8: - # pack 4 int8(int4*2) to int32, because in pytorch, we need to use int32 to represent int4 - assert ( - weight.shape[-1] % 4 == 0 - ), "the last dim of weight needs to be divided by 4" - new_weight = weight.view(torch.int32).contiguous() - else: - raise ValueError(f"{weight.dtype=} is not supported !") - return new_weight - - def unpack_from_int32( - self, - value: torch.Tensor, - num_bits: int, - shape: torch.Size = None, - packed_dim=1, - ) -> torch.Tensor: - """ - Unpacks a tensor of packed int32 weights into individual int8s, maintaining the - original bit range. - - Return tensors in int8 - - :param value: tensor to unpack - :param num_bits: number of bits to unpack each data point into - :param shape: shape to unpack into, used to remove padding - :returns: unpacked int8 tensor - """ - if value.dtype is not torch.int32: - raise ValueError( - f"Expected {torch.int32} but got {value.dtype}, Aborting unpack." - ) - - if num_bits > 8: - raise ValueError("Unpacking is only supported for less than 8 bits") - - pack_factor = 32 // num_bits - - # unpack - mask = (1 << num_bits) - 1 - - if packed_dim == 1: - unpacked = torch.zeros( - (value.shape[0], value.shape[1] * pack_factor), - device=value.device, - dtype=torch.int32, - ) - for i in range(pack_factor): - unpacked[:, i::pack_factor] = (value >> (num_bits * i)) & mask - - # remove padding - if shape is not None: - original_row_size = int(shape[1]) - unpacked = unpacked[:, :original_row_size] - else: - unpacked = torch.zeros( - (value.shape[0] * pack_factor, value.shape[1]), - device=value.device, - dtype=torch.int32, - ) - for i in range(pack_factor): - unpacked[i::pack_factor, :] = (value >> (num_bits * i)) & mask - - # remove padding - original_row_size = int(shape[0]) - unpacked = unpacked[:original_row_size, :] - - # bits are packed in unsigned format, reformat to signed - # update the value range from unsigned to signed - offset = pow(2, num_bits) // 2 - unpacked = (unpacked - offset).to(torch.int8) - - return unpacked - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - w13_weight_scale = layer.w13_weight_scale.data.transpose(-1, -2).contiguous() - w2_weight_scale = layer.w2_weight_scale.data.transpose(-1, -2).contiguous() - layer.w13_weight_scale = Parameter(w13_weight_scale, requires_grad=False) - layer.w2_weight_scale = Parameter(w2_weight_scale, requires_grad=False) - - layer.w13_weight_offset = Parameter( - layer.w13_weight_offset.data.transpose(-1, -2).contiguous(), - requires_grad=False, - ) - layer.w2_weight_offset = Parameter( - layer.w2_weight_offset.data.transpose(-1, -2).contiguous(), - requires_grad=False, - ) - - # w = [n, k // 8] --> [k, n // 8] - # w13_weight = layer.w13_weight.data.transpose(1, 2).contiguous() - # w2_weight = layer.w2_weight.data.transpose(1, 2).contiguous() - unpacked_w13_weight = ( - self.unpack_from_int32(layer.w13_weight.data.flatten(0, 1), 4) - .view(layer.w13_weight.data.shape[0], layer.w13_weight.data.shape[1], -1) - .transpose(1, 2) - .contiguous() - .int() - ) - unpacked_w2_weight = ( - self.unpack_from_int32(layer.w2_weight.data.flatten(0, 1), 4) - .view(layer.w2_weight.data.shape[0], layer.w2_weight.data.shape[1], -1) - .transpose(1, 2) - .contiguous() - .int() - ) - - w13_weight = self.pack_to_int32(unpacked_w13_weight) - w2_weight = self.pack_to_int32(unpacked_w2_weight) - - layer.w13_weight = Parameter(w13_weight, requires_grad=False) - layer.w2_weight = Parameter(w2_weight, requires_grad=False) - - def create_moe_runner( - self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig - ): - self.moe_runner_config = moe_runner_config - - def apply( - self, - layer, - dispatch_output: StandardDispatchOutput, - ) -> CombineInput: - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - - x = dispatch_output.hidden_states - topk_output = dispatch_output.topk_output - - topk_weights, topk_ids, _ = topk_output - topk_ids = topk_ids.to(torch.int32) - topk_weights = topk_weights.to(x.dtype) - output = npu_fused_experts( - hidden_states=x, - w13=layer.w13_weight, - w13_scale=layer.w13_weight_scale, - w13_offset=layer.w13_weight_offset, - w2=layer.w2_weight, - w2_scale=layer.w2_weight_scale, - w2_offset=layer.w2_weight_offset, - topk_weights=topk_weights, - topk_ids=topk_ids, - top_k=topk_ids.shape[1], - use_wna16=True, - ) - return StandardCombineInput(hidden_states=output) - - def apply_without_routing_weights( - self, - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ): - if hidden_states_scale is None: - # gmm1: gate_up_proj - hidden_states = torch_npu.npu_grouped_matmul( - x=[hidden_states], - weight=[layer.w13_weight], - antiquant_scale=[layer.w13_weight_scale], - antiquant_offset=[layer.w13_weight_offset], - split_item=2, - group_list_type=group_list_type, - group_type=0, - group_list=group_list, - output_dtype=output_dtype, - )[0] - - # act_fn: swiglu - hidden_states = torch_npu.npu_swiglu(hidden_states) - - # gmm2: down_proj - out_hidden = torch_npu.npu_grouped_matmul( - x=[hidden_states], - weight=[layer.w2_weight], - antiquant_scale=[layer.w2_weight_scale], - antiquant_offset=[layer.w2_weight_offset], - split_item=2, - group_list_type=group_list_type, - group_type=0, - group_list=group_list, - output_dtype=output_dtype, - )[0] - else: - raise ValueError( - "when weight is int4, hidden_states only supports non-quant dtype!" - ) - - return out_hidden diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index 09f082e62..fa75fe003 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/python/sglang/srt/lora/lora_manager.py @@ -37,16 +37,9 @@ from sglang.srt.lora.utils import ( from sglang.srt.managers.io_struct import LoRAUpdateOutput from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.server_args import ServerArgs -from sglang.srt.utils import is_npu, replace_submodule +from sglang.srt.utils import replace_submodule from sglang.srt.utils.hf_transformers_utils import AutoConfig -if is_npu(): - from torch_npu.contrib import transfer_to_npu # noqa: F401 - - # Re-mock torch.cuda.is_available cuz transfer_to_npu mocks it to True - torch.cuda.is_available = lambda: False - - logger = logging.getLogger(__name__) diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 07cc92563..29b5a17c3 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -384,7 +384,7 @@ class HiCacheController: self.tp_size = get_tensor_model_parallel_world_size() self.dp_rank = 0 - # Currently, AscendMLAPagedTokenToKVPool is the subclass of MLATokenToKVPool. + # Currently, NPUMLATokenToKVPool is the subclass of MLATokenToKVPool. is_mla_backend = isinstance(self.mem_pool_device, MLATokenToKVPool) return HiCacheStorageConfig( diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index a80c82fa7..c4c5a9ebb 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -83,11 +83,8 @@ from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.server_args import ServerArgs, get_global_server_args from sglang.srt.utils import flatten_nested_list -from sglang.srt.utils.common import is_npu from sglang.srt.utils.cuda_ipc_transport_utils import CudaIpcTensorTransportProxy -_is_npu = is_npu() - if TYPE_CHECKING: from sglang.srt.configs.model_config import ModelConfig from sglang.srt.speculative.eagle_info import EagleDraftInput @@ -1152,10 +1149,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): has_grammar: bool = False # Device - if not _is_npu: - device: str = "cuda" - else: - device: str = "npu" + device: str = "cuda" # Speculative decoding spec_algorithm: SpeculativeAlgorithm = None diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 64a1ba366..99550e64a 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -62,8 +62,6 @@ logger = logging.getLogger(__name__) GB = 1024 * 1024 * 1024 _is_cuda = is_cuda() _is_npu = is_npu() -if _is_npu: - import torch_npu def get_tensor_size_bytes(t: Union[torch.Tensor, List[torch.Tensor]]): @@ -1001,10 +999,16 @@ class HybridLinearKVPool(KVCache): assert not enable_kvcache_transpose self.use_mla = use_mla if not use_mla: + + TokenToKVPoolClass = MHATokenToKVPool + if _is_npu: - TokenToKVPoolClass = AscendTokenToKVPool - else: - TokenToKVPoolClass = MHATokenToKVPool + from sglang.srt.hardware_backend.npu.memory_pool_npu import ( + NPUMHATokenToKVPool, + ) + + TokenToKVPoolClass = NPUMHATokenToKVPool + self.full_kv_pool = TokenToKVPoolClass( size=size, page_size=self.page_size, @@ -1016,7 +1020,16 @@ class HybridLinearKVPool(KVCache): enable_memory_saver=enable_memory_saver, ) else: + TokenToKVPoolClass = MLATokenToKVPool + + if _is_npu: + from sglang.srt.hardware_backend.npu.memory_pool_npu import ( + NPUMLATokenToKVPool, + ) + + TokenToKVPoolClass = NPUMLATokenToKVPool + self.full_kv_pool = TokenToKVPoolClass( size=size, page_size=self.page_size, @@ -1290,95 +1303,6 @@ class SWAKVPool(KVCache): ) -class AscendTokenToKVPool(MHATokenToKVPool): - - def _create_buffers(self): - with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): - # [size, head_num, head_dim] for each layer - # The padded slot 0 is used for writing dummy outputs from padded tokens. - # Continuous memory improves the efficiency of Ascend`s transmission backend, - # while other backends remain unchanged. - self.kv_buffer = torch.zeros( - ( - 2, - self.layer_num, - self.size // self.page_size + 1, - self.page_size, - self.head_num, - self.head_dim, - ), - dtype=self.store_dtype, - device=self.device, - ) - self.k_buffer = self.kv_buffer[0] - self.v_buffer = self.kv_buffer[1] - - # for disagg - def get_contiguous_buf_infos(self): - # layer_num x [seq_len, head_num, head_dim] - # layer_num x [page_num, page_size, head_num, head_dim] - kv_data_ptrs = [ - self.get_key_buffer(i).data_ptr() - for i in range(self.start_layer, self.start_layer + self.layer_num) - ] + [ - self.get_value_buffer(i).data_ptr() - for i in range(self.start_layer, self.start_layer + self.layer_num) - ] - kv_data_lens = [ - self.get_key_buffer(i).nbytes - for i in range(self.start_layer, self.start_layer + self.layer_num) - ] + [ - self.get_value_buffer(i).nbytes - for i in range(self.start_layer, self.start_layer + self.layer_num) - ] - kv_item_lens = [ - self.get_key_buffer(i)[0].nbytes - for i in range(self.start_layer, self.start_layer + self.layer_num) - ] + [ - self.get_value_buffer(i)[0].nbytes - for i in range(self.start_layer, self.start_layer + self.layer_num) - ] - return kv_data_ptrs, kv_data_lens, kv_item_lens - - def set_kv_buffer( - self, - layer: RadixAttention, - loc: torch.Tensor, - cache_k: torch.Tensor, - cache_v: torch.Tensor, - k_scale: Optional[float] = None, - v_scale: Optional[float] = None, - layer_id_override: Optional[int] = None, - ): - if layer_id_override is not None: - layer_id = layer_id_override - else: - layer_id = layer.layer_id - if cache_k.dtype != self.dtype: - if k_scale is not None: - cache_k.div_(k_scale) - if v_scale is not None: - cache_v.div_(v_scale) - cache_k = cache_k.to(self.dtype) - cache_v = cache_v.to(self.dtype) - - if self.store_dtype != self.dtype: - cache_k = cache_k.view(self.store_dtype) - cache_v = cache_v.view(self.store_dtype) - - torch_npu._npu_reshape_and_cache( - key=cache_k, - value=cache_v, - key_cache=self.k_buffer[layer_id - self.start_layer].view( - -1, self.page_size, self.head_num, self.head_dim - ), - value_cache=self.v_buffer[layer_id - self.start_layer].view( - -1, self.page_size, self.head_num, self.head_dim - ), - slot_indices=loc, - ) - - class MLATokenToKVPool(KVCache): def __init__( self, @@ -1879,202 +1803,6 @@ class NSATokenToKVPool(MLATokenToKVPool): return kv_size_bytes -class AscendMLAPagedTokenToKVPool(MLATokenToKVPool): - def __init__( - self, - size: int, - page_size: int, - dtype: torch.dtype, - kv_lora_rank: int, - qk_rope_head_dim: int, - index_head_dim: Optional[int], - layer_num: int, - device: str, - enable_memory_saver: bool, - start_layer: Optional[int] = None, - end_layer: Optional[int] = None, - ): - super(MLATokenToKVPool, self).__init__( - size, - page_size, - dtype, - layer_num, - device, - enable_memory_saver, - start_layer, - end_layer, - ) - - self.kv_lora_rank = kv_lora_rank - self.qk_rope_head_dim = qk_rope_head_dim - self.index_head_dim = index_head_dim - - self.custom_mem_pool = None - - with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): - # The padded slot 0 is used for writing dummy outputs from padded tokens. - self.k_buffer = torch.zeros( - ( - layer_num, - self.size // self.page_size + 1, - self.page_size, - 1, - self.kv_lora_rank, - ), - dtype=self.store_dtype, - device=self.device, - ) - self.v_buffer = torch.zeros( - ( - layer_num, - self.size // self.page_size + 1, - self.page_size, - 1, - self.qk_rope_head_dim, - ), - dtype=self.store_dtype, - device=self.device, - ) - if self.index_head_dim is not None: - self.index_k_buffer = torch.zeros( - ( - layer_num, - self.size // self.page_size + 1, - self.page_size, - 1, - self.index_head_dim, - ), - dtype=self.store_dtype, - device=self.device, - ) - - self._finalize_allocation_log(size) - - def get_kv_size_bytes(self): - assert hasattr(self, "k_buffer") - assert hasattr(self, "v_buffer") - kv_size_bytes = 0 - for k_cache in self.k_buffer: - kv_size_bytes += get_tensor_size_bytes(k_cache) - for v_cache in self.v_buffer: - kv_size_bytes += get_tensor_size_bytes(v_cache) - if self.index_head_dim is not None: - assert hasattr(self, "index_k_buffer") - for index_k_cache in self.index_k_buffer: - kv_size_bytes += get_tensor_size_bytes(index_k_cache) - return kv_size_bytes - - def get_kv_buffer(self, layer_id: int): - if self.layer_transfer_counter is not None: - self.layer_transfer_counter.wait_until(layer_id - self.start_layer) - return ( - self.k_buffer[layer_id - self.start_layer], - self.v_buffer[layer_id - self.start_layer], - ) - - def get_key_buffer(self, layer_id: int): - if self.layer_transfer_counter is not None: - self.layer_transfer_counter.wait_until(layer_id - self.start_layer) - - if self.store_dtype != self.dtype: - return self.k_buffer[layer_id - self.start_layer].view(self.dtype) - return self.k_buffer[layer_id - self.start_layer] - - def get_value_buffer(self, layer_id: int): - if self.layer_transfer_counter is not None: - self.layer_transfer_counter.wait_until(layer_id - self.start_layer) - - if self.store_dtype != self.dtype: - return self.v_buffer[layer_id - self.start_layer].view(self.dtype) - return self.v_buffer[layer_id - self.start_layer] - - def get_index_k_buffer(self, layer_id: int): - if self.layer_transfer_counter is not None: - self.layer_transfer_counter.wait_until(layer_id - self.start_layer) - - if self.store_dtype != self.dtype: - return self.index_k_buffer[layer_id - self.start_layer].view(self.dtype) - return self.index_k_buffer[layer_id - self.start_layer] - - # for disagg - def get_contiguous_buf_infos(self): - # MLA has only one kv_buffer, so only the information of this buffer needs to be returned. - kv_data_ptrs = [self.k_buffer[i].data_ptr() for i in range(self.layer_num)] + [ - self.v_buffer[i].data_ptr() for i in range(self.layer_num) - ] - kv_data_lens = [self.k_buffer[i].nbytes for i in range(self.layer_num)] + [ - self.v_buffer[i].nbytes for i in range(self.layer_num) - ] - kv_item_lens = [self.k_buffer[i][0].nbytes for i in range(self.layer_num)] + [ - self.v_buffer[i][0].nbytes for i in range(self.layer_num) - ] - if self.index_head_dim is not None: - kv_data_ptrs += [ - self.index_k_buffer[i].data_ptr() for i in range(self.layer_num) - ] - kv_data_lens += [ - self.index_k_buffer[i].nbytes for i in range(self.layer_num) - ] - kv_item_lens += [ - self.index_k_buffer[i][0].nbytes for i in range(self.layer_num) - ] - return kv_data_ptrs, kv_data_lens, kv_item_lens - - def set_kv_buffer( - self, - layer: RadixAttention, - loc: torch.Tensor, - cache_k: torch.Tensor, - cache_v: torch.Tensor, - ): - layer_id = layer.layer_id - if cache_k.dtype != self.dtype: - cache_k = cache_k.to(self.dtype) - cache_v = cache_v.to(self.dtype) - - if self.store_dtype != self.dtype: - cache_k = cache_k.view(self.store_dtype) - cache_v = cache_v.view(self.store_dtype) - - if cache_v is None: - cache_k, cache_v = cache_k.split( - [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 - ) - - torch_npu.npu_scatter_nd_update_( - self.k_buffer[layer_id - self.start_layer].view(-1, 1, self.kv_lora_rank), - loc.view(-1, 1), - cache_k.view(-1, 1, self.kv_lora_rank), - ) - torch_npu.npu_scatter_nd_update_( - self.v_buffer[layer_id - self.start_layer].view( - -1, 1, self.qk_rope_head_dim - ), - loc.view(-1, 1), - cache_v.view(-1, 1, self.qk_rope_head_dim), - ) - - def set_index_k_buffer( - self, - layer_id: int, - loc: torch.Tensor, - index_k: torch.Tensor, - ): - if index_k.dtype != self.dtype: - index_k = index_k.to(self.dtype) - - if self.store_dtype != self.dtype: - index_k = index_k.view(self.store_dtype) - - torch_npu.npu_scatter_nd_update_( - self.index_k_buffer[layer_id - self.start_layer].view( - -1, 1, self.index_head_dim - ), - loc.view(-1, 1), - index_k.view(-1, 1, self.index_head_dim), - ) - - class DoubleSparseTokenToKVPool(KVCache): def __init__( self, diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index bb428fbee..d5df87c93 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -666,7 +666,7 @@ class MLATokenToKVPoolHost(HostKVCache): 1, self.kv_lora_rank + self.qk_rope_head_dim, ) - # Ascend-specific: Aligns with AscendMLAPagedTokenToKVPool layout + # Ascend-specific: Aligns with NPUMLATokenToKVPool layout # Separately allocate k_buffer and v_buffer for easier data transfer. elif self.layout == "page_first_kv_split": base_dims = ( diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index de74275d6..06f384e0e 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -78,6 +78,7 @@ from sglang.srt.eplb.expert_location import ( set_global_expert_location_metadata, ) from sglang.srt.eplb.expert_location_updater import ExpertLocationUpdater +from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import NPUGraphRunner from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.attention.attention_registry import ( ATTENTION_BACKENDS, @@ -104,10 +105,7 @@ from sglang.srt.mem_cache.allocator import ( SWATokenToKVPoolAllocator, TokenToKVPoolAllocator, ) -from sglang.srt.mem_cache.allocator_ascend import AscendPagedTokenToKVPoolAllocator from sglang.srt.mem_cache.memory_pool import ( - AscendMLAPagedTokenToKVPool, - AscendTokenToKVPool, DoubleSparseTokenToKVPool, HybridLinearKVPool, HybridReqToTokenPool, @@ -132,7 +130,6 @@ from sglang.srt.model_executor.forward_batch_info import ( ) from sglang.srt.model_executor.hook_manager import register_forward_hooks from sglang.srt.model_executor.input_buffers import GraphInputBuffers -from sglang.srt.model_executor.npu_graph_runner import NPUGraphRunner from sglang.srt.model_executor.piecewise_cuda_graph_runner import ( PiecewiseCudaGraphRunner, ) @@ -187,6 +184,17 @@ from sglang.srt.weight_sync.tensor_bucket import ( FlattenedTensorMetadata, ) +_is_cuda = is_cuda() +_is_hip = is_hip() +_is_npu = is_npu() +_is_cpu_amx_available = cpu_has_amx_support() +_is_xpu_xmx_available = xpu_has_xmx_support() + +if _is_npu: + from sglang.srt.hardware_backend.npu.utils import init_npu_backend + + init_npu_backend() + MLA_ATTENTION_BACKENDS = [ "aiter", "flashinfer", @@ -224,12 +232,6 @@ def add_chunked_prefix_cache_attention_backend(backend_name): ) -_is_cuda = is_cuda() -_is_hip = is_hip() -_is_npu = is_npu() -_is_cpu_amx_available = cpu_has_amx_support() -_is_xpu_xmx_available = xpu_has_xmx_support() - # Use a small KV cache pool size for tests in CI SGLANG_CI_SMALL_KV_SIZE = os.getenv("SGLANG_CI_SMALL_KV_SIZE", None) @@ -241,12 +243,6 @@ MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO = 3 logger = logging.getLogger(__name__) -if _is_npu: - import torch_npu - - torch.npu.config.allow_internal_format = True - torch_npu.npu.set_compile_mode(jit_compile=False) - def resolve_language_model(model: nn.Module) -> nn.Module: model_cls_name = model.__class__.__name__ @@ -1814,7 +1810,11 @@ class ModelRunner: is_nsa_model = is_deepseek_nsa(self.model_config.hf_config) if self.server_args.attention_backend == "ascend": if self.use_mla_backend: - self.token_to_kv_pool = AscendMLAPagedTokenToKVPool( + from sglang.srt.hardware_backend.npu.memory_pool_npu import ( + NPUMLATokenToKVPool, + ) + + self.token_to_kv_pool = NPUMLATokenToKVPool( self.max_total_num_tokens, page_size=self.page_size, dtype=self.kv_cache_dtype, @@ -1828,7 +1828,11 @@ class ModelRunner: end_layer=self.end_layer, ) else: - self.token_to_kv_pool = AscendTokenToKVPool( + from sglang.srt.hardware_backend.npu.memory_pool_npu import ( + NPUMHATokenToKVPool, + ) + + self.token_to_kv_pool = NPUMHATokenToKVPool( self.max_total_num_tokens, page_size=self.page_size, dtype=self.kv_cache_dtype, @@ -1986,7 +1990,11 @@ class ModelRunner: self.server_args.attention_backend == "ascend" or self.hybrid_gdn_config is not None ): - self.token_to_kv_pool_allocator = AscendPagedTokenToKVPoolAllocator( + from sglang.srt.hardware_backend.npu.allocator_npu import ( + NPUPagedTokenToKVPoolAllocator, + ) + + self.token_to_kv_pool_allocator = NPUPagedTokenToKVPoolAllocator( self.max_total_num_tokens, page_size=self.page_size, dtype=self.kv_cache_dtype, diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 29a896649..7d0a16003 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -54,10 +54,6 @@ from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.amx_utils import PackWeightMethod -from sglang.srt.layers.attention.npu_ops.mla_preprocess import ( - NPUFusedMLAPreprocess, - is_mla_preprocess_enabled, -) from sglang.srt.layers.attention.nsa.dequant_k_cache import dequantize_k_cache_paged from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer from sglang.srt.layers.attention.nsa.utils import ( @@ -214,10 +210,14 @@ elif _is_hip: awq_dequantize_triton as awq_dequantize, ) elif _is_npu: - import custom_ops # noqa: F401 - import sgl_kernel_npu # noqa: F401 - import torch_npu # noqa: F401 - + from sglang.srt.hardware_backend.npu.modules.deepseek_v2_attention_mla_npu import ( + forward_dsa_core_npu, + forward_dsa_prepare_npu, + forward_mha_core_npu, + forward_mha_prepare_npu, + forward_mla_core_npu, + forward_mla_prepare_npu, + ) from sglang.srt.layers.quantization.awq_triton import ( awq_dequantize_decomposition as awq_dequantize, ) @@ -260,9 +260,6 @@ class AttnForwardMethod(IntEnum): # Use absorbed multi-latent attention MLA = auto() - # Use Deepseek V3.2 sparse multi-latent attention - NPU_MLA_SPARSE = auto() - # Use multi-head attention, but with KV cache chunked. # This method can avoid OOM when prefix lengths are long. MHA_CHUNKED_KV = auto() @@ -277,6 +274,15 @@ class AttnForwardMethod(IntEnum): # Use MLA with fused RoPE kernel for CPU MLA_FUSED_ROPE_CPU = auto() + # Use multi-head attention for NPU + MHA_NPU = auto() + + # Use absorbed multi-latent attention for NPU + MLA_NPU = auto() + + # Use Deepseek V3.2 sparse multi-latent attention for NPU + DSA_NPU = auto() + def _dispatch_mla_subtype(attn, forward_batch): if _is_hip: @@ -311,14 +317,14 @@ def handle_attention_ascend(attn, forward_batch): and not forward_batch.forward_mode.is_draft_extend_v2() ): if hasattr(attn, "indexer"): - return AttnForwardMethod.NPU_MLA_SPARSE + return AttnForwardMethod.DSA_NPU else: - return AttnForwardMethod.MHA + return AttnForwardMethod.MHA_NPU else: if hasattr(attn, "indexer"): - return AttnForwardMethod.NPU_MLA_SPARSE + return AttnForwardMethod.DSA_NPU else: - return AttnForwardMethod.MLA + return AttnForwardMethod.MLA_NPU def _get_sum_extend_prefix_lens(forward_batch): @@ -1243,6 +1249,7 @@ class DeepseekV2AttentionMLA(nn.Module): self.v_head_dim = v_head_dim self.q_lora_rank = q_lora_rank self.kv_lora_rank = kv_lora_rank + self.quant_config = quant_config attn_tp_rank = get_attention_tp_rank() attn_tp_size = get_attention_tp_size() self.use_nsa = is_deepseek_nsa(config) @@ -1470,12 +1477,6 @@ class DeepseekV2AttentionMLA(nn.Module): self.weight_block_size = ( self.fused_qkv_a_proj_with_mqa.quant_method.quant_config.weight_block_size ) - self.is_mla_preprocess_enabled = is_mla_preprocess_enabled() - if self.is_mla_preprocess_enabled: - assert ( - quant_config is None or quant_config.get_name() == "w8a8_int8" - ), "MLA Preprocess only works with Unquant or W8A8Int8" - self.mla_preprocess = None def dispatch_attn_forward_method( self, forward_batch: ForwardBatch @@ -1571,31 +1572,7 @@ class DeepseekV2AttentionMLA(nn.Module): positions, hidden_states, forward_batch, zero_allocator ) elif attn_forward_method == AttnForwardMethod.MLA: - if not self.is_mla_preprocess_enabled: - inner_state = self.forward_absorb_prepare( - positions, hidden_states, forward_batch, zero_allocator - ) - else: - # TODO(iforgetmyname): to be separated as a standalone func - if self.mla_preprocess is None: - self.mla_preprocess = NPUFusedMLAPreprocess( - self.fused_qkv_a_proj_with_mqa, - self.q_a_layernorm, - self.kv_a_layernorm, - self.q_b_proj, - self.w_kc, - self.rotary_emb, - self.layer_id, - self.num_local_heads, - self.qk_nope_head_dim, - self.qk_rope_head_dim, - ) - inner_state = self.mla_preprocess.forward( - positions, hidden_states, forward_batch, zero_allocator - ) - inner_state = (*inner_state, None) # add a position for topk_indices - elif attn_forward_method == AttnForwardMethod.NPU_MLA_SPARSE: - inner_state = self.forward_npu_sparse_prepare( + inner_state = self.forward_absorb_prepare( positions, hidden_states, forward_batch, zero_allocator ) elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE: @@ -1606,6 +1583,18 @@ class DeepseekV2AttentionMLA(nn.Module): inner_state = self.forward_absorb_fused_mla_rope_cpu_prepare( positions, hidden_states, forward_batch, zero_allocator ) + elif attn_forward_method == AttnForwardMethod.MHA_NPU: + inner_state = forward_mha_prepare_npu( + self, positions, hidden_states, forward_batch, zero_allocator + ) + elif attn_forward_method == AttnForwardMethod.MLA_NPU: + inner_state = forward_mla_prepare_npu( + self, positions, hidden_states, forward_batch, zero_allocator + ) + elif attn_forward_method == AttnForwardMethod.DSA_NPU: + inner_state = forward_dsa_prepare_npu( + self, positions, hidden_states, forward_batch, zero_allocator + ) else: raise NotImplementedError return None, attn_forward_method, forward_batch, inner_state @@ -1625,12 +1614,16 @@ class DeepseekV2AttentionMLA(nn.Module): return self.forward_normal_one_shot_core(*inner_state) elif attn_forward_method == AttnForwardMethod.MLA: return self.forward_absorb_core(*inner_state) - elif attn_forward_method == AttnForwardMethod.NPU_MLA_SPARSE: - return self.forward_npu_sparse_core(*inner_state) elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE: return self.forward_absorb_fused_mla_rope_core(*inner_state) elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_CPU: return self.forward_absorb_fused_mla_rope_cpu_core(*inner_state) + elif attn_forward_method == AttnForwardMethod.MHA_NPU: + return forward_mha_core_npu(self, *inner_state) + elif attn_forward_method == AttnForwardMethod.MLA_NPU: + return forward_mla_core_npu(self, *inner_state) + elif attn_forward_method == AttnForwardMethod.DSA_NPU: + return forward_dsa_core_npu(self, *inner_state) else: raise NotImplementedError @@ -2147,253 +2140,6 @@ class DeepseekV2AttentionMLA(nn.Module): return output - def forward_npu_sparse_prepare( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - forward_batch: ForwardBatch, - zero_allocator: BumpAllocator, - ): - """ - Reuse `self.q_lora_rank is not None` branch from forward_absorb_prepare - """ - if self.is_mla_preprocess_enabled and forward_batch.forward_mode.is_decode(): - if self.mla_preprocess is None: - self.mla_preprocess = NPUFusedMLAPreprocess( - self.fused_qkv_a_proj_with_mqa, - self.q_a_layernorm, - self.kv_a_layernorm, - self.q_b_proj, - self.w_kc, - self.rotary_emb, - self.layer_id, - self.num_local_heads, - self.qk_nope_head_dim, - self.qk_rope_head_dim, - ) - ( - q_pe, - k_pe, - q_nope_out, - k_nope, - forward_batch, - zero_allocator, - positions, - ) = self.mla_preprocess.forward( - positions, hidden_states, forward_batch, zero_allocator - ) - - fused_qkv_a_proj_out = self.fused_qkv_a_proj_with_mqa(hidden_states)[0] - q, _ = fused_qkv_a_proj_out.split( - [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1 - ) - q_lora = self.q_a_layernorm(q) - else: - from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode - - if ( - (not isinstance(hidden_states, tuple)) - and hidden_states.shape[0] <= 16 - and self.use_min_latency_fused_a_gemm - ): - fused_qkv_a_proj_out = dsv3_fused_a_gemm( - hidden_states, self.fused_qkv_a_proj_with_mqa.weight.T - ) - else: - fused_qkv_a_proj_out = self.fused_qkv_a_proj_with_mqa(hidden_states)[0] - q, latent_cache = fused_qkv_a_proj_out.split( - [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1 - ) - k_nope = latent_cache[..., : self.kv_lora_rank] - - # overlap qk norm - if self.alt_stream is not None and get_is_capture_mode(): - current_stream = torch.cuda.current_stream() - self.alt_stream.wait_stream(current_stream) - q = self.q_a_layernorm(q) - with torch.cuda.stream(self.alt_stream): - k_nope = self.kv_a_layernorm(k_nope) - current_stream.wait_stream(self.alt_stream) - else: - if _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8: - q, k_nope, *_ = fused_rms_mxfp4_quant( - q, - self.q_a_layernorm.weight, - self.q_a_layernorm.variance_epsilon, - k_nope, - self.kv_a_layernorm.weight, - self.kv_a_layernorm.variance_epsilon, - ) - else: - if ( - _use_aiter_gfx95 - and self.q_b_proj.weight.dtype == torch.float8_e4m3fn - ): - - q, _, k_nope, _ = fused_rms_fp8_group_quant( - q, - self.q_a_layernorm.weight, - self.q_a_layernorm.variance_epsilon, - k_nope, - self.kv_a_layernorm.weight, - self.kv_a_layernorm.variance_epsilon, - group_size=128, - dtype_quant=torch.float8_e4m3fn, - res1=None, - output_unquantized_inp1=False, - ) - - else: - q = self.q_a_layernorm(q) - k_nope = self.kv_a_layernorm(k_nope) - - q_lora = q.clone() # required for topk_indices - k_nope = k_nope.unsqueeze(1) - q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) - - q_nope, q_pe = q.split( - [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 - ) - k_pe = latent_cache[..., self.kv_lora_rank :].unsqueeze(1) - - if self.use_deep_gemm_bmm: - q_nope_val, q_nope_scale, masked_m, expected_m, aligned_m = ( - per_token_group_quant_mla_deep_gemm_masked_fp8( - q_nope.transpose(0, 1) - ) - ) - q_nope_out = q_nope.new_empty( - (self.num_local_heads, aligned_m, self.kv_lora_rank) - ) - deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked( - (q_nope_val, q_nope_scale), - (self.w_kc, self.w_scale_k), - q_nope_out, - masked_m, - expected_m, - ) - q_nope_out = q_nope_out[:, :expected_m, :] - elif _is_hip: - # TODO(haishaw): add bmm_fp8 to ROCm - if _use_aiter_gfx95 and self.w_kc.dtype == torch.uint8: - x = q_nope.transpose(0, 1) - q_nope_out = torch.empty( - x.shape[0], - x.shape[1], - self.w_kc.shape[2], - device=x.device, - dtype=torch.bfloat16, - ) - batched_gemm_afp4wfp4_pre_quant( - x, - self.w_kc.transpose(-2, -1), - self.w_scale_k.transpose(-2, -1), - torch.bfloat16, - q_nope_out, - ) - else: - if _use_aiter_gfx95 and self.w_kc.dtype == torch.float8_e4m3fn: - - q_nope_out = batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant( - X=q_nope, - WQ=self.w_kc.transpose(-1, -2), - w_scale=self.w_scale, # - group_size=128, - YQ=None, # allocate (B, M, N) - transpose_bm=False, # (B, M, N) - transpose_bm_in=True, # (M, B, K) - dtype=torch.bfloat16, - ) - else: - q_nope_out = torch.bmm( - q_nope.to(torch.bfloat16).transpose(0, 1), - self.w_kc.to(torch.bfloat16) * self.w_scale, - ) - elif self.w_kc.dtype == torch.float8_e4m3fn: - q_nope_val, q_nope_scale = per_tensor_quant_mla_fp8( - q_nope.transpose(0, 1), - zero_allocator.allocate(1), - ) - q_nope_out = bmm_fp8( - q_nope_val, self.w_kc, q_nope_scale, self.w_scale, torch.bfloat16 - ) - else: - q_nope_out = torch.bmm(q_nope.transpose(0, 1), self.w_kc) - - q_nope_out = q_nope_out.transpose(0, 1) - - if not self._fuse_rope_for_trtllm_mla(forward_batch) and ( - not _use_aiter or not _is_gfx95_supported - ): - q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) - - # TODO: multi-stream indexer - topk_indices = self.indexer( - hidden_states, q_lora, positions, forward_batch, self.layer_id - ) - - return ( - q_pe, - k_pe, - q_nope_out, - k_nope, - topk_indices, - forward_batch, - zero_allocator, - positions, - ) - - def forward_npu_sparse_core( - self, - q_pe, - k_pe, - q_nope_out, - k_nope, - topk_indices, - forward_batch, - zero_allocator, - positions, - ): - attn_output = self.attn_mqa( - q_nope_out.contiguous(), - k_nope.contiguous(), - k_nope.contiguous(), - forward_batch, - save_kv_cache=True, # False if forward_batch.forward_mode.is_extend() else True, - q_rope=q_pe.contiguous(), - k_rope=k_pe.contiguous(), - topk_indices=topk_indices, - ) - attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank) - - attn_bmm_output = torch.empty( - (attn_output.shape[0], self.num_local_heads, self.v_head_dim), - dtype=attn_output.dtype, - device=attn_output.device, - ) - - if not forward_batch.forward_mode.is_decode(): - attn_output = attn_output.transpose(0, 1) - torch.bmm( - attn_output, - self.w_vc, - out=attn_bmm_output.view( - -1, self.num_local_heads, self.v_head_dim - ).transpose(0, 1), - ) - else: - attn_output = attn_output.contiguous() - torch.ops.npu.batch_matmul_transpose( - attn_output, self.w_vc, attn_bmm_output - ) - - attn_bmm_output = attn_bmm_output.reshape( - -1, self.num_local_heads * self.v_head_dim - ) - - output, _ = self.o_proj(attn_bmm_output) - return output - def forward_absorb_fused_mla_rope_prepare( self, positions: torch.Tensor, @@ -4126,12 +3872,8 @@ class DeepseekV2ForCausalLM(nn.Module): del self.lm_head.weight self.model.embed_tokens.weight = embed self.lm_head.weight = head - if not _is_npu: - torch.cuda.empty_cache() - torch.cuda.synchronize() - else: - torch.npu.empty_cache() - torch.npu.synchronize() + torch.cuda.empty_cache() + torch.cuda.synchronize() @classmethod def get_model_config_for_expert_location(cls, config): diff --git a/python/sglang/srt/models/qwen3.py b/python/sglang/srt/models/qwen3.py index 7d8603a0e..30b92acbd 100644 --- a/python/sglang/srt/models/qwen3.py +++ b/python/sglang/srt/models/qwen3.py @@ -30,13 +30,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.server_args import get_global_server_args -from sglang.srt.utils import ( - add_prefix, - get_cmo_stream, - is_cuda, - is_npu, - wait_cmo_stream, -) +from sglang.srt.utils import add_prefix, is_cuda, is_npu Qwen3Config = None @@ -47,6 +41,8 @@ _is_npu = is_npu() if _is_npu: from sgl_kernel_npu.norm.split_qkv_rmsnorm_rope import split_qkv_rmsnorm_rope + from sglang.srt.hardware_backend.npu.cmo import get_cmo_stream, wait_cmo_stream + class Qwen3Attention(nn.Module): def __init__( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 526465ee5..691cc8629 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -103,6 +103,7 @@ QUANTIZATION_CHOICES = [ "mxfp4", "auto-round", "compressed-tensors", # for Ktransformers + "modelslim", # for NPU ] ATTENTION_BACKEND_CHOICES = [ @@ -615,6 +616,7 @@ class ServerArgs: # Handle device-specific backends. self._handle_hpu_backends() self._handle_cpu_backends() + self._handle_npu_backends() # Apply model-specific adjustments. self._handle_model_specific_adjustments() @@ -928,6 +930,12 @@ class ServerArgs: self.attention_backend = "intel_amx" self.sampling_backend = "pytorch" + def _handle_npu_backends(self): + if self.device == "npu": + from sglang.srt.hardware_backend.npu.utils import set_default_server_args + + set_default_server_args(self) + def _handle_model_specific_adjustments(self): from sglang.srt.configs.model_config import is_deepseek_nsa @@ -1300,8 +1308,6 @@ class ServerArgs: self.attention_backend = "fa3" elif is_hip(): self.attention_backend = "aiter" - elif is_npu(): - self.attention_backend = "ascend" else: self.attention_backend = ( "flashinfer" if is_flashinfer_available() else "triton" @@ -1319,8 +1325,6 @@ class ServerArgs: self.attention_backend = "aiter" else: self.attention_backend = "triton" - elif is_npu(): - self.attention_backend = "ascend" else: self.attention_backend = "triton" @@ -1421,13 +1425,6 @@ class ServerArgs: if model_config.context_len > 8192: self.mem_fraction_static *= 0.85 - # NPU platforms backends - if is_npu() and self.attention_backend in ["ascend"]: - logger.warning( - "At this moment Ascend attention backend only supports a page_size of 128, change page_size to 128." - ) - self.page_size = 128 - # Other platforms backends if ( self.attention_backend == "intel_amx" @@ -1620,7 +1617,7 @@ class ServerArgs: ) if self.hicache_mem_layout == "page_first_direct": - if self.hicache_io_backend != "direct": + if self.hicache_io_backend not in ["direct", "kernel_ascend"]: self.hicache_io_backend = "direct" logger.warning( "Page first direct layout only support direct io backend" @@ -1644,20 +1641,6 @@ class ServerArgs: "Setting hicache_io_backend to vanilla I/O, which may lead to suboptimal performance with small page sizes." ) - # Below are the only parameters currently supported on Ascend - if self.enable_hierarchical_cache and is_npu(): - # FIXME(iforgetmyname) fix decode_attention_backend on ascend - self.decode_attention_backend = "ascend" - self.hicache_io_backend = "kernel_ascend" - if self.use_mla_backend(): - self.hicache_mem_layout = "page_first_kv_split" - else: - self.hicache_mem_layout = "page_first_direct" - logger.warning( - f"Ascend NPU Platform detected, change `hicache_io_backend` to `kernel_ascend` and " - f"`hicache_mem_layout` to `{self.hicache_mem_layout}`" - ) - def _handle_speculative_decoding(self): if ( self.speculative_draft_model_path is not None diff --git a/python/sglang/srt/speculative/draft_utils.py b/python/sglang/srt/speculative/draft_utils.py index 394dca74a..d3246d30b 100644 --- a/python/sglang/srt/speculative/draft_utils.py +++ b/python/sglang/srt/speculative/draft_utils.py @@ -176,7 +176,7 @@ class DraftBackendFactory: ) def _create_ascend_decode_backend(self): - from sglang.srt.layers.attention.ascend_backend import ( + from sglang.srt.hardware_backend.npu.attention.ascend_backend import ( AscendAttnMultiStepDraftBackend, ) @@ -231,7 +231,9 @@ class DraftBackendFactory: return TRTLLMMLABackend(self.draft_model_runner, skip_prefill=False) def _create_ascend_prefill_backend(self): - from sglang.srt.layers.attention.ascend_backend import AscendAttnBackend + from sglang.srt.hardware_backend.npu.attention.ascend_backend import ( + AscendAttnBackend, + ) return AscendAttnBackend(self.draft_model_runner) diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index 35f30a016..51e7bf060 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -38,9 +38,7 @@ from sglang.srt.speculative.spec_utils import ( get_src_tgt_cache_loc, get_target_cache_loc, ) -from sglang.srt.utils import is_cuda, is_npu, next_power_of_2 - -_is_npu = is_npu() +from sglang.srt.utils import is_cuda, next_power_of_2 if is_cuda(): from sgl_kernel import ( @@ -77,22 +75,18 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin): @classmethod def create_idle_input(cls, topk: int, spec_steps: int, num_verify_tokens: int): - if not _is_npu: - device = "cuda" - else: - device = "npu" return cls( - draft_token=torch.empty((0,), dtype=torch.long, device=device), - custom_mask=torch.full((0,), True, dtype=torch.bool, device=device), - positions=torch.empty((0,), dtype=torch.int64, device=device), + draft_token=torch.empty((0,), dtype=torch.long, device="cuda"), + custom_mask=torch.full((0,), True, dtype=torch.bool, device="cuda"), + positions=torch.empty((0,), dtype=torch.int64, device="cuda"), retrive_index=torch.full( - (0, num_verify_tokens), -1, dtype=torch.long, device=device + (0, num_verify_tokens), -1, dtype=torch.long, device="cuda" ), retrive_next_token=torch.full( - (0, num_verify_tokens), -1, dtype=torch.long, device=device + (0, num_verify_tokens), -1, dtype=torch.long, device="cuda" ), retrive_next_sibling=torch.full( - (0, num_verify_tokens), -1, dtype=torch.long, device=device + (0, num_verify_tokens), -1, dtype=torch.long, device="cuda" ), retrive_cum_len=None, topk=topk, @@ -282,7 +276,7 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin): "Falling back to greedy verification." ) - if is_all_greedy or not TREE_SPEC_KERNEL_AVAILABLE or _is_npu: + if is_all_greedy or not TREE_SPEC_KERNEL_AVAILABLE: target_predict = torch.argmax(logits_output.next_token_logits, dim=-1) target_predict = target_predict.reshape(bs, self.draft_token_num) predict, accept_index, accept_length = verify_tree_greedy_func( diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index fd60d358a..41879abab 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/python/sglang/srt/speculative/eagle_worker.py @@ -5,6 +5,9 @@ from typing import List, Optional, Tuple import torch from sglang.srt.distributed import get_tp_group +from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_npu_graph_runner import ( + EAGLEDraftNpuGraphRunner, +) from sglang.srt.layers.dp_attention import get_attention_tp_group from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.moe.utils import speculative_moe_backend_context @@ -31,7 +34,6 @@ from sglang.srt.speculative.eagle_draft_cuda_graph_runner import ( from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import ( EAGLEDraftExtendCudaGraphRunner, ) -from sglang.srt.speculative.eagle_draft_npu_graph_runner import EAGLEDraftNpuGraphRunner from sglang.srt.speculative.eagle_info import ( EagleDraftInput, EagleVerifyInput, diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 663c4f341..579d8f57c 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -6,6 +6,12 @@ from typing import List, Optional, Tuple import torch from sglang.srt.environ import envs +from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_extend_npu_graph_runner import ( + EAGLEDraftExtendNpuGraphRunner, +) +from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_npu_graph_runner import ( + EAGLEDraftNpuGraphRunner, +) from sglang.srt.layers.moe.utils import speculative_moe_backend_context from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput from sglang.srt.managers.schedule_batch import ModelWorkerBatch @@ -21,10 +27,6 @@ from sglang.srt.speculative.eagle_draft_cuda_graph_runner import ( from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import ( EAGLEDraftExtendCudaGraphRunner, ) -from sglang.srt.speculative.eagle_draft_extend_npu_graph_runner import ( - EAGLEDraftExtendNpuGraphRunner, -) -from sglang.srt.speculative.eagle_draft_npu_graph_runner import EAGLEDraftNpuGraphRunner from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput from sglang.srt.speculative.eagle_info_v2 import ( assign_extend_cache_locs, diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 189136568..789e918f7 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -645,60 +645,6 @@ def make_layers_non_pp( return layers -cmo_stream = None - - -def get_cmo_stream(): - """ - Cache Management Operation(CMO). - Launch a new stream to prefetch the weight of matmul when running other - AIV or communication kernels, aiming to overlap the memory access time. - """ - global cmo_stream - return cmo_stream - - -def set_cmo_stream(stream): - global cmo_stream - cmo_stream = stream - - -def prepare_weight_cache(handle, cache, PREFETCH_MAX_SIZE=1000000000): - """ - PREFETCH_MAX_SIZE: maximum size (bytes) for each prefetch operation. - This affects the time spent in prefetch: - time ≈ PREFETCH_MAX_SIZE / system_bandwidth - """ - import torch_npu - - stream = get_cmo_stream() - if stream is None: - stream = torch.get_device_module().Stream() - set_cmo_stream(stream) - stream.wait_stream(torch.get_device_module().current_stream()) - with torch.get_device_module().stream(stream): - if isinstance(cache, list): - for weight in cache: - torch_npu.npu_prefetch( - weight, - handle, - PREFETCH_MAX_SIZE, - ) - else: - torch_npu.npu_prefetch( - cache, - handle, - PREFETCH_MAX_SIZE, - ) - - -def wait_cmo_stream(): - stream = get_cmo_stream() - if stream is not None: - cur_stream = torch.get_device_module().current_stream() - cur_stream.wait_stream(stream) - - @lru_cache(maxsize=1) def get_device_module(): return torch.get_device_module() diff --git a/python/sglang/srt/utils/numa_utils.py b/python/sglang/srt/utils/numa_utils.py index a144d1002..1dac8c5fd 100644 --- a/python/sglang/srt/utils/numa_utils.py +++ b/python/sglang/srt/utils/numa_utils.py @@ -6,8 +6,8 @@ import time from contextlib import contextmanager from pathlib import Path -from sglang import ServerArgs from sglang.srt.environ import envs +from sglang.srt.server_args import ServerArgs logger = logging.getLogger(__name__) diff --git a/test/srt/ascend/test_ascend_deepep.py b/test/srt/ascend/test_ascend_deepep.py index f8cf760d5..c5f2cb6fa 100644 --- a/test/srt/ascend/test_ascend_deepep.py +++ b/test/srt/ascend/test_ascend_deepep.py @@ -35,7 +35,7 @@ class TestAscendDeepEP(CustomTestCase): "--attention-backend", "ascend", "--quantization", - "w8a8_int8", + "modelslim", "--mem-fraction-static", 0.8, "--disable-radix-cache", diff --git a/test/srt/ascend/test_ascend_deepseek_mtp.py b/test/srt/ascend/test_ascend_deepseek_mtp.py index d8f33dcd2..43089f885 100644 --- a/test/srt/ascend/test_ascend_deepseek_mtp.py +++ b/test/srt/ascend/test_ascend_deepseek_mtp.py @@ -33,7 +33,7 @@ class TestAscendDeepSeekMTP(CustomTestCase): "--attention-backend", "ascend", "--quantization", - "w8a8_int8", + "modelslim", "--mem-fraction-static", 0.8, "--disable-radix-cache", diff --git a/test/srt/ascend/test_ascend_mla_fia_w8a8int8.py b/test/srt/ascend/test_ascend_mla_fia_w8a8int8.py index 6de97b04d..1a0eb7f6d 100644 --- a/test/srt/ascend/test_ascend_mla_fia_w8a8int8.py +++ b/test/srt/ascend/test_ascend_mla_fia_w8a8int8.py @@ -38,7 +38,7 @@ class TestAscendMlaW8A8Int8(CustomTestCase): "--attention-backend", "ascend", "--quantization", - "w8a8_int8", + "modelslim", "--tp-size", 2, "--disable-radix-cache", diff --git a/test/srt/ascend/test_ascend_mla_w8a8int8.py b/test/srt/ascend/test_ascend_mla_w8a8int8.py index 70f7edab4..eddae3086 100644 --- a/test/srt/ascend/test_ascend_mla_w8a8int8.py +++ b/test/srt/ascend/test_ascend_mla_w8a8int8.py @@ -37,7 +37,7 @@ class TestAscendMlaW8A8Int8(CustomTestCase): "--attention-backend", "ascend", "--quantization", - "w8a8_int8", + "modelslim", "--tp-size", 4, "--disable-radix-cache",