chore: import upstream snapshot with attribution

python/sglang/srt/hardware_backend/npu/allocator_npu.py

@@ -0,0 +1,155 @@
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.mem_cache.allocator import (
+    PagedTokenToKVPoolAllocator,
+    alloc_extend_naive,
+)
+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
+
+
+class NPUPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
+    def __init__(
+        self,
+        size: int,
+        page_size: int,
+        dtype: torch.dtype,
+        device: str,
+        kvcache: "KVCache",
+        need_sort: bool,
+    ):
+        super().__init__(size, page_size, dtype, device, kvcache, need_sort)
+        self.roundup = page_size - 1
+
+    def alloc_extend(
+        self,
+        prefix_lens: torch.Tensor,
+        prefix_lens_cpu: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        last_loc: torch.Tensor,
+        extend_num_tokens: int,
+        num_new_pages: int = None,
+    ):
+        if self.debug_mode:
+            assert torch.all(
+                (last_loc + 1) % self.page_size == prefix_lens % self.page_size
+            )
+
+        if num_new_pages is None:
+            num_new_pages_tensor = (
+                (seq_lens + self.roundup) // self.page_size
+                - (prefix_lens + self.roundup) // self.page_size
+            ).sum()
+            num_new_pages_item = num_new_pages_tensor.item()
+        else:
+            num_new_pages_item = num_new_pages
+        if self.need_sort and num_new_pages_item > len(self.free_pages):
+            self.merge_and_sort_free()
+
+        if num_new_pages_item > len(self.free_pages):
+            return None
+
+        if num_new_pages_item < 200:
+            from sgl_kernel_npu.mem_cache.allocator import alloc_extend_kernel
+
+            out_indices = torch.empty(
+                (extend_num_tokens,),
+                dtype=torch.int64,
+                device=self.device,
+            )
+            max_num_extend_tokens = next_power_of_2(extend_num_tokens)
+            bs = prefix_lens.shape[0]
+            alloc_extend_kernel[(bs,)](
+                prefix_lens,
+                seq_lens,
+                last_loc,
+                self.free_pages,
+                out_indices,
+                next_power_of_2(bs),
+                self.page_size,
+                max_num_extend_tokens,
+            )
+
+        else:
+            out_indices = torch.empty(
+                (extend_num_tokens,),
+                dtype=torch.int32,
+                device=self.device,
+            )
+            alloc_extend_naive(
+                prefix_lens,
+                seq_lens,
+                last_loc,
+                self.free_pages,
+                out_indices,
+                self.page_size,
+                self.device,
+            )
+
+        if self.debug_mode:
+            assert len(torch.unique(out_indices)) == len(out_indices)
+
+        self.free_pages = self.free_pages[num_new_pages_item:]
+        return out_indices.int()
+
+    def alloc_decode(
+        self,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        last_loc: torch.Tensor,
+    ):
+        if self.debug_mode:
+            assert torch.all(
+                (last_loc + 2) % self.page_size == seq_lens % self.page_size
+            )
+
+        num_new_pages = get_num_new_pages(
+            seq_lens=seq_lens_cpu,
+            page_size=self.page_size,
+            decode=True,
+        )
+
+        if num_new_pages > len(self.free_pages):
+            self.merge_and_sort_free()
+
+        if num_new_pages > len(self.free_pages):
+            return None
+
+        need_new_pages = (seq_lens % self.page_size == 1).int()
+        end_new_pages = torch.cumsum(need_new_pages, 0)
+        start_new_pages = end_new_pages - need_new_pages
+        if num_new_pages == 0:
+            out_indices = last_loc + 1
+        else:
+            out_indices = (last_loc + 1) * (1 - need_new_pages) + self.free_pages[
+                start_new_pages
+            ] * self.page_size * need_new_pages
+
+        if self.debug_mode:
+            assert len(torch.unique(out_indices)) == len(out_indices)
+
+        self.free_pages = self.free_pages[num_new_pages:]
+        return out_indices.int()
+
+    def free(self, free_index: torch.Tensor):
+        if free_index.numel() == 0:
+            return
+
+        if self.is_not_in_free_group:
+            device = free_index.device
+            free_page_indices = torch.unique(free_index.cpu() // self.page_size)
+            free_page_indices = free_page_indices.to(device)
+            if self.need_sort:
+                self.release_pages = torch.cat((free_page_indices, self.release_pages))
+            else:
+                self.free_pages = torch.cat((free_page_indices, self.free_pages))
+        else:
+            self.free_group.append(free_index)
+
+        if self.debug_mode:
+            assert len(torch.unique(self.free_pages)) == len(self.free_pages)

python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py

@@ -0,0 +1,391 @@
+from typing import Optional, Tuple, Union
+
+import torch
+from sgl_kernel_npu.fla.fused_gdn_gating import (
+    fused_gdn_gating_kernel_without_sigmoid,
+    fused_gdn_gating_npu,
+)
+from sgl_kernel_npu.mamba.causal_conv1d import (
+    causal_conv1d_fn_npu,
+    causal_conv1d_update_npu,
+    causal_conv1d_update_v2,
+)
+
+from sglang.srt.hardware_backend.npu.attention.ascend_hybrid_linear_attn_backend import (
+    AscendMambaAttnBackendBase,
+)
+from sglang.srt.layers.attention.linear.gdn_backend import GDNKernelDispatcher
+from sglang.srt.layers.attention.linear.utils import (
+    get_linear_attn_decode_backend,
+    get_linear_attn_prefill_backend,
+)
+from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
+from sglang.srt.mem_cache.memory_pool import MambaPool
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.model_executor.model_runner import ModelRunner
+from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
+
+fused_gdn_gating = fused_gdn_gating_npu
+causal_conv1d_fn = causal_conv1d_fn_npu
+causal_conv1d_update = causal_conv1d_update_npu
+
+
+class AscendGDNAttnBackend(AscendMambaAttnBackendBase):
+
+    def __init__(self, model_runner: ModelRunner):
+        super().__init__(model_runner)
+        self.conv_states_shape = torch.Size(
+            (
+                *model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape[
+                    :-2
+                ],
+                model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape[-1],
+                model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape[-2],
+            )
+        )
+        decode_backend = get_linear_attn_decode_backend()
+        prefill_backend = get_linear_attn_prefill_backend()
+        self.kernel_dispatcher = GDNKernelDispatcher(decode_backend, prefill_backend)
+
+    def _prepare_mamba_track_metadata(self, forward_batch: ForwardBatch):
+        if self.forward_metadata.has_mamba_track_mask:
+            mamba_track_mask_indices = forward_batch.mamba_track_mask.nonzero(
+                as_tuple=True
+            )[0]
+            self.forward_metadata.conv_states_mask_indices = (
+                forward_batch.mamba_track_indices[mamba_track_mask_indices]
+            )
+
+    def prepare_gdn_inputs(
+        self,
+        bs: int,
+        forward_mode: ForwardMode,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        cache_indices = self.forward_metadata.mamba_cache_indices
+        self.num_accept_tokens = torch.ones(
+            [bs], dtype=torch.int32, device=cache_indices.device
+        )
+        self.actual_seq_lengths = torch.ones(
+            [bs], dtype=torch.int32, device=cache_indices.device
+        )
+        if forward_mode.is_target_verify():
+            seq_len = spec_info.draft_token_num
+            self.actual_seq_lengths = self.actual_seq_lengths * seq_len
+            # indices
+            self.ssm_state_indices = torch.arange(
+                cache_indices.shape[0] * seq_len,
+                dtype=torch.int32,
+                device=cache_indices.device,
+            )
+        else:
+            self.ssm_state_indices = cache_indices
+
+    def init_forward_metadata_out_graph(
+        self,
+        forward_batch: ForwardBatch,
+        in_capture: bool = False,
+    ):
+        if forward_batch.forward_mode.is_draft_extend_v2():
+            return
+        super().init_forward_metadata_out_graph(forward_batch, in_capture=in_capture)
+        self.prepare_gdn_inputs(
+            forward_batch.batch_size,
+            forward_batch.forward_mode,
+            forward_batch.spec_info,
+        )
+        self._prepare_mamba_track_metadata(forward_batch)
+        self.graph_mode = True
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        if forward_batch.forward_mode.is_draft_extend_v2():
+            return
+        super().init_forward_metadata(forward_batch)
+        self.prepare_gdn_inputs(
+            forward_batch.batch_size,
+            forward_batch.forward_mode,
+            forward_batch.spec_info,
+        )
+        self._prepare_mamba_track_metadata(forward_batch)
+        self.graph_mode = False
+
+    def forward_decode(
+        self,
+        layer: RadixLinearAttention,
+        forward_batch: ForwardBatch,
+        mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
+        a: torch.Tensor,
+        b: torch.Tensor,
+        **kwargs,
+    ):
+        layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
+        conv_states = layer_cache.conv[0]
+        ssm_states = layer_cache.temporal
+        query_start_loc = self.forward_metadata.query_start_loc
+        cache_indices = self.forward_metadata.mamba_cache_indices
+
+        assert isinstance(mixed_qkv, torch.Tensor)
+        conv_states_tmp = conv_states.transpose(1, 2).clone()
+        mixed_qkv = causal_conv1d_update(
+            mixed_qkv,
+            conv_states_tmp,
+            layer.conv_weights,
+            layer.bias,
+            layer.activation,
+            conv_state_indices=cache_indices,
+        )
+        conv_states[:] = conv_states_tmp.transpose(1, 2)
+
+        query, key, value = torch.split(
+            mixed_qkv,
+            [layer.q_dim, layer.k_dim, layer.v_dim],
+            dim=-1,
+        )
+        bs = forward_batch.batch_size
+        query = query.view(1, bs, layer.num_q_heads, layer.head_q_dim)
+        key = key.view(1, bs, layer.num_k_heads, layer.head_k_dim)
+        value = value.view(1, bs, layer.num_v_heads, layer.head_v_dim)
+
+        core_attn_out = self.kernel_dispatcher.decode(
+            q=query,
+            k=key,
+            v=value,
+            a=a,
+            b=b,
+            A_log=layer.A_log,
+            dt_bias=layer.dt_bias,
+            ssm_states=ssm_states,
+            cache_indices=cache_indices,
+            query_start_loc=query_start_loc,
+        )
+
+        self._track_mamba_state_decode(
+            forward_batch, conv_states, ssm_states, cache_indices
+        )
+        return core_attn_out
+
+    def forward_extend(
+        self,
+        layer: RadixLinearAttention,
+        forward_batch: ForwardBatch,
+        mixed_qkv: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
+        a: torch.Tensor,
+        b: torch.Tensor,
+        **kwargs,
+    ):
+        assert isinstance(mixed_qkv, torch.Tensor)
+        seq_len = mixed_qkv.shape[0]
+        is_target_verify = forward_batch.forward_mode.is_target_verify()
+        forward_metadata = self.forward_metadata
+
+        query_start_loc = forward_metadata.query_start_loc
+        cache_indices = forward_metadata.mamba_cache_indices
+        retrieve_next_token = forward_metadata.retrieve_next_token
+        retrieve_next_sibling = forward_metadata.retrieve_next_sibling
+        retrieve_parent_token = forward_metadata.retrieve_parent_token
+
+        mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
+        conv_states = mamba_cache_params.conv[0]
+        ssm_states = mamba_cache_params.temporal
+        if is_target_verify:
+            assert isinstance(mamba_cache_params, MambaPool.SpeculativeState)
+            intermediate_state_cache = mamba_cache_params.intermediate_ssm
+            intermediate_conv_window_cache = (
+                mamba_cache_params.intermediate_conv_window[0]
+            )
+            has_initial_states = torch.ones(
+                seq_len // forward_batch.spec_info.draft_token_num,
+                dtype=torch.bool,
+                device=forward_batch.input_ids.device,
+            )
+        else:
+            has_initial_states = forward_batch.extend_prefix_lens > 0
+        if is_target_verify:
+            num_token_padding = mixed_qkv.shape[0]
+            if (
+                not self.graph_mode
+                and forward_batch.num_token_non_padded_cpu != num_token_padding
+            ):
+                mixed_qkv = mixed_qkv[: forward_batch.num_token_non_padded_cpu]
+                a = a[: forward_batch.num_token_non_padded_cpu]
+                b = b[: forward_batch.num_token_non_padded_cpu]
+                seq_len = forward_batch.num_token_non_padded_cpu
+
+            batch_size = cache_indices.shape[0]
+            draft_token_num = forward_batch.spec_info.draft_token_num
+            num_accepted_tokens = torch.full(
+                (batch_size,),
+                draft_token_num,
+                dtype=torch.int32,
+                device=mixed_qkv.device,
+            )
+            mixed_qkv = causal_conv1d_update_v2(
+                x=mixed_qkv.view(batch_size, draft_token_num, -1).contiguous(),
+                conv_state=conv_states.contiguous(),
+                weight=layer.conv_weights.transpose(0, 1).contiguous(),
+                bias=layer.bias,
+                activation=layer.activation,
+                conv_state_indices=cache_indices,
+                num_accepted_tokens=num_accepted_tokens,
+                pad_slot_id=-1,
+                validate_data=False,
+            ).view(seq_len, -1)
+        else:
+            mixed_qkv = mixed_qkv.transpose(0, 1)
+            if forward_metadata.has_mamba_track_mask:
+                mixed_qkv_to_track = mixed_qkv[
+                    :, forward_metadata.track_conv_indices
+                ].transpose(0, 1)
+                conv_states.transpose(1, 2)[
+                    forward_metadata.conv_states_mask_indices
+                ] = mixed_qkv_to_track
+            kernel_size = layer.conv_weights.shape[-1]
+            conv_states_for_prefill = conv_states[:, -(kernel_size - 1) :, :]
+            conv_states_tmp = conv_states_for_prefill.transpose(1, 2).contiguous()
+
+            mixed_qkv = causal_conv1d_fn(
+                mixed_qkv,
+                layer.conv_weights,
+                layer.bias,
+                activation=layer.activation,
+                conv_states=conv_states_tmp,
+                has_initial_state=has_initial_states,
+                cache_indices=cache_indices,
+                query_start_loc=query_start_loc,
+                seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
+            ).transpose(0, 1)[:seq_len]
+            conv_states[:, -(kernel_size - 1) :, :] = conv_states_tmp.transpose(
+                1, 2
+            ).contiguous()
+        if is_target_verify:
+            g, beta = fused_gdn_gating_kernel_without_sigmoid(
+                layer.A_log, a, b, layer.dt_bias
+            )
+            beta = beta.unsqueeze(0)
+            num_heads, head_k_dim = layer.num_q_heads, layer.head_q_dim
+            num_value_heads, head_v_dim = layer.num_v_heads, layer.head_v_dim
+
+            mixed_qkv_last_dim = mixed_qkv.shape[-1]
+
+            mixed_qkv = mixed_qkv.view(batch_size, -1, mixed_qkv_last_dim)
+            beta = beta.view(batch_size, -1, num_value_heads)
+            g = g.view(batch_size, -1, num_value_heads)
+
+            core_attn_out = self.fused_recurrent_gated_delta_rule_update(
+                mixed_qkv,
+                num_heads,
+                num_value_heads,
+                head_k_dim,
+                head_v_dim,
+                recurrent_state=ssm_states,
+                beta=beta,
+                g=g,
+                cache_indices=cache_indices,
+                intermediate_state=intermediate_state_cache,
+            )
+            core_attn_out = core_attn_out.view(-1, num_value_heads, head_v_dim)
+            if (not self.graph_mode) and core_attn_out.shape[0] < num_token_padding:
+                core_attn_out = torch.cat(
+                    [
+                        core_attn_out,
+                        core_attn_out.new_zeros(
+                            num_token_padding - core_attn_out.shape[0],
+                            *core_attn_out.shape[1:],
+                        ),
+                    ],
+                    dim=0,
+                )
+        else:
+            query, key, value = torch.split(
+                mixed_qkv,
+                [layer.q_dim, layer.k_dim, layer.v_dim],
+                dim=-1,
+            )
+
+            actual_seq_len = query.shape[0]
+            query = query.view(1, actual_seq_len, layer.num_q_heads, layer.head_q_dim)
+            key = key.view(1, actual_seq_len, layer.num_k_heads, layer.head_k_dim)
+            value = value.view(1, actual_seq_len, layer.num_v_heads, layer.head_v_dim)
+
+            g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias)
+            core_attn_out, last_recurrent_state, h = self.kernel_dispatcher.extend(
+                q=query,
+                k=key,
+                v=value,
+                g=g,
+                beta=beta,
+                ssm_states=ssm_states,
+                cache_indices=cache_indices,
+                query_start_loc=query_start_loc,
+            )
+            if last_recurrent_state is not None:
+                last_recurrent_state = last_recurrent_state.to(
+                    ssm_states.dtype, copy=False
+                )
+                ssm_states[cache_indices] = last_recurrent_state
+            last_recurrent_state = last_recurrent_state.to(ssm_states.dtype, copy=False)
+            ssm_states[cache_indices] = last_recurrent_state
+            if h is not None:
+                self._track_mamba_state_extend(
+                    forward_batch, h, ssm_states, forward_metadata
+                )
+
+        return core_attn_out
+
+    def fused_recurrent_gated_delta_rule_update(
+        self,
+        mix_qkv: torch.Tensor,
+        num_heads,
+        num_value_heads,
+        head_k_dim,
+        head_v_dim,
+        recurrent_state: torch.Tensor,
+        beta: torch.Tensor,
+        g: torch.Tensor,
+        cache_indices: torch.Tensor,
+        intermediate_state: Optional[torch.Tensor] = None,
+    ):
+        beta = beta.to(torch.bfloat16)
+        g = g.to(torch.float32)
+        batch_size = mix_qkv.shape[0]
+        seq_len = mix_qkv.shape[1]
+        scale = 1 / (head_k_dim**0.5)
+
+        if intermediate_state is not None:
+            intermediate_state = intermediate_state.view(
+                -1, num_value_heads, head_k_dim, head_v_dim
+            )
+
+        if self.graph_mode:
+            num_accept_tokens = torch.full(
+                [batch_size], 1, dtype=torch.int32, device=cache_indices.device
+            )
+            actual_seq_lengths = torch.full(
+                [batch_size], seq_len, dtype=torch.int32, device=cache_indices.device
+            )
+            ssm_state_indices = self.forward_metadata.mamba_cache_indices_gdn
+        else:
+            num_accept_tokens = self.num_accept_tokens
+            actual_seq_lengths = self.actual_seq_lengths
+            ssm_state_indices = self.ssm_state_indices
+
+        attn_core_out = torch.ops.npu.recurrent_gated_delta_rule(
+            mix_qkv,
+            recurrent_state,
+            beta=beta,
+            scale=scale,
+            actual_seq_lengths=actual_seq_lengths,
+            ssm_state_indices=ssm_state_indices.view(batch_size, seq_len),
+            nk=num_heads,
+            nv=num_value_heads,
+            intermediate_state=intermediate_state,
+            cache_indices=cache_indices,
+            num_accepted_tokens=num_accept_tokens,
+            g=g,
+        )
+
+        if intermediate_state is not None:
+            intermediate_state = intermediate_state.view(
+                -1, seq_len, num_value_heads, head_k_dim, head_v_dim
+            )
+        return attn_core_out

python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py

@@ -0,0 +1,320 @@
+import logging
+from typing import Optional, Union
+
+import torch
+from sgl_kernel_npu.mamba.mamba_state_update_triton import (
+    conv_state_rollback,
+    move_intermediate_cache,
+)
+
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
+    HybridLinearAttnBackend,
+    MambaAttnBackendBase,
+)
+from sglang.srt.layers.attention.mamba.mamba2_metadata import (
+    ForwardMetadata,
+)
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.model_executor.model_runner import ModelRunner
+from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
+from sglang.srt.speculative.spec_info import SpecInput
+
+logger = logging.getLogger(__name__)
+
+
+class AscendMambaAttnBackendBase(MambaAttnBackendBase):
+    def __init__(self, model_runner: ModelRunner):
+        super().__init__(model_runner)
+        self.state_indices_list_gdn = []
+
+    def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
+        assert (
+            max_num_tokens % max_bs == 0
+        ), f"max_num_tokens={max_num_tokens} must be divisible by max_bs={max_bs}"
+        draft_token_num = max_num_tokens // max_bs
+        for i in range(max_bs):
+            self.state_indices_list.append(
+                torch.full(
+                    (i + 1,), self.pad_slot_id, dtype=torch.int32, device=self.device
+                )
+            )
+            self.state_indices_list_gdn.append(
+                torch.full(
+                    ((i + 1) * draft_token_num,),
+                    self.pad_slot_id,
+                    dtype=torch.int32,
+                    device=self.device,
+                )
+            )
+            self.query_start_loc_list.append(
+                torch.zeros((i + 2,), dtype=torch.int32, device=self.device)
+            )
+            self.retrieve_next_token_list.append(
+                torch.zeros(
+                    (i + 1, draft_token_num), dtype=torch.int32, device=self.device
+                )
+            )
+            self.retrieve_next_sibling_list.append(
+                torch.zeros(
+                    (i + 1, draft_token_num), dtype=torch.int32, device=self.device
+                )
+            )
+            self.retrieve_parent_token_list.append(
+                torch.zeros(
+                    (i + 1, draft_token_num), dtype=torch.int32, device=self.device
+                )
+            )
+        self.cached_cuda_graph_decode_query_start_loc = torch.arange(
+            0, max_bs + 1, dtype=torch.int32, device=self.device
+        )
+        self.cached_cuda_graph_verify_query_start_loc = torch.arange(
+            0,
+            max_bs * draft_token_num + 1,
+            step=draft_token_num,
+            dtype=torch.int32,
+            device=self.device,
+        )
+
+    def _capture_metadata(
+        self,
+        bs: int,
+        req_pool_indices: torch.Tensor,
+        forward_mode: ForwardMode,
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+    ):
+        mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices)
+        self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices)
+        if forward_mode.is_decode_or_idle():
+            self.query_start_loc_list[bs - 1].copy_(
+                self.cached_cuda_graph_decode_query_start_loc[: bs + 1]
+            )
+        elif forward_mode.is_target_verify():
+            self.query_start_loc_list[bs - 1].copy_(
+                self.cached_cuda_graph_verify_query_start_loc[: bs + 1]
+            )
+            ssm_state_indices = torch.arange(
+                mamba_indices.shape[0] * spec_info.draft_token_num,
+                dtype=torch.int32,
+                device=mamba_indices.device,
+            )
+            self.state_indices_list_gdn[bs - 1][
+                : len(mamba_indices) * spec_info.draft_token_num
+            ].copy_(ssm_state_indices)
+        else:
+            raise ValueError(f"Invalid forward mode: {forward_mode=}")
+
+        # If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask
+        if forward_mode.is_target_verify() and spec_info.topk > 1:
+            # They are None during cuda graph capture so skip the copy_...
+            # self.retrieve_next_token_list[bs - 1].copy_(spec_info.retrive_next_token)
+            # self.retrieve_next_sibling_list[bs - 1].copy_(spec_info.retrive_next_sibling)
+            return ForwardMetadata(
+                query_start_loc=self.query_start_loc_list[bs - 1],
+                mamba_cache_indices=self.state_indices_list[bs - 1],
+                retrieve_next_token=self.retrieve_next_token_list[bs - 1],
+                retrieve_next_sibling=self.retrieve_next_sibling_list[bs - 1],
+                retrieve_parent_token=self.retrieve_parent_token_list[bs - 1],
+            )
+        else:
+            return ForwardMetadata(
+                query_start_loc=self.query_start_loc_list[bs - 1],
+                mamba_cache_indices=self.state_indices_list[bs - 1],
+                mamba_cache_indices_gdn=self.state_indices_list_gdn[bs - 1],
+            )
+
+    def _replay_metadata(
+        self,
+        bs: int,
+        req_pool_indices: torch.Tensor,
+        forward_mode: ForwardMode,
+        spec_info: Optional[SpecInput],
+        seq_lens_cpu: Optional[torch.Tensor],
+        num_padding: Optional[int] = None,
+        in_capture: bool = False,
+        mamba_track_indices: Optional[torch.Tensor] = None,
+        *args,
+        **kwargs,
+    ):
+        # out_graph passes seq_lens_cpu=None at capture; mirror the base guard.
+        if seq_lens_cpu is None:
+            num_padding = 0
+        else:
+            num_padding = torch.count_nonzero(
+                seq_lens_cpu == self.get_cuda_graph_seq_len_fill_value()
+            )
+        # Make sure forward metadata is correctly handled for padding reqs
+        req_pool_indices[bs - num_padding :] = 0
+        mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices)
+        mamba_indices[bs - num_padding :] = 0
+        self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices)
+        track_buf = None
+        if mamba_track_indices is not None:
+            track_buf = mamba_track_indices
+        if forward_mode.is_decode_or_idle():
+            if num_padding == 0:
+                self.query_start_loc_list[bs - 1].copy_(
+                    self.cached_cuda_graph_decode_query_start_loc[: bs + 1]
+                )
+            else:
+                self.query_start_loc_list[bs - 1][: bs - num_padding].copy_(
+                    self.cached_cuda_graph_decode_query_start_loc[: bs - num_padding]
+                )
+                self.query_start_loc_list[bs - 1][bs - num_padding :].fill_(
+                    bs - num_padding
+                )
+        elif forward_mode.is_target_verify():
+            ssm_state_indices = torch.arange(
+                bs * spec_info.draft_token_num,
+                dtype=torch.int32,
+                device=mamba_indices.device,
+            )
+            self.state_indices_list_gdn[bs - 1].copy_(ssm_state_indices)
+            if num_padding == 0:
+                self.query_start_loc_list[bs - 1].copy_(
+                    self.cached_cuda_graph_verify_query_start_loc[: bs + 1]
+                )
+            else:
+                self.query_start_loc_list[bs - 1][: bs - num_padding].copy_(
+                    self.cached_cuda_graph_verify_query_start_loc[: bs - num_padding]
+                )
+                self.query_start_loc_list[bs - 1][bs - num_padding :].fill_(
+                    (bs - num_padding) * spec_info.draft_token_num
+                )
+        else:
+            raise ValueError(f"Invalid forward mode: {forward_mode=}")
+
+        # If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask
+        if forward_mode.is_target_verify() and spec_info.topk > 1:
+            bs_without_pad = spec_info.retrive_next_token.shape[0]
+            self.retrieve_next_token_list[bs - 1][:bs_without_pad].copy_(
+                spec_info.retrive_next_token
+            )
+            self.retrieve_next_sibling_list[bs - 1][:bs_without_pad].copy_(
+                spec_info.retrive_next_sibling
+            )
+            return ForwardMetadata(
+                query_start_loc=self.query_start_loc_list[bs - 1],
+                mamba_cache_indices=self.state_indices_list[bs - 1],
+                mamba_track_indices=track_buf,
+                retrieve_next_token=self.retrieve_next_token_list[bs - 1],
+                retrieve_next_sibling=self.retrieve_next_sibling_list[bs - 1],
+                retrieve_parent_token=self.retrieve_parent_token_list[bs - 1],
+            )
+        else:
+            return ForwardMetadata(
+                query_start_loc=self.query_start_loc_list[bs - 1],
+                mamba_cache_indices=self.state_indices_list[bs - 1],
+                mamba_cache_indices_gdn=self.state_indices_list_gdn[bs - 1],
+                mamba_track_indices=track_buf,
+            )
+
+    def get_cuda_graph_seq_len_fill_value(self):
+        return 0  # Mamba attn does not use seq lens to index kv cache
+
+
+class AscendMamba2AttnBackend(AscendMambaAttnBackendBase):
+    pass
+
+
+class AscendHybridLinearAttnBackend(HybridLinearAttnBackend):
+    def __init__(
+        self,
+        full_attn_backend: AttentionBackend,
+        linear_attn_backend: AscendMambaAttnBackendBase,
+        full_attn_layers: list[int],
+    ):
+        super().__init__(full_attn_backend, linear_attn_backend, full_attn_layers)
+
+    def update_mamba_state_after_mtp_verify(
+        self,
+        last_correct_step_indices: torch.Tensor,
+        mamba_track_indices: Optional[torch.Tensor],
+        mamba_steps_to_track: Optional[torch.Tensor],
+        model,
+    ):
+        """
+        Update mamba states after MTP verify using fully fused Triton kernel.
+
+        This replaces the original advanced indexing operations with a single fused
+        gather-scatter kernel that also handles masking internally, avoiding:
+        - index_elementwise_kernel from tensor[bool_mask]
+        - index_select kernel launches
+        - nonzero kernel launches
+        """
+        request_number = last_correct_step_indices.shape[0]
+
+        state_indices_tensor = (
+            self.linear_attn_backend.forward_metadata.mamba_cache_indices[
+                :request_number
+            ]
+        )
+
+        mamba_caches = (
+            self.linear_attn_backend.req_to_token_pool.get_speculative_mamba2_params_all_layers()
+        )
+
+        conv_states = mamba_caches.conv[0]
+        ssm_states = mamba_caches.temporal
+        intermediate_state_cache = mamba_caches.intermediate_ssm
+        dst_indices_tensor = state_indices_tensor.to(torch.int64)  # [N]
+        src_indices_tensor = torch.arange(
+            dst_indices_tensor.shape[0],
+            device=dst_indices_tensor.device,
+            dtype=torch.int64,
+        )
+        last_steps = last_correct_step_indices.to(torch.int64)  # [N]
+
+        move_intermediate_cache(
+            ssm_states,
+            intermediate_state_cache,
+            dst_indices_tensor,
+            src_indices_tensor,
+            last_steps,
+        )
+
+        draft_token_num = intermediate_state_cache.shape[2]
+        if mamba_track_indices is not None:
+            assert mamba_steps_to_track is not None
+            mamba_track_indices = mamba_track_indices.to(torch.int64)
+            mamba_steps_to_track = mamba_steps_to_track.to(torch.int64)
+
+            move_intermediate_cache(
+                ssm_states,
+                intermediate_state_cache,
+                mamba_track_indices,
+                src_indices_tensor,
+                mamba_steps_to_track,
+            )
+
+            track_mask = mamba_steps_to_track >= 0
+            # Track conv state from the verify-time window before rolling back
+            # the working slot; NPU does not keep per-step conv intermediates.
+            track_indices = mamba_track_indices[track_mask]
+            if track_indices.numel() > 0:
+                conv_states[:, track_indices] = conv_states[
+                    :, dst_indices_tensor[track_mask]
+                ]
+
+        if dst_indices_tensor.numel() > 0:
+            conv_state_rollback(
+                conv_states,
+                dst_indices_tensor,
+                last_steps,
+                draft_token_num,
+            )
+
+        if mamba_track_indices is not None and mamba_track_indices.numel() > 0:
+            conv_state_rollback(
+                conv_states,
+                mamba_track_indices,
+                mamba_steps_to_track,
+                draft_token_num,
+            )
+
+        return
+
+    def update_verify_buffers_to_fill_after_draft(
+        self, spec_info: SpecInput, cuda_graph_bs: Optional[int]
+    ):
+        pass

python/sglang/srt/hardware_backend/npu/attention/ascend_torch_native_backend.py

@@ -0,0 +1,336 @@
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+from torch.nn.functional import scaled_dot_product_attention
+
+
+class AscendTorchNativeAttnBackend:
+    def __init__(self):
+        pass
+
+    def scaled_dot_product_attention_with_softcapping(
+        self,
+        query,
+        key,
+        value,
+        attn_mask=None,
+        is_causal=False,
+        scale=None,
+        enable_gqa=False,
+        logit_cap=0.0,
+        logit_capping_method="tanh",
+    ) -> torch.Tensor:
+        L, S = query.size(-2), key.size(-2)
+        scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
+        attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device)
+        if is_causal:
+            assert attn_mask is None
+            temp_mask = torch.ones(L, S, dtype=torch.bool, device=query.device).tril(
+                diagonal=0
+            )
+            attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
+            attn_bias.to(query.dtype)
+
+        if attn_mask is not None:
+            if attn_mask.dtype == torch.bool:
+                attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
+            else:
+                attn_bias = attn_mask + attn_bias
+
+        if enable_gqa:
+            key = key.repeat_interleave(query.size(-3) // key.size(-3), -3)
+            value = value.repeat_interleave(query.size(-3) // value.size(-3), -3)
+
+        attn_weight = query @ key.transpose(-2, -1) * scale_factor
+
+        if logit_cap > 0:
+            if logit_capping_method == "tanh":
+                attn_weight = logit_cap * torch.tanh(attn_weight / logit_cap)
+
+        attn_weight += attn_bias
+        attn_weight = torch.softmax(attn_weight, dim=-1)
+        return attn_weight @ value
+
+    def run_sdpa_forward_extend(
+        self,
+        query: torch.Tensor,
+        output: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        req_to_token: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        extend_prefix_lens: torch.Tensor,
+        extend_seq_lens: torch.Tensor,
+        encoder_lens: torch.Tensor = None,
+        is_cross_attention: bool = False,
+        scaling=None,
+        enable_gqa=False,
+        causal=False,
+        sliding_window_size: int = -1,
+        full_to_swa_mapping: Optional[torch.Tensor] = None,
+        logit_cap: float = 0.0,
+        logit_capping_method: str = "tanh",
+    ):
+        """Run the extend forward by using torch native sdpa op.
+
+        Args:
+            query: [num_tokens, num_heads, head_size]
+            output: [num_tokens, num_heads, head_size]
+            k_cache: [max_total_num_tokens, num_heads, head_size]
+            v_cache: [max_total_num_tokens, num_heads, head_size]
+            req_to_token: [max_num_reqs, max_context_len]
+            req_pool_indices: [num_seqs]
+            seq_lens: [num_seqs]
+            extend_prefix_lens: [num_seqs]
+            extend_seq_lens: [num_seqs]
+            encoder_lens: [num_seqs]
+            is_cross_attention: [bool]
+            scaling: float or None
+            enable_gqa: bool
+            causal: bool
+            sliding_window_size: int, -1 means no sliding window
+            full_to_swa_mapping: mapping from full pool index to SWA pool index,
+                required for SWA layers to translate req_to_token indices
+
+        Returns:
+            output: [num_tokens, num_heads, head_size]
+        """
+
+        assert seq_lens.shape[0] == extend_prefix_lens.shape[0]
+        assert seq_lens.shape[0] == extend_seq_lens.shape[0]
+
+        # [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
+        query = query.movedim(0, query.dim() - 2)
+
+        start_q, start_kv = 0, 0
+        for seq_idx in range(seq_lens.shape[0]):
+            # Need optimize the performance later.
+
+            extend_seq_len_q = int(extend_seq_lens[seq_idx].item())
+            prefill_seq_len_q = int(extend_prefix_lens[seq_idx].item())
+
+            seq_len_kv = int(seq_lens[seq_idx].item())
+            end_q = start_q + extend_seq_len_q
+            end_kv = start_kv + seq_len_kv
+            atten_start_kv = 0
+            atten_end_kv = seq_len_kv
+            # support cross attention
+            if encoder_lens is not None:
+                if is_cross_attention:
+                    atten_end_kv = int(encoder_lens[seq_idx].item())
+                else:
+                    atten_start_kv = int(encoder_lens[seq_idx].item())
+                    atten_end_kv = atten_start_kv + extend_seq_len_q
+
+            is_swa_self_attn = (
+                sliding_window_size is not None
+                and sliding_window_size > -1
+                and encoder_lens is None
+            )
+            if is_swa_self_attn:
+                # For extend, the sliding window must be anchored at the first
+                # query token in this chunk rather than the final sequence
+                # length. Otherwise a large extend chunk can no longer fit in
+                # the cropped query suffix, which breaks the native fallback.
+                atten_start_kv = max(
+                    prefill_seq_len_q - sliding_window_size, atten_start_kv
+                )
+
+            per_req_query = query[:, start_q:end_q, :]
+
+            # SWA crops the front of the KV window, so the redundant query
+            # tensor must match the cropped window to keep Q.len == K.len for
+            # the causal mask. In cross-attention (and non-SWA self-attention)
+            # the original sizing — text seq len — must be preserved, since Q
+            # (text) and KV (encoder) lengths legitimately differ there.
+            if is_swa_self_attn:
+                redundant_len = atten_end_kv - atten_start_kv
+                query_start_idx = max(prefill_seq_len_q - atten_start_kv, 0)
+            else:
+                redundant_len = int(seq_lens[seq_idx].item())
+                query_start_idx = prefill_seq_len_q
+
+            per_req_query_redundant = torch.zeros(
+                (per_req_query.shape[0], redundant_len, per_req_query.shape[2]),
+                dtype=per_req_query.dtype,
+                device=per_req_query.device,
+            )
+
+            per_req_query_redundant[:, query_start_idx:, :] = per_req_query
+
+            # get key and value from cache. per_req_tokens contains the kv cache
+            # index for each token in the sequence.
+            req_pool_idx = req_pool_indices[seq_idx]
+            per_req_tokens = req_to_token[req_pool_idx, atten_start_kv:atten_end_kv]
+            # For SWA layers, k_cache/v_cache are from the SWA pool but
+            # req_to_token stores full pool indices. Translate before indexing.
+            if full_to_swa_mapping is not None:
+                per_req_tokens = full_to_swa_mapping[per_req_tokens]
+            per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
+            per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
+
+            if not (per_req_query.dtype == per_req_key.dtype == per_req_value.dtype):
+                # scaled_dot_product_attention() expects query, key, and value to have the same dtype
+                per_req_key = per_req_key.to(per_req_query.dtype)
+                per_req_value = per_req_value.to(per_req_query.dtype)
+
+            if logit_cap > 0:
+                per_req_out_redundant = (
+                    self.scaled_dot_product_attention_with_softcapping(
+                        per_req_query_redundant.unsqueeze(0),
+                        per_req_key.unsqueeze(0),
+                        per_req_value.unsqueeze(0),
+                        enable_gqa=enable_gqa,
+                        scale=scaling,
+                        is_causal=causal,
+                        logit_cap=logit_cap,
+                        logit_capping_method=logit_capping_method,
+                    )
+                    .squeeze(0)
+                    .movedim(query.dim() - 2, 0)
+                )
+            else:
+                per_req_out_redundant = (
+                    scaled_dot_product_attention(
+                        per_req_query_redundant.unsqueeze(0),
+                        per_req_key.unsqueeze(0),
+                        per_req_value.unsqueeze(0),
+                        enable_gqa=enable_gqa,
+                        scale=scaling,
+                        is_causal=causal,
+                    )
+                    .squeeze(0)
+                    .movedim(query.dim() - 2, 0)
+                )
+            output[start_q:end_q, :, :] = per_req_out_redundant[
+                query_start_idx : query_start_idx + extend_seq_len_q, :, :
+            ]
+            start_q, start_kv = end_q, end_kv
+        return output
+
+    def run_sdpa_forward_decode(
+        self,
+        query: torch.Tensor,
+        output: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        req_to_token: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        encoder_lens: torch.Tensor = None,
+        is_cross_attention: bool = False,
+        scaling=None,
+        enable_gqa=False,
+        causal=False,
+        sliding_window_size: int = -1,
+        full_to_swa_mapping: Optional[torch.Tensor] = None,
+        logit_cap: float = 0.0,
+        logit_capping_method: str = "tanh",
+    ):
+        """Run the decode forward by using torch native sdpa op.
+
+        Args:
+            query: [num_tokens, num_heads, head_size]
+            output: [num_tokens, num_heads, head_size]
+            k_cache: [max_total_num_tokens, num_heads, head_size]
+            v_cache: [max_total_num_tokens, num_heads, head_size]
+            req_to_token: [max_num_reqs, max_context_len]
+            req_pool_indices: [num_seqs]
+            seq_lens: [num_seqs]
+            encoder_lens: [num_seqs]
+            is_cross_attention: [bool]
+            scaling: float or None
+            enable_gqa: bool
+            causal: bool
+
+        Returns:
+            output: [num_tokens, num_heads, head_size]
+        """
+
+        # [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
+        query = query.movedim(0, query.dim() - 2)
+
+        start_q, start_kv = 0, 0
+        for seq_idx in range(seq_lens.shape[0]):
+            # Need optimize the performance later.
+
+            seq_len_q = 1
+            seq_len_kv = int(seq_lens[seq_idx].item())
+            end_q = start_q + seq_len_q
+            end_kv = start_kv + seq_len_kv
+            atten_start_kv = 0
+            atten_end_kv = seq_len_kv
+            # support cross attention
+            if encoder_lens is not None:
+                if is_cross_attention:
+                    atten_end_kv = int(encoder_lens[seq_idx].item())
+                else:
+                    atten_start_kv = int(encoder_lens[seq_idx].item())
+                    atten_end_kv = atten_start_kv + seq_len_kv
+
+            if (
+                sliding_window_size is not None
+                and sliding_window_size > -1
+                and encoder_lens is None
+            ):
+                atten_start_kv = max(
+                    atten_end_kv - (sliding_window_size + 1), atten_start_kv
+                )
+
+            per_req_query = query[:, start_q:end_q, :]
+
+            # get key and value from cache. per_req_tokens contains the kv cache
+            # index for each token in the sequence.
+            req_pool_idx = req_pool_indices[seq_idx]
+            per_req_tokens = req_to_token[req_pool_idx, atten_start_kv:atten_end_kv]
+            # For SWA layers, k_cache/v_cache are from the SWA pool but
+            # req_to_token stores full pool indices. Translate before indexing.
+            if full_to_swa_mapping is not None:
+                per_req_tokens = full_to_swa_mapping[per_req_tokens]
+            per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
+            per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
+
+            if not (per_req_query.dtype == per_req_key.dtype == per_req_value.dtype):
+                # scaled_dot_product_attention() expects query, key, and value to have the same dtype
+                per_req_key = per_req_key.to(per_req_query.dtype)
+                per_req_value = per_req_value.to(per_req_query.dtype)
+
+            if logit_cap > 0:
+                per_req_out = (
+                    self.scaled_dot_product_attention_with_softcapping(
+                        per_req_query.unsqueeze(0),
+                        per_req_key.unsqueeze(0),
+                        per_req_value.unsqueeze(0),
+                        enable_gqa=enable_gqa,
+                        scale=scaling,
+                        is_causal=causal,
+                        logit_cap=logit_cap,
+                        logit_capping_method=logit_capping_method,
+                    )
+                    .squeeze(0)
+                    .movedim(query.dim() - 2, 0)
+                )
+            else:
+                per_req_out = (
+                    scaled_dot_product_attention(
+                        per_req_query.unsqueeze(0),
+                        per_req_key.unsqueeze(0),
+                        per_req_value.unsqueeze(0),
+                        enable_gqa=enable_gqa,
+                        scale=scaling,
+                        is_causal=causal,
+                    )
+                    .squeeze(0)
+                    .movedim(query.dim() - 2, 0)
+                )
+            output[start_q:end_q, :, :] = per_req_out
+            start_q, start_kv = end_q, end_kv
+
+        return output
+
+    def support_triton(self):
+        return False

python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py

@@ -0,0 +1,493 @@
+import re
+from functools import lru_cache
+from typing import TYPE_CHECKING, Optional
+
+import torch
+import torch.nn.functional as F
+
+from sglang.srt.hardware_backend.npu.utils import npu_format_cast
+from sglang.srt.model_executor.forward_context import (
+    get_attn_backend,
+    get_token_to_kv_pool,
+)
+from sglang.srt.utils import get_bool_env_var
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.quantization.base_config import QuantizationConfig
+
+
+@lru_cache(maxsize=1)
+def is_mla_preprocess_enabled() -> bool:
+    return get_bool_env_var("SGLANG_NPU_USE_MLAPO")
+
+
+@lru_cache(maxsize=1)
+def is_fia_nz() -> bool:
+    is_fia_nz_ = get_bool_env_var("SGLANG_USE_FIA_NZ")
+    if is_fia_nz_:
+        assert (
+            is_mla_preprocess_enabled()
+        ), "SGLANG_USE_FIA_NZ must be enable with SGLANG_NPU_USE_MLAPO"
+    return is_fia_nz_
+
+
+def round_up(val: int, align: int) -> int:
+    if align == 0:
+        return 0
+    return -(val // -align) * align
+
+
+def transdata(nd_mat, block_size: tuple = (16, 16)):
+    r = round_up(nd_mat.shape[0], block_size[0])
+    c = round_up(nd_mat.shape[1], block_size[1])
+    r_pad = r - nd_mat.shape[0]
+    c_pad = c - nd_mat.shape[1]
+    nd_mat = F.pad(nd_mat, ((0, r_pad, 0, c_pad)))
+    nz_mat = torch.permute(
+        torch.reshape(
+            nd_mat,
+            (r // block_size[0], block_size[0], c // block_size[1], block_size[1]),
+        ),
+        [2, 0, 1, 3],
+    )
+    nz_mat = torch.reshape(
+        nz_mat, (nz_mat.shape[0], nz_mat.shape[1] * nz_mat.shape[2], nz_mat.shape[3])
+    )
+    return nz_mat
+
+
+def trans_rope_weight(weight, rope_dim):
+    weight_1 = weight[..., -rope_dim::2, :].contiguous()
+    weight_2 = weight[..., -rope_dim + 1 :: 2, :].contiguous()
+    weight[..., -rope_dim:, :] = torch.cat([weight_1, weight_2], dim=-2)
+
+    return weight.contiguous()
+
+
+class NPUFusedMLAPreprocess(torch.nn.Module):
+    def __init__(
+        self,
+        fused_qkv_a_proj_with_mqa,
+        q_a_layernorm,
+        kv_a_layernorm,
+        q_b_proj,
+        w_kc,
+        rotary_emb,
+        layer_id,
+        num_local_heads,
+        qk_nope_head_dim,
+        qk_rope_head_dim,
+        v_head_dim,
+        quant_config: Optional["QuantizationConfig"] = None,
+    ):
+        super().__init__()
+        self.qkv_a_proj = fused_qkv_a_proj_with_mqa
+        self.q_a_layernorm = q_a_layernorm
+        self.kv_a_layernorm = kv_a_layernorm
+        self.q_b_proj = q_b_proj
+        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
+
+        self.q_lora_rank = self.q_b_proj.input_size  # 1536
+        self.kv_lora_rank = self.kv_a_layernorm.hidden_size  # 512
+        self.num_local_heads = num_local_heads  # tp
+        self.qk_nope_head_dim = qk_nope_head_dim  # 128
+        self.qk_rope_head_dim = qk_rope_head_dim  # 64
+        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
+        self.v_head_dim = v_head_dim
+        self.q_b_proj_weight_scale = self.q_b_proj.weight_scale.view(1, -1).to(
+            torch.float
+        )
+
+    def preprocess_weights(self, hidden_states):
+        self.dummy = torch.zeros(
+            (hidden_states.shape[-1]),
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+        self.qkv_a_proj_input_offset = self.qkv_a_proj.input_offset.to(dtype=torch.int8)
+        self.q_b_proj_input_offset = self.q_b_proj.input_offset.to(dtype=torch.int8)
+
+        # matmul_0 weight [7168, 2112]
+        fused_qkv_a_proj_with_mqa_weight_q = self.qkv_a_proj.weight.data[
+            :, : self.q_lora_rank
+        ].clone()  # [7168, 1536]
+        fused_qkv_a_proj_with_mqa_weight_kv = self.qkv_a_proj.weight.data[
+            :, self.q_lora_rank :
+        ].clone()  # [7168, 576]
+        # rope fit
+        fused_qkv_a_proj_with_mqa_weight_kv_t = (
+            fused_qkv_a_proj_with_mqa_weight_kv.t().contiguous()
+        )
+        fused_qkv_a_proj_with_mqa_weight_kv_t = trans_rope_weight(
+            fused_qkv_a_proj_with_mqa_weight_kv_t, self.qk_rope_head_dim
+        )
+        fused_qkv_a_proj_with_mqa_weight_kv = (
+            fused_qkv_a_proj_with_mqa_weight_kv_t.t().contiguous()
+        )
+        # cat nz
+        fused_qkv_a_proj_with_mqa_weight_new = torch.cat(
+            (fused_qkv_a_proj_with_mqa_weight_kv, fused_qkv_a_proj_with_mqa_weight_q),
+            dim=-1,
+        )
+        fused_qkv_a_proj_with_mqa_weight = (
+            fused_qkv_a_proj_with_mqa_weight_new.t().contiguous()
+        )
+        fused_qkv_a_proj_with_mqa_weight_nz = (
+            transdata(fused_qkv_a_proj_with_mqa_weight, block_size=(16, 32))
+            .unsqueeze(0)
+            .contiguous()
+        )
+        self.qkv_a_proj_weight_nz = npu_format_cast(fused_qkv_a_proj_with_mqa_weight_nz)
+
+        # matmul_0 deq_scale [2112]
+        fused_qkv_a_proj_with_mqa_deq_scale_q = self.qkv_a_proj.deq_scale.data[
+            : self.q_lora_rank
+        ].clone()  # [7168, 1536]
+        fused_qkv_a_proj_with_mqa_deq_scale_kv = self.qkv_a_proj.deq_scale.data[
+            self.q_lora_rank :
+        ].clone()  # [7168, 576]
+        # rope fit
+        fused_qkv_a_proj_with_mqa_deq_scale_kv = (
+            fused_qkv_a_proj_with_mqa_deq_scale_kv.reshape(
+                self.kv_lora_rank + self.qk_rope_head_dim, -1
+            ).contiguous()
+        )
+        fused_qkv_a_proj_with_mqa_deq_scale_kv = trans_rope_weight(
+            fused_qkv_a_proj_with_mqa_deq_scale_kv, self.qk_rope_head_dim
+        )
+        fused_qkv_a_proj_with_mqa_deq_scale_kv = (
+            fused_qkv_a_proj_with_mqa_deq_scale_kv.view(
+                self.kv_lora_rank + self.qk_rope_head_dim
+            ).contiguous()
+        )
+        self.qkv_a_proj_deq_scale_kvq = torch.cat(
+            (
+                fused_qkv_a_proj_with_mqa_deq_scale_kv,
+                fused_qkv_a_proj_with_mqa_deq_scale_q,
+            ),
+            dim=-1,
+        )
+
+        # matmul_0 quant_bias [2112]
+        fused_qkv_a_proj_with_mqa_quant_bias_q = self.qkv_a_proj.quant_bias.data[
+            : self.q_lora_rank
+        ].clone()  # [7168, 1536]
+        fused_qkv_a_proj_with_mqa_quant_bias_kv = self.qkv_a_proj.quant_bias.data[
+            self.q_lora_rank :
+        ].clone()  # [7168, 576]
+        # rope fit
+        fused_qkv_a_proj_with_mqa_quant_bias_kv = (
+            fused_qkv_a_proj_with_mqa_quant_bias_kv.reshape(
+                self.kv_lora_rank + self.qk_rope_head_dim, -1
+            ).contiguous()
+        )
+        fused_qkv_a_proj_with_mqa_quant_bias_kv = trans_rope_weight(
+            fused_qkv_a_proj_with_mqa_quant_bias_kv, self.qk_rope_head_dim
+        )
+        fused_qkv_a_proj_with_mqa_quant_bias_kv = (
+            fused_qkv_a_proj_with_mqa_quant_bias_kv.view(
+                self.kv_lora_rank + self.qk_rope_head_dim
+            ).contiguous()
+        )
+        self.qkv_a_proj_quant_bias_kvq = torch.cat(
+            (
+                fused_qkv_a_proj_with_mqa_quant_bias_kv,
+                fused_qkv_a_proj_with_mqa_quant_bias_q,
+            ),
+            dim=-1,
+        )
+
+        # matmul_1 weight [1536, num_head * 192]
+        q_b_proj_weight = self.q_b_proj.weight.data.clone()
+        q_b_proj_weight = q_b_proj_weight.t().reshape(
+            self.num_local_heads, self.qk_nope_head_dim + self.qk_rope_head_dim, -1
+        )
+        q_b_proj_weight = trans_rope_weight(q_b_proj_weight, self.qk_rope_head_dim)
+        q_b_proj_weight = q_b_proj_weight.reshape(
+            self.num_local_heads * (self.qk_nope_head_dim + self.qk_rope_head_dim), -1
+        )
+        q_b_proj_weight_nz = (
+            transdata(q_b_proj_weight, block_size=(16, 32)).unsqueeze(0).contiguous()
+        )
+        self.q_b_proj_weight_nz = npu_format_cast(q_b_proj_weight_nz)
+
+        # matmul_1 deq_scale [num_head * 192]
+        q_b_proj_deq_scale = self.q_b_proj.deq_scale.data.clone()
+        q_b_proj_deq_scale = q_b_proj_deq_scale.reshape(
+            self.num_local_heads, self.qk_nope_head_dim + self.qk_rope_head_dim, -1
+        )
+        q_b_proj_deq_scale = trans_rope_weight(
+            q_b_proj_deq_scale, self.qk_rope_head_dim
+        )
+        self.q_b_proj_deq_scale = q_b_proj_deq_scale.reshape(
+            self.num_local_heads * (self.qk_nope_head_dim + self.qk_rope_head_dim)
+        )
+
+        # matmul_1 quant_bias [num_head * 192]
+        q_b_proj_quant_bias = self.q_b_proj.quant_bias.data.clone()
+        q_b_proj_quant_bias = q_b_proj_quant_bias.reshape(
+            self.num_local_heads, self.qk_nope_head_dim + self.qk_rope_head_dim, -1
+        )
+        q_b_proj_quant_bias = trans_rope_weight(
+            q_b_proj_quant_bias, self.qk_rope_head_dim
+        )
+        self.q_b_proj_quant_bias = q_b_proj_quant_bias.reshape(
+            self.num_local_heads * (self.qk_nope_head_dim + self.qk_rope_head_dim)
+        )
+
+    def mlaprolog_preprocess_weight(self):
+        self.qkv_a_proj.weight.data = self.qkv_a_proj.weight.data.transpose(0, 1)
+        qkv_a_proj_weight_q = self.qkv_a_proj.weight.data[:, : self.q_lora_rank].clone()
+        qkv_a_proj_weight_kv = self.qkv_a_proj.weight.data[
+            :, self.q_lora_rank :
+        ].clone()
+        self.q_a_proj_weight = npu_format_cast(qkv_a_proj_weight_q)
+        self.kv_a_proj_weight = npu_format_cast(qkv_a_proj_weight_kv)
+
+    def get_sin_cos(self, positions):
+        cos_sin = self.rotary_emb.cos_sin_cache[positions]
+        cos, sin = cos_sin.chunk(2, dim=-1)
+        cos = cos.repeat(1, 2)
+        sin = sin.repeat(1, 2)
+        return cos, sin
+
+    def get_kv_cache_and_cache_idx(self, forward_batch):
+        k_cache, v_cache = get_token_to_kv_pool().get_kv_buffer(self.layer_id)
+        slot_mapping = forward_batch.out_cache_loc.to(dtype=torch.int32)
+        return k_cache, v_cache, slot_mapping
+
+    def forward_absorb_prepare_npu_rms_norm_cache(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        forward_batch,
+        zero_allocator,
+    ):
+        bsz, _ = hidden_states.view(-1, hidden_states.shape[-1]).shape
+        self.dtype = hidden_states.dtype
+        if self.layer_id == 0:
+            self.cos, self.sin = self.get_sin_cos(positions)
+            self.rotary_emb.cos_cached, self.rotary_emb.sin_cache = self.cos, self.sin
+        else:
+            self.cos, self.sin = self.rotary_emb.cos_cached, self.rotary_emb.sin_cache
+
+        self.kvCache, self.kvCacheRope, self.slotmapping = (
+            self.get_kv_cache_and_cache_idx(forward_batch)
+        )
+
+        if not self.has_preprocess_weights:
+            self.has_preprocess_weights = True
+
+        cos, sin = self.cos, self.sin
+
+        if self.q_lora_rank is not None:
+            fused_qkv_a_proj_out = self.qkv_a_proj(hidden_states)[0]
+            q_lowrank, latent_cache = fused_qkv_a_proj_out.split(
+                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1
+            )
+            q = self.q_a_layernorm(q_lowrank)
+            q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
+        else:
+            q = self.q_proj(hidden_states)[0].view(
+                -1, self.num_local_heads, self.qk_head_dim
+            )
+            latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
+
+        q_nope, q_pe = torch.split(
+            q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
+        )  # b*s,n,d
+
+        q_nope = q_nope.view(-1, self.num_local_heads, self.qk_nope_head_dim)
+        q_nope = torch.matmul(q_nope.transpose(0, 1), self.w_kc).transpose(0, 1)
+
+        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.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(
+            -1, 1, 1, self.kv_lora_rank + self.qk_rope_head_dim
+        )  # (B*S,N,1,D)
+
+        cache_mode = "PA_NZ" if is_fia_nz() else "PA_BNSD"
+        self.kvCache = self.kvCache.view(
+            -1,
+            get_attn_backend().page_size,
+            1,
+            get_attn_backend().kv_lora_rank,
+        )
+        self.kvCacheRope = self.kvCacheRope.view(
+            -1,
+            get_attn_backend().page_size,
+            1,
+            get_attn_backend().qk_rope_head_dim,
+        )
+        k_rope, k_nope, _, _ = torch.ops.npu.npu_kv_rmsnorm_rope_cache(
+            latent_cache,
+            self.kv_a_layernorm.weight,
+            cos,
+            sin,
+            self.slotmapping.to(torch.int64),
+            self.kvCacheRope,
+            self.kvCache,
+            epsilon=self.kv_a_layernorm.variance_epsilon,
+            cache_mode=cache_mode,
+        )
+
+        return (q_pe, k_rope, q_nope, k_nope, forward_batch, zero_allocator, positions)
+
+    def forward_mlapo(self, positions, hidden_states, forward_batch, zero_allocator):
+        input_dtype = hidden_states.dtype
+        if not self.has_preprocess_weights:
+            self.preprocess_weights(hidden_states)
+            self.has_preprocess_weights = True
+            self.dtype = hidden_states.dtype
+
+        if self.layer_id == 0:
+            cos, sin = self.get_sin_cos(positions)
+            self.rotary_emb.cos_cached, self.rotary_emb.sin_cache = cos, sin
+        else:
+            cos, sin = self.rotary_emb.cos_cached, self.rotary_emb.sin_cache
+
+        k_cache, v_cache, slot_mapping = self.get_kv_cache_and_cache_idx(forward_batch)
+
+        q_nope_out = torch.empty(
+            (hidden_states.shape[0], self.w_kc.shape[0], k_cache.shape[-1]),
+            dtype=input_dtype,
+            device=hidden_states.device,
+        )
+        q_rope_out = torch.empty(
+            (hidden_states.shape[0], self.w_kc.shape[0], v_cache.shape[-1]),
+            dtype=input_dtype,
+            device=hidden_states.device,
+        )
+        if is_fia_nz():
+            kv_shape, kv_rope_shape = k_cache.shape, v_cache.shape
+            num_blocks, block_size, num_heads, _ = kv_shape
+            k_cache = k_cache.view(
+                num_blocks, num_heads * self.kv_lora_rank // 16, block_size, 16
+            )
+            v_cache = v_cache.view(
+                num_blocks, num_heads * self.qk_rope_head_dim // 16, block_size, 16
+            )
+        # TODO: dummy inputs to be removed
+        # https://github.com/sgl-project/sgl-kernel-npu/issues/78
+        if hasattr(self.q_a_layernorm, "bias"):
+            q_a_layernorm_bias = self.q_a_layernorm.bias
+        else:
+            q_a_layernorm_bias = self.dummy
+
+        torch.ops.npu.mla_preprocess(
+            hidden_states,
+            self.dummy,
+            self.dummy,
+            self.qkv_a_proj_weight_nz,
+            self.qkv_a_proj_deq_scale_kvq,
+            self.q_a_layernorm.weight,
+            q_a_layernorm_bias,
+            self.q_b_proj_weight_nz,
+            self.q_b_proj_deq_scale,
+            self.kv_a_layernorm.weight,
+            cos,
+            sin,
+            self.w_kc,
+            k_cache,
+            v_cache,
+            slot_mapping,
+            quant_scale0=self.qkv_a_proj.input_scale,
+            quant_offset0=self.qkv_a_proj_input_offset,
+            bias0=self.qkv_a_proj_quant_bias_kvq,
+            quant_scale1=self.q_b_proj.input_scale,
+            quant_offset1=self.q_b_proj_input_offset,
+            bias1=self.q_b_proj_quant_bias,
+            cache_mode="nzcache" if is_fia_nz() else "krope_ctkv",
+            quant_mode="per_tensor_quant_asymm",
+            q_out0=q_nope_out,
+            kv_cache_out0=k_cache,
+            q_out1=q_rope_out,
+            kv_cache_out1=v_cache,
+        )
+
+        if is_fia_nz():
+            k_cache = k_cache.view(kv_shape)
+            v_cache = v_cache.view(kv_rope_shape)
+
+        return (
+            q_rope_out,
+            v_cache,
+            q_nope_out,
+            k_cache,
+            forward_batch,
+            zero_allocator,
+            positions,
+        )
+
+    def forward_mlaprolog(self, positions, hidden_states, forward_batch):
+        if not self.has_preprocess_weights:
+            self.mlaprolog_preprocess_weight()
+            self.has_preprocess_weights = True
+        self.cos, self.sin = self.get_sin_cos(positions)
+        k_cache, v_cache, slot_mapping = self.get_kv_cache_and_cache_idx(forward_batch)
+        mla_prolog_input_args = {
+            "token_x": hidden_states,
+            "weight_dq": self.q_a_proj_weight,
+            "weight_uq_qr": self.q_b_proj.weight,
+            "weight_uk": self.w_kc,
+            "weight_dkv_kr": self.kv_a_proj_weight,
+            "rmsnorm_gamma_cq": self.q_a_layernorm.weight,
+            "rmsnorm_gamma_ckv": self.kv_a_layernorm.weight,
+            "rope_sin": self.sin,
+            "rope_cos": self.cos,
+            "kv_cache": k_cache,
+            "kr_cache": v_cache,
+            "cache_index": slot_mapping.to(dtype=torch.int64),
+            "dequant_scale_w_uq_qr": self.q_b_proj_weight_scale,
+            "rmsnorm_epsilon_cq": self.q_a_layernorm.variance_epsilon,
+            "rmsnorm_epsilon_ckv": self.kv_a_layernorm.variance_epsilon,
+            "cache_mode": "PA_BSND",
+            "query_norm_flag": True,
+            "weight_quant_mode": 1,  # 0:no quant; 1:uq_qr: quant; 2: weight_dq,weight_uq_qr,weight_dkv_kr: quant
+        }
+        q_nope, q_pe, dequant_scale_q_nope, qr, dequant_q_norm = (
+            torch.ops.custom.npu_mla_prolog_v3(**mla_prolog_input_args)
+        )
+        dequant_q_norm = dequant_q_norm.view(hidden_states.shape[0])
+        return (
+            q_pe,
+            v_cache,
+            q_nope,
+            k_cache,
+            qr,
+            forward_batch,
+            positions,
+            dequant_q_norm,
+        )
+
+    def forward(self, positions, hidden_states, forward_batch, zero_allocator):
+        # assert self.quant_config and self.quant_config.get_name() == "modelslim"
+        # route by `qkv_a_proj` quant type as MTP layers can be unquantized
+        _is_w8a8 = (
+            hasattr(self.qkv_a_proj.quant_method, "quantization_config")
+            and self.qkv_a_proj.quant_method.quantization_config.get_name()
+            == "modelslim"
+        )
+        # with the mlaprolog enabled, the kv_b_proj layers are unquantized
+        _is_mlaprolog = hasattr(self.quant_config, "ignore") and any(
+            re.fullmatch(r".*kv_b_proj", l) for l in self.quant_config.ignore
+        )
+        if _is_w8a8:
+            return self.forward_mlapo(
+                positions, hidden_states, forward_batch, zero_allocator
+            )
+        elif _is_mlaprolog:
+            return self.forward_mlaprolog(positions, hidden_states, forward_batch)
+        else:
+            return self.forward_absorb_prepare_npu_rms_norm_cache(
+                positions, hidden_states, forward_batch, zero_allocator
+            )

python/sglang/srt/hardware_backend/npu/batch_invariant_ops/npu_batch_invariant_ops.py

@@ -0,0 +1,62 @@
+# Adapted from https://github.com/thinking-machines-lab/batch_invariant_ops/blob/main/batch_invariant_ops/batch_invariant_ops.py
+
+import batch_invariant_ops  # noqa: F401
+import torch
+import torch_npu
+
+
+def npu_mm_batch_invariant(a, b):
+    return torch.ops.batch_invariant_ops.npu_mm_batch_invariant(a, b)
+
+
+def npu_matmul_batch_invariant(a, b):
+    return torch.ops.batch_invariant_ops.npu_matmul_batch_invariant(a, b)
+
+
+def npu_mean_batch_invariant(
+    input, dim, keepdim=False, dtype: torch.dtype | None = None
+):
+    assert dtype is None or dtype == torch.float32, f"unsupported dtype: {dtype}"
+    if len(dim) == 1:
+        return torch.ops.batch_invariant_ops.npu_reduce_mean_batch_invariant(
+            input, dim[0], keepdim=keepdim
+        )
+    else:
+        assert input.dtype in {
+            torch.float16,
+            torch.bfloat16,
+            torch.float32,
+        }, "only float types supported for now"
+        n_elems = 1
+        for d in dim:
+            n_elems *= input.shape[d]
+        return torch.sum(input, dim=dim, keepdim=keepdim, dtype=torch.float32) / n_elems
+
+
+def npu_log_softmax_batch_invariant(input, dim, _half_to_float):
+    assert not _half_to_float, "not implemented"
+    return torch.ops.batch_invariant_ops.npu_log_softmax_batch_invariant(input, dim=dim)
+
+
+def npu_fused_infer_attention_score_batch_invariant(*args, **kwargs):
+    return (
+        torch.ops.batch_invariant_ops.npu_fused_infer_attention_score_batch_invariant(
+            *args, **kwargs
+        )
+    )
+
+
+def npu_add_rms_norm_batch_invariant(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+):
+    """
+    AclnnAddRmsNorm can't ensure batch invariant,
+    so we need to split it into add and rms_norm.
+    """
+    x_ = x + residual
+    residual_ = x_
+    x_, _ = torch_npu.npu_rms_norm(x_, weight, eps)
+    return x_, None, residual_

python/sglang/srt/hardware_backend/npu/cmo.py

@@ -0,0 +1,83 @@
+import torch
+
+cmo_stream = None
+share_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)
+
+
+def get_share_stream():
+    global share_stream
+    return share_stream
+
+
+def set_share_stream(stream):
+    global share_stream
+    share_stream = stream
+
+
+def wait_share_stream():
+    stream = get_share_stream()
+    if stream is not None:
+        cur_stream = torch.npu.current_stream()
+        cur_stream.wait_stream(stream)
+
+
+def shared_expert_on_independent_stream(hidden_states, forward_func):
+    stream = get_share_stream()
+    if stream is None:
+        stream = torch.npu.Stream()
+        set_share_stream(stream)
+    stream.wait_stream(torch.npu.current_stream())
+    with torch.npu.stream(stream):
+        shared_output = forward_func(hidden_states)
+        return shared_output

python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py

@@ -0,0 +1,749 @@
+"""DSV4-NPU SWA + c4/c128 paged allocator.
+
+Subclasses :class:`SWATokenToKVPoolAllocator` and adds paged allocation for the
+c4/c128 compressed-KV pools and their tail-only compress-state pools, alongside
+the parent's full + SWA pools.
+
+Per ``alloc_extend`` / ``alloc_decode``:
+  1. super() allocates the full + SWA slots (``out_full_loc``).
+  2. Allocate c4/c128 KV slots — one compressed token per ``ratio`` raw tokens
+     (``seq_len // ratio - prefix_len // ratio``) — via the standard
+     :class:`NPUPagedTokenToKVPoolAllocator` over the pool's c4/c128 KV buffers.
+  3. Allocate the c4/c128 compress-state slots the same way, tail-only per req,
+     using the per-req lens the scheduler packed into ``DSV4StateLens``.
+  4. Return a :class:`DSV4OutCacheLoc` bundling all five slot families.
+
+State slots are paged because the NPU fused compressor runs ``cache_mode=1``; the
+base class' ``translate_kv_loc_to_compress_state_loc`` ring-hash is the CUDA-only
+path and is unused on NPU. The bundle is the explicit return value:
+mem_cache/common.py unpacks ``out_full_loc`` and stashes the bundle on
+``batch.out_cache_loc_dsv4``; ``DSV4NPUReqToTokenPool`` writes the per-req
+``req_to_token_c{4,128}[_state]`` tables that :meth:`free` and the last_loc
+lookups read back.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Optional
+
+import torch
+
+from sglang.srt.configs.model_config import is_deepseek_v4
+from sglang.srt.hardware_backend.npu.allocator_npu import NPUPagedTokenToKVPoolAllocator
+from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import (
+    maybe_write_dsv4_extend,
+)
+from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
+from sglang.srt.mem_cache.common import alloc_paged_token_slots_extend
+from sglang.srt.model_executor.forward_batch_info import DSV4OutCacheLoc, DSV4StateLens
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req
+
+
+def get_last_loc(
+    req_to_token: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    prefix_lens: torch.Tensor,
+) -> torch.Tensor:
+    """Slot id of each req's last already-allocated token, or -1 when
+    ``prefix_lens[i] == 0`` (fresh req).
+
+    Looks up ``req_to_token[req, prefix_lens - 1]`` to anchor the paged
+    allocator's ``alloc_extend`` on the real previous tail slot, preserving the
+    intra-page slot continuity the kernel's ``cmp_block_table`` relies on (the
+    allocator debug-asserts ``(last_loc + 1) % page_size == prefix_lens %
+    page_size``). Result dtype matches ``prefix_lens``.
+    """
+    req_pool_indices = req_pool_indices.to(torch.int64)
+    safe_idx = (prefix_lens.to(torch.int64) - 1).clamp(min=0)
+    looked_up = req_to_token[req_pool_indices, safe_idx].to(prefix_lens.dtype)
+    return torch.where(
+        prefix_lens > 0,
+        looked_up,
+        torch.full_like(prefix_lens, -1),
+    )
+
+
+def alloc_paged_token_slots_extend_npu(*args, batch=None, **kwargs):
+    if batch is not None and is_deepseek_v4(batch.model_config.hf_config):
+        return alloc_paged_token_slots_reserve_extend(*args, batch=batch, **kwargs)
+    return alloc_paged_token_slots_extend(*args, batch=batch, **kwargs)
+
+
+def alloc_paged_token_slots_reserve_extend(
+    tree_cache,
+    prefix_lens: torch.Tensor,
+    prefix_lens_cpu: torch.Tensor,
+    seq_lens: torch.Tensor,
+    seq_lens_cpu: torch.Tensor,
+    last_loc: torch.Tensor,
+    extend_num_tokens: int,
+    *,
+    req_pool_indices: Optional[torch.Tensor] = None,
+    dsv4_state_lens: Optional[DSV4StateLens] = None,
+    batch=None,
+):
+    """Allocate reserved draft slots and update DSV4 per-request tables."""
+    if dsv4_state_lens is None and batch is not None:
+        allocator = batch.token_to_kv_pool_allocator
+        dsv4_state_lens = (
+            allocator.compute_dsv4_state_lens_reserve(
+                batch.reqs, prefix_lens_cpu, seq_lens_cpu
+            )
+            if hasattr(allocator, "compute_dsv4_state_lens_reserve")
+            else None
+        )
+
+    out_cache_loc = alloc_paged_token_slots_extend(
+        tree_cache,
+        prefix_lens,
+        prefix_lens_cpu,
+        seq_lens,
+        seq_lens_cpu,
+        last_loc,
+        extend_num_tokens,
+        req_pool_indices=req_pool_indices,
+        dsv4_state_lens=dsv4_state_lens,
+        batch=batch,
+    )
+    if batch is not None:
+        maybe_write_dsv4_extend(
+            batch,
+            batch.req_pool_indices_cpu,
+            prefix_lens_cpu,
+            seq_lens_cpu,
+            c4_state_alloc_offsets=prefix_lens_cpu,
+            c128_state_alloc_offsets=prefix_lens_cpu,
+        )
+    return out_cache_loc
+
+
+class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
+    """SWA allocator + c4/c128 KV and compress-state paged allocators for DSV4 on NPU."""
+
+    def __init__(
+        self,
+        size: int,
+        size_swa: int,
+        page_size: int,
+        dtype: torch.dtype,
+        device: str,
+        kvcache,
+        need_sort: bool,
+    ):
+        super().__init__(
+            size=size,
+            size_swa=size_swa,
+            page_size=page_size,
+            dtype=dtype,
+            device=device,
+            kvcache=kvcache,
+            need_sort=need_sort,
+        )
+
+        def mk(pool_size, pool):
+            # c4/c128 KV and state sub-pools implement KVCache, so they drop into
+            # the standard paged allocator. pool_size is in compressed-token units.
+            return NPUPagedTokenToKVPoolAllocator(
+                pool_size,
+                page_size=page_size,
+                dtype=dtype,
+                device=device,
+                kvcache=pool,
+                need_sort=need_sort,
+            )
+
+        self.c4_attn_allocator = mk(kvcache.c4_size, kvcache.c4_kv_pool)
+        self.c128_attn_allocator = mk(kvcache.c128_size, kvcache.c128_kv_pool)
+
+        # State allocators (paged, NPU-only). Any layer's pool works as KVCache
+        # pointer (slot alloc is layer-agnostic); None when no c{ratio} layers or
+        # zero budget.
+        self.c4_state_attn_allocator: Optional[NPUPagedTokenToKVPoolAllocator] = None
+        self.c128_state_attn_allocator: Optional[NPUPagedTokenToKVPoolAllocator] = None
+        state_pools = getattr(kvcache, "compress_state_pools", None)
+        if state_pools:
+
+            def first_state_pool(want_ratio):
+                return next(
+                    (
+                        p
+                        for r, p in zip(kvcache.compression_ratios, state_pools)
+                        if r == want_ratio and p is not None
+                    ),
+                    None,
+                )
+
+            c4_state_pool = first_state_pool(4)
+            c128_state_pool = first_state_pool(128)
+            if c4_state_pool is not None and kvcache.c4_state_pool_size > 0:
+                self.c4_state_attn_allocator = mk(
+                    kvcache.c4_state_pool_size, c4_state_pool
+                )
+            if c128_state_pool is not None and kvcache.c128_state_pool_size > 0:
+                self.c128_state_attn_allocator = mk(
+                    kvcache.c128_state_pool_size, c128_state_pool
+                )
+
+        # Returned by the c-pool helpers when a step adds no compressed tokens.
+        self._empty_loc = torch.empty((0,), dtype=torch.int64, device=device)
+
+        # Per-call handle to the DSV4NPUReqToTokenPool, stashed by alloc_extend/
+        # alloc_decode for last_loc lookups; avoids a permanent allocator->pool ref.
+        self._cur_req_to_token_pool = None
+
+    @staticmethod
+    def _compute_c_extend_counts(
+        prefix_lens_cpu: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        ratio: int,
+    ) -> int:
+        """New compressed-K tokens this extend produces across the batch:
+        ``sum_i (seq_lens[i] // ratio - prefix_lens[i] // ratio)``."""
+        if prefix_lens_cpu is None or seq_lens_cpu is None:
+            return 0
+        diff = ((seq_lens_cpu // ratio) - (prefix_lens_cpu // ratio)).clamp(min=0)
+        return int(diff.sum().item())
+
+    @staticmethod
+    def _pool_exhausted(
+        ratio: int, kind: str, need: int, available: int
+    ) -> RuntimeError:
+        return RuntimeError(
+            f"DSV4 c{ratio} {kind} pool exhausted: need {need} new slots, "
+            f"available={available}. Raise --mem-fraction-static, lower "
+            f"--max-running-requests, or check that "
+            f"DSV4NPUTokenToKVPoolAllocator.free(req=...) releases {kind} slots "
+            f"on req finish."
+        )
+
+    def _alloc_state_extend(
+        self,
+        allocator: Optional[NPUPagedTokenToKVPoolAllocator],
+        raw_prefix_lens: torch.Tensor,
+        state_prefix_lens: torch.Tensor,
+        state_prefix_lens_cpu: torch.Tensor,
+        state_seq_lens: torch.Tensor,
+        state_seq_lens_cpu: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        last_loc_dtype: torch.dtype,
+        state_extend_num_tokens: int,
+        ratio: int,
+    ) -> torch.Tensor:
+        """Allocate tail-only state-pool slots for an extend at ``ratio``.
+
+        The state pool is a separate paged slot space; each req allocates only
+        its trailing window (cumulative lens precomputed by
+        ``ScheduleBatch._compute_dsv4_state_lens_*`` and passed via
+        ``DSV4StateLens``). ``state_last_loc`` is looked up from
+        ``req_to_token_c{ratio}_state`` at the RAW position
+        ``raw_prefix_lens - 1`` (the last position the previous extend/decode
+        populated). Returns ``_empty_loc`` when the allocator is absent (no
+        c{ratio} layers) or there is nothing to add.
+        """
+        if allocator is None or state_extend_num_tokens == 0:
+            return self._empty_loc
+
+        assert self._cur_req_to_token_pool is not None, (
+            "alloc_extend/alloc_decode must be called with req_to_token_pool= "
+            "for the state-pool last_loc lookup."
+        )
+        state_table = (
+            self._cur_req_to_token_pool.req_to_token_c4_state
+            if ratio == 4
+            else self._cur_req_to_token_pool.req_to_token_c128_state
+        )
+        state_last_loc = get_last_loc(
+            state_table, req_pool_indices, raw_prefix_lens
+        ).to(last_loc_dtype)
+
+        result = allocator.alloc_extend(
+            state_prefix_lens,
+            state_prefix_lens_cpu,
+            state_seq_lens,
+            state_seq_lens_cpu,
+            state_last_loc,
+            state_extend_num_tokens,
+        )
+        if result is None:
+            raise self._pool_exhausted(
+                ratio, "state", state_extend_num_tokens, allocator.available_size()
+            )
+        return result
+
+    def _alloc_c_extend(
+        self,
+        allocator: NPUPagedTokenToKVPoolAllocator,
+        prefix_lens: torch.Tensor,
+        prefix_lens_cpu: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        last_loc_dtype: torch.dtype,
+        ratio: int,
+    ) -> torch.Tensor:
+        """Allocate compressed-KV slots for an extend at ``ratio``.
+
+        Prefix/seq lens are translated to compressed units (``// ratio``); the
+        c-pool last_loc comes from ``req_to_token_c{ratio}`` via
+        :func:`get_last_loc` so the paged allocator continues in-page (or opens
+        a fresh page at a ratio boundary), keeping the intra-page continuity the
+        ``cmp_block_table`` reader relies on. Returns ``_empty_loc`` when this
+        step closes no compressed token.
+        """
+        c_extend = self._compute_c_extend_counts(prefix_lens_cpu, seq_lens_cpu, ratio)
+        if c_extend == 0:
+            return self._empty_loc
+
+        assert self._cur_req_to_token_pool is not None, (
+            "alloc_extend/alloc_decode must be called with req_to_token_pool= "
+            "for the c-pool last_loc lookup."
+        )
+        c_table = (
+            self._cur_req_to_token_pool.req_to_token_c4
+            if ratio == 4
+            else self._cur_req_to_token_pool.req_to_token_c128
+        )
+        c_prefix = (prefix_lens // ratio).to(prefix_lens.dtype)
+        c_seq = (seq_lens // ratio).to(seq_lens.dtype)
+        c_last_loc = get_last_loc(c_table, req_pool_indices, c_prefix).to(
+            last_loc_dtype
+        )
+
+        result = allocator.alloc_extend(
+            c_prefix,
+            prefix_lens_cpu // ratio,
+            c_seq,
+            seq_lens_cpu // ratio,
+            c_last_loc,
+            c_extend,
+        )
+        if result is None:
+            raise self._pool_exhausted(
+                ratio, "KV", c_extend, allocator.available_size()
+            )
+        return result
+
+    def _alloc_c_and_state(
+        self,
+        out_full_loc: torch.Tensor,
+        out_swa_loc: torch.Tensor,
+        prefix_lens: torch.Tensor,
+        prefix_lens_cpu: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        last_loc_dtype: torch.dtype,
+        req_pool_indices: Optional[torch.Tensor],
+        dsv4_state_lens: Optional[DSV4StateLens],
+    ) -> DSV4OutCacheLoc:
+        """Allocate c4/c128 KV + state slots and bundle them with full/swa loc.
+
+        Shared by alloc_extend / alloc_decode (which differ only in how
+        prefix_lens is derived). State lens are tail-only, precomputed by
+        ScheduleBatch._compute_dsv4_state_lens_*; raw prefix_lens drives the
+        state last_loc lookup.
+        """
+        assert req_pool_indices is not None, (
+            "DSV4NPUTokenToKVPoolAllocator requires req_pool_indices "
+            "(forwarded from batch.req_pool_indices)."
+        )
+        if dsv4_state_lens is not None:
+            out_c4_state_loc = self._alloc_state_extend(
+                self.c4_state_attn_allocator,
+                prefix_lens,
+                dsv4_state_lens.c4_prefix_lens,
+                dsv4_state_lens.c4_prefix_lens_cpu,
+                dsv4_state_lens.c4_seq_lens,
+                dsv4_state_lens.c4_seq_lens_cpu,
+                req_pool_indices,
+                last_loc_dtype,
+                dsv4_state_lens.c4_extend_num_tokens,
+                ratio=4,
+            )
+            out_c128_state_loc = self._alloc_state_extend(
+                self.c128_state_attn_allocator,
+                prefix_lens,
+                dsv4_state_lens.c128_prefix_lens,
+                dsv4_state_lens.c128_prefix_lens_cpu,
+                dsv4_state_lens.c128_seq_lens,
+                dsv4_state_lens.c128_seq_lens_cpu,
+                req_pool_indices,
+                last_loc_dtype,
+                dsv4_state_lens.c128_extend_num_tokens,
+                ratio=128,
+            )
+        else:
+            out_c4_state_loc = self._empty_loc
+            out_c128_state_loc = self._empty_loc
+        out_c4_loc = self._alloc_c_extend(
+            self.c4_attn_allocator,
+            prefix_lens,
+            prefix_lens_cpu,
+            seq_lens,
+            seq_lens_cpu,
+            req_pool_indices,
+            last_loc_dtype,
+            ratio=4,
+        )
+        out_c128_loc = self._alloc_c_extend(
+            self.c128_attn_allocator,
+            prefix_lens,
+            prefix_lens_cpu,
+            seq_lens,
+            seq_lens_cpu,
+            req_pool_indices,
+            last_loc_dtype,
+            ratio=128,
+        )
+        return DSV4OutCacheLoc(
+            out_full_loc=out_full_loc,
+            out_swa_loc=out_swa_loc,
+            out_c4_loc=out_c4_loc,
+            out_c128_loc=out_c128_loc,
+            out_c4_state_loc=out_c4_state_loc,
+            out_c128_state_loc=out_c128_state_loc,
+        )
+
+    def compute_dsv4_state_lens_extend(
+        self, reqs: List[Req], seq_lens: List[int]
+    ) -> Optional[DSV4StateLens]:
+        """Per-req c{4,128}_state pool alloc lens for extend (tail-only).
+
+        State pool stores only the trailing portion of each sequence (the c{N}
+        compressor's read/write window); the tail length depends on raw
+        seq_len's alignment to the SWA page boundary (128)::
+
+            c4_alloc_len  = tail + 128 if (tail <= 3 and seq_len >= 128) else tail
+            c128_alloc_len = tail                  where tail = seq_len % 128
+
+        Long prefills allocate only the trailing partial window, not slots for
+        already-compressed positions, so the small paged state pool (~256
+        slots/req) stays sufficient even for 28k-token prompts.
+
+        Mutates per-req cumulative state via getattr/setattr so the community
+        ``Req`` needs no DSV4 field declarations:
+          * ``req.c{4,128}_state_kv_len`` — cumulative slot count (prefix for
+            the paged allocator; never decreases on eviction).
+          * ``req.c{4,128}_state_alloc_offset`` — low-water raw-position mark
+            for eviction (see ``dsv4_common_hooks.maybe_evict_dsv4_state``).
+
+        Returns None when this model has no paged state pools (CUDA / non-V4 /
+        zero budget) — callers pass that straight through as ``dsv4_state_lens``.
+        """
+        if self.c4_state_attn_allocator is None:
+            return None
+        c4_prefix: List[int] = []
+        c4_seq: List[int] = []
+        c128_prefix: List[int] = []
+        c128_seq: List[int] = []
+        for req, seq_len in zip(reqs, seq_lens):
+            tail = seq_len % 128
+            c4_alloc_len = tail + 128 if (tail <= 3 and seq_len >= 128) else tail
+            c128_alloc_len = tail
+
+            prev_c4 = getattr(req, "c4_state_kv_len", 0)
+            prev_c128 = getattr(req, "c128_state_kv_len", 0)
+            new_c4 = prev_c4 + c4_alloc_len
+            new_c128 = prev_c128 + c128_alloc_len
+
+            c4_prefix.append(prev_c4)
+            c4_seq.append(new_c4)
+            c128_prefix.append(prev_c128)
+            c128_seq.append(new_c128)
+
+            req.c4_state_kv_len = new_c4
+            req.c128_state_kv_len = new_c128
+            req.c4_state_alloc_offset = seq_len - c4_alloc_len
+            req.c128_state_alloc_offset = seq_len - c128_alloc_len
+
+        return self._pack_state_lens(
+            c4_prefix,
+            c4_seq,
+            c128_prefix,
+            c128_seq,
+            c4_extend_num_tokens=int(sum(s - p for s, p in zip(c4_seq, c4_prefix))),
+            c128_extend_num_tokens=int(
+                sum(s - p for s, p in zip(c128_seq, c128_prefix))
+            ),
+        )
+
+    def compute_dsv4_state_lens_decode(
+        self, reqs: List[Req]
+    ) -> Optional[DSV4StateLens]:
+        """Per-req c{4,128}_state pool alloc lens for decode: exactly 1 new
+        state slot per req per pool. ``c{N}_state_alloc_offset`` does NOT
+        advance here (only eviction advances it). Returns None when there are
+        no paged state pools."""
+        if self.c4_state_attn_allocator is None:
+            return None
+        c4_prefix: List[int] = []
+        c4_seq: List[int] = []
+        c128_prefix: List[int] = []
+        c128_seq: List[int] = []
+        for req in reqs:
+            prev_c4 = getattr(req, "c4_state_kv_len", 0)
+            prev_c128 = getattr(req, "c128_state_kv_len", 0)
+            c4_prefix.append(prev_c4)
+            c4_seq.append(prev_c4 + 1)
+            c128_prefix.append(prev_c128)
+            c128_seq.append(prev_c128 + 1)
+            req.c4_state_kv_len = prev_c4 + 1
+            req.c128_state_kv_len = prev_c128 + 1
+
+        bs = len(reqs)
+        return self._pack_state_lens(
+            c4_prefix,
+            c4_seq,
+            c128_prefix,
+            c128_seq,
+            c4_extend_num_tokens=bs,
+            c128_extend_num_tokens=bs,
+        )
+
+    def compute_dsv4_state_lens_reserve(
+        self, reqs: List[Req], prefix_lens: List[int], seq_lens: List[int]
+    ) -> Optional[DSV4StateLens]:
+        """Allocate state slots for a speculative pre-reserved raw interval."""
+        if self.c4_state_attn_allocator is None:
+            return None
+
+        c4_prefix: List[int] = []
+        c4_seq: List[int] = []
+        c128_prefix: List[int] = []
+        c128_seq: List[int] = []
+        for req, prefix_len, seq_len in zip(reqs, prefix_lens, seq_lens):
+            reserve = max(0, int(seq_len) - int(prefix_len))
+            prev_c4 = getattr(req, "c4_state_kv_len", 0)
+            prev_c128 = getattr(req, "c128_state_kv_len", 0)
+            c4_prefix.append(prev_c4)
+            c4_seq.append(prev_c4 + reserve)
+            c128_prefix.append(prev_c128)
+            c128_seq.append(prev_c128 + reserve)
+            req.c4_state_kv_len = prev_c4 + reserve
+            req.c128_state_kv_len = prev_c128 + reserve
+
+        total = sum(max(0, int(s) - int(p)) for p, s in zip(prefix_lens, seq_lens))
+        return self._pack_state_lens(
+            c4_prefix,
+            c4_seq,
+            c128_prefix,
+            c128_seq,
+            c4_extend_num_tokens=total,
+            c128_extend_num_tokens=total,
+        )
+
+    def _pack_state_lens(
+        self,
+        c4_prefix: List[int],
+        c4_seq: List[int],
+        c128_prefix: List[int],
+        c128_seq: List[int],
+        *,
+        c4_extend_num_tokens: int,
+        c128_extend_num_tokens: int,
+    ) -> DSV4StateLens:
+        c4_prefix_cpu = torch.tensor(c4_prefix, dtype=torch.int64)
+        c4_seq_cpu = torch.tensor(c4_seq, dtype=torch.int64)
+        c128_prefix_cpu = torch.tensor(c128_prefix, dtype=torch.int64)
+        c128_seq_cpu = torch.tensor(c128_seq, dtype=torch.int64)
+        return DSV4StateLens(
+            c4_prefix_lens=c4_prefix_cpu.to(self.device, non_blocking=True),
+            c4_prefix_lens_cpu=c4_prefix_cpu,
+            c4_seq_lens=c4_seq_cpu.to(self.device, non_blocking=True),
+            c4_seq_lens_cpu=c4_seq_cpu,
+            c4_extend_num_tokens=c4_extend_num_tokens,
+            c128_prefix_lens=c128_prefix_cpu.to(self.device, non_blocking=True),
+            c128_prefix_lens_cpu=c128_prefix_cpu,
+            c128_seq_lens=c128_seq_cpu.to(self.device, non_blocking=True),
+            c128_seq_lens_cpu=c128_seq_cpu,
+            c128_extend_num_tokens=c128_extend_num_tokens,
+        )
+
+    def alloc_extend(
+        self,
+        prefix_lens: torch.Tensor,
+        prefix_lens_cpu: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        last_loc: torch.Tensor,
+        extend_num_tokens: int,
+        *,
+        req_pool_indices: Optional[torch.Tensor] = None,
+        dsv4_state_lens: Optional[DSV4StateLens] = None,
+        req_to_token_pool=None,
+    ) -> Optional[DSV4OutCacheLoc]:
+        # Stash per-req tables for this call's last_loc lookups (read by
+        # _alloc_c_extend / _alloc_state_extend); no permanent allocator->pool ref.
+        self._cur_req_to_token_pool = req_to_token_pool
+        out_full_loc = super().alloc_extend(
+            prefix_lens,
+            prefix_lens_cpu,
+            seq_lens,
+            seq_lens_cpu,
+            last_loc,
+            extend_num_tokens,
+        )
+        if out_full_loc is None:
+            return None
+
+        out_swa_loc = self.translate_loc_from_full_to_swa(out_full_loc)
+        assert out_swa_loc is not None, (
+            "translate_loc_from_full_to_swa returned None — "
+            "full_to_swa_index_mapping not initialized?"
+        )
+        return self._alloc_c_and_state(
+            out_full_loc,
+            out_swa_loc,
+            prefix_lens,
+            prefix_lens_cpu,
+            seq_lens,
+            seq_lens_cpu,
+            last_loc.dtype,
+            req_pool_indices,
+            dsv4_state_lens,
+        )
+
+    def alloc_decode(
+        self,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        last_loc: torch.Tensor,
+        *,
+        req_pool_indices: Optional[torch.Tensor] = None,
+        dsv4_state_lens: Optional[DSV4StateLens] = None,
+        req_to_token_pool=None,
+    ) -> Optional[DSV4OutCacheLoc]:
+        self._cur_req_to_token_pool = req_to_token_pool
+        out_full_loc = super().alloc_decode(seq_lens, seq_lens_cpu, last_loc)
+        if out_full_loc is None:
+            return None
+
+        out_swa_loc = self.translate_loc_from_full_to_swa(out_full_loc)
+        # One new token per req. Model as an extend from (seq_len-1)//ratio to
+        # seq_len//ratio so _alloc_c_extend anchors on the real c-pool last_loc.
+        prefix_lens = (seq_lens - 1).clamp(min=0)
+        prefix_lens_cpu = (seq_lens_cpu - 1).clamp(min=0)
+        return self._alloc_c_and_state(
+            out_full_loc,
+            out_swa_loc,
+            prefix_lens,
+            prefix_lens_cpu,
+            seq_lens,
+            seq_lens_cpu,
+            last_loc.dtype,
+            req_pool_indices,
+            dsv4_state_lens,
+        )
+
+    def free(
+        self,
+        free_index: Optional[torch.Tensor] = None,
+        *,
+        req=None,
+        req_to_token_pool=None,
+    ):
+        """Unified free for full/swa/c4/c128 pools. Two forms (may co-fire):
+
+          * ``free(free_index)`` — full + SWA only (tail/radix eviction; no req
+            identity, so c-pool free can't run).
+          * ``free(req=, req_to_token_pool=)`` — from DSV4NPUReqToTokenPool.free
+            on req finish: reads the per-req slot lists from
+            ``req_to_token_c{4,128}[_state]`` and returns them to the c-pools
+            (the paged allocator dedupes by page).
+
+        KV pools free ``[0, kv_len // ratio)``. State pools are 1-per-raw-token
+        and free only the tail ``[c{N}_state_alloc_offset, kv_len)`` — the prefix
+        was already returned by ScheduleBatch._evict_swa (state rides SWA
+        eviction); freeing it again would double-free (caught by the paged
+        allocator's debug_mode assert, corrupts the free list otherwise).
+        """
+        if free_index is not None:
+            super().free(free_index)
+
+        if req is None or req_to_token_pool is None:
+            return
+        kv_len = max(req.kv_committed_len, req.kv_allocated_len)
+        req_pool_idx = req.req_pool_idx
+        if kv_len <= 0 or req_pool_idx is None:
+            return
+
+        # KV pools: free the leading [0, kv_len // ratio) compressed slots.
+        for ratio, allocator, table_attr in (
+            (4, self.c4_attn_allocator, "req_to_token_c4"),
+            (128, self.c128_attn_allocator, "req_to_token_c128"),
+        ):
+            n = kv_len // ratio
+            if n > 0 and hasattr(req_to_token_pool, table_attr):
+                slots = getattr(req_to_token_pool, table_attr)[req_pool_idx, :n]
+                # to int64 — paged allocator's free does cpu()//page_size on it.
+                allocator.free(slots.to(torch.int64))
+
+        # State pools: free only the tail [c{N}_state_alloc_offset, kv_len).
+        for ratio, allocator, table_attr, off_attr in (
+            (
+                4,
+                self.c4_state_attn_allocator,
+                "req_to_token_c4_state",
+                "c4_state_alloc_offset",
+            ),
+            (
+                128,
+                self.c128_state_attn_allocator,
+                "req_to_token_c128_state",
+                "c128_state_alloc_offset",
+            ),
+        ):
+            if allocator is None or not hasattr(req_to_token_pool, table_attr):
+                continue
+            off = getattr(req, off_attr, 0)
+            if kv_len > off:
+                slots = getattr(req_to_token_pool, table_attr)[req_pool_idx, off:kv_len]
+                allocator.free(slots.to(torch.int64))
+
+    def backup_state(self):
+        # EAGLE/NEXTN draft preprocess allocates speculative c{4,128} KV via
+        # alloc_extend(backup_state=True) and rolls it back with restore_state.
+        # The base SWATokenToKVPoolAllocator only snapshots the full + SWA pools,
+        # so without this override the draft's c{4,128} (+ state) slots are never
+        # rolled back -> they leak every draft step until the c4 pool exhausts.
+        # Snapshot the sub-allocators alongside the base pools.
+        return (
+            super().backup_state(),
+            self.c4_attn_allocator.backup_state(),
+            self.c128_attn_allocator.backup_state(),
+            (
+                self.c4_state_attn_allocator.backup_state()
+                if self.c4_state_attn_allocator is not None
+                else None
+            ),
+            (
+                self.c128_state_attn_allocator.backup_state()
+                if self.c128_state_attn_allocator is not None
+                else None
+            ),
+        )
+
+    def restore_state(self, state):
+        base, c4, c128, c4_state, c128_state = state
+        super().restore_state(base)
+        self.c4_attn_allocator.restore_state(c4)
+        self.c128_attn_allocator.restore_state(c128)
+        if self.c4_state_attn_allocator is not None and c4_state is not None:
+            self.c4_state_attn_allocator.restore_state(c4_state)
+        if self.c128_state_attn_allocator is not None and c128_state is not None:
+            self.c128_state_attn_allocator.restore_state(c128_state)
+
+    def clear(self):
+        super().clear()
+        # super().__init__ calls clear() before our sub-allocators exist;
+        # getattr(..., None) tolerates that and the always-None state allocators.
+        for attr in (
+            "c4_attn_allocator",
+            "c128_attn_allocator",
+            "c4_state_attn_allocator",
+            "c128_state_attn_allocator",
+        ):
+            allocator = getattr(self, attr, None)
+            if allocator is not None:
+                allocator.clear()

python/sglang/srt/hardware_backend/npu/dsv4/dsv4_common_hooks.py

@@ -0,0 +1,404 @@
+"""Helpers used by mem_cache/common.py to wire DSV4-NPU per-req tables.
+
+mem_cache/common.py runs platform-agnostic alloc flow. When the model is
+DSV4 on NPU, ``alloc_paged_token_slots_{extend,decode}`` already stashed the
+:class:`DSV4OutCacheLoc` the allocator returned onto
+``batch.out_cache_loc_dsv4``. After each ``alloc_extend`` / ``alloc_decode``
+these hooks then:
+
+  1. Read the bundle from ``batch.out_cache_loc_dsv4``.
+  2. Write the per-pool slot ids into the per-req tables on the
+     :class:`DSV4NPUReqToTokenPool`.
+
+Non-DSV4 paths leave ``batch.out_cache_loc_dsv4`` None, so this module is a
+no-op for them.
+
+TODO: the disagg DSV4 path bypasses these hooks — it calls
+``allocator.alloc_extend`` directly then ``req_to_token_pool.write`` without
+going through ``mem_cache/common.py`` (see ``disaggregation/decode.py``). The
+DSV4OutCacheLoc bundle is still produced but never written into the per-req
+tables, so disagg + DSV4 is unsupported here (c-pages leak). Fixing requires
+calling these hooks from disagg's per-req alloc loop, or moving the write
+into the allocator itself.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Sequence
+
+import torch
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
+
+
+def maybe_write_dsv4_extend(
+    batch: ScheduleBatch,
+    req_pool_indices_cpu: torch.Tensor,
+    prefix_lens_cpu: torch.Tensor,
+    seq_lens_cpu: torch.Tensor,
+    *,
+    c4_state_alloc_offsets: Sequence[int] | torch.Tensor | None = None,
+    c128_state_alloc_offsets: Sequence[int] | torch.Tensor | None = None,
+) -> None:
+    """Post-alloc_extend hook for DSV4. No-op when allocator/pool is not DSV4.
+
+    For each compressed pool (c4 / c128), spreads the flat
+    ``out_c{4,128}_loc`` tensor across requests using per-req extend
+    counts (``seq_lens[i] // ratio - prefix_lens[i] // ratio``) and writes
+    the resulting slot ids into ``req_to_token_c{4,128}[req, prefix:seq]``.
+
+    Also writes ``req_to_token_swa[req, prefix:seq]`` with the swa slots
+    derived from out_full_loc via the SWA index mapping.
+    """
+    # Bundle stashed on batch.out_cache_loc_dsv4 by mem_cache/common.py;
+    # None on CUDA / non-V4 paths → no-op.
+    bundle = batch.out_cache_loc_dsv4
+    if bundle is None:
+        return
+
+    req_to_token_pool = batch.req_to_token_pool
+    if not hasattr(req_to_token_pool, "write_c4"):
+        return  # non-DSV4 pool; skip defensively (shouldn't happen)
+
+    # SWA writes: prefix..seq token positions, one slot per raw token.
+    _write_per_req_slice(
+        req_to_token_pool.write_swa,
+        req_pool_indices_cpu,
+        prefix_lens_cpu,
+        seq_lens_cpu,
+        bundle.out_swa_loc,
+        ratio=1,
+    )
+
+    # c4 / c128 writes: prefix//ratio .. seq//ratio compressed positions.
+    _write_per_req_slice(
+        req_to_token_pool.write_c4,
+        req_pool_indices_cpu,
+        prefix_lens_cpu,
+        seq_lens_cpu,
+        bundle.out_c4_loc,
+        ratio=4,
+    )
+    _write_per_req_slice(
+        req_to_token_pool.write_c128,
+        req_pool_indices_cpu,
+        prefix_lens_cpu,
+        seq_lens_cpu,
+        bundle.out_c128_loc,
+        ratio=128,
+    )
+
+    # c4_state / c128_state writes: tail-only. Bundle length is
+    # sum(c{N}_state_alloc_len_i), NOT total raw extend tokens. Normal extend
+    # uses the per-Req low-water marks; reserve callers can pass explicit raw
+    # offsets for the pre-reserved interval.
+    if c4_state_alloc_offsets is None:
+        c4_state_alloc_offsets = [
+            getattr(r, "c4_state_alloc_offset", 0) for r in batch.reqs
+        ]
+    if c128_state_alloc_offsets is None:
+        c128_state_alloc_offsets = [
+            getattr(r, "c128_state_alloc_offset", 0) for r in batch.reqs
+        ]
+    if bundle.out_c4_state_loc is not None and hasattr(
+        req_to_token_pool, "write_c4_state"
+    ):
+        _write_state_tail_per_req(
+            req_to_token_pool.write_c4_state,
+            req_pool_indices_cpu,
+            c4_state_alloc_offsets,
+            seq_lens_cpu,
+            bundle.out_c4_state_loc,
+        )
+    if bundle.out_c128_state_loc is not None and hasattr(
+        req_to_token_pool, "write_c128_state"
+    ):
+        _write_state_tail_per_req(
+            req_to_token_pool.write_c128_state,
+            req_pool_indices_cpu,
+            c128_state_alloc_offsets,
+            seq_lens_cpu,
+            bundle.out_c128_state_loc,
+        )
+
+
+def maybe_write_dsv4_decode(
+    batch: ScheduleBatch,
+    seq_lens_cpu: torch.Tensor,
+    token_per_req: int,
+) -> None:
+    """Post-alloc_decode hook for DSV4. Spreads the new token slot ids
+    (one per req for swa, gated by ratio boundary for c4/c128) into the
+    per-req tables on DSV4NPUReqToTokenPool.
+
+    ``seq_lens_cpu`` is the POST-decode seq len (already incremented by
+    ``token_per_req``); the new compressed tokens go at positions
+    ``[(old_seq) // ratio, (new_seq) // ratio)``.
+    """
+    # Bundle stashed on batch.out_cache_loc_dsv4 by mem_cache/common.py;
+    # None on CUDA / non-V4 paths → no-op.
+    bundle = batch.out_cache_loc_dsv4
+    if bundle is None:
+        return
+
+    req_to_token_pool = batch.req_to_token_pool
+    if not hasattr(req_to_token_pool, "write_c4"):
+        return
+
+    prefix_lens_cpu = (seq_lens_cpu - token_per_req).clamp(min=0)
+    req_pool_indices_cpu = batch.req_pool_indices.cpu()
+
+    _write_per_req_slice(
+        req_to_token_pool.write_swa,
+        req_pool_indices_cpu,
+        prefix_lens_cpu,
+        seq_lens_cpu,
+        bundle.out_swa_loc,
+        ratio=1,
+    )
+    _write_per_req_slice(
+        req_to_token_pool.write_c4,
+        req_pool_indices_cpu,
+        prefix_lens_cpu,
+        seq_lens_cpu,
+        bundle.out_c4_loc,
+        ratio=4,
+    )
+    _write_per_req_slice(
+        req_to_token_pool.write_c128,
+        req_pool_indices_cpu,
+        prefix_lens_cpu,
+        seq_lens_cpu,
+        bundle.out_c128_loc,
+        ratio=128,
+    )
+
+    # State table decode writes: one slot per raw decode token (ratio=1).
+    if bundle.out_c4_state_loc is not None and hasattr(
+        req_to_token_pool, "write_c4_state"
+    ):
+        _write_per_req_slice(
+            req_to_token_pool.write_c4_state,
+            req_pool_indices_cpu,
+            prefix_lens_cpu,
+            seq_lens_cpu,
+            bundle.out_c4_state_loc,
+            ratio=1,
+        )
+    if bundle.out_c128_state_loc is not None and hasattr(
+        req_to_token_pool, "write_c128_state"
+    ):
+        _write_per_req_slice(
+            req_to_token_pool.write_c128_state,
+            req_pool_indices_cpu,
+            prefix_lens_cpu,
+            seq_lens_cpu,
+            bundle.out_c128_state_loc,
+            ratio=1,
+        )
+
+
+def maybe_build_dsv4_verify_bundle(batch: ScheduleBatch, draft_token_num: int):
+    """Build the DSV4 cache-location view for one target-verify pass.
+
+    Spec-v2 reserves cache ahead of time, so target verify must select only the
+    current draft interval from the per-request DSV4 tables instead of reusing
+    the larger allocation bundle produced during decode preparation.
+    """
+    pool = batch.req_to_token_pool
+    if not hasattr(pool, "req_to_token_c4"):
+        return None
+    reserve_bundle = batch.out_cache_loc_dsv4
+    if reserve_bundle is None:
+        return None
+
+    req_indices = batch.req_pool_indices_cpu.tolist()
+    seq_lens = batch.seq_lens_cpu.tolist()
+
+    def flatten_interval(table: torch.Tensor, ratio: int) -> torch.Tensor:
+        chunks = []
+        for req_idx, seq_len in zip(req_indices, seq_lens):
+            start = int(seq_len) // ratio
+            end = (int(seq_len) + draft_token_num) // ratio
+            if end > start:
+                chunks.append(table[int(req_idx), start:end])
+        return torch.cat(chunks) if chunks else table.new_empty((0,))
+
+    return type(reserve_bundle)(
+        out_full_loc=batch.out_cache_loc,
+        out_swa_loc=flatten_interval(pool.req_to_token_swa, 1),
+        out_c4_loc=flatten_interval(pool.req_to_token_c4, 4),
+        out_c128_loc=flatten_interval(pool.req_to_token_c128, 128),
+        out_c4_state_loc=flatten_interval(pool.req_to_token_c4_state, 1),
+        out_c128_state_loc=flatten_interval(pool.req_to_token_c128_state, 1),
+    )
+
+
+def _write_per_req(
+    write_fn,
+    req_pool_indices_cpu: torch.Tensor,
+    flat_loc: torch.Tensor,
+    bounds_fn,
+) -> None:
+    """Distribute a flat ``[total_alloc]`` slot tensor across reqs.
+
+    ``bounds_fn(i) -> (lo, hi)`` gives req i's write window; the matching
+    ``hi - lo`` slots are sliced off ``flat_loc`` in order and written via
+    ``write_fn((req_idx, slice(lo, hi)), values)``. flat_loc may be None /
+    empty when the alloc path bypassed DSV4NPUTokenToKVPoolAllocator (e.g.
+    page_size=1 or HiSparse wrapper); skip then.
+    """
+    if flat_loc is None or flat_loc.numel() == 0:
+        return
+    pt = 0
+    for i in range(req_pool_indices_cpu.shape[0]):
+        lo, hi = bounds_fn(i)
+        alloc_len = max(0, hi - lo)
+        if alloc_len == 0:
+            continue
+        req_idx = int(req_pool_indices_cpu[i].item())
+        chunk = flat_loc[pt : pt + alloc_len].to(torch.int32)
+        write_fn((req_idx, slice(lo, hi)), chunk)
+        pt += alloc_len
+
+
+def _write_state_tail_per_req(
+    write_fn,
+    req_pool_indices_cpu: torch.Tensor,
+    state_alloc_offsets: list,
+    seq_lens_cpu: torch.Tensor,
+    flat_loc: torch.Tensor,
+) -> None:
+    """Tail-only state write: req i's slots go at ``[state_alloc_offsets[i],
+    seq_lens[i])`` in ``req_to_token_c{N}_state``."""
+    _write_per_req(
+        write_fn,
+        req_pool_indices_cpu,
+        flat_loc,
+        lambda i: (int(state_alloc_offsets[i]), int(seq_lens_cpu[i].item())),
+    )
+
+
+def _write_per_req_slice(
+    write_fn,
+    req_pool_indices_cpu: torch.Tensor,
+    prefix_lens_cpu: torch.Tensor,
+    seq_lens_cpu: torch.Tensor,
+    flat_loc: torch.Tensor,
+    ratio: int,
+) -> None:
+    """Compressed-position write: req i's slots go at
+    ``[prefix_lens[i] // ratio, seq_lens[i] // ratio)``."""
+    _write_per_req(
+        write_fn,
+        req_pool_indices_cpu,
+        flat_loc,
+        lambda i: (
+            int(prefix_lens_cpu[i].item()) // ratio,
+            int(seq_lens_cpu[i].item()) // ratio,
+        ),
+    )
+
+
+def maybe_evict_dsv4_state(batch: ScheduleBatch, req: Req, pre_len: int) -> None:
+    """Per-decode evict for the DSV4-NPU compress-state pools, independent of
+    SWA evict cadence. Called every decode step from ``ScheduleBatch``.
+
+    The state pool is small (~2 pages c4 / ~3 pages c128 of raw positions per
+    req) — with a large sliding_window (SWA evict fires every
+    ``eviction_interval`` and needs ``pre_len > sliding_window + page_size`` to
+    free anything) the pool exhausts before the first SWA frontier advance, so
+    we drain it here on its own cadence.
+
+    Retention windows (kernel read window + decode lookahead margin):
+    c4 = 8 + 16, c128 = 128 + 64 raw positions — intentionally smaller than one
+    SWA page so the first eviction fires before the small pool fills. Watermarks
+    are page-aligned so freed slots are whole pages reclaimable by the paged
+    allocator. ``req.c{4,128}_state_alloc_offset`` (read/written via getattr/
+    setattr) is the low-water mark. No-op on non-DSV4-NPU paths.
+    """
+    allocator = batch.token_to_kv_pool_allocator
+    pool = batch.req_to_token_pool
+    if not hasattr(allocator, "c4_state_attn_allocator") or (
+        allocator.c4_state_attn_allocator is None
+        and allocator.c128_state_attn_allocator is None
+    ):
+        return
+
+    page_size = batch.tree_cache.page_size
+    c4_watermark = ((max(0, pre_len - (8 + 16))) // page_size) * page_size
+    c128_watermark = ((max(0, pre_len - (128 + 64))) // page_size) * page_size
+
+    _free_state_range(
+        allocator.c4_state_attn_allocator,
+        pool,
+        "req_to_token_c4_state",
+        req,
+        "c4_state_alloc_offset",
+        c4_watermark,
+    )
+    _free_state_range(
+        allocator.c128_state_attn_allocator,
+        pool,
+        "req_to_token_c128_state",
+        req,
+        "c128_state_alloc_offset",
+        c128_watermark,
+    )
+
+
+def maybe_evict_dsv4_state_on_swa(
+    allocator, pool, req: Req, new_swa_evicted_seqlen: int
+) -> None:
+    """Free compress-state slots that ride along with SWA eviction.
+
+    State at raw positions < ``swa_evicted_seqlen`` is no longer readable (the
+    compressor only reads the trailing ``2*ratio`` window) and is returned to
+    its paged allocator to keep the small state pool from exhausting on long
+    generations. No-op when the DSV4-NPU state allocators are absent.
+
+    This path is needed for small-sliding-window models where
+    ``sliding_window < retention`` (e.g. c128 retention 192 > window 128):
+    in that case the watermark-based eviction alone may not free slots
+    fast enough, and the SWA-ride eviction is the primary reclaim mechanism.
+    For typical large-window models (DS-V4 with window >> 192), the
+    watermark eviction always runs first, making this path a no-op.
+    """
+    if not hasattr(allocator, "c4_state_attn_allocator"):
+        return
+    _free_state_range(
+        allocator.c4_state_attn_allocator,
+        pool,
+        "req_to_token_c4_state",
+        req,
+        "c4_state_alloc_offset",
+        new_swa_evicted_seqlen,
+    )
+    _free_state_range(
+        allocator.c128_state_attn_allocator,
+        pool,
+        "req_to_token_c128_state",
+        req,
+        "c128_state_alloc_offset",
+        new_swa_evicted_seqlen,
+    )
+
+
+def _free_state_range(
+    state_allocator,
+    pool,
+    table_attr: str,
+    req: Req,
+    offset_attr: str,
+    watermark: int,
+) -> None:
+    """Free ``[alloc_offset, watermark)`` raw-position state slots for ``req``
+    and advance its low-water mark. No-op when the allocator/table is absent or
+    the watermark hasn't advanced past the current offset."""
+    offset = getattr(req, offset_attr, 0)
+    if state_allocator is None or not hasattr(pool, table_attr) or watermark <= offset:
+        return
+    free_slots = getattr(pool, table_attr)[req.req_pool_idx, offset:watermark]
+    state_allocator.free(free_slots.to(torch.int64))
+    setattr(req, offset_attr, watermark)

python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py

@@ -0,0 +1,593 @@
+"""NPU-only KV pool variant for DeepSeek-V4.
+
+Subclasses :class:`DeepSeekV4TokenToKVPool` to swap the ring-buffered
+:class:`CompressStatePool` for the paged :class:`NPUCompressStatePool` that
+the on-NPU fused compressor kernel (``torch.ops.custom.compressor`` with
+``cache_mode=1``) requires. Atlas A3 rejects ``cache_mode=2`` (ring) entirely,
+so this is the only valid layout on that hardware.
+
+Selected at pool construction time by
+:meth:`ModelRunnerKVCacheMixin._init_pools` when the model is DSV4 AND the
+device is NPU. CUDA continues to use the unchanged base class.
+
+The subclass overrides only:
+
+  * ``_make_attn_state_pool`` / ``_make_indexer_state_pool`` — the per-ratio
+    state-pool factories the base ``_init_paged_compress_states`` loop calls.
+    Both return :class:`NPUCompressStatePool` (paged, ``cache_mode=1``)
+    instead of the base's ring-buffered :class:`CompressStatePool`.
+  * ``translate_kv_loc_to_compress_state_loc`` — raise loudly. The ring
+    hash this method implements is meaningless on the paged kernel; callers
+    must consume ``out_cache_loc_dsv4.out_c{4,128}_state_loc`` from the
+    allocator bundle instead. Currently the only NPU caller that still
+    invokes translate is the unfused Python compressor decode path
+    (``layers/attention/dsv4/compressor.py``); with USE_FUSED_COMPRESSOR=1
+    that path is dead. If someone disables the fused compressor, they hit
+    the raise with a clear message.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch_npu
+
+from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
+from sglang.srt.mem_cache.deepseek_v4_compress_state import CompressStatePool
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
+    ONLINE_C128,
+    DeepSeekV4IndexerPool,
+    DeepSeekV4SingleKVPool,
+    DeepSeekV4TokenToKVPool,
+)
+
+
+class NPUDeepSeekV4SingleKVPool(DeepSeekV4SingleKVPool):
+    """NPU bf16 variant of the full / SWA / c4 / c128 single-KV pool.
+
+    ``npu_sparse_attn_sharedkv`` reads KV in PA_ND layout
+    ``(num_pages, kernel_page_size, num_kv_heads=1, dim)`` with ``dim`` packing
+    K_nope + K_rope as bf16, and requires ``cmp_kv.shape[1] == ori_kv.shape[1]``.
+    So the c4/c128 pools (whose token-level page_size is ``page_size // ratio``)
+    are allocated at the GLOBAL ``kernel_page_size`` rather than their own
+    per-ratio page_size; the SWA pool uses ``kernel_page_size == page_size``.
+    The CUDA fp8-packed-bytes layout (the base ``create_buffer``) is untouched.
+    """
+
+    def __init__(self, *args, kernel_page_size: int, **kwargs):
+        # Set before super().__init__ — it calls _create_buffers() ->
+        # create_buffer(), which reads self.kernel_page_size.
+        self.kernel_page_size = kernel_page_size
+        super().__init__(*args, **kwargs)
+
+    def create_buffer(self, *, num_pages: int):
+        # Non-bf16 store dtype (shouldn't happen here) falls back to base layout.
+        if self.store_dtype != torch.bfloat16:
+            return super().create_buffer(num_pages=num_pages)
+        kv_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+        self.kv_cache_total_dim = kv_dim
+        # GLOBAL kernel_page_size keeps cmp_kv.shape[1] == ori_kv.shape[1]; writes
+        # are flat-indexed by loc, so page granularity affects shape not location.
+        npu_num_pages = (self.size + self.kernel_page_size + 1) // self.kernel_page_size
+        return torch.zeros(
+            npu_num_pages,
+            self.kernel_page_size,
+            1,
+            kv_dim,
+            dtype=torch.bfloat16,
+            device=self.device,
+        )
+
+
+def npu_state_pool_size(
+    *,
+    ratio: int,
+    page_size: int,
+    max_num_reqs: int,
+) -> int:
+    """Per-pool state slot count for the NPU paged state pool's
+    :class:`NPUPagedTokenToKVPoolAllocator`.
+
+    Sizing formula::
+
+        max(2, ceil(1.8 * ratio / page_size) + 1) * max_num_reqs * page_size
+
+    Sized for steady-state during decode: each req keeps roughly the trailing
+    ``sliding_window_size`` worth of state slots live at any one time (SWA
+    eviction in :meth:`ScheduleBatch._evict_swa` frees state slots as it
+    advances), and the 1.8x factor adds headroom for the tail-only allocation
+    pattern across page boundaries.
+
+    Prefill no longer drives sizing because allocation is tail-only — long
+    prompts only allocate ``c{ratio}_alloc_len`` slots (``≤ tail + 128`` for
+    c4, ``≤ tail`` for c128, where ``tail = seq_len % 128``), not the full raw
+    seqlen. See :meth:`ScheduleBatch._compute_dsv4_state_lens_extend` for the
+    per-req formula.
+
+    Result is in TOKEN units (matches the SGLang allocator
+    ``PagedTokenToKVPoolAllocator(size, ...)`` convention where
+    ``num_pages = size // page_size`` is the count of USABLE pages handed out
+    by ``free_pages = arange(1, num_pages+1)``). The BUFFER allocates one extra
+    page (see :class:`NPUCompressStatePool`, sized ``(num_pages + 1) *
+    page_size`` — page 0 is the kernel's skip-sentinel).
+    """
+    blocks_per_req = max(2, math.ceil(1.8 * ratio / page_size) + 1)
+    num_usable_pages = blocks_per_req * max_num_reqs
+    return num_usable_pages * page_size
+
+
+class NPUCompressStatePool(CompressStatePool):
+    """Paged compress-state pool for the NPU fused compressor kernel.
+
+    ``torch.ops.custom.compressor`` (cache_mode=1) reads/writes the compress
+    state via ``state_cache`` shape ``(block_num, page_size, 2*coff*head_dim)``
+    indexed by a paged ``state_block_table`` (block ids from 1; value 0 means
+    "skip this slot"). The CUDA :class:`CompressStatePool` sizes itself
+    ring-style, which misaddresses slots under cache_mode=1 (ring is also
+    unsupported on Atlas A3). This subclass keeps the parent's buffer layout
+    (``(self._size, 2*coff*head_dim)`` flat; ``state_cache_3d`` reshapes to
+    ``(num_blocks, page_size, 2*coff*head_dim)``) but replaces the size formula
+    with a paged one derived from ``max_num_reqs``. Block 0 is reserved as the
+    kernel's skip-sentinel (zero kv / -inf score) so any ``state_block_table``
+    entry defaulting to 0 lands in a deterministic, attention-neutral place.
+
+    NPU-only; CUDA keeps using the unchanged :class:`CompressStatePool`.
+    """
+
+    def __init__(
+        self,
+        *,
+        size: int,
+        overlap: bool,
+        head_dim: int,
+        dtype: torch.dtype,
+        device: str,
+        enable_memory_saver: bool,
+        ratio: int,
+        page_size: int,
+    ):
+        # Bypass parent __init__ — its ring-based sizing is incompatible with the
+        # kernel's paged block-id contract. We redo buffer alloc and set the same
+        # fields so the parent API (state_cache_3d, kv_score_buffer) stays intact.
+        assert ratio in (
+            4,
+            128,
+        ), f"NPUCompressStatePool only supports ratio in (4, 128); got {ratio}"
+        assert page_size > 1, (
+            "NPUCompressStatePool requires page_size>1 (kernel's "
+            "state_cache_3d view is (block_num, page_size, slot_dim)). "
+            "Got page_size=%d." % page_size
+        )
+
+        # ``size`` is the ALLOCATOR's size (npu_state_pool_size output). Buffer
+        # needs one EXTRA page so the free list arange(1, num_pages+1) indexes it
+        # without OOB (page 0 = skip sentinel; pages 1..num_pages handed out).
+        num_usable_pages = (size + page_size - 1) // page_size
+        num_buffer_pages = num_usable_pages + 1
+        self._size = num_buffer_pages * page_size
+        self.page_size = page_size
+        # ring_size=0 marks "not ring-buffered" (paged allocator replaces the
+        # parent's ring hashing); kept so downstream hasattr probes don't break.
+        self.ring_size = 0
+        # online compress is a CUDA-only opt with no NPU fused-compressor support;
+        # force off so layout matches kernel expectations.
+        self.online = False
+
+        # Slot dim = 2 * coff * head_dim = [kv | score]; coff = 1 (no overlap) or
+        # 2 (overlap). Matches CompressStatePool non-online layout.
+        self.last_dim = 2 * (1 + int(overlap)) * head_dim
+
+        # Reuse parent's buffer-alloc helper; only self._size differs from the
+        # ring-based parent path.
+        self._alloc_kv_score_buffer(
+            dtype=dtype, device=device, enable_memory_saver=enable_memory_saver
+        )
+
+        # Block 0 = kernel skip-sentinel: kv zeroed, score -inf (softmax → 0).
+        # The free list excludes it; only stale state_block_table entries land here.
+        self.kv_score_buffer.kv[:page_size].zero_()
+        self.kv_score_buffer.score[:page_size].fill_(float("-inf"))
+
+
+class NPUDeepSeekV4IndexerPool(DeepSeekV4IndexerPool):
+    """NPU c4-indexer pool. Keeps the base packed CUDA buffer (read by
+    get_contiguous_buf_infos / NSA) and ADDS dedicated int8 K + float16 scale
+    buffers in PA_ND layout at the global ``kernel_page_size``, written by
+    ``torch_npu.npu_scatter_nd_update_`` and read by
+    ``torch.ops.custom.npu_quant_lightning_indexer``.
+    """
+
+    def __init__(self, *args, kernel_page_size: int, **kwargs):
+        # Set before super().__init__ — it calls _create_buffer().
+        self._kernel_page_size = kernel_page_size
+        super().__init__(*args, **kwargs)
+
+    def _create_buffer(self):
+        # Base allocates the packed CUDA index_k_with_scale_buffer (kept for
+        # get_contiguous_buf_infos / NSA compat); then add the NPU buffers.
+        super()._create_buffer()
+        kp = self._kernel_page_size
+        npu_num_pages = (self.size + kp + 1) // kp
+        with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
+            self.index_k_buffer = [
+                torch.zeros(
+                    npu_num_pages,
+                    kp,
+                    1,
+                    self.index_head_dim,
+                    dtype=torch.int8,
+                    device=self.device,
+                )
+                for _ in range(self.layer_num)
+            ]
+            self.index_scale_buffer = [
+                torch.zeros(
+                    npu_num_pages,
+                    kp,
+                    1,
+                    1,
+                    dtype=torch.float16,
+                    device=self.device,
+                )
+                for _ in range(self.layer_num)
+            ]
+
+    @property
+    def has_npu_storage(self) -> bool:
+        return True
+
+    def get_index_k(self, layer_id: int) -> torch.Tensor:
+        return self.index_k_buffer[layer_id]
+
+    def get_index_scale(self, layer_id: int) -> torch.Tensor:
+        return self.index_scale_buffer[layer_id]
+
+    def set_index_k_scale(
+        self,
+        layer_id: int,
+        loc: torch.Tensor,
+        index_k: torch.Tensor,
+        index_k_scale: Optional[torch.Tensor],
+    ) -> None:
+        # int8 K + fp16 scale come from _compressor_epilog_npu's npu_dynamic_quant
+        # output (index_k: int8 [T, D], index_k_scale: fp16 [T, 1]).
+        d = self.index_head_dim
+        loc_long = loc.view(-1, 1).long()
+        torch_npu.npu_scatter_nd_update_(
+            self.index_k_buffer[layer_id].view(-1, 1, d),
+            loc_long,
+            index_k.to(torch.int8).view(-1, 1, d),
+        )
+        if index_k_scale is not None:
+            torch_npu.npu_scatter_nd_update_(
+                self.index_scale_buffer[layer_id].view(-1, 1, 1),
+                loc_long,
+                index_k_scale.to(torch.float16).view(-1, 1, 1),
+            )
+
+
+class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool):
+    """NPU-only DSV4 KV pool with paged compress-state buffers.
+
+    The full / SWA / c4 / c128 KV pools use the NPU bf16 PA_ND layout
+    (:class:`NPUDeepSeekV4SingleKVPool`); the compress-state pool is paged
+    (:class:`NPUCompressStatePool`) rather than ring-buffered; and the indexer
+    pool adds dedicated int8 K + fp16 scale buffers
+    (:class:`NPUDeepSeekV4IndexerPool`). The generic-accessor / port-hook
+    methods at the bottom of this class are the NPU equivalents of the CUDA
+    DSV4 store-cache chain — kept here, not in the community base, which raises
+    ``NotImplementedError`` for them (CUDA goes through the radix / store_cache
+    accessors instead).
+    """
+
+    def _make_kv_pool(
+        self,
+        *,
+        size: int,
+        page_size: int,
+        dtype: torch.dtype,
+        layer_num: int,
+        device: str,
+        enable_memory_saver: bool,
+        global_page_size: int,
+        cls: type = DeepSeekV4SingleKVPool,
+    ) -> NPUDeepSeekV4SingleKVPool:
+        # NPU does not use the HiSparse c4 device pool; fail loud if someone
+        # enables it so the silent layout mismatch surfaces at init.
+        assert cls is DeepSeekV4SingleKVPool, (
+            "enable_hisparse is not supported on the NPU DSV4 KV pool "
+            f"(got c4 pool class {cls.__name__})."
+        )
+        return NPUDeepSeekV4SingleKVPool(
+            size,
+            page_size,
+            dtype,
+            self.qk_nope_head_dim,
+            self.qk_rope_head_dim,
+            layer_num,
+            device,
+            enable_memory_saver,
+            kernel_page_size=global_page_size,
+        )
+
+    def _get_state_pool(self, layer_id: int, from_indexer: bool) -> CompressStatePool:
+        """Select this layer's attention vs c4-indexer compress-state pool.
+        Wraps the community getters so the NPU port hooks below don't index the
+        pool lists directly."""
+        if from_indexer:
+            return self.get_indexer_compress_states(layer_id)
+        return self.get_attention_compress_states(layer_id)
+
+    def _make_attn_state_pool(
+        self, ratio: int, enable_memory_saver: bool
+    ) -> NPUCompressStatePool:
+        # ONLINE_C128 (CUDA-only) collapses the c128 ring to size 1; the NPU fused
+        # compressor has no online mode, so assert the config mismatch early.
+        assert not (ratio == 128 and ONLINE_C128), (
+            "SGLANG_OPT_USE_ONLINE_COMPRESS is incompatible with the "
+            "NPU fused compressor (no online mode in the kernel)."
+        )
+        return NPUCompressStatePool(
+            size=self._state_pool_size(ratio),
+            overlap=ratio == 4,
+            head_dim=self.qk_nope_head_dim + self.qk_rope_head_dim,
+            dtype=self.c4_state_dtype if ratio == 4 else self.c128_state_dtype,
+            device=self.device,
+            enable_memory_saver=enable_memory_saver,
+            ratio=ratio,
+            page_size=self.swa_page_size,
+        )
+
+    def _make_indexer_state_pool(
+        self, ratio: int, enable_memory_saver: bool
+    ) -> NPUCompressStatePool:
+        # c4 indexer shares the c4 state pool size budget but has its own
+        # slot_dim (indexer_head_dim vs attention head_dim).
+        return NPUCompressStatePool(
+            size=self.c4_state_pool_size,
+            overlap=ratio == 4,
+            head_dim=self.indexer_head_dim,
+            device=self.device,
+            dtype=self.c4_state_dtype,
+            enable_memory_saver=enable_memory_saver,
+            ratio=ratio,
+            page_size=self.swa_page_size,
+        )
+
+    def clear_unaccepted_c128_draft_states(
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        accept_lens: torch.Tensor,
+        num_draft_tokens: int,
+    ) -> None:
+        pass
+
+    def _make_indexer_pool(
+        self,
+        size: int,
+        page_size: int,
+        dtype: torch.dtype,
+        index_head_dim: int,
+        layer_num: int,
+        device: str,
+        enable_memory_saver: bool,
+    ) -> NPUDeepSeekV4IndexerPool:
+        # NPU dedicated int8 K + fp16 scale buffers use the GLOBAL page_size
+        # (= self.page_size) as kernel_page_size, matching ori_kv for the kernel.
+        return NPUDeepSeekV4IndexerPool(
+            size,
+            page_size,
+            dtype,
+            index_head_dim,
+            layer_num,
+            device,
+            enable_memory_saver,
+            kernel_page_size=self.page_size,
+        )
+
+    def get_state_cache(self, layer_id: int, from_indexer: bool) -> torch.Tensor:
+        """fp32 ``[block_num, page_size, 2*coff*D]`` view of this layer's
+        kv+score buffer — the fused compressor op
+        (``torch.ops.custom.compressor``)'s ``state_cache`` argument."""
+        return self._get_state_pool(layer_id, from_indexer).state_cache_3d
+
+    # ------------------------------------------------------------------
+    # Generic KV accessors (community base raises NotImplementedError; CUDA uses
+    # store_cache). AscendAttnBackend reads KV through these, routed to the right
+    # sub-pool by compression ratio.
+    # ------------------------------------------------------------------
+
+    def get_key_buffer(self, layer_id: int) -> torch.Tensor:
+        item = self.layer_mapping[layer_id]
+        ratio = item.compress_ratio
+        if ratio == 0:
+            return self.swa_kv_pool.kv_buffer[item.compress_layer_id]
+        if ratio == 4:
+            return self.c4_kv_pool.kv_buffer[item.compress_layer_id]
+        if ratio == 128:
+            return self.c128_kv_pool.kv_buffer[item.compress_layer_id]
+        raise ValueError(f"unsupported compress_ratio={ratio} for get_key_buffer")
+
+    def get_value_buffer(self, layer_id: int) -> torch.Tensor:
+        # V4 uses MQA / latent attention — the K buffer doubles as V.
+        return self.get_key_buffer(layer_id)
+
+    def get_kv_buffer(self, layer_id: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        buf = self.get_key_buffer(layer_id)
+        return buf, buf
+
+    def get_swa_buffer(
+        self, layer_id: int, loc: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """Return the SWA layer's KV cache in PA_ND layout
+        (num_pages, page_size, num_kv_heads=1, dim). When ``loc`` is given,
+        flatten across (num_pages, page_size) and gather the matching tokens —
+        shape becomes (num_tokens, 1, dim).
+        """
+        # Index by RAW layer_id, not compress_layer_id (a per-bucket counter that
+        # would collide across ratios). swa_kv_pool is sized layer_num=total_layers.
+        kv = self.swa_kv_pool.kv_buffer[layer_id]
+        if loc is not None:
+            kv = kv.flatten(0, 1)[loc]
+        return kv
+
+    def get_compress_buffer(
+        self,
+        layer_id: int,
+        from_indexer: bool = False,
+        loc: Optional[torch.Tensor] = None,
+    ) -> Optional[torch.Tensor]:
+        """Return the compressed KV buffer for a c4 / c128 layer.
+
+        Routes to c4 / c128 kv_pool by layer compression ratio. Returns
+        ``None`` for ratio == 0 (no compress KV exists). The
+        from_indexer=True branch returns the dedicated int8 K buffer that
+        ``torch.ops.custom.npu_quant_lightning_indexer`` consumes.
+        """
+        item = self.layer_mapping[layer_id]
+        if item.compress_ratio == 4:
+            if from_indexer:
+                kv = self.c4_indexer_kv_pool.get_index_k(item.compress_layer_id)
+            else:
+                kv = self.c4_kv_pool.kv_buffer[item.compress_layer_id]
+        elif item.compress_ratio == 128:
+            assert not from_indexer, "c128 has no indexer pool"
+            kv = self.c128_kv_pool.kv_buffer[item.compress_layer_id]
+        else:
+            return None
+        if loc is not None:
+            kv = kv.flatten(0, 1)[loc]
+        return kv
+
+    def set_swa_buffer(
+        self,
+        layer_id: int,
+        loc: torch.Tensor,
+        cache: torch.Tensor,
+    ) -> None:
+        """Write ``cache`` into the SWA pool at flat token positions ``loc``.
+
+        ``cache`` shape: (num_tokens, num_kv_heads=1, dim). The buffer view is
+        (num_pages, page_size, 1, dim) so we flatten the first two dims and
+        index_put.
+        """
+        # Index by raw layer_id (see get_swa_buffer) to avoid bucket collision.
+        buf = self.swa_kv_pool.kv_buffer[layer_id]
+        buf_flat = buf.flatten(0, 1)  # (num_pages * page_size, 1, dim)
+        # Caller (V4 MQALayer) may hand us cache shaped (T, dim); the buffer has
+        # an explicit num_kv_heads=1 axis, so insert it.
+        if cache.ndim == buf_flat.ndim - 1:
+            cache = cache.unsqueeze(1)
+        buf_flat[loc] = cache.to(buf_flat.dtype)
+
+    # ------------------------------------------------------------------
+    # NPU port hooks — used by dsv4/{compressor,indexer}.py forward_npu.
+    # CompressStatePool stores a fused [kv | score] tensor; split is a last-dim slice.
+    # ------------------------------------------------------------------
+
+    def set_state_buffer(
+        self,
+        layer_id: int,
+        loc: torch.Tensor,
+        kv: torch.Tensor,
+        score: torch.Tensor,
+        from_indexer: bool,
+    ) -> None:
+        # KVAndScore.kv_score is [..., 2*coff*head_dim] = [kv | score].
+        kv_score = self._get_state_pool(layer_id, from_indexer).kv_score_buffer.kv_score
+        last_dim = kv_score.shape[-1]
+        half = last_dim // 2
+        kv_view = kv.reshape(-1, half).to(kv_score.dtype)
+        score_view = score.reshape(-1, half).to(kv_score.dtype)
+        kv_score[loc, :half] = kv_view
+        kv_score[loc, half:] = score_view
+
+    def get_state_buffer(
+        self,
+        layer_id: int,
+        from_indexer: bool,
+        kv_indices: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        kv_score = self._get_state_pool(layer_id, from_indexer).kv_score_buffer.kv_score
+        if kv_indices is not None:
+            kv_score = kv_score[kv_indices]
+        last_dim = kv_score.shape[-1]
+        half = last_dim // 2
+        kv = kv_score[..., :half].unsqueeze(-2)  # add num_kv_heads=1 axis
+        score = kv_score[..., half:].unsqueeze(-2)
+        return kv, score
+
+    def set_compress_buffer(
+        self,
+        layer_id: int,
+        loc: torch.Tensor,
+        kv: torch.Tensor,
+        kv_scale: Optional[torch.Tensor],
+        from_indexer: bool,
+    ) -> None:
+        # Routes to c4_indexer (from_indexer) / c4_kv (ratio 4) / c128_kv (ratio
+        # 128). NPU bypasses CUDA fused_store_cache with direct bf16 writes.
+        ratio, compress_layer_id, _ = self.layer_mapping[layer_id]
+        device_type = kv.device.type
+        if from_indexer:
+            assert ratio == 4, f"indexer only on c4 layers, got ratio={ratio}"
+            if device_type == "npu":
+                assert (
+                    self.c4_indexer_kv_pool.has_npu_storage
+                ), "NPU index buffers not allocated — pool was init'd on CUDA?"
+                self.c4_indexer_kv_pool.set_index_k_scale(
+                    compress_layer_id, loc, kv, kv_scale
+                )
+                return
+            if kv_scale is None:
+                self.c4_indexer_kv_pool.set_index_fused(compress_layer_id, loc, kv)
+                return
+            self.c4_indexer_kv_pool.set_index_k_scale_buffer(
+                compress_layer_id, loc, kv, kv_scale
+            )
+            return
+        compress_pool = self.c4_kv_pool if ratio == 4 else self.c128_kv_pool
+        if device_type == "npu":
+            # PA_ND layout: kv_buffer[layer_id] shape = (num_pages, page_size,
+            # 1, kv_dim). Flatten (num_pages, page_size) and index by `loc`.
+            buf = compress_pool.kv_buffer[compress_layer_id]
+            buf_flat = buf.flatten(0, 1)
+            kv_view = kv.to(buf_flat.dtype)
+            if kv_view.ndim == buf_flat.ndim - 1:
+                kv_view = kv_view.unsqueeze(1)
+            buf_flat[loc] = kv_view
+            return
+        compress_pool.set_key_buffer_fused(compress_layer_id, loc, kv)
+
+    def get_compress_dequant_scale_buffer(
+        self,
+        layer_id: int,
+        from_indexer: bool,
+    ) -> torch.Tensor:
+        # Returns the float16 dequant scale buffer (NPU indexer pool's dedicated
+        # scale buffer alongside the int8 K buffer).
+        assert from_indexer, "only indexer compress pool has dequant scale"
+        compress_layer_id = self.layer_mapping[layer_id].compress_layer_id
+        return self.c4_indexer_kv_pool.get_index_scale(compress_layer_id)
+
+    def translate_kv_loc_to_compress_state_loc(
+        self,
+        kv_loc: torch.Tensor,
+        compress_ratio: int,
+    ) -> torch.Tensor:
+        # Parent's ring-buffer hash is meaningless under the paged cache_mode=1
+        # contract; returning a stale value would silently corrupt state. Fail loud.
+        raise RuntimeError(
+            "DSV4NPUTokenToKVPool.translate_kv_loc_to_compress_state_loc was "
+            "called, but the NPU fused compressor kernel uses a paged state "
+            "pool (cache_mode=1) and does not support ring-buffer state "
+            "addressing (cache_mode=2 is explicitly unsupported on Atlas A3). "
+            "Callers must consume out_cache_loc_dsv4.out_c{4,128}_state_loc "
+            "from the allocator bundle (set during alloc_extend/alloc_decode) "
+            "and read state_page_table from req_to_token_c{4,128}_state on "
+            "the DSV4NPUReqToTokenPool instead. See "
+            "hardware_backend/npu/dsv4_memory_pool.py for the rationale."
+        )

python/sglang/srt/hardware_backend/npu/dsv4/dsv4_req_to_token_pool.py

@@ -0,0 +1,128 @@
+"""DSV4-NPU per-request mapping pool.
+
+Subclass of ``ReqToTokenPool`` that adds five auxiliary per-request tables
+needed by the DSV4 attention backend:
+
+  * ``req_to_token_swa``        — slot ids in the SWA full-pool view
+  * ``req_to_token_c4``         — slot ids in the c4 compressed-KV pool
+  * ``req_to_token_c128``       — slot ids in the c128 compressed-KV pool
+  * ``req_to_token_c4_state``   — c4 state-pool slot ids, 1 per raw token
+  * ``req_to_token_c128_state`` — c128 state-pool slot ids, 1 per raw token
+
+Compressed KV pools store 1 slot per ``ratio`` raw tokens, so their per-req
+table column count is ``max_context_len // ratio``. swa mirrors the raw
+token count. Elements are token-level slot ids; the attention backend
+converts to page ids via ``// page_size`` when constructing PA_ND block
+tables.
+
+The c4/c128 STATE pools also have per-req tables here: the NPU fused
+compressor uses a paged state pool (``cache_mode=1``), so each raw token's
+state slot id is recorded (1 column per raw token) and the backend builds
+``state_block_table = req_to_token_c{N}_state[req, ::page_size] // page_size``
+to feed the kernel. (The base class' ``translate_kv_loc_to_compress_state_loc``
+ring-hash is the CUDA-only path; it is disabled on NPU.)
+
+Memory cost example (size=64, max_context_len=32K): swa 8MB + c4 2MB +
+c128 64KB ≈ 10MB extra on top of the base req_to_token (8MB).
+
+The tables are populated by the ``dsv4_common_hooks`` writers (driven from
+``mem_cache/common.py``) immediately after a successful alloc_extend /
+alloc_decode, using the per-pool slot indices returned in ``DSV4OutCacheLoc``.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
+
+
+class DSV4NPUReqToTokenPool(ReqToTokenPool):
+    """ReqToTokenPool extended with DSV4 SWA + c4/c128 per-req tables.
+
+    Drop-in replacement for ReqToTokenPool when the model is DeepSeek-V4 on
+    NPU. Selected by ``model_runner_kv_cache_mixin`` based on model arch +
+    device. Non-DSV4 and non-NPU paths continue to use the base class.
+
+    The auxiliary tables are intentionally NOT zeroed on ``clear()``: they are
+    indexed only by active rows (via req_pool_idx) and only each row's
+    ``[:seq_len]`` prefix is read, so stale entries past kv_committed_len are
+    unreachable by the attention metadata builder.
+    """
+
+    def __init__(
+        self,
+        size: int,
+        max_context_len: int,
+        device: str,
+        enable_memory_saver: bool,
+    ):
+        super().__init__(size, max_context_len, device, enable_memory_saver)
+
+        memory_saver_adapter = TorchMemorySaverAdapter.create(
+            enable=enable_memory_saver
+        )
+
+        # Back-ref to DSV4NPUTokenToKVPoolAllocator, wired via
+        # register_dsv4_allocator after both exist, so free(req) can release
+        # c4/c128 pages. None at construction so base clear() runs safely.
+        self._dsv4_allocator = None
+
+        # (name, columns). swa + state tables: 1 slot per raw token; c4/c128:
+        # 1 slot per `ratio` raw tokens. Init zero so unallocated columns map to
+        # block 0 (kernel skip sentinel cleared by NPUCompressStatePool).
+        with memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
+            for name, cols in (
+                ("req_to_token_swa", max_context_len),
+                ("req_to_token_c4", max(1, max_context_len // 4)),
+                ("req_to_token_c128", max(1, max_context_len // 128)),
+                ("req_to_token_c4_state", max_context_len),
+                ("req_to_token_c128_state", max_context_len),
+            ):
+                setattr(
+                    self,
+                    name,
+                    torch.zeros(
+                        (self._alloc_size, cols),
+                        dtype=torch.int32,
+                        device=device,
+                    ),
+                )
+
+    # ------------------------------------------------------------------
+    # Per-pool write helpers, called by mem_cache/common.py after alloc, using
+    # slot indices from DSV4OutCacheLoc. Args: (req_pool_idx, token_offset), slot.
+    # ------------------------------------------------------------------
+
+    def write_swa(self, indices, values: torch.Tensor) -> None:
+        self.req_to_token_swa[indices] = values
+
+    def write_c4(self, indices, values: torch.Tensor) -> None:
+        self.req_to_token_c4[indices] = values
+
+    def write_c128(self, indices, values: torch.Tensor) -> None:
+        self.req_to_token_c128[indices] = values
+
+    def write_c4_state(self, indices, values: torch.Tensor) -> None:
+        self.req_to_token_c4_state[indices] = values
+
+    def write_c128_state(self, indices, values: torch.Tensor) -> None:
+        self.req_to_token_c128_state[indices] = values
+
+    def register_dsv4_allocator(self, allocator) -> None:
+        """Wire the DSV4NPUTokenToKVPoolAllocator ref so ``free(req)`` can
+        release c4/c128 pool pages alongside the req_pool_idx slot. This is a
+        one-way ref (pool -> allocator). The reverse direction (the allocator
+        reading these per-req tables for its c-pool / state last_loc lookup) is
+        no longer a stored back-ref: mem_cache/common.py passes this pool into
+        ``alloc_extend`` / ``alloc_decode`` per call instead."""
+        self._dsv4_allocator = allocator
+
+    def free(self, req):
+        # Trigger c4/c128 free via the allocator's unified free path. May be None
+        # between __init__ and register_dsv4_allocator — defensive None check.
+        if self._dsv4_allocator is not None:
+            self._dsv4_allocator.free(req=req, req_to_token_pool=self)
+        super().free(req)

python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py

@@ -0,0 +1,49 @@
+# Copyright 2024-2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.configs.model_config import is_deepseek_dsa
+from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
+    EAGLEDraftExtendCudaGraphRunner,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker
+
+
+class EAGLEDraftExtendNpuGraphRunner(EAGLEDraftExtendCudaGraphRunner):
+    def __init__(self, eagle_worker: EagleDraftWorker):
+        super().__init__(eagle_worker)
+
+    def _cache_loc_dtype(self):
+        return torch.int32
+
+    def _replay_graph(self, shape_key, forward_batch):
+        if not is_deepseek_dsa(self.model_runner.model_config.hf_config):
+            seq_lens = forward_batch.seq_lens_cpu.tolist() + [0] * (
+                self.bs - self.raw_bs
+            )
+            return self.backend.replay_with_input_update(
+                shape_key,
+                seq_lens=seq_lens,
+                attr_name="actual_seq_lengths_kv",
+                attr_type=[],
+            )
+        else:
+            return self.backend.replay(shape_key, forward_batch)

python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py

@@ -0,0 +1,69 @@
+# Copyright 2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Dict, Union
+
+import torch
+
+from sglang.srt.configs.model_config import AttentionArch, is_deepseek_dsa
+from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
+    EAGLEDraftCudaGraphRunner,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker
+
+
+class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner):
+    def __init__(self, eagle_worker: EagleDraftWorker):
+        self._init_arch_map()
+        super().__init__(eagle_worker)
+
+    def _init_arch_map(self):
+        self.attr_name: Dict[str, str] = {
+            AttentionArch.MLA: "actual_seq_lengths_kv",
+            AttentionArch.MHA: "context_lens",
+        }
+        self.attr_type: Dict[str, Union[list, torch.Tensor]] = {
+            AttentionArch.MLA: [],
+            AttentionArch.MHA: torch.Tensor(),
+        }
+
+    def _cache_loc_dtype(self):
+        return torch.int32
+
+    def _get_update_attr_name(self):
+        return self.attr_name[AttentionArch.MLA]
+
+    def _get_update_attr_type(self):
+        return self.attr_type[AttentionArch.MLA]
+
+    def _replay_graph(self, shape_key, forward_batch):
+        if not is_deepseek_dsa(self.model_runner.model_config.hf_config):
+            seq_lens_for_each_draft_step = []
+            for speculative_step_id in range(self.speculative_num_steps - 1):
+                seq_lens_cpu = (
+                    forward_batch.seq_lens_cpu[: self.raw_bs] + speculative_step_id + 1
+                )
+                seq_lens = seq_lens_cpu.tolist() + [0] * (self.bs - self.raw_bs)
+                seq_lens_for_each_draft_step.append(seq_lens)
+            attr_name = self._get_update_attr_name()
+            cpu_update_input = [{attr_name: sl} for sl in seq_lens_for_each_draft_step]
+            return self.backend.replay_with_input_update(
+                shape_key, seq_lens=None, cpu_update_input=cpu_update_input
+            )
+        else:
+            return self.backend.replay(shape_key, forward_batch)

python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py

@@ -0,0 +1,54 @@
+# Copyright 2024-2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
+from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import (
+    MultiLayerEagleDraftExtendCudaGraphRunner,
+    MultiLayerEagleMultiStepDraftExtendCudaGraphRunner,
+)
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    pass
+
+
+class MultiLayerEagleDraftExtendNpuGraphRunner(
+    MultiLayerEagleDraftExtendCudaGraphRunner
+):
+    def _replay_graph(self, shape_key, forward_batch):
+        seq_lens = self.buffers.seq_lens_cpu[: self.raw_bs].tolist() + [0] * (
+            self.bs - self.raw_bs
+        )
+        return self.backend.replay_with_input_update(
+            shape_key,
+            seq_lens=seq_lens,
+            attr_name="actual_seq_kvlen",
+            attr_type=[],
+        )
+
+
+class MultiLayerEagleMultiStepDraftExtendNpuGraphRunner(
+    MultiLayerEagleMultiStepDraftExtendCudaGraphRunner
+):
+    def _create_runner(self, step: int) -> MultiLayerEagleDraftExtendNpuGraphRunner:
+        return MultiLayerEagleDraftExtendNpuGraphRunner(self.eagle_worker, step)
+
+    def _cuda_graph_disabled(self) -> bool:
+        return cuda_graph_fully_disabled()

python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py

@@ -0,0 +1,177 @@
+"""NPUCudaGraphBackend — Ascend NPU full-graph capture (torch.npu.NPUGraph).
+
+Mirrors FullCudaGraphBackend with two differences:
+  - Captures via torch.npu.graph(...) into torch.npu.NPUGraph.
+  - replay_with_input_update(shape_key, seq_lens, attr_name) rebinds
+    the recorded graph's input bindings for variable seq_lens at replay
+    time (NPU's NPUGraph.update(...) API).
+
+torch.npu is imported lazily inside methods so the module loads on
+non-NPU hosts.
+"""
+
+from __future__ import annotations
+
+import threading
+from contextlib import AbstractContextManager, contextmanager
+from functools import partial
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+
+import numpy as np
+import torch
+
+from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    set_graph_pool_id,
+)
+from sglang.srt.model_executor.runner.shape_key import ShapeKey
+from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
+    BaseCudaGraphBackend,
+)
+from sglang.srt.utils import empty_context, get_bool_env_var
+from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
+        BaseCudaGraphRunner,
+    )
+
+
+class NPUCudaGraphBackend(BaseCudaGraphBackend):
+    """One torch.npu.NPUGraph per shape; attention metadata captured
+    inside the graph. replay_with_input_update substitutes fresh
+    seq_lens without re-recording."""
+
+    def __init__(
+        self,
+        cuda_graph_runner: BaseCudaGraphRunner,
+        *,
+        enable_memory_saver: bool = False,
+    ) -> None:
+        self._graphs: Dict[Any, Any] = {}
+        self._outputs: Dict[Any, Any] = {}
+        self._pool = None
+        self._device_module = cuda_graph_runner.device_module
+        self._tp_group = cuda_graph_runner.model_runner.tp_group
+        self._capture_stream = None
+        self._memory_saver_adapter: Optional[Any] = TorchMemorySaverAdapter.create(
+            enable=enable_memory_saver
+            and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH")
+        )
+        self._enable_torch_compile = getattr(
+            cuda_graph_runner, "enable_torch_compile", False
+        )
+
+    @contextmanager
+    def capture_session(self, stream):
+        if self._pool is None:
+            self._pool = self._device_module.graph_pool_handle()
+        set_graph_pool_id(self._pool)
+        self._capture_stream = stream
+        try:
+            yield
+        finally:
+            self._capture_stream = None
+
+    def capture_one(
+        self,
+        shape_key: ShapeKey,
+        forward_fn: Callable[[], Any],
+        dummies: Optional[Any] = None,
+        post_warmup_hook: Optional[Callable[[], None]] = None,
+    ) -> None:
+        import torch_npu  # noqa: F401  (verifies NPU availability)
+
+        # Two warmups so kernels are loaded and one-time setup is paid before capture.
+        # post_warmup_hook lets the attention backend reset state that warmup mutated.
+        for _ in range(2):
+            self._device_module.synchronize()
+            self._tp_group.barrier()
+            forward_fn()
+            if post_warmup_hook is not None:
+                post_warmup_hook()
+
+        graph = torch.npu.NPUGraph()
+
+        if self._enable_torch_compile:
+            skip_guard_context = torch.compiler.set_stance(skip_guard_eval_unsafe=True)
+        else:
+            skip_guard_context = empty_context()
+
+        graph_ctx: Callable[..., AbstractContextManager]
+        if (
+            self._memory_saver_adapter is not None
+            and self._memory_saver_adapter.enabled
+        ):
+            graph_ctx = partial(
+                self._memory_saver_adapter.cuda_graph,
+                tag=GPU_MEMORY_TYPE_CUDA_GRAPH,
+            )
+        else:
+            graph_ctx = torch.npu.graph
+
+        with skip_guard_context, graph_ctx(
+            graph,
+            pool=self._pool,
+            stream=self._capture_stream,
+            auto_dispatch_capture=True,
+        ):
+            out = forward_fn()
+
+        self._graphs[shape_key] = graph
+        self._outputs[shape_key] = out
+
+    def can_run(self, forward_batch: ForwardBatch, shape_key: ShapeKey) -> bool:
+        return shape_key in self._graphs
+
+    @contextmanager
+    def replay_session(self):
+        yield
+
+    def replay(
+        self,
+        shape_key: ShapeKey,
+        static_forward_batch: ForwardBatch,
+        **kwargs,
+    ) -> Any:
+        self._graphs[shape_key].replay()
+        return self._outputs[shape_key]
+
+    def replay_with_input_update(
+        self,
+        shape_key: ShapeKey,
+        seq_lens: Any,
+        attr_name: str = None,
+        attr_type: Any = None,
+        cpu_update_input: list = None,
+    ) -> Any:
+        """Rebind seq_lens on the recorded NPU graph in a background
+        thread, then replay. Used when the model is not deepseek-nsa.
+
+        Two calling conventions:
+        1. (legacy) seq_lens + attr_name + attr_type:
+           Constructs cpu_update_input=[{attr_name: seq_lens}] internally.
+        2. cpu_update_input: A list of {attr_name: seq_lens} dicts,
+           one per speculative step.  Used by EAGLE draft runners.
+        """
+        if cpu_update_input is None:
+            if isinstance(attr_type, torch.Tensor):
+                seq_lens = torch.from_numpy(np.array(seq_lens).astype(np.int32))
+            cpu_update_input = [{attr_name: seq_lens}]
+
+        graph = self._graphs[shape_key]
+
+        def _update():
+            graph.update(cpu_update_input=cpu_update_input)
+
+        thread = threading.Thread(target=_update)
+        thread.start()
+        graph.replay()
+        thread.join()
+        return self._outputs[shape_key]
+
+    def cleanup(self) -> None:
+        self._graphs.clear()
+        self._outputs.clear()
+        self._pool = None

python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py

@@ -0,0 +1,284 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Run the model with NPU graph and torch.compile.
+
+NPUGraphRunner is a thin subclass of DecodeCudaGraphRunner: the
+factory returns NPUCudaGraphBackend for NPU devices, so all
+capture/replay mechanics live in the backend. This class adds:
+  - NPU-specific patch_model monkey-patch for the decode-Full +
+    torch.compile path.
+  - Profile context override (NPU profiler emits to disk, not in-mem).
+  - Replay override that issues an async NPUGraph.update for
+    seq_lens before replay (skipped for deepseek-nsa).
+  - Smaller cache_loc dtype (int32 instead of int64).
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from contextlib import contextmanager
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, Optional, Union
+
+import numpy as np
+import torch
+
+from sglang.srt.configs.model_config import (
+    AttentionArch,
+    is_deepseek_dsa,
+    is_deepseek_v4,
+)
+from sglang.srt.distributed.parallel_state import GroupCoordinator
+from sglang.srt.environ import envs
+from sglang.srt.model_executor.runner import DecodeCudaGraphRunner
+from sglang.srt.utils import (
+    empty_context,
+    get_bool_env_var,
+    get_compiler_backend,
+    is_npu,
+)
+
+is_npu = is_npu()
+
+if is_npu:
+    import torch_npu
+    from torch_npu.profiler import ProfilerActivity, profile
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+
+
+@contextmanager
+def patch_model_npu(
+    model: torch.nn.Module,
+    enable_compile: bool,
+    num_tokens: int,
+    tp_group: GroupCoordinator,
+):
+    if enable_compile:
+        backend = get_compiler_backend("npugraph_ex")
+        yield torch.compile(
+            torch.no_grad()(model.forward),
+            fullgraph=True,
+            dynamic=False,
+            backend=backend,
+        )
+    else:
+        yield model.forward
+
+
+class NPUGraphRunner(DecodeCudaGraphRunner):
+    """A NPUGraphRunner runs the forward pass of a model with NPU graph and torch.compile."""
+
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        *,
+        attn_backend=None,
+        speculative_num_steps: Optional[int] = None,
+        speculative_num_draft_tokens: Optional[int] = None,
+    ):
+        # NPU patch_model override: monkey-patch torch_compile_decoration's
+        # patch_model with the NPU-specific version.
+        from sglang.srt.compilation import torch_compile_decoration
+
+        torch_compile_decoration.patch_model = patch_model_npu
+        super().__init__(
+            model_runner,
+            attn_backend=attn_backend,
+            speculative_num_steps=speculative_num_steps,
+            speculative_num_draft_tokens=speculative_num_draft_tokens,
+        )
+        self.update_attr_name = None
+        self.update_attr_type = None
+        self.model_runner = model_runner
+        self._init_arch_map()
+        self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False")
+        self.if_use_v2 = any(
+            arch
+            in ("MiMoV2ForCausalLM", "MiMoV2FlashForCausalLM", "Step3p5ForCausalLM")
+            for arch in (model_runner.model_config.hf_config.architectures or [])
+        )
+
+    def _init_arch_map(self):
+        if self.is_dllm:
+            self.attr_name: Dict[str, str] = {
+                AttentionArch.MLA: "actual_seq_lengths_kv",
+                AttentionArch.MHA: "actual_seq_lengths_kv",
+                "TARGET_VERIFY": "actual_seq_kvlen",
+            }
+        else:
+            self.attr_name: Dict[str, str] = {
+                AttentionArch.MLA: "actual_seq_lengths_kv",
+                AttentionArch.MHA: "context_lens",
+                "TARGET_VERIFY": "actual_seq_kvlen",
+            }
+        self.attr_type: Dict[str, Union[list, torch.Tensor]] = {
+            AttentionArch.MLA: [],
+            AttentionArch.MHA: torch.Tensor(),
+            "TARGET_VERIFY": [],
+        }
+
+    def _create_device_graph(self):
+        return torch.npu.NPUGraph()
+
+    def _capture_graph(self, graph, pool, stream, run_once_fn):
+        if self.enable_torch_compile:
+            skip_guard_context = torch.compiler.set_stance(skip_guard_eval_unsafe=True)
+        else:
+            skip_guard_context = empty_context()
+
+        with (
+            skip_guard_context,
+            torch.npu.graph(
+                graph,
+                pool=pool,
+                stream=stream,
+                auto_dispatch_capture=True,
+            ),
+        ):
+            out = run_once_fn()
+        return out
+
+    def _get_update_attr_name(self):
+        if self.if_use_v2:
+            return self.attr_name["TARGET_VERIFY"]
+        return self.attr_name[AttentionArch.MLA]
+
+    def _get_update_attr_type(self):
+        if self.if_use_v2:
+            return self.attr_type["TARGET_VERIFY"]
+        return self.attr_type[AttentionArch.MLA]
+
+    def _update_inputs(self, seq_lens):
+        if isinstance(self.update_attr_type, torch.Tensor):
+            seq_lens = torch.from_numpy(np.array(seq_lens).astype(np.int32))
+
+        self.graphs[self.bs].update(
+            cpu_update_input=[{self.update_attr_name: seq_lens}]
+        )
+
+    def _cache_loc_dtype(self):
+        return torch.int32
+
+    def _init_profile_context_and_memory_record(self):
+        output_dir = os.path.join(
+            os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp"), "graph_capture_profile"
+        )
+        if not Path(output_dir).exists():
+            Path(output_dir).mkdir(parents=True, exist_ok=True)
+        logger.info(
+            f"Profiling starts for graph capture for NPU. Traces will be saved to: {output_dir}"
+        )
+        experimental_config = torch_npu.profiler._ExperimentalConfig(
+            export_type=[torch_npu.profiler.ExportType.Text],
+            profiler_level=torch_npu.profiler.ProfilerLevel.Level1,
+        )
+        profile_context = profile(
+            activities=[ProfilerActivity.CPU, ProfilerActivity.NPU],
+            record_shapes=True,
+            profile_memory=True,
+            on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(
+                output_dir, async_mode=True
+            ),
+            experimental_config=experimental_config,
+        )
+        return profile_context
+
+    def _post_process_after_profile(self, prof_context):
+        # for NPU, profile data will be saved to disk for further analysis.
+        pass
+
+    def execute(
+        self,
+        forward_batch: ForwardBatch,
+        pp_proxy_tensors: Optional[PPProxyTensors] = None,
+    ) -> Union[LogitsProcessorOutput, PPProxyTensors]:
+        if forward_batch.needs_forward_metadata_init():
+            self.load_batch(forward_batch, pp_proxy_tensors)
+        else:
+            # In speculative decoding, these two fields are still needed.
+            self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids)
+            self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions)
+            if (
+                self.model_runner.spec_algorithm.is_dflash()
+                and self.model_runner.is_draft_worker
+                and forward_batch.input_embeds is not None
+            ):
+                self.buffers.input_embeds[: self.raw_num_token].copy_(
+                    forward_batch.input_embeds
+                )
+            if (
+                envs.SGLANG_ENABLE_OVERLAP_PLAN_STREAM.get()
+                and forward_batch.mrope_positions is not None
+            ):
+                self.buffers.mrope_positions[:, : self.raw_num_token].copy_(
+                    forward_batch.mrope_positions
+                )
+
+        graph_key = self._make_graph_key(self.bs)
+
+        if not (
+            is_deepseek_dsa(self.model_runner.model_config.hf_config)
+            or is_deepseek_v4(self.model_runner.model_config.hf_config)
+        ):
+            if forward_batch.forward_mode.is_target_verify():
+                seq_lens_cpu = forward_batch.seq_lens.cpu() + self.num_tokens_per_bs
+                seq_lens = seq_lens_cpu.tolist() + [0] * (self.bs - self.raw_bs)
+            else:
+                seq_lens = forward_batch.seq_lens.cpu().tolist() + [0] * (
+                    self.bs - self.raw_bs
+                )
+            output = self.backend.replay_with_input_update(
+                graph_key,
+                seq_lens=seq_lens,
+                attr_name=self._get_update_attr_name(),
+                attr_type=self._get_update_attr_type(),
+            )
+        else:
+            output = self.backend.replay(graph_key, forward_batch)
+
+        if isinstance(output, LogitsProcessorOutput):
+            if self.is_dllm:
+                next_token_logits = None
+                full_logits = (
+                    output.full_logits[: self.raw_num_token]
+                    if output.full_logits is not None
+                    else None
+                )
+            else:
+                full_logits = None
+                next_token_logits = (
+                    output.next_token_logits[: self.raw_num_token]
+                    if output.next_token_logits is not None
+                    else None
+                )
+            return LogitsProcessorOutput(
+                next_token_logits=next_token_logits,
+                full_logits=full_logits,
+                hidden_states=(
+                    output.hidden_states[: self.raw_num_token]
+                    if output.hidden_states is not None
+                    else None
+                ),
+            )
+        else:
+            assert isinstance(output, PPProxyTensors)
+            return PPProxyTensors({k: v[: self.bs] for k, v in output.tensors.items()})

python/sglang/srt/hardware_backend/npu/graph_runner/vit_npu_graph_runner.py

@@ -0,0 +1,229 @@
+# Copyright 2023-2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""ViT NPU Graph Runner class."""
+
+from __future__ import annotations
+
+from typing import Dict, Hashable, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch_npu
+
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    set_graph_pool_id,
+)
+from sglang.srt.layers.attention.vision import VisionAttention
+from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner
+from sglang.srt.runtime_context import get_server_args
+
+
+class ViTNpuGraphRunner(ViTCudaGraphRunner):
+    """Generic ViT NPU Graph Runner.
+
+    This runner captures the "blocks + merger + deepstack merger (optional)" part
+    of a vision transformer into a NPU graph and replays it for identical shapes.
+
+    Optional for Qwen3 deepstack:
+      - vit.deepstack_vision_indexes: Sequence[int]
+      - vit.deepstack_merger_list: nn.ModuleList (same length as deepstack_vision_indexes)
+    """
+
+    _graph_memory_pool = None
+
+    def __init__(
+        self,
+        vit: nn.Module,
+    ) -> None:
+        super().__init__(vit)
+        self.device_module = torch.get_device_module(self.device)
+        self.cu_seq_lens: Dict[Hashable, torch.Tensor] = {}
+
+        # rotary position buffers shared across graphs
+        self.sin_cos_ws: Dict[Hashable, Tuple[torch.Tensor, torch.Tensor]] = {}
+
+    @property
+    def device(self) -> torch.device:
+        return self.vit.device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.vit.dtype
+
+    def _create_graph(
+        self,
+        graph_key: int,
+    ):
+
+        graph = torch_npu.npu.NPUGraph()
+        vit = self.vit
+
+        override_backend = get_server_args().mm_attention_backend
+        with torch_npu.npu.graph(graph, pool=ViTNpuGraphRunner._graph_memory_pool):
+            y = None
+            deepstack_outs: List[torch.Tensor] = []
+            deepstack_capture_idx = 0
+
+            for layer_num, blk in enumerate(vit.blocks):
+                if override_backend == "ascend_attn":
+                    cu_seq_lens = self.cu_seq_lens[graph_key]
+                else:
+                    raise RuntimeError("Not supported ViT attention backend")
+
+                if layer_num == 0:
+                    y = blk(
+                        self.block_input[graph_key],
+                        cu_seqlens=cu_seq_lens,
+                        rotary_pos_emb_cos=self.sin_cos_ws[graph_key][0],
+                        rotary_pos_emb_sin=self.sin_cos_ws[graph_key][1],
+                        output_ws=self.block_ws[graph_key],
+                    )
+                else:
+                    y = blk(
+                        y,
+                        cu_seqlens=cu_seq_lens,
+                        rotary_pos_emb_cos=self.sin_cos_ws[graph_key][0],
+                        rotary_pos_emb_sin=self.sin_cos_ws[graph_key][1],
+                        output_ws=self.block_ws[graph_key],
+                    )
+
+                # Optional deepstack support (Qwen3-VL)
+                if (
+                    self._deepstack_visual_indexes
+                    and layer_num in self._deepstack_visual_indexes
+                ):
+                    if self._deepstack_merger_list is None:
+                        raise RuntimeError(
+                            "deepstack_visual_indexes exists but deepstack_merger_list is missing."
+                        )
+                    deepstack_out = self._deepstack_merger_list[deepstack_capture_idx](
+                        y
+                    )
+                    deepstack_outs.append(deepstack_out)
+                    deepstack_capture_idx += 1
+
+            main_out = vit.merger(y)
+
+            if deepstack_outs:
+                self.block_output[graph_key] = torch.cat(
+                    [main_out] + deepstack_outs, dim=1
+                )
+            else:
+                self.block_output[graph_key] = main_out
+
+        self.block_graphs[graph_key] = graph
+
+    def create_graph(
+        self,
+        x_3d: torch.Tensor,  # [S, 1, H]
+        cu_seqlens: torch.Tensor,
+        rotary_pos_emb_cos: Optional[torch.Tensor] = None,
+        rotary_pos_emb_sin: Optional[torch.Tensor] = None,
+    ) -> int:
+        vit = self.vit
+        graph_key = self._get_graph_key(x_3d)
+
+        if graph_key in self.block_graphs:
+            return graph_key
+
+        if ViTNpuGraphRunner._graph_memory_pool is None:
+            ViTNpuGraphRunner._graph_memory_pool = (
+                self.device_module.graph_pool_handle()
+            )
+        # Set graph pool id globally to be able to use symmetric memory
+        set_graph_pool_id(ViTNpuGraphRunner._graph_memory_pool)
+
+        # pre-allocate workspace
+        attn_module: VisionAttention = vit.blocks[0].attn
+        num_heads = attn_module.num_attention_heads_per_partition
+        attn_head_dim = attn_module.head_size
+
+        if graph_key not in self.block_output:
+            self.block_output[graph_key] = x_3d
+            self.block_input[graph_key] = x_3d
+            self.block_ws[graph_key] = torch.empty(
+                graph_key,
+                num_heads,
+                attn_head_dim,
+                device=self.device,
+                dtype=self.dtype,
+            )
+            if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None:
+                self.sin_cos_ws[graph_key] = (rotary_pos_emb_cos, rotary_pos_emb_sin)
+
+        if graph_key not in self.cu_seq_lens:
+            seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+            self.cu_seq_lens[graph_key] = seq_lens.to("cpu").to(torch.int32)
+
+        if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None:
+            self._create_graph(
+                graph_key=graph_key,
+            )
+
+        return graph_key
+
+    def replay(
+        self,
+        graph_key: int,
+        x_3d: torch.Tensor,
+        rotary_pos_emb_cos: Optional[torch.Tensor] = None,
+        rotary_pos_emb_sin: Optional[torch.Tensor] = None,
+        output_indices: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None:
+            # update rotary workspace content
+            self.sin_cos_ws[graph_key][0].copy_(rotary_pos_emb_cos)
+            self.sin_cos_ws[graph_key][1].copy_(rotary_pos_emb_sin)
+
+        # copy input
+        self.block_input[graph_key].copy_(x_3d)
+
+        # replay
+        self.block_graphs[graph_key].replay()
+
+        out = self.block_output[graph_key]
+
+        # Optional output reordering (Qwen2.5-VL window permutation inverse)
+        if output_indices is not None:
+            out = out.index_select(0, output_indices)
+
+        return out
+
+    def run(
+        self,
+        x: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        rotary_pos_emb_cos: Optional[torch.Tensor] = None,
+        rotary_pos_emb_sin: Optional[torch.Tensor] = None,
+        output_indices: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        # x: [seq_len, hidden] -> [S, B=1, H]
+        x_3d = x.unsqueeze(1)
+        graph_key = self._get_graph_key(x_3d)
+        if graph_key not in self.block_graphs:
+            self.create_graph(
+                x_3d=x_3d,
+                cu_seqlens=cu_seqlens,
+                rotary_pos_emb_cos=rotary_pos_emb_cos,
+                rotary_pos_emb_sin=rotary_pos_emb_sin,
+            )
+
+        return self.replay(
+            graph_key=graph_key,
+            x_3d=x_3d,
+            rotary_pos_emb_cos=rotary_pos_emb_cos,
+            rotary_pos_emb_sin=rotary_pos_emb_sin,
+            output_indices=output_indices,
+        )

python/sglang/srt/hardware_backend/npu/memory_pool_npu.py

@@ -0,0 +1,553 @@
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
+from sglang.srt.mem_cache.memory_pool import (
+    MHATokenToKVPool,
+    MLATokenToKVPool,
+    get_tensor_size_bytes,
+    unwrap_write_loc,
+)
+from sglang.srt.utils import get_bool_env_var
+from sglang.srt.utils.common import is_npu
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.radix_attention import RadixAttention
+
+if is_npu():
+    import torch_npu
+
+
+def _init_npu_conv_state(
+    conv_state_in, conv_state_shape, speculative_num_draft_tokens: Optional[int] = None
+):
+    extra_conv_len = 0
+    if speculative_num_draft_tokens is not None:
+        extra_conv_len = speculative_num_draft_tokens - 1
+
+    # conv_state shape (layers, pool_size, conv_wind + draft_step, dim) for conv1d ascendc ops require dim as last dim
+    conv_state = [
+        torch.zeros(
+            size=(
+                conv_state_in.shape[0],
+                conv_state_in.shape[1],
+                conv_shape[1] + extra_conv_len,
+                conv_shape[0],
+            ),
+            dtype=conv_state_in.dtype,
+            device=conv_state_in.device,
+        )
+        for conv_shape in conv_state_shape
+    ]
+    return conv_state
+
+
+class NPUMHATokenToKVPool(MHATokenToKVPool):
+
+    def __init__(
+        self,
+        size: int,
+        page_size: int,
+        dtype: torch.dtype,
+        head_num: int,
+        head_dim: int,
+        layer_num: int,
+        device: str,
+        enable_memory_saver: bool,
+        v_head_dim: Optional[int] = None,
+        swa_head_num: Optional[int] = None,
+        swa_head_dim: Optional[int] = None,
+        swa_v_head_dim: Optional[int] = None,
+        start_layer: Optional[int] = None,
+        end_layer: Optional[int] = None,
+        enable_alt_stream: bool = True,
+        enable_kv_cache_copy: bool = False,
+        **kwargs,
+    ):
+        self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False")
+        super().__init__(
+            size=size,
+            page_size=page_size,
+            dtype=dtype,
+            head_num=head_num,
+            head_dim=head_dim,
+            layer_num=layer_num,
+            device=device,
+            enable_memory_saver=enable_memory_saver,
+            v_head_dim=v_head_dim,
+            swa_head_num=swa_head_num,
+            swa_head_dim=swa_head_dim,
+            swa_v_head_dim=swa_v_head_dim,
+            start_layer=start_layer,
+            end_layer=end_layer,
+            enable_alt_stream=enable_alt_stream,
+            enable_kv_cache_copy=enable_kv_cache_copy,
+            **kwargs,
+        )
+
+    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.k_buffer = torch.zeros(
+                (
+                    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.v_buffer = torch.zeros(
+                (
+                    self.layer_num,
+                    self.size // self.page_size + 1,
+                    self.page_size,
+                    self.head_num,
+                    self.v_head_dim,
+                ),
+                dtype=self.store_dtype,
+                device=self.device,
+            )
+
+            if self.use_fia:
+                # Use per-layer Python lists to avoid torch.compile capturing
+                # the entire multi-layer tensor (OOM during graph capture).
+                # Each layer view: [P*ps, 1, H, D], sharing the contiguous
+                # storage allocated above.
+                self.k_buffer = [
+                    self.k_buffer[i].view(-1, 1, self.head_num, self.head_dim)
+                    for i in range(self.layer_num)
+                ]
+                self.v_buffer = [
+                    self.v_buffer[i].view(-1, 1, self.head_num, self.v_head_dim)
+                    for i in range(self.layer_num)
+                ]
+
+    def _init_kv_copy_and_warmup(self):
+        # implementation relies on self.data_strides / self.data_ptrs, which the
+        # NPU paged buffer layout never builds.
+        self._kv_copy_config = None
+
+    # 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)
+        ]
+        if self.use_fia:
+            kv_item_lens = [
+                self.get_key_buffer(i)[0].nbytes * self.page_size
+                for i in range(self.start_layer, self.start_layer + self.layer_num)
+            ] + [
+                self.get_value_buffer(i)[0].nbytes * self.page_size
+                for i in range(self.start_layer, self.start_layer + self.layer_num)
+            ]
+        else:
+            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_info,
+        cache_k: torch.Tensor,
+        cache_v: torch.Tensor,
+        k_scale: Optional[float] = None,
+        v_scale: Optional[float] = None,
+        layer_id_override: Optional[int] = None,
+        dcp_kv_mask: Optional[torch.Tensor] = None,
+    ):
+        loc, _, _ = unwrap_write_loc(loc_info)
+        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)
+
+        if self.use_fia:
+            k_buffer_layer = self.k_buffer[layer_id - self.start_layer]
+            v_buffer_layer = self.v_buffer[layer_id - self.start_layer]
+
+            torch_npu.npu_scatter_nd_update_(
+                k_buffer_layer,
+                loc.view(-1, 1),
+                cache_k.view(-1, 1, self.head_num, self.head_dim),
+            )
+            torch_npu.npu_scatter_nd_update_(
+                v_buffer_layer,
+                loc.view(-1, 1),
+                cache_v.view(-1, 1, self.head_num, self.v_head_dim),
+            )
+        else:
+            loc = loc.to(torch.int32)
+            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.v_head_dim
+                ),
+                slot_indices=loc,
+            )
+
+    def _chunk_copy_npu_to_cpu(self, buf_of_layers, indices):
+        chunk_size = self.cpu_offloading_chunk_size
+        out = []
+        for tensors_per_layer in buf_of_layers:  # [k_buf, v_buf]
+            layer_chunks = []
+            for i in range(0, len(indices), chunk_size):
+                ci = indices[i : i + chunk_size]
+                layer_chunks.append(
+                    [
+                        t[ci].to("cpu", non_blocking=True)
+                        for t in tensors_per_layer
+                        if t is not None
+                    ]
+                )
+            out.append(layer_chunks)
+        return out
+
+    # Parent MHATokenToKVPool.get_cpu_copy / load_cpu_copy use
+    # `self.k_buffer[layer_id][chunk_indices]` which indexes the first dim.
+    # NPUMHATokenToKVPool stores buffers as
+    #   (num_pages, page_size, head_num, head_dim)            # use_fia=False
+    #   (num_pages*page_size, 1, head_num, head_dim)          # use_fia=True
+    def get_cpu_copy(self, indices, mamba_indices=None):
+        torch.npu.synchronize()
+        buf_of_layers = []
+        for local_layer_id in range(self.layer_num):
+            k_layer = self.k_buffer[local_layer_id].view(
+                -1, self.head_num, self.head_dim
+            )
+            v_layer = self.v_buffer[local_layer_id].view(
+                -1, self.head_num, self.head_dim
+            )
+            buf_of_layers.append([k_layer, v_layer])
+        kv_cache_cpu = self._chunk_copy_npu_to_cpu(buf_of_layers, indices)
+        torch.npu.synchronize()
+        return kv_cache_cpu
+
+    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+        torch.npu.synchronize()
+        chunk_size = self.cpu_offloading_chunk_size
+        for local_layer_id in range(self.layer_num):
+            k_layer = self.k_buffer[local_layer_id].view(
+                -1, self.head_num, self.head_dim
+            )
+            v_layer = self.v_buffer[local_layer_id].view(
+                -1, self.head_num, self.head_dim
+            )
+            for i in range(0, len(indices), chunk_size):
+                chunk_indices = indices[i : i + chunk_size]
+                k_cpu, v_cpu = (
+                    kv_cache_cpu[local_layer_id][i // chunk_size][0],
+                    kv_cache_cpu[local_layer_id][i // chunk_size][1],
+                )
+                assert k_cpu.shape[0] == v_cpu.shape[0] == len(chunk_indices)
+                k_layer[chunk_indices] = k_cpu.to(k_layer.device, non_blocking=True)
+                v_layer[chunk_indices] = v_cpu.to(v_layer.device, non_blocking=True)
+        torch.npu.synchronize()
+
+
+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=size,
+            page_size=page_size,
+            dtype=dtype,
+            layer_num=layer_num,
+            device=device,
+            enable_memory_saver=enable_memory_saver,
+            start_layer=start_layer,
+            end_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,
+            )
+            self.index_k_buffer = None
+            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_state_buf_infos(self):
+        if self.index_head_dim is None:
+            return [], [], []
+        data_ptrs = [self.index_k_buffer[i].data_ptr() for i in range(self.layer_num)]
+        data_lens = [self.index_k_buffer[i].nbytes for i in range(self.layer_num)]
+        item_lens = [self.index_k_buffer[i][0].nbytes for i in range(self.layer_num)]
+        return data_ptrs, data_lens, item_lens
+
+    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_info,
+        cache_k: torch.Tensor,
+        cache_v: torch.Tensor,
+    ):
+        loc, _, _ = unwrap_write_loc(loc_info)
+        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),
+        )
+
+    def _chunk_copy_npu_to_cpu(self, buf_of_layers, indices):
+        chunk_size = self.cpu_offloading_chunk_size
+        out = []
+        for tensors_per_layer in buf_of_layers:  # [k_buf, v_buf, ik_buf/None]
+            layer_chunks = []
+            for i in range(0, len(indices), chunk_size):
+                ci = indices[i : i + chunk_size]
+                layer_chunks.append(
+                    [
+                        t[ci].to("cpu", non_blocking=True)
+                        for t in tensors_per_layer
+                        if t is not None
+                    ]
+                )
+            out.append(layer_chunks)
+        return out
+
+    def get_cpu_copy(self, indices, mamba_indices=None):
+        torch.npu.synchronize()
+        buf_of_layers = []
+        has_ik = self.index_head_dim is not None
+        for local_layer_id in range(self.layer_num):
+            k_layer = self.k_buffer[local_layer_id].view(-1, 1, self.kv_lora_rank)
+            v_layer = self.v_buffer[local_layer_id].view(-1, 1, self.qk_rope_head_dim)
+            ik_layer = (
+                self.index_k_buffer[local_layer_id].view(-1, 1, self.index_head_dim)
+                if has_ik
+                else None
+            )
+            buf_of_layers.append([k_layer, v_layer, ik_layer])
+
+        kv_cache_cpu = self._chunk_copy_npu_to_cpu(buf_of_layers, indices)
+        torch.npu.synchronize()
+        return kv_cache_cpu
+
+    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
+        torch.npu.synchronize()
+        chunk_size = self.cpu_offloading_chunk_size
+        has_ik = self.index_head_dim is not None
+        for local_layer_id in range(self.layer_num):
+            k_layer = self.k_buffer[local_layer_id].view(-1, 1, self.kv_lora_rank)
+            v_layer = self.v_buffer[local_layer_id].view(-1, 1, self.qk_rope_head_dim)
+            ik_layer = (
+                self.index_k_buffer[local_layer_id].view(-1, 1, self.index_head_dim)
+                if has_ik
+                else None
+            )
+            for i in range(0, len(indices), chunk_size):
+                chunk_indices = indices[i : i + chunk_size]
+                chunk = kv_cache_cpu[local_layer_id][i // chunk_size]
+                k_cpu, v_cpu = chunk[0], chunk[1]
+                assert k_cpu.shape[0] == len(chunk_indices)
+                k_layer[chunk_indices] = k_cpu.to(k_layer.device, non_blocking=True)
+                v_layer[chunk_indices] = v_cpu.to(v_layer.device, non_blocking=True)
+                if has_ik:
+                    ik_cpu = chunk[2]
+                    ik_layer[chunk_indices] = ik_cpu.to(
+                        ik_layer.device, non_blocking=True
+                    )
+        torch.npu.synchronize()

python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py

@@ -0,0 +1,604 @@
+import re
+from typing import TYPE_CHECKING
+
+import torch
+import torch_npu
+from sgl_kernel_npu.norm.fused_split_qk_norm import fused_split_qk_norm
+
+from sglang.srt.environ import envs
+from sglang.srt.hardware_backend.npu.attention.mla_preprocess import (
+    NPUFusedMLAPreprocess,
+    is_fia_nz,
+    is_mla_preprocess_enabled,
+)
+from sglang.srt.layers.attention.dsa.dsa_indexer import scattered_to_tp_attn_full
+from sglang.srt.layers.attention.dsa.utils import (
+    dsa_use_prefill_cp,
+)
+from sglang.srt.layers.communicator import ScatterMode, get_attn_tp_context
+from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
+
+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
+_use_ag_after_qlora = envs.SGLANG_USE_AG_AFTER_QLORA.get()
+
+
+# region MHA
+def forward_mha_prepare_npu(
+    m: "DeepseekV2AttentionMLA",
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    forward_batch: "ForwardBatch",
+    zero_allocator: "BumpAllocator",
+    layer_scatter_modes,
+):
+    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,
+            )
+        )
+
+        # DSA Indexer: cache quantized keys, auto-skip topk for sequences <= dsa_index_topk
+
+        if m.use_dsa:
+            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)
+            if (
+                _use_ag_after_qlora
+                and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED
+                and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL
+            ):
+                q = scattered_to_tp_attn_full(q, forward_batch)
+                latent_cache = scattered_to_tp_attn_full(latent_cache, forward_batch)
+            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)
+
+    if m.use_deepseek_yarn_rope:
+        B, S = q.shape[0], 1
+        cos, sin = m.rotary_emb.get_cos_sin_cache(
+            positions, hidden_states.dtype, offsets=None
+        )
+        q_pe = torch_npu.npu_interleave_rope(
+            q_pe.reshape(B, -1, S, m.qk_rope_head_dim),
+            cos,
+            sin,
+        )
+        q_pe = q_pe.reshape(B, -1, m.qk_rope_head_dim)
+
+        ckv_cache, k_rope_cache = get_token_to_kv_pool().get_kv_buffer(m.layer_id)
+        _, _, k_pe, kv_a = torch_npu.npu_kv_rmsnorm_rope_cache(
+            latent_cache.view(-1, 1, 1, m.kv_lora_rank + m.qk_rope_head_dim),  # bnsd
+            m.kv_a_layernorm.weight,
+            cos,
+            sin,
+            forward_batch.out_cache_loc.to(torch.int64),
+            k_rope_cache,
+            ckv_cache,
+            k_rope_scale=None,
+            c_kv_scale=None,
+            k_rope_offset=None,
+            c_kv_offset=None,
+            epsilon=m.kv_a_layernorm.variance_epsilon,
+            cache_mode="PA_NZ" if is_fia_nz() else "PA_BNSD",
+            is_output_kv=True,
+        )  # adapter NZ
+
+        k_pe = k_pe.reshape(B, -1, m.qk_rope_head_dim)
+    else:
+        kv_a = m.kv_a_layernorm(kv_a)
+        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)
+        # this is for model kimi-vl-a3B-instruct
+        get_token_to_kv_pool().set_kv_buffer(
+            m, forward_batch.out_cache_loc, kv_a.unsqueeze(1), k_pe
+        )
+
+    q[..., m.qk_nope_head_dim :] = q_pe
+
+    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",
+    layer_scatter_modes,
+):
+    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:
+            qkv_latent = get_attn_tp_context().fetch_qkv_latent()
+            if (
+                _use_ag_after_qlora
+                and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED
+                and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL
+            ):
+                q, latent_cache = 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)
+                q = scattered_to_tp_attn_full(q, forward_batch)
+                latent_cache = scattered_to_tp_attn_full(latent_cache, forward_batch)
+
+                k_nope = m.kv_a_layernorm(k_nope).unsqueeze(1)
+                k_pe = latent_cache[..., m.kv_lora_rank :].unsqueeze(1)
+            else:
+                if qkv_latent.shape[0] < 65536 and not dsa_use_prefill_cp(
+                    forward_batch
+                ):
+                    q, k_nope, k_pe = fused_split_qk_norm(
+                        qkv_latent,
+                        m.q_a_layernorm,
+                        m.kv_a_layernorm,
+                        m.q_lora_rank,
+                        m.kv_lora_rank,
+                        m.qk_rope_head_dim,
+                        eps=m.q_a_layernorm.variance_epsilon,
+                    )
+                else:
+                    q, latent_cache = 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).unsqueeze(1)
+                    k_pe = latent_cache[..., m.kv_lora_rank :].unsqueeze(1)
+
+            # q_lora needed by indexer
+            if m.use_dsa:
+                q_lora = 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]
+            k_nope = latent_cache[..., : m.kv_lora_rank]
+            k_nope = m.kv_a_layernorm(k_nope).unsqueeze(1)
+            k_pe = latent_cache[..., m.kv_lora_rank :].unsqueeze(1)
+
+        q_nope, q_pe = q.split([m.qk_nope_head_dim, m.qk_rope_head_dim], dim=-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)
+
+        if dsa_use_prefill_cp(forward_batch):
+            # 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,
+    )
+
+    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
+
+
+# region DSA
+def forward_dsa_prepare_npu(
+    m: "DeepseekV2AttentionMLA",
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    forward_batch: "ForwardBatch",
+    zero_allocator: "BumpAllocator",
+    layer_scatter_modes,
+    prev_topk_indices: torch.Tensor = None,
+):
+    dynamic_scale = None
+    if is_mla_preprocess_enabled() and forward_batch.forward_mode.is_decode():
+        (
+            q_pe,
+            k_pe,
+            q_nope_out,
+            k_nope,
+            q_lora,
+            forward_batch,
+            zero_allocator,
+            positions,
+            dynamic_scale,
+        ) = npu_mla_preprocess(
+            m,
+            hidden_states,
+            positions,
+            forward_batch,
+            zero_allocator,
+        )
+    else:
+        fused_qkv_a_proj_out = m.fused_qkv_a_proj_with_mqa(hidden_states)[0]
+        if m.rotary_emb.is_neox_style:
+            q, latent_cache = fused_qkv_a_proj_out.split(
+                [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1
+            )
+            # overlap qk norm
+            q = m.q_a_layernorm(q)
+            if (
+                _use_ag_after_qlora
+                and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED
+                and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL
+            ):
+                q = scattered_to_tp_attn_full(q, forward_batch)
+                latent_cache = scattered_to_tp_attn_full(latent_cache, forward_batch)
+            q_lora = q.clone()  # required for topk_indices
+
+            q_event = None
+            if m.alt_stream is not None:
+                m.alt_stream.wait_stream(torch.npu.current_stream())
+                with torch.npu.stream(m.alt_stream):
+                    q = m.q_b_proj(q_lora)[0].view(-1, m.num_local_heads, m.qk_head_dim)
+                    # record q to ensure memory space will not be released
+                    q.record_stream(m.alt_stream)
+                    q_event = m.alt_stream.record_event()
+            else:
+                q = m.q_b_proj(q_lora)[0].view(-1, m.num_local_heads, m.qk_head_dim)
+
+            k_nope, k_pe = latent_cache.unsqueeze(1).split(
+                [m.kv_lora_rank, m.qk_rope_head_dim], dim=-1
+            )
+            k_nope = m.kv_a_layernorm(k_nope)
+            # main stream waits for the completion of the event on the alt stream to ensure data dependency is complete
+            if q_event is not None:
+                torch.npu.current_stream().wait_event(q_event)
+        else:
+            if fused_qkv_a_proj_out.shape[0] < 65535 and not dsa_use_prefill_cp(
+                forward_batch
+            ):
+                q_lora, k_nope, k_pe = fused_split_qk_norm(
+                    fused_qkv_a_proj_out,
+                    m.q_a_layernorm,
+                    m.kv_a_layernorm,
+                    m.q_lora_rank,
+                    m.kv_lora_rank,
+                    m.qk_rope_head_dim,
+                    eps=m.q_a_layernorm.variance_epsilon,
+                )
+            else:
+                q, latent_cache = fused_qkv_a_proj_out.split(
+                    [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1
+                )
+                # overlap qk norm
+                q = m.q_a_layernorm(q)
+
+                q_lora = q.clone()  # required for topk_indices
+                k_nope, k_pe = latent_cache.unsqueeze(1).split(
+                    [m.kv_lora_rank, m.qk_rope_head_dim], dim=-1
+                )
+                k_nope = m.kv_a_layernorm(k_nope)
+            q = m.q_b_proj(q_lora)[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)
+
+        q_nope_out = torch.bmm(q_nope.transpose(0, 1), m.w_kc)
+
+        q_nope_out = q_nope_out.transpose(0, 1)
+
+        if m.layer_id == 0:
+            m.rotary_emb.sin_cos_cache = m.rotary_emb.cos_sin_cache.index_select(
+                0, positions
+            )
+
+        q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
+
+        if dsa_use_prefill_cp(forward_batch):
+            # support allgather+rerrange
+            k_nope, k_pe = m.rebuild_cp_kv_cache(
+                latent_cache, forward_batch, k_nope, k_pe
+            )
+
+    if not m.skip_topk or (m.is_nextn and prev_topk_indices is None):
+        topk_indices = m.indexer(
+            hidden_states,
+            q_lora,
+            positions,
+            forward_batch,
+            m.layer_id,
+            layer_scatter_modes,
+            dynamic_scale,
+        )
+    else:
+        topk_indices = prev_topk_indices
+
+    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 (
+        forward_batch.forward_mode.is_extend()
+        and not forward_batch.forward_mode.is_draft_extend_v2()
+        and not forward_batch.forward_mode.is_target_verify()
+    ):
+        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)
+    if not m.next_skip_topk:
+        return output, None
+    else:
+        return output, topk_indices
+
+
+def npu_mla_preprocess(
+    m: "DeepseekV2AttentionMLA",
+    hidden_states: torch.Tensor,
+    positions: torch.Tensor,
+    forward_batch: "ForwardBatch",
+    zero_allocator: "BumpAllocator",
+):
+    dynamic_scale = None
+    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.v_head_dim,
+            m.quant_config,
+        )
+    # mlaprolog does not require additional calculation of q_lora
+    _is_mlaprolog = hasattr(m.quant_config, "ignore") and any(
+        re.fullmatch(r".*kv_b_proj", l) for l in m.quant_config.ignore
+    )
+    if _is_mlaprolog:
+        (
+            q_pe,
+            k_pe,
+            q_nope_out,
+            k_nope,
+            q_lora,
+            forward_batch,
+            positions,
+            dynamic_scale,
+        ) = m.mla_preprocess.forward(
+            positions, hidden_states, forward_batch, zero_allocator
+        )
+    else:
+        if m.alt_stream is not None:
+            mla_event = torch.npu.Event()
+            mla_event.record()
+            with torch.npu.stream(m.alt_stream):
+                # alt stream waits for the completion of the event on the main stream to ensure data dependency is complete
+                torch.npu.current_stream().wait_event(mla_event)
+                (
+                    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)
+            torch.npu.current_stream().wait_event(m.alt_stream)
+        else:
+            (
+                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)
+
+    return (
+        q_pe,
+        k_pe,
+        q_nope_out,
+        k_nope,
+        q_lora,
+        forward_batch,
+        zero_allocator,
+        positions,
+        dynamic_scale,
+    )
+
+
+# endregion

python/sglang/srt/hardware_backend/npu/modules/glm46v_processor.py

@@ -0,0 +1,285 @@
+"""NPU patch for GLM-4.6V image and video preprocessing.
+
+The GLM-4.6V image processor (Glm46VImageProcessorFast) and video processor
+(Glm46VVideoProcessor) create 10-dimensional tensors during patch extraction,
+which exceeds Ascend NPU's 8-dimension limit.
+
+This patch restructures the computation to stay within 8 dimensions, following
+the same pattern as the Qwen VL NPU patch.
+"""
+
+import torch
+import torchvision.transforms.v2.functional as tvF
+from transformers.image_processing_utils import BatchFeature
+from transformers.image_processing_utils_fast import (
+    group_images_by_shape,
+    reorder_images,
+)
+from transformers.image_utils import (
+    ChannelDimension,
+    PILImageResampling,
+    SizeDict,
+    get_image_size,
+)
+from transformers.models.glm46v.image_processing_glm46v import smart_resize
+from transformers.utils import TensorType
+from transformers.video_utils import group_videos_by_shape, reorder_videos
+
+from sglang.srt.hardware_backend.npu.modules.qwen_vl_processor import (
+    transform_patches_to_flatten,
+)
+from sglang.srt.utils import apply_module_patch
+
+
+# Func refers to transformers.models.glm46v.image_processing_glm46v_fast.py
+# Glm46VImageProcessorFast._preprocess
+def npu_wrapper_glm46v_preprocess(func):
+
+    def _preprocess(
+        self,
+        images: list["torch.Tensor"],
+        do_resize: bool,
+        size: SizeDict,
+        resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+        do_rescale: bool,
+        rescale_factor: float,
+        do_normalize: bool,
+        image_mean: float | list[float] | None,
+        image_std: float | list[float] | None,
+        patch_size: int,
+        temporal_patch_size: int,
+        merge_size: int,
+        disable_grouping: bool | None,
+        return_tensors: str | TensorType | None,
+        **kwargs,
+    ):
+        grouped_images, grouped_images_index = group_images_by_shape(
+            images, disable_grouping=disable_grouping
+        )
+        resized_images_grouped = {}
+        for shape, stacked_images in grouped_images.items():
+            height, width = stacked_images.shape[-2:]
+            if do_resize:
+                resized_height, resized_width = smart_resize(
+                    num_frames=temporal_patch_size,
+                    height=height,
+                    width=width,
+                    temporal_factor=temporal_patch_size,
+                    factor=patch_size * merge_size,
+                    min_pixels=size.shortest_edge,
+                    max_pixels=size.longest_edge,
+                )
+                stacked_images = self.resize(
+                    stacked_images,
+                    size=SizeDict(height=resized_height, width=resized_width),
+                    resample=resample,
+                )
+            resized_images_grouped[shape] = stacked_images
+
+        resized_images = reorder_images(resized_images_grouped, grouped_images_index)
+
+        grouped_images, grouped_images_index = group_images_by_shape(
+            resized_images, disable_grouping=disable_grouping
+        )
+        processed_images_grouped = {}
+        processed_grids = {}
+
+        for shape, stacked_images in grouped_images.items():
+            resized_height, resized_width = stacked_images.shape[-2:]
+
+            patches = self.rescale_and_normalize(
+                stacked_images,
+                do_rescale,
+                rescale_factor,
+                do_normalize,
+                image_mean,
+                image_std,
+            )
+            if patches.ndim == 4:
+                patches = patches.unsqueeze(1)
+
+            if patches.shape[1] % temporal_patch_size != 0:
+                repeats = patches[:, -1:].repeat(
+                    1,
+                    temporal_patch_size - (patches.shape[1] % temporal_patch_size),
+                    1,
+                    1,
+                    1,
+                )
+                patches = torch.cat([patches, repeats], dim=1)
+
+            batch_size, t_len, channel = patches.shape[:3]
+            grid_t = t_len // temporal_patch_size
+            grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+
+            ######################################
+            # Start of modifications for sglang  #
+            ######################################
+            flatten_patches = transform_patches_to_flatten(
+                patches,
+                batch_size,
+                grid_t,
+                temporal_patch_size,
+                channel,
+                grid_h,
+                grid_w,
+                patch_size,
+                merge_size,
+            )
+            ######################################
+            #  End of modifications for sglang   #
+            ######################################
+
+            processed_images_grouped[shape] = flatten_patches
+            processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
+
+        processed_images = reorder_images(
+            processed_images_grouped, grouped_images_index
+        )
+        processed_grids = reorder_images(processed_grids, grouped_images_index)
+
+        pixel_values = torch.cat(processed_images, dim=0)
+        image_grid_thw = torch.tensor(processed_grids)
+
+        return BatchFeature(
+            data={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw},
+            tensor_type=return_tensors,
+        )
+
+    return _preprocess
+
+
+# Func refers to transformers.models.glm46v.video_processing_glm46v.py
+# Glm46VVideoProcessor._preprocess
+def npu_wrapper_glm46v_video_preprocess(func):
+
+    def _preprocess(
+        self,
+        videos: list[torch.Tensor],
+        do_convert_rgb: bool = True,
+        do_resize: bool = True,
+        size: SizeDict | None = None,
+        resample: "PILImageResampling | tvF.InterpolationMode | int | None" = PILImageResampling.BICUBIC,
+        do_rescale: bool = True,
+        rescale_factor: float = 1 / 255.0,
+        do_normalize: bool = True,
+        image_mean: float | list[float] | None = None,
+        image_std: float | list[float] | None = None,
+        patch_size: int | None = None,
+        temporal_patch_size: int | None = None,
+        merge_size: int | None = None,
+        return_tensors: str | TensorType | None = None,
+        **kwargs,
+    ):
+        grouped_videos, grouped_videos_index = group_videos_by_shape(videos)
+        resized_videos_grouped = {}
+
+        for shape, stacked_videos in grouped_videos.items():
+            B, T, C, H, W = stacked_videos.shape
+            num_frames, height, width = T, H, W
+            if do_resize:
+                resized_height, resized_width = smart_resize(
+                    num_frames=num_frames,
+                    height=height,
+                    width=width,
+                    temporal_factor=temporal_patch_size,
+                    factor=patch_size * merge_size,
+                    min_pixels=size.shortest_edge,
+                    max_pixels=size.longest_edge,
+                )
+                stacked_videos = stacked_videos.view(B * T, C, H, W)
+                stacked_videos = self.resize(
+                    stacked_videos,
+                    size=SizeDict(height=resized_height, width=resized_width),
+                    resample=resample,
+                )
+                stacked_videos = stacked_videos.view(
+                    B, T, C, resized_height, resized_width
+                )
+            resized_videos_grouped[shape] = stacked_videos
+        resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index)
+
+        # Group videos by size for further processing
+        # Needed in case do_resize is False, or resize returns videos with different sizes
+        grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos)
+        processed_videos_grouped = {}
+        processed_grids = {}
+        for shape, stacked_videos in grouped_videos.items():
+            resized_height, resized_width = get_image_size(
+                stacked_videos[0], channel_dim=ChannelDimension.FIRST
+            )
+
+            # Fused rescale and normalize
+            stacked_videos = self.rescale_and_normalize(
+                stacked_videos,
+                do_rescale,
+                rescale_factor,
+                do_normalize,
+                image_mean,
+                image_std,
+            )
+            patches = stacked_videos
+
+            # Check that videos have `num_frames` divisible by `temporal_patch_size`
+            if patches.shape[1] % temporal_patch_size != 0:
+                repeats = patches[:, -1:].repeat(1, temporal_patch_size - 1, 1, 1, 1)
+                patches = torch.cat([patches, repeats], dim=1)
+            batch_size, grid_t, channel = patches.shape[:3]
+            grid_t = grid_t // temporal_patch_size
+            grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+
+            ######################################
+            # Start of modifications for sglang  #
+            ######################################
+            flatten_patches = transform_patches_to_flatten(
+                patches,
+                batch_size,
+                grid_t,
+                temporal_patch_size,
+                channel,
+                grid_h,
+                grid_w,
+                patch_size,
+                merge_size,
+            )
+            ######################################
+            #  End of modifications for sglang   #
+            ######################################
+
+            processed_videos_grouped[shape] = flatten_patches
+            processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
+
+        processed_videos = reorder_videos(
+            processed_videos_grouped, grouped_videos_index
+        )
+        processed_grids = reorder_videos(processed_grids, grouped_videos_index)
+        pixel_values_videos = torch.cat(processed_videos, dim=0)
+        video_grid_thw = torch.tensor(processed_grids)
+        data = {
+            "pixel_values_videos": pixel_values_videos,
+            "video_grid_thw": video_grid_thw,
+        }
+
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+    return _preprocess
+
+
+_npu_glm46v_preprocess_patched = False
+
+
+def npu_apply_glm46v_image_preprocess_patch():
+    global _npu_glm46v_preprocess_patched
+    if _npu_glm46v_preprocess_patched:
+        return
+    apply_module_patch(
+        "transformers.models.glm46v.image_processing_glm46v_fast.Glm46VImageProcessorFast",
+        "_preprocess",
+        [npu_wrapper_glm46v_preprocess],
+    )
+    apply_module_patch(
+        "transformers.models.glm46v.video_processing_glm46v.Glm46VVideoProcessor",
+        "_preprocess",
+        [npu_wrapper_glm46v_video_preprocess],
+    )
+    _npu_glm46v_preprocess_patched = True

python/sglang/srt/hardware_backend/npu/modules/qwen_vl_processor.py

@@ -0,0 +1,304 @@
+import torch
+import torchvision.transforms.v2.functional as tvF
+from transformers.image_processing_utils import BatchFeature
+from transformers.image_transforms import group_images_by_shape, reorder_images
+from transformers.image_utils import (
+    ChannelDimension,
+    PILImageResampling,
+    SizeDict,
+    get_image_size,
+)
+from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
+from transformers.models.qwen3_vl.video_processing_qwen3_vl import (
+    smart_resize as smart_resize_video,
+)
+from transformers.utils import TensorType
+from transformers.video_utils import group_videos_by_shape, reorder_videos
+
+from sglang.srt.utils import apply_module_patch
+
+
+def transform_patches_to_flatten(
+    patches: torch.Tensor,
+    batch_size: int,
+    grid_t: int,
+    temporal_patch_size: int,
+    channel: int,
+    grid_h: int,
+    grid_w: int,
+    patch_size: int,
+    merge_size: int,
+) -> torch.Tensor:
+    patches = patches.view(
+        batch_size * grid_t,
+        temporal_patch_size * channel,
+        grid_h // merge_size,
+        merge_size,
+        patch_size,
+        grid_w // merge_size,
+        merge_size,
+        patch_size,
+    )
+    patches = patches.permute(0, 1, 2, 5, 3, 6, 4, 7)
+    patches = patches.reshape(
+        batch_size,
+        grid_t,
+        temporal_patch_size,
+        channel,
+        grid_h * grid_w,
+        patch_size,
+        patch_size,
+    )
+    patches = patches.permute(0, 1, 4, 3, 2, 5, 6)
+    flatten_patches = patches.reshape(
+        batch_size,
+        grid_t * grid_h * grid_w,
+        -1,
+    )
+    return flatten_patches
+
+
+# Func refers to transformers.models.qwen2_vl.image_processing_qwen2_vl.py
+# Qwen2VLImageProcessor._preprocess
+def npu_wrapper_preprocess(func):
+
+    def _preprocess(
+        self,
+        images: list["torch.Tensor"],
+        do_resize: bool,
+        size: SizeDict,
+        resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+        do_rescale: bool,
+        rescale_factor: float,
+        do_normalize: bool,
+        image_mean: float | list[float] | None,
+        image_std: float | list[float] | None,
+        patch_size: int,
+        temporal_patch_size: int,
+        merge_size: int,
+        disable_grouping: bool | None,
+        return_tensors: str | TensorType | None,
+        **kwargs,
+    ):
+        # Group images by size for batched resizing
+        grouped_images, grouped_images_index = group_images_by_shape(
+            images, disable_grouping=disable_grouping
+        )
+        resized_images_grouped = {}
+        for shape, stacked_images in grouped_images.items():
+            height, width = stacked_images.shape[-2:]
+            if do_resize:
+                resized_height, resized_width = smart_resize(
+                    height,
+                    width,
+                    factor=patch_size * merge_size,
+                    min_pixels=size.shortest_edge,
+                    max_pixels=size.longest_edge,
+                )
+                stacked_images = self.resize(
+                    image=stacked_images,
+                    size=SizeDict(height=resized_height, width=resized_width),
+                    resample=resample,
+                )
+            resized_images_grouped[shape] = stacked_images
+        resized_images = reorder_images(resized_images_grouped, grouped_images_index)
+
+        # Group images by size for further processing
+        # Needed in case do_resize is False, or resize returns images with different sizes
+        grouped_images, grouped_images_index = group_images_by_shape(
+            resized_images, disable_grouping=disable_grouping
+        )
+        processed_images_grouped = {}
+        processed_grids = {}
+        for shape, stacked_images in grouped_images.items():
+            resized_height, resized_width = stacked_images.shape[-2:]
+            # Fused rescale and normalize
+            patches = self.rescale_and_normalize(
+                stacked_images,
+                do_rescale,
+                rescale_factor,
+                do_normalize,
+                image_mean,
+                image_std,
+            )
+            if patches.ndim == 4:
+                # add a temporal dimension if we have images
+                patches = patches.unsqueeze(1)
+            if patches.shape[1] % temporal_patch_size != 0:
+                repeats = patches[:, -1:].repeat(1, temporal_patch_size - 1, 1, 1, 1)
+                patches = torch.cat([patches, repeats], dim=1)
+            batch_size, grid_t, channel = patches.shape[:3]
+            grid_t = grid_t // temporal_patch_size
+            grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+
+            ######################################
+            # Start of modifications for sglang  #
+            ######################################
+            flatten_patches = transform_patches_to_flatten(
+                patches,
+                batch_size,
+                grid_t,
+                temporal_patch_size,
+                channel,
+                grid_h,
+                grid_w,
+                patch_size,
+                merge_size,
+            )
+            ######################################
+            #  End of modifications for sglang   #
+            ######################################
+
+            processed_images_grouped[shape] = flatten_patches
+            processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
+
+        processed_images = reorder_images(
+            processed_images_grouped, grouped_images_index
+        )
+        processed_grids = reorder_images(processed_grids, grouped_images_index)
+        pixel_values = torch.cat(processed_images, dim=0)
+        image_grid_thw = torch.tensor(processed_grids)
+
+        return BatchFeature(
+            data={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw},
+            tensor_type=return_tensors,
+        )
+
+    return _preprocess
+
+
+# Func refers to transformers.models.qwen3_vl.video_processing_qwen3_vl.py
+# Qwen3VLVideoProcessor._preprocess
+def npu_wrapper_video_preprocess(func):
+
+    def _preprocess(
+        self,
+        videos: list[torch.Tensor],
+        do_convert_rgb: bool = True,
+        do_resize: bool = True,
+        size: SizeDict | None = None,
+        resample: "PILImageResampling | tvF.InterpolationMode | int | None" = PILImageResampling.BICUBIC,
+        do_rescale: bool = True,
+        rescale_factor: float = 1 / 255.0,
+        do_normalize: bool = True,
+        image_mean: float | list[float] | None = None,
+        image_std: float | list[float] | None = None,
+        patch_size: int | None = None,
+        temporal_patch_size: int | None = None,
+        merge_size: int | None = None,
+        return_tensors: str | TensorType | None = None,
+        **kwargs,
+    ):
+        grouped_videos, grouped_videos_index = group_videos_by_shape(videos)
+        resized_videos_grouped = {}
+
+        for shape, stacked_videos in grouped_videos.items():
+            B, T, C, H, W = stacked_videos.shape
+            num_frames, height, width = T, H, W
+            if do_resize:
+                resized_height, resized_width = smart_resize_video(
+                    num_frames=num_frames,
+                    height=height,
+                    width=width,
+                    temporal_factor=temporal_patch_size,
+                    factor=patch_size * merge_size,
+                    min_pixels=size.shortest_edge,
+                    max_pixels=size.longest_edge,
+                )
+                stacked_videos = stacked_videos.view(B * T, C, H, W)
+                stacked_videos = self.resize(
+                    stacked_videos,
+                    size=SizeDict(height=resized_height, width=resized_width),
+                    resample=resample,
+                )
+                stacked_videos = stacked_videos.view(
+                    B, T, C, resized_height, resized_width
+                )
+            resized_videos_grouped[shape] = stacked_videos
+        resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index)
+
+        # Group videos by size for further processing
+        # Needed in case do_resize is False, or resize returns videos with different sizes
+        grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos)
+        processed_videos_grouped = {}
+        processed_grids = {}
+        for shape, stacked_videos in grouped_videos.items():
+            resized_height, resized_width = get_image_size(
+                stacked_videos[0], channel_dim=ChannelDimension.FIRST
+            )
+
+            # Fused rescale and normalize
+            stacked_videos = self.rescale_and_normalize(
+                stacked_videos,
+                do_rescale,
+                rescale_factor,
+                do_normalize,
+                image_mean,
+                image_std,
+            )
+            patches = stacked_videos
+
+            # Check that videos have `num_frames` divisible by `temporal_patch_size`
+            T = patches.shape[1]
+            if pad := -T % temporal_patch_size:
+                repeats = patches[:, -1:].expand(-1, pad, -1, -1, -1)
+                patches = torch.cat((patches, repeats), dim=1)
+            batch_size, grid_t, channel = patches.shape[:3]
+            grid_t = grid_t // temporal_patch_size
+            grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+
+            ######################################
+            # Start of modifications for sglang  #
+            ######################################
+            flatten_patches = transform_patches_to_flatten(
+                patches,
+                batch_size,
+                grid_t,
+                temporal_patch_size,
+                channel,
+                grid_h,
+                grid_w,
+                patch_size,
+                merge_size,
+            )
+            ######################################
+            #  End of modifications for sglang   #
+            ######################################
+
+            processed_videos_grouped[shape] = flatten_patches
+            processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
+
+        processed_videos = reorder_videos(
+            processed_videos_grouped, grouped_videos_index
+        )
+        processed_grids = reorder_videos(processed_grids, grouped_videos_index)
+        pixel_values_videos = torch.cat(processed_videos, dim=0)
+        video_grid_thw = torch.tensor(processed_grids)
+        data = {
+            "pixel_values_videos": pixel_values_videos,
+            "video_grid_thw": video_grid_thw,
+        }
+
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+    return _preprocess
+
+
+_npu_preprocess_patched = False
+
+
+def npu_apply_qwen_image_preprocess_patch():
+    global _npu_preprocess_patched
+    if _npu_preprocess_patched:
+        return
+    apply_module_patch(
+        "transformers.models.qwen2_vl.image_processing_qwen2_vl.Qwen2VLImageProcessor",
+        "_preprocess",
+        [npu_wrapper_preprocess],
+    )
+    apply_module_patch(
+        "transformers.models.qwen3_vl.video_processing_qwen3_vl.Qwen3VLVideoProcessor",
+        "_preprocess",
+        [npu_wrapper_video_preprocess],
+    )
+    _npu_preprocess_patched = True

python/sglang/srt/hardware_backend/npu/moe/fuseep.py

@@ -0,0 +1,171 @@
+"""Ascend FuseEP fused dispatch+GEMM+combine forward path.
+
+Follows the mega_moe shape: a free-function bypass invoked from
+``FusedMoE.forward`` when ``--moe-a2a-backend ascend_fuseep`` is set, plus a
+weight-postprocess helper that NPU quant_methods call from their
+``process_weights_after_loading`` when the same backend is selected.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.distributed import get_tp_group
+from sglang.srt.environ import envs
+from sglang.srt.hardware_backend.npu.utils import FusedMoEMode, npu_format_cast
+from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPBuffer
+from sglang.srt.layers.moe.utils import DeepEPMode
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
+    from sglang.srt.layers.moe.topk import TopKOutput
+
+
+_PARAMS_BYTES = 2  # bf16 — Ascend's Dispatch & Combine does not support fp16
+
+
+def _get_fuseep_buffer(layer: FusedMoE):
+    DeepEPBuffer.set_dispatch_mode_as_low_latency()
+    return DeepEPBuffer.get_deepep_buffer(
+        get_tp_group().device_group,
+        layer.hidden_size,
+        _PARAMS_BYTES,
+        DeepEPMode.LOW_LATENCY,
+        envs.SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get(),
+        layer.num_experts,
+    )
+
+
+def forward_fuseep(
+    layer: FusedMoE,
+    hidden_states: torch.Tensor,
+    topk_output: TopKOutput,
+) -> torch.Tensor:
+    buf = _get_fuseep_buffer(layer)
+    hidden_states, _ = buf.fused_deep_moe(
+        hidden_states,
+        topk_idx=topk_output.topk_ids,
+        topk_weights=topk_output.topk_weights,
+        gmm1_permuted_weight=layer.w13_weight,
+        gmm1_permuted_weight_scale=layer.w13_weight_scale,
+        gmm2_weight=layer.w2_weight,
+        gmm2_weight_scale=layer.w2_weight_scale,
+        num_max_dispatch_tokens_per_rank=(
+            envs.SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get()
+        ),
+        num_experts=layer.num_experts,
+        fuse_mode=envs.SGLANG_NPU_FUSED_MOE_MODE.get(),
+    )
+    return hidden_states
+
+
+def _permute_w13_weight_scale(w: torch.Tensor, tile_n: int) -> torch.Tensor:
+    if tile_n % 2 != 0:
+        raise ValueError(f"tile_n must be even, got {tile_n}")
+
+    *dims, n = w.shape
+    if n % tile_n != 0:
+        raise ValueError(f"Last dimension {n} must be divisible by tile_n {tile_n}")
+
+    w_reshaped = w.reshape(*dims, 2, n // tile_n, tile_n // 2)
+    perm_order = list(range(len(dims))) + [-2, -3, -1]
+    return w_reshaped.permute(perm_order).reshape(*dims, n)
+
+
+def _reshape_w13_weight(
+    weight: torch.Tensor, dim: int, chunk_size: int = 64
+) -> torch.Tensor:
+    # Achieving greater computing power through reshape on Ascend.
+    original_shape = weight.shape
+    if dim < 0:
+        dim += len(original_shape)
+
+    if original_shape[dim] % (2 * chunk_size) != 0:
+        raise ValueError(
+            f"Dimension {dim} size {original_shape[dim]} must be divisible by "
+            f"{2 * chunk_size}"
+        )
+
+    new_shape = (
+        *original_shape[:dim],
+        2,
+        original_shape[dim] // (2 * chunk_size),
+        chunk_size,
+        *original_shape[dim + 1 :],
+    )
+
+    weight = weight.view(new_shape)
+    weight = weight.transpose(dim, dim + 1).contiguous()
+    return weight.view(*original_shape[:dim], -1, *original_shape[dim + 1 :])
+
+
+def _release_weight_cache(weight: torch.Tensor) -> torch.Tensor:
+    # .contiguous() introduces additional memory overhead; release with resize_(0)
+    origin_weight = weight.data.transpose(1, 2)
+    new_weight = origin_weight.contiguous()
+    origin_weight.untyped_storage().resize_(0)
+    return new_weight
+
+
+def _scale_from_float_to_int64(scale: torch.Tensor) -> torch.nn.Parameter:
+    import numpy as np
+
+    converted = torch.from_numpy(
+        np.frombuffer(
+            scale.cpu().to(torch.float32).numpy().tobytes(), dtype=np.int32
+        ).astype(np.int64)
+    ).to(scale.device)
+    return torch.nn.Parameter(converted, requires_grad=False)
+
+
+def process_fuseep_weights(layer: torch.nn.Module) -> None:
+    """Apply the Ascend FuseEP-specific weight layout.
+
+    Replaces NPU quant_method weight layouts with the form required by the
+    fused_deep_moe op. Invoked from NPU ``process_weights_after_loading``
+    when ``--moe-a2a-backend ascend_fuseep`` is set.
+    """
+    if envs.SGLANG_NPU_FUSED_MOE_MODE.get() == FusedMoEMode.DISPATCH_FFN_COMBINE.value:
+        w13_weight = _release_weight_cache(layer.w13_weight)
+        layer.w13_weight.data = npu_format_cast(w13_weight)
+        w2_weight = _release_weight_cache(layer.w2_weight)
+        layer.w2_weight.data = npu_format_cast(w2_weight)
+
+        layer.w13_weight_scale.data = layer.w13_weight_scale.data.view(
+            layer.w13_weight_scale.data.shape[0], -1
+        )
+        w2_scale = layer.w2_weight_scale.data.squeeze(-1).contiguous()
+        layer.w2_weight_scale = torch.nn.Parameter(
+            w2_scale.to(torch.float32), requires_grad=False
+        )
+
+        layer.w13_weight_scale = _scale_from_float_to_int64(layer.w13_weight_scale.data)
+        layer.w2_weight_scale = _scale_from_float_to_int64(layer.w2_weight_scale.data)
+    else:
+        cpu_w13 = layer.w13_weight.data.transpose(1, 2).cpu()
+        layer.w13_weight.data = _reshape_w13_weight(cpu_w13, -1).npu()
+        w13_scale = layer.w13_weight_scale.data.squeeze(-1).contiguous()
+        w13_scale = _permute_w13_weight_scale(w13_scale, 128)
+        layer.w13_weight_scale = torch.nn.Parameter(
+            w13_scale.to(torch.float32), 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)
+
+        w2_scale = layer.w2_weight_scale.data.squeeze(-1).contiguous()
+        layer.w2_weight_scale = torch.nn.Parameter(
+            w2_scale.to(torch.float32), requires_grad=False
+        )
+
+    if hasattr(layer, "w13_weight_offset"):
+        layer.w13_weight_offset = torch.nn.Parameter(
+            layer.w13_weight_offset.data.squeeze(-1).contiguous(),
+            requires_grad=False,
+        )
+    if hasattr(layer, "w2_weight_offset"):
+        layer.w2_weight_offset = torch.nn.Parameter(
+            layer.w2_weight_offset.data.squeeze(-1).contiguous(),
+            requires_grad=False,
+        )

python/sglang/srt/hardware_backend/npu/moe/topk.py

@@ -0,0 +1,113 @@
+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,
+    capture_routed_experts_if_allowed,
+    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,
+    layer_id: Optional[int] = None,
+) -> "TopKOutput":
+
+    use_grouped_topk = topk_config.use_grouped_topk
+    renormalize = topk_config.renormalize
+    correction_bias = topk_config.correction_bias
+
+    # Fast path: simple top-k without grouped routing and bias
+    if not use_grouped_topk and correction_bias is None:
+        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)
+
+    # sqrtsoftplus (DSV4 noaux_tc): the NPU op only scores sigmoid/softmax, so use
+    # a torch path. top-k over (scores + bias); weights from un-biased scores.
+    elif topk_config.scoring_func == "sqrtsoftplus":
+        scores = torch.nn.functional.softplus(router_logits.float()).sqrt()
+        scores_for_choice = (
+            scores + correction_bias.unsqueeze(0).float()
+            if correction_bias is not None
+            else scores
+        )
+        _, topk_ids = torch.topk(
+            scores_for_choice, k=topk_config.top_k, dim=-1, sorted=False
+        )
+        topk_ids = topk_ids.to(torch.int32)
+        topk_weights = scores.gather(1, topk_ids)
+        if renormalize:
+            topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+        else:
+            topk_weights = topk_weights * topk_config.routed_scaling_factor
+        topk_weights = topk_weights.to(torch.float32)
+
+    # Support grouped top-k or correction bias or sigmoid or routed_scaling_factor
+    elif (
+        correction_bias is not None
+        or topk_config.scoring_func == "sigmoid"
+        or num_token_non_padded is not None
+    ):
+        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)
+                if correction_bias is not None
+                else None
+            ),
+            # num_expert_group and topk_group in some topk_config without group is None, (not supported by this ops)
+            k_group=topk_config.topk_group if use_grouped_topk else 1,
+            group_count=topk_config.num_expert_group if use_grouped_topk else 1,
+            group_select_mode=(1 if use_grouped_topk else 0),
+            renorm=0,
+            # 1 for sigmoid, 0 for softmax
+            norm_type=1,
+            routed_scaling_factor=(
+                1 if renormalize else topk_config.routed_scaling_factor
+            ),
+            eps=float(1e-20),
+        )
+        topk_weights = topk_weights.to(torch.float32)
+
+    # torch native is not yet supported num_token_non_padded
+    # Fallback to torch native implementation
+    else:
+        topk_config.torch_native = True
+        return select_experts(
+            hidden_states=hidden_states,
+            layer_id=layer_id,
+            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)
+    capture_routed_experts_if_allowed(topk_config, layer_id, topk_ids)
+
+    return StandardTopKOutput(topk_weights, topk_ids, router_logits)

python/sglang/srt/hardware_backend/npu/quantization/awq_kernels.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
+    NPUW4A16Int4DynamicMoEMethod,
+)
+from sglang.srt.layers.quantization.utils import replace_parameter
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
+    from sglang.srt.layers.quantization.base_config import QuantizationConfig
+
+import torch_npu
+
+
+class AWQAscendLinearKernel:
+    def __init__(self, quant_config: Optional[QuantizationConfig] = None):
+        self.quant_config = quant_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False)
+        qweight_tmp = torch.zeros_like(layer.qweight.data)
+        qzeros_tmp = layer.qzeros.data
+        qzeros_list = []
+        shifts = [0, 4, 1, 5, 2, 6, 3, 7]
+
+        for i in range(0, self.quant_config.pack_factor):
+            shift_num = shifts[i] * 4
+            qzeros_list.append((qzeros_tmp.reshape(-1, 1) >> shift_num) & 0xF)
+            qweight_tmp.bitwise_or_(
+                ((layer.qweight.data >> shift_num) & 0xF) << (4 * i)
+            )
+
+        qweight_tmp.bitwise_xor_(0x88888888)
+
+        qzeros_tmp = torch.cat(qzeros_list, dim=-1).reshape(qzeros_tmp.shape[0], -1)
+        qzeros_tmp = -(qzeros_tmp - 8)
+        qzeros_tmp = qzeros_tmp.to(layer.scales.data.dtype)
+
+        layer.zeros = torch.nn.Parameter(qzeros_tmp, requires_grad=False)
+        layer.weight = torch.nn.Parameter(qweight_tmp, requires_grad=False)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        qweight = layer.weight
+        scales = layer.scales
+        qzeros = layer.zeros
+        pack_factor = self.quant_config.pack_factor
+        out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,)
+        reshaped_x = x.reshape(-1, x.shape[-1])
+
+        if bias is not None and bias.dtype == torch.bfloat16:
+            bias = bias.float()
+
+        out = torch_npu.npu_weight_quant_batchmatmul(
+            reshaped_x,
+            qweight,
+            antiquant_scale=scales,
+            antiquant_offset=qzeros,
+            antiquant_group_size=self.quant_config.group_size,
+            bias=bias,
+        )
+
+        return out.reshape(out_shape)
+
+
+class AWQAscendMoEKernel:
+    def __init__(self, quant_config: Optional[QuantizationConfig] = None):
+        self.quant_config = quant_config
+        self.kernel = NPUW4A16Int4DynamicMoEMethod()
+
+    @staticmethod
+    def _register_or_replace_parameter(
+        layer: torch.nn.Module, name: str, tensor: torch.Tensor
+    ) -> None:
+        if hasattr(layer, name):
+            replace_parameter(layer, name, tensor)
+        else:
+            layer.register_parameter(
+                name, torch.nn.Parameter(tensor, requires_grad=False)
+            )
+
+    def _convert_awq_weight_to_npu_layout(self, qweight: torch.Tensor) -> torch.Tensor:
+        num_experts, input_size, _ = qweight.shape
+        unpacked_weight = (
+            self.kernel._unpack_from_int32(qweight.flatten(0, 1), 4)
+            .view(num_experts, input_size, -1)
+            .transpose(1, 2)
+            .contiguous()
+            .int()
+        )
+        return self.kernel._pack_to_int32(unpacked_weight)
+
+    def _convert_awq_qzeros_to_npu_offset(
+        self, qzeros: torch.Tensor, dtype: torch.dtype
+    ) -> torch.Tensor:
+        num_experts, num_groups, _ = qzeros.shape
+        offset = (
+            -self.kernel._unpack_from_int32(qzeros.flatten(0, 1), 4)
+            .view(num_experts, num_groups, -1)
+            .transpose(1, 2)
+            .contiguous()
+        )
+        return offset.to(dtype)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        self._register_or_replace_parameter(
+            layer,
+            "w13_weight",
+            self._convert_awq_weight_to_npu_layout(layer.w13_qweight.data),
+        )
+        self._register_or_replace_parameter(
+            layer,
+            "w2_weight",
+            self._convert_awq_weight_to_npu_layout(layer.w2_qweight.data),
+        )
+        self._register_or_replace_parameter(
+            layer,
+            "w13_weight_scale",
+            layer.w13_scales.data.transpose(1, 2).contiguous(),
+        )
+        self._register_or_replace_parameter(
+            layer,
+            "w2_weight_scale",
+            layer.w2_scales.data.transpose(1, 2).contiguous(),
+        )
+        self._register_or_replace_parameter(
+            layer,
+            "w13_weight_offset",
+            self._convert_awq_qzeros_to_npu_offset(
+                layer.w13_qzeros.data, layer.w13_scales.data.dtype
+            ),
+        )
+        self._register_or_replace_parameter(
+            layer,
+            "w2_weight_offset",
+            self._convert_awq_qzeros_to_npu_offset(
+                layer.w2_qzeros.data, layer.w2_scales.data.dtype
+            ),
+        )
+
+        self.kernel.process_weights_after_loading(layer)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: StandardDispatchOutput,
+    ) -> torch.Tensor:
+        return self.kernel.apply(layer, dispatch_output)
+
+    def apply_without_routing_weights(
+        self,
+        layer,
+        hidden_states,
+        hidden_states_scale,
+        group_list_type,
+        group_list,
+        output_dtype,
+    ):
+        return self.kernel.apply_without_routing_weights(
+            layer,
+            hidden_states,
+            hidden_states_scale,
+            group_list_type,
+            group_list,
+            output_dtype,
+        )

python/sglang/srt/hardware_backend/npu/quantization/gptq_kernels.py

@@ -0,0 +1,315 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+import torch_npu
+
+from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
+    npu_fused_experts,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe import MoeRunnerConfig
+    from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
+    from sglang.srt.layers.quantization.base_config import QuantizationConfig
+
+
+def unpack_from_int32(
+    weight: torch.Tensor,
+    num_bits: int,
+    packed_dim: int = 1,
+) -> torch.Tensor:
+    """
+    Unpacks quantized weights from int32 format back to original bits.
+
+    :param weight: The packed int32 tensor containing quantized weights
+    :param num_bits: The number of bits used for quantization (<= 8)
+    :param packed_dim: Dimension along which weights are packed (0 or 1), defaults to 1
+    :return: Unpacked tensor with int8 dtype after applying offset correction
+    """
+    assert (
+        weight.dtype == torch.int32
+    ), f"Expecting `weight.dtype` is torch.int32 but got {weight.dtype}."
+    assert (
+        num_bits <= 8
+    ), f"Expecting `num_bits` should not be larger than 8 but got {num_bits}."
+
+    pack_factor = 32 // num_bits
+    mask = (1 << num_bits) - 1
+
+    if packed_dim == 1:
+        unpacked_weight = torch.zeros(
+            (weight.shape[0], weight.shape[1] * pack_factor),
+            device=weight.device,
+            dtype=torch.int32,
+        )
+        for i in range(pack_factor):
+            unpacked_weight[:, i::pack_factor] = (weight >> (num_bits * i)) & mask
+    else:
+        unpacked_weight = torch.zeros(
+            (weight.shape[0] * pack_factor, weight.shape[1]),
+            device=weight.device,
+            dtype=torch.int32,
+        )
+        for i in range(pack_factor):
+            unpacked_weight[i::pack_factor, :] = (weight >> (num_bits * i)) & mask
+    offset = pow(2, num_bits) // 2
+    unpacked_weight = (unpacked_weight - offset).to(torch.int8)
+    return unpacked_weight
+
+
+class GPTQLinearAscendKernel:
+    def __init__(self, quant_config: Optional[QuantizationConfig] = None):
+        self.quant_config = quant_config
+        self.use_v2_format = quant_config.checkpoint_format == "gptq_v2"
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+
+        layer.qzeros = torch.nn.Parameter(
+            unpack_from_int32(
+                layer.qzeros.data.contiguous(),
+                self.quant_config.weight_bits,
+                packed_dim=1,
+            ).to(layer.scales.dtype),
+            requires_grad=False,
+        )
+        if not self.use_v2_format:
+            layer.qzeros += 1
+
+        qweight_tmp = unpack_from_int32(
+            layer.qweight.data.contiguous(), self.quant_config.weight_bits, packed_dim=0
+        )
+        # use int8 to store weight by default
+        if self.quant_config.weight_bits != 4:
+            layer.qweight = torch.nn.Parameter(
+                qweight_tmp,
+                requires_grad=False,
+            )
+            return
+
+        # for 4bit case we need to pack 4bit weight to int32 to save memory
+        layer.qweight = torch.nn.Parameter(
+            torch_npu.npu_convert_weight_to_int4pack(qweight_tmp.to(torch.int32)),
+            requires_grad=False,
+        )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        qweight = layer.qweight
+        scales = layer.scales
+        qzeros = layer.qzeros
+
+        reshaped_x = x.reshape(-1, x.shape[-1])
+
+        if bias is not None and bias.dtype == torch.bfloat16:
+            bias = bias.float()
+
+        # 4bit weight is packed to int32(8 x int4)
+        if self.quant_config.weight_bits == 4:
+            out_shape = x.shape[:-1] + (qweight.shape[-1] * 8,)
+        else:
+            out_shape = x.shape[:-1] + (qweight.shape[-1],)
+
+        out = torch_npu.npu_weight_quant_batchmatmul(
+            reshaped_x,
+            qweight,
+            antiquant_scale=scales,
+            antiquant_offset=qzeros,
+            antiquant_group_size=self.quant_config.group_size,
+            bias=bias,
+        )
+
+        return out.reshape(out_shape)
+
+
+class GPTQMoEAscendKernel:
+    def __init__(self, quant_config: Optional[QuantizationConfig] = None):
+        self.quant_config = quant_config
+        self.use_v2_format = quant_config.checkpoint_format == "gptq_v2"
+        self.moe_runner_config: Optional[MoeRunnerConfig] = None
+
+    def create_moe_runner(
+        self,
+        layer: torch.nn.Module,
+        moe_runner_config: MoeRunnerConfig,
+        **extra_weight_attrs,
+    ):
+        self.moe_runner_config = moe_runner_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        w13_qzeros_2d = layer.w13_qzeros.data.contiguous().reshape(
+            -1, layer.w13_qzeros.shape[-1]
+        )
+        layer.w13_qzeros = torch.nn.Parameter(
+            unpack_from_int32(
+                w13_qzeros_2d,
+                self.quant_config.weight_bits,
+                packed_dim=1,
+            )
+            .reshape(layer.w13_qzeros.shape[0], layer.w13_qzeros.shape[1], -1)
+            .to(layer.w13_scales.dtype),
+            requires_grad=False,
+        )
+        if not self.use_v2_format:
+            layer.w13_qzeros += 1
+
+        w2_qzeros_2d = layer.w2_qzeros.data.contiguous().reshape(
+            -1, layer.w2_qzeros.shape[-1]
+        )
+        layer.w2_qzeros = torch.nn.Parameter(
+            unpack_from_int32(
+                w2_qzeros_2d,
+                self.quant_config.weight_bits,
+                packed_dim=1,
+            )
+            .reshape(layer.w2_qzeros.shape[0], layer.w2_qzeros.shape[1], -1)
+            .to(layer.w2_scales.dtype),
+            requires_grad=False,
+        )
+        if not self.use_v2_format:
+            layer.w2_qzeros += 1
+
+        w13_qweight_2d = (
+            layer.w13_qweight.data.transpose(-1, -2)
+            .contiguous()
+            .reshape(-1, layer.w13_qweight.shape[-2])
+        )
+        w13_qweight_tmp = unpack_from_int32(
+            w13_qweight_2d, self.quant_config.weight_bits, packed_dim=1
+        )
+
+        if self.quant_config.weight_bits == 4:
+            group_size = self.quant_config.group_size
+            scale_expanded = layer.w13_scales.data.repeat_interleave(group_size, dim=1)
+
+            neg_mask = scale_expanded < 0
+
+            if neg_mask.any():
+                neg_mask = neg_mask.transpose(-1, -2)
+                neg_mask = neg_mask.contiguous().reshape(w13_qweight_tmp.shape)
+                w13_qweight_tmp[neg_mask] = -w13_qweight_tmp[neg_mask]
+
+                if w13_qweight_tmp.max() > 7:
+                    w13_qweight_tmp.clamp_(max=7)
+
+                layer.w13_scales.data.abs_()
+
+            layer.w13_qweight = torch.nn.Parameter(
+                torch_npu.npu_convert_weight_to_int4pack(
+                    w13_qweight_tmp.reshape(
+                        layer.w13_qweight.shape[0], layer.w13_qweight.shape[2], -1
+                    )
+                    .transpose(-1, -2)
+                    .contiguous()
+                    .reshape(-1, layer.w13_qweight.shape[2])
+                    .to(torch.int32)
+                )
+                .reshape(layer.w13_qweight.shape[0], layer.w13_qweight.shape[1] * 8, -1)
+                .contiguous(),
+                requires_grad=False,
+            )
+        # use int8 to store weight by default
+        else:
+            layer.w13_qweight = torch.nn.Parameter(
+                w13_qweight_tmp.reshape(
+                    layer.w13_qweight.shape[0], layer.w13_qweight.shape[2], -1
+                )
+                .transpose(-1, -2)
+                .contiguous(),
+                requires_grad=False,
+            )
+
+        w2_qweight_2d = (
+            layer.w2_qweight.data.transpose(-1, -2)
+            .contiguous()
+            .reshape(-1, layer.w2_qweight.shape[-2])
+        )
+        w2_qweight_tmp = unpack_from_int32(
+            w2_qweight_2d, self.quant_config.weight_bits, packed_dim=1
+        )
+
+        if self.quant_config.weight_bits == 4:
+            group_size = self.quant_config.group_size
+            scale_expanded = layer.w2_scales.data.repeat_interleave(group_size, dim=1)
+
+            neg_mask = scale_expanded < 0
+
+            if neg_mask.any():
+                neg_mask = neg_mask.transpose(-1, -2)
+                neg_mask = neg_mask.contiguous().reshape(w2_qweight_tmp.shape)
+                w2_qweight_tmp[neg_mask] = -w2_qweight_tmp[neg_mask]
+
+                if w2_qweight_tmp.max() > 7:
+                    w2_qweight_tmp.clamp_(max=7)
+
+                layer.w2_scales.data.abs_()
+
+            layer.w2_qweight = torch.nn.Parameter(
+                torch_npu.npu_convert_weight_to_int4pack(
+                    w2_qweight_tmp.reshape(
+                        layer.w2_qweight.shape[0], layer.w2_qweight.shape[2], -1
+                    )
+                    .transpose(-1, -2)
+                    .contiguous()
+                    .reshape(-1, layer.w2_qweight.shape[2])
+                    .to(torch.int32)
+                )
+                .reshape(layer.w2_qweight.shape[0], layer.w2_qweight.shape[1] * 8, -1)
+                .contiguous(),
+                requires_grad=False,
+            )
+        # use int8 to store weight by default
+        else:
+            layer.w2_qweight = torch.nn.Parameter(
+                w2_qweight_tmp.reshape(
+                    layer.w2_qweight.shape[0], layer.w2_qweight.shape[2], -1
+                )
+                .transpose(-1, -2)
+                .contiguous(),
+                requires_grad=False,
+            )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: StandardDispatchOutput,
+    ) -> torch.Tensor:
+        from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
+
+        assert self.moe_runner_config is not None, (
+            "moe_runner_config is not set. "
+            "Did you forget to call create_weights/create_moe_runner?"
+        )
+
+        assert self.moe_runner_config.activation in ("silu", "swiglu"), (
+            f"Only SiLU/Swiglu activation is supported, "
+            f"got {self.moe_runner_config.activation!r}."
+        )
+
+        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_qweight,
+            w13_scale=layer.w13_scales,
+            w13_offset=layer.w13_qzeros,
+            w2=layer.w2_qweight,
+            w2_scale=layer.w2_scales,
+            w2_offset=layer.w2_qzeros,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            top_k=topk_ids.shape[1],
+            use_wna16=True,
+        )
+
+        return StandardCombineInput(hidden_states=output)

python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py

@@ -0,0 +1,622 @@
+import logging
+from typing import TYPE_CHECKING, Optional
+
+import torch
+from torch.nn.parameter import Parameter
+
+from sglang.srt.hardware_backend.npu.utils import NPUACLFormat, npu_format_cast
+from sglang.srt.layers.quantization.base_config import LinearMethodBase
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.quantization.base_config import QuantizationConfig
+
+logger = logging.getLogger(__name__)
+
+MXFP8_BLOCK_SIZE = 32
+# W4A8_MXFP block (group) size — fixed at 32 by the msmodelslim export format.
+MXFP4_BLOCK_SIZE = 32
+
+
+# NPU ops are reached via torch.ops.npu.* (registered when torch_npu is imported
+# by the runtime), so this module needs no top-level `import torch_npu` and stays
+# importable on CUDA/CPU/AMD/XPU CI.
+def _get_float8_e8m0fnu_dtype():
+    # Resolve lazily rather than as a module-level constant: this module is
+    # imported early (during quant-scheme registration), so reading the dtype at
+    # call time keeps it correct regardless of import order / platform.
+    return getattr(torch, "float8_e8m0fnu", None)
+
+
+def _get_float4_e2m1fn_x2_dtype():
+    # The packed-FP4 dtype MUST come from torch_npu (an int enum, e.g. 296), not
+    # from torch. The NPU ops that consume it -- npu_dynamic_mx_quant(dst_type=),
+    # npu_quant_matmul(x2_dtype=), npu_format_cast(input_dtype=) -- REJECT the
+    # torch dtype object torch.float4_e2m1fn_x2 in op-plugin on recent torch_npu
+    # builds (it raises, or with None gives "output y must be same shape as input
+    # x"), even though torch.float4_e2m1fn_x2 exists. This is fp4-specific: fp8 /
+    # float8_e8m0fnu is accepted from torch either way. Verified on A5 /
+    # torch_npu 2.10.0.post2.dev20260704 (see llm/probe_fp4_w4a8_chain.py: dst=296
+    # passes the full quant->format_cast->matmul chain, dst=torch dtype fails).
+    #
+    # Lazy import so this NPU-only path keeps the module importable on
+    # CUDA/CPU/AMD/XPU CI (no top-level torch_npu; see AGENTS.md known pitfalls).
+    from sglang.srt.utils import is_npu
+
+    if is_npu():
+        import torch_npu
+
+        npu_dtype = getattr(torch_npu, "float4_e2m1fn_x2", None)
+        if npu_dtype is not None:
+            return npu_dtype
+    return getattr(torch, "float4_e2m1fn_x2", None)
+
+
+class _NPULinearMethodBase(LinearMethodBase):
+
+    def __init__(
+        self,
+        quant_config: Optional["QuantizationConfig"] = None,
+    ):
+        self.quant_config = quant_config
+
+
+class NPUW8A8Int8LinearMethod(_NPULinearMethodBase):
+
+    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()
+        # Compressed-tensors format doesn't have this field
+        if hasattr(layer, "weight_offset"):
+            layer.weight_offset.data = layer.weight_offset.data.flatten()
+
+        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,
+        )
+
+    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,
+        )
+
+
+class NPUW8A8Int8DynamicLinearMethod(_NPULinearMethodBase):
+
+    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()
+        # Compressed-tensors format doesn't have this field
+        if hasattr(layer, "weight_offset"):
+            layer.weight_offset.data = layer.weight_offset.data.flatten()
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        if isinstance(x, tuple):
+            """dynamic_scale is calculated in malprolog kernel"""
+            original_dtype = torch.bfloat16
+            quant_out, dynamic_scale = x
+        else:
+            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.flatten(),
+            bias=bias,
+            output_dtype=original_dtype,
+        )
+
+
+class NPUMXFP8LinearMethod(_NPULinearMethodBase):
+    """Ascend NPU MXFP8 linear method for LLM (SRT) models.
+
+    Shared kernel for both the online config path (``--quantization mxfp8``) and
+    the offline ModelSlimMXFP8Scheme (which delegates to this as ``self.kernel``).
+    process_weights_after_loading branches on weight dtype: FP16/BF16 weights are
+    quantised to MXFP8 at load time (online); pre-quantised float8_e4m3fn weights
+    are only re-laid-out (offline). Inference: dynamic MXFP8 activation quant +
+    MXFP8 matmul (block_size=32).
+    """
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes,
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        from sglang.srt.layers.parameter import ModelWeightParameter
+
+        output_size_per_partition = sum(output_partition_sizes)
+        weight_loader = extra_weight_attrs.get("weight_loader")
+
+        layer.logical_widths = output_partition_sizes
+        layer.input_size_per_partition = input_size_per_partition
+        layer.output_size_per_partition = output_size_per_partition
+        layer.orig_dtype = params_dtype
+
+        # Load weights in original dtype; quantise later in process_weights_after_loading
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition,
+                dtype=params_dtype,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight = layer.weight.data
+        if weight.dtype == torch.float8_e4m3fn:
+            # Offline (ModelSlim) path: weight is already MXFP8-quantised and
+            # layer.weight_scale holds the uint8 block scales [out, in/32]. Only
+            # re-layout to [in, out] / [in//64, out, 2] strided views below.
+            n_dim, k_dim = layer.weight_scale.data.shape
+            scale = layer.weight_scale.data.reshape(n_dim, k_dim // 2, 2)
+            layer.weight = Parameter(weight.transpose(0, 1), requires_grad=False)
+            layer.weight_scale_inv = Parameter(
+                scale.transpose(0, 1), requires_grad=False
+            )
+            # weight_scale is now folded into weight_scale_inv (which keeps the
+            # underlying storage alive via its view); drop the stale parameter so
+            # it doesn't linger in named_parameters() / state_dict().
+            del layer.weight_scale
+        else:
+            # Online path: quantise FP16/BF16 weights to MXFP8 at load time.
+            if weight.dtype not in (torch.float16, torch.bfloat16):
+                logger.warning(
+                    "NPUMXFP8LinearMethod: weight dtype %s is not float16/bfloat16; "
+                    "casting to bfloat16 before MXFP8 quantisation.",
+                    weight.dtype,
+                )
+                weight = weight.to(torch.bfloat16)
+            # Move weight to NPU if needed (cpu offload may move it back to CPU).
+            if not weight.is_npu:
+                weight = weight.to(f"npu:{torch.npu.current_device()}")
+            # Online MXFP8 quantisation of weights (block_size=32).
+            # qw: [out, in] float8_e4m3fn, w_scale: [out, in//64, 2] uint8.
+            qw, w_scale = torch.ops.npu.npu_dynamic_mx_quant(
+                weight, dst_type=torch.float8_e4m3fn
+            )
+            layer.weight = Parameter(qw.transpose(0, 1), requires_grad=False)
+            layer.weight_scale_inv = Parameter(
+                w_scale.transpose(0, 1), requires_grad=False
+            )
+
+        # Both paths produce weight [in, out] and weight_scale_inv [in//64, out,
+        # 2] as strided transpose views — DO NOT call .contiguous(). The matmul
+        # reduction loop scans the in-dim per output column; the [out, in]
+        # row-major source gives stride-1 access for that scan via the transpose
+        # view (matches msmodelslim's offline layout and vllm-ascend's
+        # AscendW8A8MXFP8DynamicLinearMethod). Calling .contiguous() physically
+        # reorders to [in, out] row-major, making the inner-loop stride = out and
+        # tanking HBM bandwidth.
+
+        # Cache FP32 bias once to avoid a per-forward dtype conversion + alloc.
+        if (
+            getattr(layer, "bias", None) is not None
+            and layer.bias.dtype != torch.float32
+        ):
+            layer.bias_fp32 = Parameter(
+                layer.bias.data.to(torch.float32), requires_grad=False
+            )
+        else:
+            layer.bias_fp32 = None
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        original_dtype = x.dtype
+        if original_dtype not in (torch.float16, torch.bfloat16):
+            x = x.to(torch.bfloat16)
+            original_dtype = torch.bfloat16
+
+        # Flatten to 2D [tokens, hidden] for npu_dynamic_mx_quant
+        input_shape = x.shape
+        x_2d = x.reshape(-1, x.shape[-1])
+
+        # Dynamic MXFP8 activation quantisation
+        qx, input_scale = torch.ops.npu.npu_dynamic_mx_quant(
+            x_2d, dst_type=torch.float8_e4m3fn
+        )
+
+        # MXFP8 matmul (weight & scale already transposed at load time)
+        # Use the cached FP32 bias from process_weights_after_loading; fall back
+        # to per-call conversion if the cache was bypassed (e.g. dynamic bias).
+        if bias is None:
+            quant_bias = None
+        elif (
+            bias is getattr(layer, "bias", None)
+            and getattr(layer, "bias_fp32", None) is not None
+        ):
+            quant_bias = layer.bias_fp32
+        else:
+            quant_bias = bias.to(torch.float32)
+
+        e8m0_dtype = _get_float8_e8m0fnu_dtype()
+        output = torch.ops.npu.npu_quant_matmul(
+            qx,
+            layer.weight,
+            layer.weight_scale_inv,
+            scale_dtype=e8m0_dtype,
+            pertoken_scale=input_scale,
+            pertoken_scale_dtype=e8m0_dtype,
+            bias=quant_bias,
+            output_dtype=original_dtype,
+            group_sizes=[1, 1, MXFP8_BLOCK_SIZE],
+        )
+
+        # Restore original shape (replace last dim with output features)
+        output_shape = list(input_shape[:-1]) + [output.shape[-1]]
+        return output.reshape(output_shape)
+
+
+class NPU_W4A4DynamicLinearMethod(_NPULinearMethodBase):
+
+    def process_weights_after_loading(self, layer):
+        layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
+        layer.weight_scale.data = layer.weight_scale.data.flatten()
+        layer.weight_scale_fp32 = layer.weight_scale.data.to(torch.float32)
+        layer.weight_offset.data = layer.weight_offset.data.flatten()
+        layer.weight.data = torch.ops.npu.npu_convert_weight_to_int4pack(
+            layer.weight.data.to(torch.int32)
+        )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+        tp_rank: Optional[int] = 0,
+    ) -> torch.Tensor:
+        original_dtype = x.dtype
+        quant_out, dynamic_scale = torch.ops.npu.npu_dynamic_quant(
+            x, dst_type=torch.quint4x2
+        )
+        return torch.ops.npu.npu_quant_matmul(
+            quant_out,
+            layer.weight,
+            layer.weight_scale,
+            pertoken_scale=dynamic_scale.flatten(),
+            bias=bias,
+            output_dtype=original_dtype,
+        )
+
+
+class NPUMXFP4W4A8LinearMethod(_NPULinearMethodBase):
+    """Ascend NPU W4A8 online quantization: MXFP4 weights + MXFP8 activations.
+
+    This is a *true* W4(weight) A8(activation) path: it mirrors the offline
+    ``W4A8_MXFP`` kernel (``NPUMXFP4W4A8OfflineLinearMethod``) exactly — the only
+    difference is that the FP4 weights are produced online from BF16/FP16
+    (round-to-nearest, no calibration) instead of being loaded from a msmodelslim
+    checkpoint. An earlier version of this method ran a *dual-level* scheme that
+    also compressed the activation to FP4 (W4A4 compute via
+    ``npu_dual_level_quant_matmul``); that was a large accuracy regression — 4-bit
+    activations — so it was replaced with the single-level FP8-activation path
+    below, aligned with the offline W4A8 implementation.
+
+    Weight quantization (process_weights_after_loading):
+        BF16/FP16 weight → npu_dynamic_mx_quant(dst=float4_e2m1fn_x2) → packed FP4
+        + UE8M0 block scale → npu_format_cast to FRACTAL_NZ → transpose [in//2, out]
+
+    Inference (apply):
+        BF16/FP16 activation → npu_dynamic_mx_quant(dst=float8_e4m3fn)  (A8, FP8)
+        → npu_quant_matmul(x2_dtype=float4_e2m1fn_x2, group_sizes=[0, 0, block])
+
+    Hardware: Ascend 950 (A5) + a recent torch_npu with the FP4 npu_quant_matmul
+    (same requirement as the offline W4A8 path — see that class's docstring).
+    """
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes,
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        """Register an unquantized (``params_dtype``) weight placeholder.
+
+        Online quantization needs its own ``create_weights`` because the
+        checkpoint still holds full-precision BF16/FP16 weights: the loader
+        fills this buffer, then ``process_weights_after_loading`` quantizes it to
+        MXFP4 in place. This differs from the offline/int8 methods, whose weights
+        are created by the scheme's own ``create_weights`` to match the
+        already-quantized (FP8 / uint8-packed) layout the checkpoint provides.
+        """
+        from sglang.srt.layers.parameter import ModelWeightParameter
+
+        output_size_per_partition = sum(output_partition_sizes)
+        weight_loader = extra_weight_attrs.get("weight_loader")
+
+        layer.logical_widths = output_partition_sizes
+        layer.input_size_per_partition = input_size_per_partition
+        layer.output_size_per_partition = output_size_per_partition
+        layer.orig_dtype = params_dtype
+
+        # Load weights in original dtype; quantise to MXFP4 in
+        # process_weights_after_loading.
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition,
+                dtype=params_dtype,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        # Online single-level MXFP4 weight quant, then lay the weight out exactly
+        # like the offline W4A8 path so the same npu_quant_matmul(x2_dtype=fp4)
+        # kernel accepts it. All NPU ops go through torch.ops.npu.* (no torch_npu).
+        fp4_dtype = _get_float4_e2m1fn_x2_dtype()
+
+        weight_fp = layer.weight.data
+        if weight_fp.dtype not in (torch.float16, torch.bfloat16):
+            weight_fp = weight_fp.to(torch.bfloat16)
+        # Move to NPU if needed (cpu offload may have put it on CPU).
+        if not weight_fp.is_npu:
+            weight_fp = weight_fp.to(f"npu:{torch.npu.current_device()}")
+
+        # BF16 -> packed FP4 (float4_e2m1fn_x2, [out, in//2]) + UE8M0 block scale.
+        # npu_dynamic_mx_quant returns the scale as [out, in//64, 2] (3D); older
+        # builds may return [out, in//32] (2D) — handle both before the transpose.
+        qw, w_scale = torch.ops.npu.npu_dynamic_mx_quant(
+            weight_fp, dst_type=fp4_dtype, round_mode="round"
+        )
+
+        # weight: packed FP4 -> FRACTAL_NZ (float8_e4m3fn view) -> transpose
+        # [in//2, out]. Mirror the offline path (no .contiguous() on the NZ view);
+        # view as uint8 first because npu_format_cast only accepts int-dtype tensors.
+        qw_nz = npu_format_cast(
+            qw.view(torch.uint8),
+            NPUACLFormat.ACL_FORMAT_FRACTAL_NZ,
+            customize_dtype=torch.float8_e4m3fn,
+            input_dtype=fp4_dtype,
+        )
+        layer.weight = Parameter(qw_nz.transpose(-1, -2), requires_grad=False)
+
+        # weight_scale -> [in//64, out, 2] to match npu_quant_matmul.
+        if w_scale.dim() == 2:
+            n, k = w_scale.shape
+            w_scale = w_scale.reshape(n, k // 2, 2)
+        layer.weight_scale = Parameter(w_scale.transpose(-3, -2), requires_grad=False)
+
+        # Cache FP32 bias once to avoid a per-forward dtype conversion + alloc.
+        if (
+            getattr(layer, "bias", None) is not None
+            and layer.bias.dtype != torch.float32
+        ):
+            layer.bias_fp32 = Parameter(
+                layer.bias.data.to(torch.float32), requires_grad=False
+            )
+        else:
+            layer.bias_fp32 = None
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        e8m0_dtype = _get_float8_e8m0fnu_dtype()
+        fp4_dtype = _get_float4_e2m1fn_x2_dtype()
+
+        original_dtype = x.dtype
+        if original_dtype not in (torch.float16, torch.bfloat16):
+            x = x.to(torch.bfloat16)
+            original_dtype = torch.bfloat16
+
+        # Flatten to 2D [tokens, hidden] for npu_dynamic_mx_quant.
+        input_shape = x.shape
+        x_2d = x.reshape(-1, x.shape[-1])
+
+        # Dynamic MXFP8 activation quantisation (A8 — FP8, not FP4).
+        quantized_x, dynamic_scale = torch.ops.npu.npu_dynamic_mx_quant(
+            x_2d, dst_type=torch.float8_e4m3fn
+        )
+
+        # Use the cached FP32 bias from process_weights_after_loading; fall back
+        # to per-call conversion if the cache was bypassed (e.g. dynamic bias).
+        if bias is None:
+            quant_bias = None
+        elif (
+            bias is getattr(layer, "bias", None)
+            and getattr(layer, "bias_fp32", None) is not None
+        ):
+            quant_bias = layer.bias_fp32
+        else:
+            quant_bias = bias.to(torch.float32)
+
+        # True W4(weight)A8(activation) matmul, identical to the offline path.
+        output = torch.ops.npu.npu_quant_matmul(
+            quantized_x,
+            layer.weight,
+            layer.weight_scale,
+            scale_dtype=e8m0_dtype,
+            pertoken_scale=dynamic_scale,
+            pertoken_scale_dtype=e8m0_dtype,
+            bias=quant_bias,
+            output_dtype=original_dtype,
+            x2_dtype=fp4_dtype,
+            group_sizes=[0, 0, MXFP4_BLOCK_SIZE],
+        )
+
+        # Restore original shape (replace last dim with output features).
+        output_shape = list(input_shape[:-1]) + [output.shape[-1]]
+        return output.reshape(output_shape)
+
+
+class NPUMXFP4W4A8OfflineLinearMethod(_NPULinearMethodBase):
+    """Ascend NPU offline W4A8 (ModelSlim ``W4A8_MXFP``): packed-FP4 weights + MXFP8 activations.
+
+    Kernel for the offline ModelSlimMXFP4W4A8Scheme (delegated as ``self.kernel``).
+    The msmodelslim ``W4A8_MXFP`` checkpoint stores weights as *packed FP4*
+    (``pack_fp4_to_uint8`` → ``uint8`` shape ``[out, in//2]``) plus UE8M0 block
+    scales (``uint8`` shape ``[out, in//group_size]``):
+
+      process_weights_after_loading:
+        weight (uint8 packed FP4 [out, in//2]) → npu_format_cast(29,
+            customize_dtype=float8_e4m3fn, input_dtype=float4_e2m1fn_x2) → FRACTAL_NZ
+            → transpose [in//2, out]
+        weight_scale [out, in/32] → reshape [out, in/64, 2] → transpose → [in/64, out, 2]
+
+      apply:
+        BF16/FP16 activation → npu_dynamic_mx_quant(dst=float8_e4m3fn)  (A8, MXFP8)
+        → npu_quant_matmul(x2_dtype=float4_e2m1fn_x2, group_sizes=[0, 0, block])
+
+    Mirrors vllm-ascend ``AscendW4A8MXFPDynamicLinearMethod`` exactly (Ascend 950/A5).
+    The weight is cast to FRACTAL_NZ then transposed; ``npu_dynamic_mx_quant`` already
+    returns a 3D ``[tokens, in//64, 2]`` block scale so the matmul needs no extra
+    scale-layout normalization.
+
+    ⚠️ REQUIRES a recent torch_npu build for the FP4 ``npu_quant_matmul``. On the
+    A5 this device forces ``allow_internal_format=False`` (the NZ cast still produces
+    a ``FRACTAL_NZ_C0_16`` tensor, which is fine). Older torch_npu (e.g.
+    ``2.10.0.dev20260320``) had a broken FP4 matmul that rejected the NZ weight in
+    *prefill* with ``x2 should be in ... nz format, but it is 2``;
+    ``2.10.0.post1.dev20260624`` (and later) runs the vllm-aligned NZ path
+    correctly. If you hit ``it is 2``, update torch_npu — do NOT "fix" it by
+    switching the weight to ND.
+
+    ⚠️ A ``atb::OperationSetup`` *segfault during decode* (not prefill) is a
+    DIFFERENT, unrelated issue: it is the eager-decode ``ascend`` attention
+    backend, NOT this matmul (verified by stage-sync bisection — qkv's matmul
+    syncs clean, the fault surfaces at the entry-sync of the next layer, i.e. the
+    decode attention between qkv and o_proj). Run with the NPU decode graph (do
+    NOT pass ``--disable-cuda-graph``); graph mode is the NPU default and what
+    vllm uses. This attention issue is model-agnostic and out of scope for W4A8.
+
+    This is a true W4(weight) A8(activation) single-level matmul. The *online*
+    ``NPUMXFP4W4A8LinearMethod`` now uses this exact apply path — the only
+    difference is that it quantizes BF16/FP16 weights to FP4 at load time instead
+    of loading them from a msmodelslim checkpoint. ``group_size`` is fixed at 32
+    by the ``W4A8_MXFP`` export format.
+    """
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        # Mirror vllm-ascend AscendW4A8MXFPDynamicLinearMethod: cast the packed-FP4
+        # weight to FRACTAL_NZ then transpose. All NPU ops go through
+        # torch.ops.npu.* (no torch_npu). Requires a recent torch_npu build (see
+        # class docstring): older builds reject the NZ weight ("x2 ... it is 2").
+        fp4_dtype = _get_float4_e2m1fn_x2_dtype()
+
+        # weight: packed-FP4 uint8 [out, in//2] -> FRACTAL_NZ (float8_e4m3fn view)
+        # -> transpose to [in//2, out].
+        layer.weight.data = npu_format_cast(
+            layer.weight.data,
+            NPUACLFormat.ACL_FORMAT_FRACTAL_NZ,
+            customize_dtype=torch.float8_e4m3fn,
+            input_dtype=fp4_dtype,
+        )
+        layer.weight.data = layer.weight.data.transpose(-1, -2)
+        # weight_scale: [out, in/32] uint8 -> [in/64, out, 2].
+        n, k = layer.weight_scale.data.shape
+        layer.weight_scale.data = layer.weight_scale.data.reshape(
+            n, k // 2, 2
+        ).transpose(-3, -2)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        e8m0_dtype = _get_float8_e8m0fnu_dtype()
+        fp4_dtype = _get_float4_e2m1fn_x2_dtype()
+
+        original_dtype = x.dtype
+        if original_dtype not in (torch.float16, torch.bfloat16):
+            x = x.to(torch.bfloat16)
+            original_dtype = torch.bfloat16
+
+        # Flatten to 2D [tokens, hidden] for npu_dynamic_mx_quant.
+        input_shape = x.shape
+        x_2d = x.reshape(-1, x.shape[-1])
+
+        # Dynamic MXFP8 activation quantisation (A8).
+        quantized_x, dynamic_scale = torch.ops.npu.npu_dynamic_mx_quant(
+            x_2d, dst_type=torch.float8_e4m3fn
+        )
+
+        if bias is not None and bias.dtype != torch.float32:
+            bias = bias.to(torch.float32)
+
+        # W4(weight)A8(activation) matmul, mirroring vllm-ascend exactly.
+        output = torch.ops.npu.npu_quant_matmul(
+            quantized_x,
+            layer.weight,
+            layer.weight_scale,
+            scale_dtype=e8m0_dtype,
+            pertoken_scale=dynamic_scale,
+            pertoken_scale_dtype=e8m0_dtype,
+            bias=bias,
+            output_dtype=original_dtype,
+            x2_dtype=fp4_dtype,
+            group_sizes=[0, 0, MXFP4_BLOCK_SIZE],
+        )
+
+        # Restore original shape (replace last dim with output features).
+        output_shape = list(input_shape[:-1]) + [output.shape[-1]]
+        return output.reshape(output_shape)

python/sglang/srt/hardware_backend/npu/utils.py

@@ -0,0 +1,366 @@
+import functools
+import logging
+import sys
+from enum import IntEnum
+from typing import TYPE_CHECKING, Callable
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.utils import get_npu_memory_capacity, is_npu
+
+if TYPE_CHECKING:
+    from sglang.srt.server_args import ServerArgs
+
+logger = logging.getLogger(__name__)
+_is_npu = is_npu()
+indexer_weight_stream = None
+gva_is_inited = False
+
+
+class NPUACLFormat(IntEnum):
+    ACL_FORMAT_UNDEFINED = -1
+    ACL_FORMAT_ND = 2
+    ACL_FORMAT_FRACTAL_NZ = 29
+
+
+class FusedMoEMode(IntEnum):
+    FUSED_DEEP_MOE = 1
+    DISPATCH_FFN_COMBINE = 2
+
+
+def _call_once(fn: Callable):
+
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        if getattr(fn, "_has_been_called", False):
+            logger.debug("Function {} has already been called.", fn.__name__)
+            return
+
+        fn._has_been_called = True
+        return fn(*args, **kwargs)
+
+    return wrapper
+
+
+def set_default_server_args(args: "ServerArgs"):
+    """
+    Set default server arguments for NPU backend.
+    """
+
+    # NPU only works with "ascend" attention backend for now
+    args.attention_backend = "ascend"
+    args.prefill_attention_backend = "ascend"
+    args.decode_attention_backend = "ascend"
+    if args.page_size is None:
+        args.page_size = 128
+
+    # NPU memory settings
+    decode = args.cuda_graph_config.decode
+    npu_mem = get_npu_memory_capacity()
+    if npu_mem <= 32 * 1024:
+        # Ascend 910B4,910B4_1
+        # (chunked_prefill_size 4k, max_bs 16 if tp < 4 else 64)
+        if args.chunked_prefill_size is None:
+            args.chunked_prefill_size = 4 * 1024
+        if decode.max_bs is None:
+            if args.tp_size < 4:
+                decode.max_bs = 16
+            else:
+                decode.max_bs = 64
+    elif npu_mem <= 64 * 1024:
+        # Ascend 910B1,910B2,910B2C,910B3,910_9391,910_9392,910_9381,910_9382,910_9372,910_9362
+        # (chunked_prefill_size 8k, max_bs 64 if tp < 4 else 256)
+        if args.chunked_prefill_size is None:
+            args.chunked_prefill_size = 8 * 1024
+        if decode.max_bs is None:
+            if args.tp_size < 4:
+                decode.max_bs = 64
+            else:
+                decode.max_bs = 256
+
+    # NPU does not support CustomAllReduce
+    args.disable_custom_all_reduce = True
+
+    # handles hierarchical cache configs
+    if args.enable_hierarchical_cache:
+        args.hicache_io_backend = "kernel_ascend"
+        if args.use_mla_backend():
+            args.hicache_mem_layout = "page_first_kv_split"
+        else:
+            args.hicache_mem_layout = "page_first_direct"
+
+
+@_call_once
+def init_npu_backend():
+    """
+    Initialize NPU backend. This function should be called only once.
+    """
+
+    assert _is_npu, "NPU backend initialization called on non-NPU device."
+
+    try:
+        import custom_ops  # noqa: F401
+        import sgl_kernel_npu  # noqa: F401
+    except ImportError as e:
+        logger.warning("NPU custom kernel packages unavailable: %s", e)
+
+    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 _is_nz_aligned(tensor: torch.Tensor) -> bool:
+    """Check whether the last two dims satisfy FRACTAL_NZ alignment rules.
+
+    Ascend FRACTAL_NZ requires:
+      BF16 / FP16 : both dims divisible by 16
+      INT8         : k % 16 == 0  and  n % 32 == 0
+      INT4         : k % 16 == 0  and  n % 64 == 0
+      FP4          : both dims divisible by 64
+    """
+    if tensor.dim() < 2:
+        return False
+    k, n = tensor.shape[-2], tensor.shape[-1]
+    if tensor.dtype in (torch.bfloat16, torch.float16):
+        return k % 16 == 0 and n % 16 == 0
+    if tensor.dtype == torch.int8:
+        return k % 16 == 0 and n % 32 == 0
+    if tensor.dtype in (torch.uint8, torch.int32):
+        # INT4 is typically packed into uint8/int32; be conservative
+        return k % 16 == 0 and n % 64 == 0
+    return True
+
+
+def npu_format_cast(
+    tensor: torch.Tensor,
+    acl_format: NPUACLFormat = NPUACLFormat.ACL_FORMAT_FRACTAL_NZ,
+    *,
+    customize_dtype=None,
+    input_dtype=None,
+) -> 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.
+        customize_dtype / input_dtype: packed-FP4 unpack kwargs (e.g.
+            ``customize_dtype=torch.float8_e4m3fn``,
+            ``input_dtype=torch.float4_e2m1fn_x2``). When either is set the unpack
+            kwargs are forwarded to the op and the ``_is_nz_aligned`` ND fallback
+            is skipped: the FP4 matmul strictly requires FRACTAL_NZ, so a silent
+            ND fallback would corrupt results.
+
+    Returns:
+        torch.Tensor: The tensor cast to the specified NPU ACL format.
+    """
+
+    if not _is_npu:
+        return tensor
+
+    if envs.SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT.get():
+        return tensor
+
+    if tensor.device == torch.device("cpu"):
+        logger.warning_once(
+            "Warning: The conversion from 'ND' to 'NZ' does not work on the CPU. "
+            "Please disable offloading, otherwise the performance will be "
+            "significantly reduced. --dit-cpu-offload false"
+        )
+        return tensor
+
+    # Skip format cast for meta tensors (used in offloader)
+    if tensor.device.type == "meta":
+        return tensor
+
+    # Packed-FP4 → FRACTAL_NZ: forward the unpack kwargs to the op, and skip the
+    # _is_nz_aligned ND fallback — the FP4 matmul strictly requires NZ, so a
+    # silent ND fallback would corrupt results.
+    if customize_dtype is not None or input_dtype is not None:
+        return torch.ops.npu.npu_format_cast(
+            tensor,
+            int(acl_format),
+            customize_dtype=customize_dtype,
+            input_dtype=input_dtype,
+        )
+
+    if acl_format == NPUACLFormat.ACL_FORMAT_FRACTAL_NZ and not _is_nz_aligned(tensor):
+        k, n = tensor.shape[-2], tensor.shape[-1]
+        logger.warning_once(
+            "Skipping FRACTAL_NZ format cast: tensor shape (%d, %d) dtype %s "
+            "is not aligned to NZ requirements. Falling back to 'ND' format, "
+            "which may reduce NPU performance.",
+            k,
+            n,
+            tensor.dtype,
+        )
+        return tensor
+
+    return torch.ops.npu.npu_format_cast(tensor, acl_format.value)
+
+
+def get_indexer_weight_stream():
+    global indexer_weight_stream
+    if indexer_weight_stream is None:
+        indexer_weight_stream = torch.npu.Stream()
+    return indexer_weight_stream
+
+
+def init_zbal(world_size, gpu_id, world_rank, do_check=True):
+    """
+    init zbal, if is mix alloc mode, only register for sma & comm
+    """
+    zbal_mem_size = envs.SGLANG_ZBAL_LOCAL_MEM_SIZE.get()
+    if not zbal_mem_size > 0:
+        return 1
+
+    global gva_is_inited
+    from zbal import is_mix_alloc, switch_to_allocator, zbal_init
+
+    if is_mix_alloc():
+        switch_to_allocator()
+        # use lazy init for mix alloc
+        return 1
+    else:
+        if envs.SGLANG_ZBAL_BOOTSTRAP_URL.get():
+            ret = zbal_init(
+                world_size,
+                gpu_id,
+                world_rank,
+                zbal_mem_size * (1024**2),
+                ip_port=envs.SGLANG_ZBAL_BOOTSTRAP_URL.get(),
+            )
+        else:
+            ret = zbal_init(world_size, gpu_id, world_rank, zbal_mem_size * (1024**2))
+
+        gva_is_inited = True
+
+        if do_check and not ret:
+            logger.error("[ZBAL] zbal init failed!")
+            sys.exit(-1)
+
+        return ret
+
+
+def lazy_init_zbal_gva_mem(
+    device, gpu_id, world_rank, world_size, cpu_group=None, do_check=True
+):
+    """
+    lazy init zbal gva mem, keep weights and kv remains alloc by dma vmm to avoid memory fragment
+    """
+    from zbal import is_mix_alloc, zbal_init
+
+    if not is_mix_alloc():
+        logger.info(
+            "lazy init is supported only in mix alloc mode, this action will be passed"
+        )
+        return 1
+
+    global gva_is_inited
+    from sglang.srt.utils.common import get_available_gpu_memory
+
+    # TODO need to use allgather if you want use total_memory stats from mem_get_info as unbalance os
+    total_memory = 61.2  # 2.5GB for other (workspace & os) outside torch
+    free_gpu_memory = get_available_gpu_memory(
+        device,
+        gpu_id,
+        distributed=world_size > 1,
+        cpu_group=cpu_group,
+        empty_cache=True,
+    )
+
+    used_memory = total_memory - free_gpu_memory
+
+    used_memory_in_mb = int(used_memory * 1024)
+    gva_in_mb = envs.SGLANG_ZBAL_LOCAL_MEM_SIZE.get() - used_memory_in_mb
+    gva_in_mb = gva_in_mb - gva_in_mb % 128  # align to 128MB
+    print(f"[ZBAL] rank {world_rank} allocated {gva_in_mb} MB gva space.")
+
+    assert not gva_is_inited, "zbal gva should be inited only once"
+    # zbal_set_logger_level(0)
+    if envs.SGLANG_ZBAL_BOOTSTRAP_URL.get():
+        res = zbal_init(
+            world_size,
+            gpu_id,
+            world_rank,
+            gva_in_mb * (1024**2),
+            ip_port=envs.SGLANG_ZBAL_BOOTSTRAP_URL.get(),
+        )
+    else:
+        res = zbal_init(world_size, gpu_id, world_rank, gva_in_mb * (1024**2))
+
+    gva_is_inited = True
+    if do_check and not res:
+        logger.error("[ZBAL] zbal lazy init failed!")
+        sys.exit(-1)
+    return res
+
+
+share_stream = None
+routed_stream = None
+
+
+def get_share_stream():
+    global share_stream
+    return share_stream
+
+
+def set_share_stream(stream):
+    global share_stream
+    share_stream = stream
+    # TODO LKL: set stream limit has impact on precision
+    # torch.npu.set_stream_limit(share_stream, 8, 16)
+
+
+def get_routed_stream():
+    global routed_stream
+    return routed_stream
+
+
+def set_routed_stream(stream):
+    global routed_stream
+    routed_stream = stream
+    # TODO LKL: set stream limit has impact on precision
+    # torch.npu.set_stream_limit(routed_stream, 16, 32)
+
+
+def wait_share_stream():
+    stream = get_share_stream()
+    if stream is not None:
+        cur_stream = torch.get_device_module().current_stream()
+        cur_stream.wait_stream(stream)
+
+
+def wait_routed_stream():
+    stream = get_routed_stream()
+    if stream is not None:
+        cur_stream = torch.get_device_module().current_stream()
+        cur_stream.wait_stream(stream)
+
+
+def process_shared_expert(hidden_states, forward_func):
+    stream = get_share_stream()
+    if stream is None:
+        stream = torch.get_device_module().Stream()
+        set_share_stream(stream)
+    stream.wait_stream(torch.get_device_module().current_stream())
+    with torch.get_device_module().stream(stream):
+        shared_output = forward_func(hidden_states)
+    return shared_output
+
+
+def process_routed_expert(hidden_states, topk_output, forward_func):
+    stream = get_routed_stream()
+    if stream is None:
+        stream = torch.get_device_module().Stream()
+        set_routed_stream(stream)
+    stream.wait_stream(torch.get_device_module().current_stream())
+    with torch.get_device_module().stream(stream):
+        shared_output = forward_func(hidden_states, topk_output)
+    return shared_output