[NPU][1/N] NPU basic functions refactor and new modelslim quant type (#13359)

.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

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)

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
+from sglang.srt.utils import Withable, get_int_env_var
-
-_is_npu = is_npu()
 
 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):

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,

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

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()
+if TYPE_CHECKING:
-_ENABLE_MLA_PREPROCESS_FLAG = get_bool_env_var("SGLANG_NPU_USE_MLAPO")
+    from sglang.srt.layers.quantization.base_config import QuantizationConfig
-_NPU_FORMAT_NZ = 29
 
 
+@lru_cache(maxsize=1)
 def is_mla_preprocess_enabled() -> bool:
-    return _is_npu and _ENABLE_MLA_PREPROCESS_FLAG
+    return get_bool_env_var("SGLANG_NPU_USE_MLAPO")
-
-
-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)
 
 
 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(
+        self.qkv_a_proj_weight_nz = npu_format_cast(fused_qkv_a_proj_with_mqa_weight_nz)
-            fused_qkv_a_proj_with_mqa_weight_nz, _NPU_FORMAT_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(
+        self.q_b_proj_weight_nz = npu_format_cast(q_b_proj_weight_nz)
-            q_b_proj_weight_nz, _NPU_FORMAT_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")
+            hasattr(self.qkv_a_proj.quant_method, "quant_config")
-            and self.qkv_a_proj.quant_method.quantization_config.get_name()
+            and self.qkv_a_proj.quant_method.quant_config.get_name() == "modelslim"
-            == "w8a8_int8"
         )
         if _is_w8a8:
             return self.forward_mlapo(

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)

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),
+        )

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

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)

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

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()

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

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)

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)
 

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

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,

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
+        from sglang.srt.hardware_backend.npu.moe.topk import fused_topk_npu
-        torch_native = self.topk_config.torch_native
-        renormalize = self.topk_config.renormalize
 
-        if not use_grouped_topk and not torch_native:
+        return fused_topk_npu(
-            topk_weights, topk_ids, _ = torch_npu.npu_moe_gating_top_k_softmax(
+            hidden_states=hidden_states,
-                router_logits,
+            router_logits=router_logits,
-                k=self.topk_config.top_k,
+            topk_config=self.topk_config,
-            )
+            num_token_non_padded=num_token_non_padded,
-
+            expert_location_dispatch_info=expert_location_dispatch_info,
-            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,
-            )
 
     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,

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}
 
 

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:

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__)
 
 

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(

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"
-        device: str = "cuda"
-    else:
-        device: str = "npu"
 
     # Speculative decoding
     spec_algorithm: SpeculativeAlgorithm = None

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
+                from sglang.srt.hardware_backend.npu.memory_pool_npu import (
-            else:
+                    NPUMHATokenToKVPool,
-                TokenToKVPoolClass = MHATokenToKVPool
+                )
+
+                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,

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 = (

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,

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
+    from sglang.srt.hardware_backend.npu.modules.deepseek_v2_attention_mla_npu import (
-    import sgl_kernel_npu  # noqa: F401
+        forward_dsa_core_npu,
-    import torch_npu  # noqa: F401
+        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(
-                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(
                 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.empty_cache()
+        torch.cuda.synchronize()
-            torch.cuda.synchronize()
-        else:
-            torch.npu.empty_cache()
-            torch.npu.synchronize()
 
     @classmethod
     def get_model_config_for_expert_location(cls, config):

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 (
+from sglang.srt.utils import add_prefix, is_cuda, is_npu
-    add_prefix,
-    get_cmo_stream,
-    is_cuda,
-    is_npu,
-    wait_cmo_stream,
-)
 
 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__(

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

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)
 

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
+from sglang.srt.utils import is_cuda, next_power_of_2
-
-_is_npu = is_npu()
 
 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),
+            draft_token=torch.empty((0,), dtype=torch.long, device="cuda"),
-            custom_mask=torch.full((0,), True, dtype=torch.bool, device=device),
+            custom_mask=torch.full((0,), True, dtype=torch.bool, device="cuda"),
-            positions=torch.empty((0,), dtype=torch.int64, device=device),
+            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(

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,

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,

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()

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__)
 

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",

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",

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",

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",