[NPU] Add sparsity-driven KV offload for DeepSeek DSA on Ascend (#33089)

Co-authored-by: itcj <itcj@foxmail.com> Co-authored-by: chenyh <chenyh18@mail.ustc.edu.cn> Co-authored-by: hibikid <7514545+hibikid@user.noreply.gitee.com> Co-authored-by: Chengjie Tang <tangcj@sxu.edu.cn> Co-authored-by: wangbiao0814 <wangbiao0814@outlook.com>
2026-09-09 09:13:22 +08:00
parent a8e45f16cc
commit 295132c4a5
9 changed files with 1737 additions and 30 deletions
@@ -923,6 +923,7 @@ class Envs:
    SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT = EnvBool(False)
    SGLANG_NPU_USE_MULTI_STREAM = EnvBool(False)
    SGLANG_NPU_USE_MLAPO = EnvBool(False)
+    SGLANG_NPU_ENABLE_SPARSE_KV_OFFLOAD = EnvBool(False)
    # Forward native implementation for activation gelu tanh for model Skywork-Reward-Gemma-2-27B-v0.2
    SGLANG_NPU_FORWARD_NATIVE_GELUTANH = EnvBool(False)
    # Forward native implementation for gemma rms norm for model Skywork-Reward-Gemma-2-27B-v0.2
@@ -19,6 +19,10 @@ from sglang.srt.hardware_backend.npu.attention.mla_preprocess import (
    is_fia_nz,
    is_mla_preprocess_enabled,
 )
+from sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config import (
+    get_sparsity_driven_kv_offload_sparse_context_len,
+    is_sparsity_driven_kv_offload_enabled,
+)
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
 from sglang.srt.layers.radix_attention import AttentionType
@@ -341,6 +345,29 @@ class AscendAttnBackend(AttentionBackend):
        self.req_to_token = model_runner.req_to_token_pool.req_to_token
        self.graph_mode = False
        self.use_fa = get_bool_env_var("ASCEND_USE_FA", "False")
+        self.enable_sparsity_driven_kv_offload = is_sparsity_driven_kv_offload_enabled(
+            model_config=model_runner.model_config,
+            use_mla_backend=model_runner.use_mla_backend,
+        )
+        self.sparse_kv_manager = None
+        if self.enable_sparsity_driven_kv_offload:
+            from sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.manager import (
+                SparseKVCacheManager,
+                register_sparse_kv_manager,
+            )
+
+            self.sparse_kv_manager = SparseKVCacheManager(
+                model_runner.req_to_token_pool,
+                model_runner.token_to_kv_pool_allocator,
+                sparse_context_len=get_sparsity_driven_kv_offload_sparse_context_len(
+                    model_config=model_runner.model_config
+                ),
+            )
+            register_sparse_kv_manager(self.sparse_kv_manager)
+            logger.info(
+                "Sparsity-driven KV offload is enabled with manager %s.",
+                self.sparse_kv_manager,
+            )
        self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False")
        self.enable_torch_compile = get_flags().capture.enable_torch_compile
        self.speculative_num_draft_tokens = get_spec().speculative_num_draft_tokens
@@ -635,6 +662,14 @@ class AscendAttnBackend(AttentionBackend):
    ) -> ForwardMetadata:
        """Create and store the per-bs ForwardMetadata for CUDA graph capture."""
        metadata = ForwardMetadata()
+        if self.enable_sparsity_driven_kv_offload:
+            from sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.host_callback import (
+                register_npu_host_callback_stream,
+            )
+
+            register_npu_host_callback_stream(
+                torch.npu.current_stream(self.device), self.device
+            )
        metadata.block_tables = self.graph_metadata["block_tables"][:bs, :]
        if self.is_hybrid_swa:
            metadata.block_tables_swa = self.graph_metadata["block_tables_swa"][:bs, :]
@@ -1198,6 +1233,23 @@ class AscendAttnBackend(AttentionBackend):
                k_rope=k_rope,
            )
        if topk_indices is not None:
+            if self.enable_sparsity_driven_kv_offload:
+                from sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.attention import (
+                    forward_sparsity_driven_kv_offload,
+                )
+
+                return forward_sparsity_driven_kv_offload(
+                    self,
+                    q,
+                    k,
+                    v,
+                    layer,
+                    forward_batch,
+                    save_kv_cache,
+                    q_rope,
+                    k_rope,
+                    topk_indices,
+                )
            return self.forward_sparse(
                q,
                k,
@@ -2551,6 +2603,23 @@ class AscendAttnBackend(AttentionBackend):
            # MLAPO does saving kv_cache
            save_kv_cache = False
        if topk_indices is not None:
+            if self.enable_sparsity_driven_kv_offload:
+                from sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.attention import (
+                    forward_sparsity_driven_kv_offload,
+                )
+
+                return forward_sparsity_driven_kv_offload(
+                    self,
+                    q,
+                    k,
+                    v,
+                    layer,
+                    forward_batch,
+                    save_kv_cache,
+                    q_rope,
+                    k_rope,
+                    topk_indices,
+                )
            return self.forward_sparse(
                q,
                k,
@@ -552,33 +552,42 @@ class NPUMLATokenToKVPool(MLATokenToKVPool):
        self.kv_lora_rank = kv_lora_rank
        self.qk_rope_head_dim = qk_rope_head_dim
        self.index_head_dim = index_head_dim
+        self.enable_sparsity_driven_kv_offload = (
+            envs.SGLANG_NPU_ENABLE_SPARSE_KV_OFFLOAD.get()
+        )
+        if self.enable_sparsity_driven_kv_offload and self.index_head_dim is None:
+            raise ValueError("Sparsity-driven KV offload requires an index KV cache.")

        self.custom_mem_pool = None

        with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
            # The padded slot 0 is used for writing dummy outputs from padded tokens.
-            self.k_buffer = torch.zeros(
-                (
-                    layer_num,
-                    self.size // self.page_size + 1,
-                    self.page_size,
-                    1,
-                    self.kv_lora_rank,
-                ),
-                dtype=self.store_dtype,
-                device=self.device,
-            )
-            self.v_buffer = torch.zeros(
-                (
-                    layer_num,
-                    self.size // self.page_size + 1,
-                    self.page_size,
-                    1,
-                    self.qk_rope_head_dim,
-                ),
-                dtype=self.store_dtype,
-                device=self.device,
-            )
+            if self.enable_sparsity_driven_kv_offload:
+                self.k_buffer = None
+                self.v_buffer = None
+            else:
+                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(
@@ -596,22 +605,34 @@ class NPUMLATokenToKVPool(MLATokenToKVPool):
        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")
+
+        if getattr(self, "k_buffer", None) is not None:
+            for k_cache in self.k_buffer:
+                kv_size_bytes += get_tensor_size_bytes(k_cache)
+        if getattr(self, "v_buffer", None) is not None:
+            for v_cache in self.v_buffer:
+                kv_size_bytes += get_tensor_size_bytes(v_cache)
+        if getattr(self, "index_k_buffer", None) is not None:
            for index_k_cache in self.index_k_buffer:
                kv_size_bytes += get_tensor_size_bytes(index_k_cache)
        return kv_size_bytes

+    def _raise_if_native_kv_cache_disabled(self):
+        if (
+            getattr(self, "k_buffer", None) is None
+            or getattr(self, "v_buffer", None) is None
+        ):
+            raise RuntimeError(
+                "Native NPU MLA device KV cache is disabled; "
+                "k_buffer/v_buffer are not available. Use the sparse KV manager "
+                "path, or re-enable native device KV cache for this code path."
+            )
+
    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)
+        self._raise_if_native_kv_cache_disabled()
        return (
            self.k_buffer[layer_id - self.start_layer],
            self.v_buffer[layer_id - self.start_layer],
@@ -628,6 +649,7 @@ class NPUMLATokenToKVPool(MLATokenToKVPool):
    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)
+        self._raise_if_native_kv_cache_disabled()

        if self.store_dtype != self.dtype:
            return self.k_buffer[layer_id - self.start_layer].view(self.dtype)
@@ -636,6 +658,7 @@ class NPUMLATokenToKVPool(MLATokenToKVPool):
    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)
+        self._raise_if_native_kv_cache_disabled()

        if self.store_dtype != self.dtype:
            return self.v_buffer[layer_id - self.start_layer].view(self.dtype)
@@ -644,6 +667,8 @@ class NPUMLATokenToKVPool(MLATokenToKVPool):
    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 getattr(self, "index_k_buffer", None) is None:
+            raise RuntimeError("NPU MLA index KV cache is not allocated.")

        if self.store_dtype != self.dtype:
            return self.index_k_buffer[layer_id - self.start_layer].view(self.dtype)
@@ -651,6 +676,7 @@ class NPUMLATokenToKVPool(MLATokenToKVPool):

    # for disagg
    def get_contiguous_buf_infos(self):
+        self._raise_if_native_kv_cache_disabled()
        # 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)
@@ -681,6 +707,7 @@ class NPUMLATokenToKVPool(MLATokenToKVPool):
        cache_v: torch.Tensor,
    ):
        loc, _, _ = unwrap_write_loc(loc_info)
+        self._raise_if_native_kv_cache_disabled()
        layer_id = layer.layer_id
        if cache_k.dtype != self.dtype:
            cache_k = cache_k.to(self.dtype)
@@ -0,0 +1,318 @@
+"""Ascend attention path backed by sparsity-driven KV offload."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+import torch_npu
+
+from sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.manager import (
+    _wait_stream_event,
+    normalize_batch_topk_indices,
+)
+from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
+
+if TYPE_CHECKING:
+    from sglang.srt.hardware_backend.npu.attention.ascend_backend import (
+        AscendAttnBackend,
+    )
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+def _get_sparse_kv_manager(backend: AscendAttnBackend):
+    if backend.sparse_kv_manager is None:
+        raise RuntimeError(
+            "Sparsity-driven KV offload is disabled or was not initialized."
+        )
+    return backend.sparse_kv_manager
+
+
+def _expand_dsa_sparse_indices(topk_indices: torch.Tensor) -> torch.Tensor:
+    """Expand [T, K] to [T, 1, K] for NPU sparse attention."""
+    if topk_indices.dim() == 2:
+        return topk_indices.unsqueeze(-2)
+    return topk_indices
+
+
+def forward_sparsity_driven_kv_offload(
+    backend: AscendAttnBackend,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    layer: RadixAttention,
+    forward_batch: ForwardBatch,
+    save_kv_cache: bool = True,
+    q_rope: Optional[torch.Tensor] = None,
+    k_rope: Optional[torch.Tensor] = None,
+    topk_indices: Optional[torch.Tensor] = None,
+):
+    """Run sparse attention using host-offloaded compact MLA KV."""
+    del v
+    if q_rope is None or k_rope is None or topk_indices is None:
+        raise ValueError(
+            "Sparsity-driven KV offload requires q_rope, k_rope, and topk_indices."
+        )
+
+    is_prefill = forward_batch.forward_mode.is_extend_without_speculative()
+
+    q_nope, q_pe = q, q_rope
+    k_nope = k.view(-1, layer.tp_k_head_num, backend.kv_lora_rank).contiguous()
+    k_pe = k_rope.view(-1, layer.tp_k_head_num, backend.qk_rope_head_dim).contiguous()
+    sparse_kv_manager = _get_sparse_kv_manager(backend)
+    stream = torch.npu.current_stream(backend.device)
+
+    if save_kv_cache:
+        sparse_kv_manager.offload_v2(k_nope, k_pe, layer, forward_batch, stream)
+
+    if is_prefill:
+        if backend.forward_metadata.actual_seq_lengths_q is not None:
+            actual_seq_qlen = backend.forward_metadata.actual_seq_lengths_q
+        else:
+            actual_seq_qlen = torch.cumsum(forward_batch.extend_seq_lens, dim=0)
+    elif backend.forward_metadata.actual_seq_lengths_q is None:
+        if (
+            forward_batch.forward_mode.is_draft_extend_v2()
+            or forward_batch.forward_mode.is_target_verify()
+        ):
+            actual_seq_qlen = (
+                torch.arange(
+                    backend.speculative_num_draft_tokens,
+                    backend.speculative_num_draft_tokens + q.shape[0],
+                    backend.speculative_num_draft_tokens,
+                    dtype=torch.int32,
+                )
+                .to(q.device)
+                .to(torch.int32)
+            )
+        else:
+            actual_seq_qlen = (
+                torch.arange(1, q.shape[0] + 1).to(q.device).to(torch.int32)
+            )
+    else:
+        actual_seq_qlen = backend.forward_metadata.actual_seq_lengths_q
+
+    if backend.forward_metadata.actual_seq_lengths_kv is not None:
+        actual_seq_lengths_kv = backend.forward_metadata.actual_seq_lengths_kv
+    elif backend.forward_metadata.seq_lens_cpu_int is not None:
+        actual_seq_lengths_kv = backend.forward_metadata.seq_lens_cpu_int
+    else:
+        actual_seq_lengths_kv = backend.forward_metadata.seq_lens
+
+    if (
+        is_prefill
+        and is_dsa_enable_prefill_cp()
+        and forward_batch.attn_cp_metadata is not None
+    ):
+        attn_out = backend.do_cp_balance_attn(
+            q_nope,
+            k_nope,
+            q_pe,
+            k_pe,
+            topk_indices,
+            layer,
+            actual_seq_qlen,
+            actual_seq_lengths_kv,
+        )
+    elif forward_batch.forward_mode.is_decode():
+        batch_size = forward_batch.batch_size
+        num_kv_heads = layer.tp_k_head_num
+        num_query_heads = layer.tp_q_head_num
+        nope_head_dim = backend.kv_lora_rank
+        rope_head_dim = backend.qk_rope_head_dim
+
+        topk_2d_input = normalize_batch_topk_indices(topk_indices)
+        effective_topk_length = topk_2d_input.shape[1]
+        selected_kv_length = sparse_kv_manager.sparse_context_len
+        if effective_topk_length <= 0:
+            raise RuntimeError("SFA BSND compact path expects a positive top-k length.")
+        if effective_topk_length > selected_kv_length:
+            raise RuntimeError(
+                "DSA top-k length exceeds sparse KV device cache capacity: "
+                f"topk_len={effective_topk_length}, "
+                f"sparse_context_len={selected_kv_length}."
+            )
+        if effective_topk_length == selected_kv_length:
+            topk_2d = topk_2d_input.contiguous()
+        else:
+            topk_2d = torch.full(
+                (batch_size, selected_kv_length),
+                -1,
+                dtype=topk_2d_input.dtype,
+                device=topk_2d_input.device,
+            )
+            topk_2d[:, :effective_topk_length] = topk_2d_input
+            topk_2d = topk_2d.contiguous()
+
+        assert num_kv_heads == 1, (
+            "FIA_v2 MLA selected KV path expects KV_N == 1, "
+            f"got num_kv_heads={num_kv_heads}"
+        )
+
+        padded_query_heads = q_nope.numel() // (batch_size * nope_head_dim)
+        assert padded_query_heads >= num_query_heads, (
+            "query head count mismatch: "
+            f"padded_query_heads={padded_query_heads}, "
+            f"num_query_heads={num_query_heads}"
+        )
+
+        # Materialize the compact top-k KV for the current attention step:
+        # device-cache hits and host misses are copied into this buffer, the
+        # device cache is refilled from it, and sparse attention consumes it
+        # directly.
+        selected_kv_buffer = torch.zeros(
+            (
+                batch_size,
+                selected_kv_length,
+                num_kv_heads,
+                nope_head_dim + rope_head_dim,
+            ),
+            dtype=k.dtype,
+            device=backend.device,
+        )
+        sparse_kv_manager.materialize_selected_kv(
+            layer, forward_batch, topk_indices, selected_kv_buffer, stream
+        )
+
+        _wait_stream_event(stream, sparse_kv_manager.hit_done)
+        _wait_stream_event(stream, sparse_kv_manager.miss_done)
+
+        selected_k_nope, selected_k_rope = selected_kv_buffer.split(
+            [nope_head_dim, rope_head_dim], dim=-1
+        )
+
+        topk_valid = topk_2d >= 0
+        if forward_batch.seq_lens is not None:
+            valid_rows = (forward_batch.seq_lens[:batch_size] > 0).view(batch_size, 1)
+            topk_valid = topk_valid & valid_rows
+
+        actual_seq_lengths_kv = (
+            topk_valid.sum(dim=1)
+            .clamp(min=1, max=selected_kv_length)
+            .to(device=q_nope.device, dtype=torch.int32)
+            .contiguous()
+        )
+        actual_seq_lengths_query = torch.ones(
+            batch_size, dtype=torch.int32, device=q_nope.device
+        ).contiguous()
+
+        compact_indices = (
+            torch.arange(selected_kv_length, device=q_nope.device, dtype=torch.int32)
+            .view(1, 1, 1, selected_kv_length)
+            .expand(batch_size, 1, num_kv_heads, selected_kv_length)
+            .clone()
+        )
+        compact_valid = topk_valid.view(batch_size, 1, 1, selected_kv_length).expand(
+            batch_size, 1, num_kv_heads, selected_kv_length
+        )
+        sparse_indices = torch.where(
+            compact_valid,
+            compact_indices,
+            torch.full_like(compact_indices, -1),
+        ).contiguous()
+
+        empty_rows = (topk_valid.sum(dim=1) == 0).view(batch_size, 1, 1)
+        sparse_indices[:, :, :, 0] = torch.where(
+            empty_rows.expand(batch_size, 1, num_kv_heads),
+            torch.zeros(
+                (batch_size, 1, num_kv_heads),
+                dtype=torch.int32,
+                device=q_nope.device,
+            ),
+            sparse_indices[:, :, :, 0],
+        )
+
+        q_nope_sfa = q_nope.view(
+            batch_size, 1, padded_query_heads, nope_head_dim
+        ).contiguous()
+        q_rope_sfa = q_pe.view(
+            batch_size, 1, padded_query_heads, rope_head_dim
+        ).contiguous()
+        k_nope_sfa = selected_k_nope.contiguous()
+        k_rope_sfa = selected_k_rope.contiguous()
+
+        assert q_nope_sfa.shape == (
+            batch_size,
+            1,
+            padded_query_heads,
+            nope_head_dim,
+        )
+        assert q_rope_sfa.shape == (
+            batch_size,
+            1,
+            padded_query_heads,
+            rope_head_dim,
+        )
+        assert k_nope_sfa.shape == (
+            batch_size,
+            selected_kv_length,
+            num_kv_heads,
+            nope_head_dim,
+        )
+        assert k_rope_sfa.shape == (
+            batch_size,
+            selected_kv_length,
+            num_kv_heads,
+            rope_head_dim,
+        )
+
+        ret = torch_npu.npu_sparse_flash_attention(
+            q_nope_sfa,
+            k_nope_sfa,
+            k_nope_sfa,
+            sparse_indices,
+            layer.scaling,
+            actual_seq_lengths_query=actual_seq_lengths_query,
+            actual_seq_lengths_kv=actual_seq_lengths_kv,
+            query_rope=q_rope_sfa,
+            key_rope=k_rope_sfa,
+            sparse_block_size=1,
+            layout_query="BSND",
+            layout_kv="BSND",
+            sparse_mode=0,
+            attention_mode=2,
+            return_softmax_lse=False,
+        )
+
+        _wait_stream_event(stream, sparse_kv_manager.refill_done)
+        _wait_stream_event(stream, sparse_kv_manager.slot_map_done)
+
+        attn_out = ret[0] if isinstance(ret, tuple) else ret
+        attn_out = attn_out[:, :, :num_query_heads, :].reshape(
+            batch_size, num_query_heads * nope_head_dim
+        )
+    else:
+        if is_prefill:
+            k_nope_sfa, k_pe_sfa = sparse_kv_manager.get_forward_kv(
+                layer, forward_batch, stream
+            )
+            forward_actual_seq_lengths_kv = torch.cumsum(forward_batch.seq_lens, dim=0)
+        else:
+            k_nope_sfa, k_pe_sfa = k_nope, k_pe
+            forward_actual_seq_lengths_kv = actual_seq_lengths_kv
+
+        topk_indices = _expand_dsa_sparse_indices(topk_indices)
+        attn_out, _, _ = torch_npu.npu_sparse_flash_attention(
+            query=q_nope,
+            key=k_nope_sfa,
+            value=k_nope_sfa,
+            query_rope=q_pe,
+            key_rope=k_pe_sfa,
+            sparse_indices=topk_indices,
+            scale_value=layer.scaling,
+            actual_seq_lengths_query=actual_seq_qlen.to(
+                device=q_nope.device, dtype=torch.int32
+            ),
+            actual_seq_lengths_kv=forward_actual_seq_lengths_kv.to(
+                device=q_nope.device, dtype=torch.int32
+            ),
+            sparse_block_size=1,
+            layout_query="TND",
+            layout_kv="TND",
+            sparse_mode=3,
+            attention_mode=2,
+            return_softmax_lse=False,
+        )
+
+    return attn_out
@@ -0,0 +1,96 @@
+"""Configuration and validation for sparsity-driven KV offload."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+from sglang.srt.configs.model_config import (
+    get_dsa_index_head_dim,
+    get_dsa_index_topk,
+    is_deepseek_dsa,
+)
+from sglang.srt.environ import envs
+from sglang.srt.runtime_context import attention_backends, get_schedule
+from sglang.srt.utils.common import is_npu
+
+if TYPE_CHECKING:
+    from sglang.srt.configs.model_config import ModelConfig
+
+
+def is_sparsity_driven_kv_offload_enabled(
+    *,
+    model_config: ModelConfig,
+    use_mla_backend: bool,
+) -> bool:
+    if not envs.SGLANG_NPU_ENABLE_SPARSE_KV_OFFLOAD.get():
+        return False
+
+    prefill_attention_backend, decode_attention_backend = attention_backends()
+    if not (
+        is_npu()
+        and prefill_attention_backend == "ascend"
+        and decode_attention_backend == "ascend"
+        and use_mla_backend
+        and is_deepseek_dsa(model_config.hf_config)
+    ):
+        raise ValueError(
+            "SGLANG_NPU_ENABLE_SPARSE_KV_OFFLOAD requires an NPU "
+            "DSA-family MLA model "
+            "(for example DeepSeek V3.2 or GLM-5.x) using the Ascend MLA "
+            "attention backend."
+        )
+    if get_schedule().max_running_requests is None:
+        raise ValueError(
+            "SGLANG_NPU_ENABLE_SPARSE_KV_OFFLOAD requires max_running_requests "
+            "to be set to bound the per-process host KV allocation."
+        )
+    return True
+
+
+def get_sparsity_driven_kv_offload_sparse_context_len(
+    *,
+    model_config: ModelConfig,
+) -> int:
+    """Return the per-request on-device sparse KV window size."""
+    sparse_context_len = int(get_dsa_index_topk(model_config.hf_config))
+    if sparse_context_len <= 0:
+        raise ValueError(
+            "Sparsity-driven KV offload requires a positive DSA index_topk, "
+            f"got {sparse_context_len}."
+        )
+    return sparse_context_len
+
+
+def get_sparsity_driven_kv_offload_index_head_dim(
+    *,
+    model_config: ModelConfig,
+) -> int:
+    index_head_dim = getattr(model_config, "index_head_dim", None)
+    if index_head_dim is None:
+        index_head_dim = get_dsa_index_head_dim(model_config.hf_config)
+    index_head_dim = int(index_head_dim)
+    if index_head_dim <= 0:
+        raise ValueError(
+            "Sparsity-driven KV offload requires a positive DSA index_head_dim, "
+            f"got {index_head_dim}."
+        )
+    return index_head_dim
+
+
+def get_sparsity_driven_kv_offload_cell_size(
+    *,
+    model_config: ModelConfig,
+    use_mla_backend: bool,
+    num_layers: int,
+    element_size: int,
+) -> Optional[int]:
+    if not is_sparsity_driven_kv_offload_enabled(
+        model_config=model_config,
+        use_mla_backend=use_mla_backend,
+    ):
+        return None
+
+    index_head_dim = get_sparsity_driven_kv_offload_index_head_dim(
+        model_config=model_config
+    )
+    return index_head_dim * num_layers * element_size
@@ -0,0 +1,80 @@
+"""ACL host-callback support for sparsity-driven KV offload."""
+
+import atexit
+import threading
+
+import acl
+import torch
+
+_REPORTERS = {}
+
+
+def _get_stream_ptr(stream) -> int:
+    for attr in ("npu_stream", "stream_ptr", "cuda_stream"):
+        if hasattr(stream, attr):
+            value = getattr(stream, attr)
+            value = value() if callable(value) else value
+            return int(value)
+    raise RuntimeError("cannot get raw NPU stream ptr")
+
+
+class _AclReportThread:
+    def __init__(self, device_index: int):
+        self.device_index = device_index
+        self.ready = threading.Event()
+        self.stop = threading.Event()
+        self.thread_id = None
+        self.streams = set()
+        self.thread = threading.Thread(target=self._loop, daemon=True)
+        self.thread.start()
+        self.ready.wait()
+        atexit.register(self.close)
+
+    def _loop(self):
+        torch.npu.set_device(self.device_index)
+        try:
+            acl.rt.set_device(self.device_index)
+        except Exception:
+            pass
+
+        self.thread_id = threading.current_thread().ident
+        self.ready.set()
+
+        while not self.stop.is_set():
+            acl.rt.process_report(100)
+
+    def subscribe(self, stream):
+        stream_ptr = _get_stream_ptr(stream)
+        if stream_ptr in self.streams:
+            return
+
+        ret = acl.rt.subscribe_report(self.thread_id, stream_ptr)
+        if ret != 0:
+            raise RuntimeError(
+                f"acl.rt.subscribe_report failed, ret={ret}, "
+                f"thread_id={self.thread_id}, stream_ptr={stream_ptr}"
+            )
+        self.streams.add(stream_ptr)
+
+    def close(self):
+        for stream_ptr in list(self.streams):
+            try:
+                acl.rt.unsubscribe_report(self.thread_id, stream_ptr)
+            except Exception:
+                pass
+        self.streams.clear()
+        self.stop.set()
+        self.thread.join(timeout=1)
+
+
+def register_npu_host_callback_stream(stream, device):
+    device_index = device.index
+    if device_index is None:
+        device_index = torch.npu.current_device()
+
+    reporter = _REPORTERS.get(device_index)
+    if reporter is None:
+        reporter = _AclReportThread(device_index)
+        _REPORTERS[device_index] = reporter
+
+    reporter.subscribe(stream)
@@ -0,0 +1,980 @@
+"""Sparsity-driven KV offload manager for the Ascend backend."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, List, Optional, Union
+
+import torch
+from sgl_kernel_npu.sparsity_driven_kv_offload import (
+    create_shm_tensor,
+    slot_map_lookup,
+    unidex_copy_inplace,
+)
+
+from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
+from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+from sglang.srt.mem_cache.memory_pool import (
+    MLATokenToKVPool,
+    ReqToTokenPool,
+)
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
+
+if TYPE_CHECKING:
+    import torch.npu
+
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.managers.schedule_batch import Req
+
+logger = logging.getLogger(__name__)
+
+
+def _record_stream_event(stream, event) -> None:
+    if hasattr(stream, "record_event"):
+        stream.record_event(event)
+    else:
+        event.record(stream)
+
+
+def _wait_stream_event(stream, event) -> None:
+    if hasattr(stream, "wait_event"):
+        stream.wait_event(event)
+    else:
+        event.wait(stream)
+
+
+def normalize_batch_topk_indices(topk_indices: torch.Tensor) -> torch.Tensor:
+    """Normalize DSA top-k indices to [batch, topk] for compact KV copies."""
+    if topk_indices.dim() == 2:
+        return topk_indices
+    if topk_indices.dim() == 3 and topk_indices.shape[1] == 1:
+        return topk_indices[:, 0, :]
+    if (
+        topk_indices.dim() == 4
+        and topk_indices.shape[1] == 1
+        and topk_indices.shape[2] == 1
+    ):
+        return topk_indices[:, 0, 0, :]
+    raise RuntimeError(
+        "Sparsity-driven KV offload expects DSA top-k indices with shape "
+        f"[batch, topk], [batch, 1, topk], or [batch, 1, 1, topk], got "
+        f"{tuple(topk_indices.shape)}."
+    )
+
+
+class SparseKVCacheManager:
+    copy_stream = None
+    miss_shm_cpu_tensor: list = []
+    miss_shm_dev_ptr: Optional[int] = None
+    miss_shm_shape: list = []
+    miss_shm_dtype: list = []
+
+    def __init__(
+        self,
+        req_to_token_pool: ReqToTokenPool,
+        token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
+        sparse_context_len: int,
+    ) -> None:
+        enable_memory_saver = False
+        memory_saver_adapter = TorchMemorySaverAdapter.create(
+            enable=enable_memory_saver
+        )
+
+        # Include the padding row because real request IDs can equal the
+        # configured capacity when row 0 is reserved for graph padding.
+        self.size = int(req_to_token_pool.req_to_token.shape[0])
+        self.max_context_len = req_to_token_pool.max_context_len
+        self.sparse_context_len = int(sparse_context_len)
+        if self.sparse_context_len <= 0:
+            raise ValueError(
+                "SparseKVCacheManager requires a positive sparse_context_len, "
+                f"got {self.sparse_context_len}."
+            )
+        self.device = req_to_token_pool.device
+        paged_kv_cache = token_to_kv_pool_allocator.get_kvcache()
+        if not isinstance(paged_kv_cache, MLATokenToKVPool):
+            raise TypeError(
+                "SparseKVCacheManager requires an MLATokenToKVPool, "
+                f"got {type(paged_kv_cache).__name__}"
+            )
+        self.paged_kv_cache = paged_kv_cache
+        self.start_layer = paged_kv_cache.start_layer
+        # MLA params
+        self.head_num = 1
+        self.kv_lora_rank = self.paged_kv_cache.kv_lora_rank
+        self.qk_rope_head_dim = self.paged_kv_cache.qk_rope_head_dim
+        # kv_cache_dim = kv_lora_rank + qk_rope_head_dim
+        self.head_dim = (
+            self.paged_kv_cache.kv_lora_rank + self.paged_kv_cache.qk_rope_head_dim
+        )
+        self.store_dtype = self.paged_kv_cache.store_dtype
+        self.layer_num = self.paged_kv_cache.layer_num
+        self._materialize_d2d_hit_stream = torch.npu.Stream()
+        self._materialize_h2d_miss_stream = torch.npu.Stream()
+        self._materialize_refill_stream = torch.npu.Stream()
+        self._materialize_slot_map_stream = torch.npu.Stream()
+
+        self.hit_done = torch.npu.Event()
+        self.miss_done = torch.npu.Event()
+        self.refill_done = torch.npu.Event()
+        self.slot_map_done = torch.npu.Event()
+
+        # device KV buffer
+        try:
+            with memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
+                # [bs, ctx_len, head_num, head_dim] for each layer
+                # The padded slot 0 is used for writing dummy outputs from padded tokens.
+                self.device_kv_buffer: list[torch.Tensor] = [
+                    torch.empty(
+                        (
+                            self.size,
+                            self.sparse_context_len,
+                            self.head_num,
+                            self.head_dim,
+                        ),
+                        dtype=self.store_dtype,
+                        device=self.device,
+                    )
+                    for _ in range(self.layer_num)
+                ]
+        except Exception as e:
+            self._raise_buffer_allocation_error("device_kv_buffer", e)
+
+        try:
+            with memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
+                # Reserve the last row for padded requests and ensure token index
+                # `max_context_len` is a valid sentinel column for masked writes.
+                # The row width is also aligned to eight int32 values (32 bytes).
+                self.device_slot_map: list[torch.Tensor] = [
+                    torch.full(
+                        (
+                            self.size + 1,
+                            (self.max_context_len // 8 + 1) * 8,
+                        ),
+                        -1,
+                        dtype=torch.int32,
+                        device=self.device,
+                    )
+                    for _ in range(self.layer_num)
+                ]
+                # Pre-filled all-(-1) template used to reset the slot map with a
+                # fast device-to-device copy instead of an indexed fill.
+                self._device_slot_map_minus_one = torch.full_like(
+                    self.device_slot_map[0], -1
+                )
+        except Exception as e:
+            self._raise_buffer_allocation_error("device_slot_map", e)
+
+        # Host KV buffer
+        # [bs, ctx_len, head_num, head_dim] for each layer
+        # The padded slot 0 is used for writing dummy outputs from padded tokens.
+        self.host_kv_buffer: list[torch.Tensor] = []
+        self.host_ptr_list: list[int] = []
+        self.dev_ptr_list: list[int] = []
+
+        host_kv_shape = (
+            self.size,
+            self.max_context_len,
+            self.head_num,
+            self.head_dim,
+        )
+        logger.info("Sparse KV host buffer shape: %s", host_kv_shape)
+        device_id = torch.npu.current_device()
+
+        try:
+            for layer_idx in range(self.layer_num):
+                shm_cpu_tensor, host_ptr, dev_ptr = create_shm_tensor(
+                    shape=host_kv_shape,
+                    dtype=self.store_dtype,
+                    device_id=device_id,
+                    name=f"host_kv_layer_{layer_idx}_rank_{device_id}",
+                )
+                self.host_kv_buffer.append(shm_cpu_tensor)
+                self.host_ptr_list.append(host_ptr)
+                self.dev_ptr_list.append(dev_ptr)
+        except Exception as e:
+            self._raise_buffer_allocation_error("host_kv_buffer", e)
+        self.host_kv_ctx_len = torch.zeros(
+            (self.size, self.max_context_len), dtype=torch.int32, device="cpu"
+        )
+        self.topk_indices_cpu = None
+        self.token_on_device_cpu = None
+        self.device_token_pos_cpu = None
+        self.current_req_indices_cpu = None
+
+        self._device_cache_slot_ids = torch.arange(
+            self.sparse_context_len, dtype=torch.long, device=self.device
+        )
+        self._slot_map_width = (self.max_context_len // 8 + 1) * 8
+
+        self._install_req_alloc_hook(req_to_token_pool)
+
+    def _raise_buffer_allocation_error(
+        self,
+        buffer_name: str,
+        exc: Exception,
+    ) -> None:
+        raise RuntimeError(
+            "Failed to allocate sparse KV buffer "
+            f"{buffer_name}: req_capacity={self.size}, "
+            f"max_context_len={self.max_context_len}, "
+            f"sparse_context_len={self.sparse_context_len}. "
+            "The sparse KV request capacity may be too large; set a smaller "
+            "--max-running-requests for sparse KV offload."
+        ) from exc
+
+    def init_req(self, req: Req) -> None:
+        if req.is_chunked > 0:
+            return
+        rid = req.req_pool_idx
+        if rid is None:
+            raise RuntimeError(
+                "Cannot initialize sparse KV state before allocating a request pool slot"
+            )
+        current_len = len(req.origin_input_ids)
+        self.host_kv_ctx_len[rid] = current_len
+        self.reset_requests([rid])
+
+    def reset_requests(self, req_ids: List[int]) -> None:
+        if not req_ids:
+            return
+
+        req_ids_tensor = torch.tensor(
+            req_ids, dtype=torch.long, device=self.device
+        ).contiguous()
+        for layer_idx in range(self.layer_num):
+            self.device_slot_map[layer_idx].index_fill_(0, req_ids_tensor, -1)
+
+    def _install_req_alloc_hook(self, req_to_token_pool: ReqToTokenPool) -> None:
+        original_alloc = getattr(
+            req_to_token_pool, "_sparse_kv_original_alloc", req_to_token_pool.alloc
+        )
+        setattr(req_to_token_pool, "_sparse_kv_original_alloc", original_alloc)
+
+        def alloc_with_sparse_reset(reqs: list[Req]) -> Optional[List[int]]:
+            newly_allocated = [req.req_pool_idx is None for req in reqs]
+            req_pool_indices = original_alloc(reqs)
+            if req_pool_indices is not None:
+                self.reset_requests(
+                    [
+                        req_pool_indices[i]
+                        for i, is_new in enumerate(newly_allocated)
+                        if is_new
+                    ]
+                )
+            return req_pool_indices
+
+        setattr(req_to_token_pool, "alloc", alloc_with_sparse_reset)
+
+    def offload(
+        self,
+        k: torch.Tensor,
+        k_rope: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        stream: torch.npu.Stream,
+    ):
+        layer_idx = layer.layer_id - self.start_layer
+        device = k.device
+
+        # k:        [total_token_slots, nhead, dim]
+        # k_rope:   [total_token_slots, nhead, dim]
+        # kv_device: [total_token_slots, nhead, 2*dim]
+        kv_device = torch.cat([k, k_rope], dim=-1)
+
+        # Source row indices into kv_device.
+        # In graph mode this tensor is expected to have a static shape.
+        src_token_indices = forward_batch.out_cache_loc.to(torch.long).contiguous()
+        static_token_slots = int(src_token_indices.numel())
+
+        if forward_batch.forward_mode.is_decode():
+            # decode graph mode:
+            # one token slot per request, padded requests are masked out
+            req_ids = forward_batch.req_pool_indices.to(torch.long)
+            token_pos = (forward_batch.seq_lens - 1).to(torch.long)
+
+            dst_token_indices = (
+                req_ids * self.max_context_len + token_pos
+            ).contiguous()
+
+            # Existing graph decode convention:
+            # padded decode requests carry seq_len == 1.
+            valid_mask = (
+                (forward_batch.seq_lens != 1)
+                & (src_token_indices >= 0)
+                & (req_ids >= 0)
+            ).contiguous()
+
+        else:
+            # prefill graph mode:
+            # assume out_cache_loc is laid out as [B, TOKENS_PER_REQ] flattened row-major
+            if (
+                forward_batch.extend_seq_lens is None
+                or forward_batch.extend_prefix_lens is None
+            ):
+                raise RuntimeError(
+                    "Sparse graph prefill offload requires extend_seq_lens and "
+                    "extend_prefix_lens in ForwardBatch."
+                )
+
+            batch_size = int(forward_batch.req_pool_indices.shape[0])
+            if batch_size <= 0:
+                return
+
+            if static_token_slots % batch_size != 0:
+                raise RuntimeError(
+                    f"out_cache_loc length {static_token_slots} is not divisible by "
+                    f"batch size {batch_size}. Cannot infer graph static token layout."
+                )
+
+            tokens_per_req = static_token_slots // batch_size
+
+            req_ids = forward_batch.req_pool_indices.to(torch.long)
+            extend_seq_lens = forward_batch.extend_seq_lens.to(torch.long)
+            extend_prefix_lens = forward_batch.extend_prefix_lens.to(torch.long)
+
+            local_offsets = (
+                torch.arange(
+                    tokens_per_req,
+                    device=device,
+                    dtype=torch.long,
+                )
+                .unsqueeze(0)
+                .expand(batch_size, tokens_per_req)
+            )
+
+            req_ids_2d = req_ids.unsqueeze(1).expand(batch_size, tokens_per_req)
+            dst_pos_2d = extend_prefix_lens.unsqueeze(1) + local_offsets
+
+            dst_token_indices = (
+                (req_ids_2d * self.max_context_len + dst_pos_2d)
+                .reshape(-1)
+                .contiguous()
+            )
+
+            valid_mask = (
+                (local_offsets < extend_seq_lens.unsqueeze(1)).reshape(-1)
+                & (src_token_indices >= 0)
+                & (req_ids_2d.reshape(-1) >= 0)
+            ).contiguous()
+
+        # Layout check:
+        # kv_device rows:                 [token_slot]
+        # host_kv_buffer[layer] rows:     [req_id, seq_pos]
+        # block dims must match on [nhead, 2*dim]
+        assert kv_device.shape[1:] == self.host_kv_buffer[layer_idx].shape[2:]
+        # torch.npu.synchronize()
+        actual_stream = stream if stream is not None else torch.npu.current_stream()
+        with torch.npu.stream(actual_stream):
+            unidex_copy_inplace(
+                kv_device,
+                self.host_kv_buffer[layer_idx],
+                src_token_indices,
+                dst_token_indices,
+                valid_mask,
+                1,  # kv_device: [token_slot, nhead, 2*dim]
+                2,  # host_kv_buffer: [num_req, max_context_len, nhead, 2*dim]
+                block_dim=48,
+                dst_ptr=self.dev_ptr_list[layer_idx],
+            )
+        # torch.npu.synchronize()
+
+    def offload_v2(
+        self,
+        k: torch.Tensor,
+        k_rope: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        stream: torch.npu.Stream,
+    ):
+        """Offload compact per-forward KV rows into the sparse host KV buffer.
+
+        v1 expects k/k_rope to be full native KV-cache views and therefore uses
+        forward_batch.out_cache_loc as source rows. v2 expects k/k_rope to be
+        compact rows produced by the current forward pass, so source rows are
+        simply [0, num_new_tokens). The native cache slot is kept only as
+        validity metadata.
+
+        Keep src_tensor/dst_tensor/src_index/dst_index/valid_mask explicit so
+        the final copy can be swapped to a custom kernel without changing the
+        graph-friendly index construction.
+        """
+        layer_idx = layer.layer_id - self.start_layer
+        device = k.device
+
+        # k:         [num_new_tokens, nhead, kv_lora_rank]
+        # k_rope:    [num_new_tokens, nhead, qk_rope_head_dim]
+        # kv_device: [num_new_tokens, nhead, kv_lora_rank + qk_rope_head_dim]
+        src_tensor = torch.cat([k, k_rope], dim=-1).contiguous()
+        dst_tensor = self.host_kv_buffer[layer_idx]
+        src_index = torch.arange(src_tensor.shape[0], device=device, dtype=torch.long)
+        num_src_rows = int(src_tensor.shape[0])
+
+        if forward_batch.forward_mode.is_decode():
+            req_ids = forward_batch.req_pool_indices.to(torch.long)
+            token_pos = (forward_batch.seq_lens - 1).to(torch.long)
+            cache_loc = forward_batch.out_cache_loc.to(torch.long)
+
+            if int(req_ids.shape[0]) != num_src_rows:
+                raise RuntimeError(
+                    "Sparse v2 decode offload expects compact KV rows to match "
+                    f"batch size, got {num_src_rows} and {int(req_ids.shape[0])}."
+                )
+            if int(cache_loc.shape[0]) != num_src_rows:
+                raise RuntimeError(
+                    "Sparse v2 decode offload expects out_cache_loc rows to match "
+                    f"compact KV rows, got {int(cache_loc.shape[0])} and "
+                    f"{num_src_rows}."
+                )
+
+            dst_index = (req_ids * self.max_context_len + token_pos).contiguous()
+            valid_mask = (
+                (forward_batch.seq_lens != 1)
+                & (cache_loc >= 0)
+                & (req_ids >= 0)
+                & (token_pos >= 0)
+                & (token_pos < self.max_context_len)
+            ).contiguous()
+        else:
+            if (
+                forward_batch.extend_seq_lens is None
+                or forward_batch.extend_prefix_lens is None
+            ):
+                raise RuntimeError(
+                    "Sparse v2 prefill offload requires extend_seq_lens and "
+                    "extend_prefix_lens in ForwardBatch."
+                )
+
+            req_ids = forward_batch.req_pool_indices.to(torch.long)
+            extend_seq_lens = forward_batch.extend_seq_lens.to(torch.long)
+            extend_prefix_lens = forward_batch.extend_prefix_lens.to(torch.long)
+
+            batch_size = int(req_ids.shape[0])
+            if batch_size <= 0:
+                return
+
+            if int(extend_seq_lens.shape[0]) != batch_size:
+                raise RuntimeError(
+                    "Sparse v2 prefill offload expects extend_seq_lens to be "
+                    f"padded to batch size {batch_size}, got "
+                    f"{int(extend_seq_lens.shape[0])}."
+                )
+
+            prefix_len_size = int(extend_prefix_lens.shape[0])
+            if prefix_len_size < batch_size:
+                extend_prefix_lens = torch.cat(
+                    [
+                        extend_prefix_lens,
+                        torch.zeros(
+                            batch_size - prefix_len_size,
+                            device=device,
+                            dtype=torch.long,
+                        ),
+                    ],
+                    dim=0,
+                )
+            elif prefix_len_size > batch_size:
+                raise RuntimeError(
+                    "Sparse v2 prefill offload expects extend_prefix_lens length "
+                    f"<= batch size {batch_size}, got {prefix_len_size}."
+                )
+
+            if forward_batch.extend_seq_lens_cpu is not None:
+                extend_seq_lens_sum = int(
+                    sum(forward_batch.extend_seq_lens_cpu[:batch_size])
+                )
+            else:
+                extend_seq_lens_sum = int(extend_seq_lens.sum().item())
+            if extend_seq_lens_sum == num_src_rows:
+                # Chunk prefill emits compact rows as [req0 tokens][req1 tokens]...
+                # instead of graph-captured padded [B, tokens_per_req] rows.
+                seq_starts = torch.cumsum(extend_seq_lens, dim=0) - extend_seq_lens
+                flat_req_ids = torch.repeat_interleave(
+                    req_ids, extend_seq_lens, output_size=num_src_rows
+                )
+                flat_seq_starts = torch.repeat_interleave(
+                    seq_starts, extend_seq_lens, output_size=num_src_rows
+                )
+                flat_prefix_lens = torch.repeat_interleave(
+                    extend_prefix_lens, extend_seq_lens, output_size=num_src_rows
+                )
+                token_pos = (
+                    flat_prefix_lens
+                    + torch.arange(num_src_rows, device=device, dtype=torch.long)
+                    - flat_seq_starts
+                )
+                dst_index = (
+                    flat_req_ids * self.max_context_len + token_pos
+                ).contiguous()
+                valid_mask = (
+                    (flat_req_ids >= 0)
+                    & (token_pos >= 0)
+                    & (token_pos < self.max_context_len)
+                )
+            else:
+                if num_src_rows % batch_size != 0:
+                    raise RuntimeError(
+                        "Sparse v2 prefill offload expects either compact ragged "
+                        "layout with rows=sum(extend_seq_lens) or graph-style "
+                        f"row-major layout [B, tokens_per_req], got rows={num_src_rows}, "
+                        f"batch={batch_size}, extend_seq_lens_sum={extend_seq_lens_sum}."
+                    )
+
+                # Graph-friendly static layout:
+                # compact rows are interpreted as [batch_size, tokens_per_req].
+                # Invalid padded columns are masked by local_offsets < extend_seq_lens.
+                tokens_per_req = num_src_rows // batch_size
+
+                local_offsets = (
+                    torch.arange(tokens_per_req, device=device, dtype=torch.long)
+                    .unsqueeze(0)
+                    .expand(batch_size, tokens_per_req)
+                )
+                req_ids_2d = req_ids.unsqueeze(1).expand(batch_size, tokens_per_req)
+                token_pos_2d = extend_prefix_lens.unsqueeze(1) + local_offsets
+
+                dst_index = (
+                    (req_ids_2d * self.max_context_len + token_pos_2d)
+                    .reshape(-1)
+                    .contiguous()
+                )
+                valid_mask = (
+                    (local_offsets < extend_seq_lens.unsqueeze(1))
+                    & (req_ids_2d >= 0)
+                    & (token_pos_2d >= 0)
+                    & (token_pos_2d < self.max_context_len)
+                ).reshape(-1)
+
+            if (
+                forward_batch.out_cache_loc is not None
+                and int(forward_batch.out_cache_loc.numel()) == num_src_rows
+            ):
+                valid_mask = valid_mask & (
+                    forward_batch.out_cache_loc.to(torch.long) >= 0
+                )
+            valid_mask = valid_mask.contiguous()
+
+        assert src_tensor.shape[1:] == dst_tensor.shape[2:]
+        assert src_index.shape == dst_index.shape == valid_mask.shape
+
+        actual_stream = stream if stream is not None else torch.npu.current_stream()
+        with torch.npu.stream(actual_stream):
+            unidex_copy_inplace(
+                src_tensor,
+                dst_tensor,
+                src_index,
+                dst_index,
+                valid_mask,
+                1,  # src_tensor: [num_new_tokens, nhead, head_dim]
+                2,  # dst_tensor: [num_req, max_context_len, nhead, head_dim]
+                block_dim=48,
+                dst_ptr=self.dev_ptr_list[layer_idx],
+            )
+
+    def get_forward_kv(
+        self,
+        layer: Union[RadixAttention, int],
+        forward_batch: ForwardBatch,
+        stream: Optional[torch.npu.Stream] = None,
+    ):
+        """Gather full request KV from sparse storage as compact TND tensors.
+
+        This helper is intended for the prefill/extend sparse path. It returns
+        KV in the same request order as forward_batch.req_pool_indices:
+        [req0 tokens][req1 tokens]... . The SFA caller should pair this with
+        actual_seq_lengths_kv = cumsum(forward_batch.seq_lens).
+
+        TODO: Add a prefill-resident device KV cache shaped like
+        [max_prefill_parallel_reqs, max_prefill_len, nhead, head_dim]. Reuse a
+        slot while the same req_id continues chunked prefill, and fully
+        overwrite it when a new req_id takes that slot. This mirrors the decode
+        device cache idea and avoids repeatedly copying the full prefix KV from
+        host for every chunk.
+        """
+        layer_id = layer.layer_id if hasattr(layer, "layer_id") else int(layer)
+        layer_idx = layer_id - self.start_layer
+        if layer_idx < 0 or layer_idx >= self.layer_num:
+            raise RuntimeError(
+                f"Invalid sparse KV layer id {layer_id}; start_layer="
+                f"{self.start_layer}, layer_num={self.layer_num}."
+            )
+
+        if forward_batch.req_pool_indices is None or forward_batch.seq_lens is None:
+            raise RuntimeError(
+                "get_forward_kv requires req_pool_indices and seq_lens in ForwardBatch."
+            )
+
+        device = (
+            forward_batch.req_pool_indices.device
+            if forward_batch.req_pool_indices.device.type == "npu"
+            else self.device
+        )
+        req_ids = forward_batch.req_pool_indices.to(device=device, dtype=torch.long)
+        seq_lens = forward_batch.seq_lens.to(device=device, dtype=torch.long)
+
+        if int(req_ids.numel()) != int(seq_lens.numel()):
+            raise RuntimeError(
+                "get_forward_kv expects req_pool_indices and seq_lens to have the "
+                f"same length, got {int(req_ids.numel())} and "
+                f"{int(seq_lens.numel())}."
+            )
+
+        valid_reqs = (seq_lens > 0) & (req_ids >= 0)
+        req_ids = req_ids[valid_reqs].contiguous()
+        seq_lens = seq_lens[valid_reqs].contiguous()
+
+        if int(seq_lens.numel()) == 0:
+            empty_nope = torch.empty(
+                (0, self.head_num, self.kv_lora_rank),
+                dtype=self.store_dtype,
+                device=device,
+            )
+            empty_pe = torch.empty(
+                (0, self.head_num, self.qk_rope_head_dim),
+                dtype=self.store_dtype,
+                device=device,
+            )
+            return empty_nope, empty_pe
+
+        if bool((req_ids >= self.size).any().item()):
+            raise RuntimeError(
+                f"get_forward_kv got req id outside sparse pool size {self.size}."
+            )
+        if bool((seq_lens > self.max_context_len).any().item()):
+            raise RuntimeError(
+                "get_forward_kv got seq_len larger than max_context_len "
+                f"{self.max_context_len}."
+            )
+
+        total_tokens = int(seq_lens.sum().item())
+        kv_cat = torch.empty(
+            (total_tokens, self.head_num, self.head_dim),
+            dtype=self.store_dtype,
+            device=device,
+        )
+
+        seq_starts = torch.cumsum(seq_lens, dim=0) - seq_lens
+        src_req_ids = torch.repeat_interleave(
+            req_ids, seq_lens, output_size=total_tokens
+        )
+        src_seq_starts = torch.repeat_interleave(
+            seq_starts, seq_lens, output_size=total_tokens
+        )
+        token_pos = (
+            torch.arange(total_tokens, device=device, dtype=torch.long) - src_seq_starts
+        )
+        src_index = (src_req_ids * self.max_context_len + token_pos).contiguous()
+        dst_index = torch.arange(total_tokens, device=device, dtype=torch.long)
+        valid_mask = torch.ones(total_tokens, device=device, dtype=torch.bool)
+
+        actual_stream = stream if stream is not None else torch.npu.current_stream()
+        with torch.npu.stream(actual_stream):
+            unidex_copy_inplace(
+                self.host_kv_buffer[layer_idx],
+                kv_cat,
+                src_index,
+                dst_index,
+                valid_mask,
+                2,  # host_kv_buffer: [num_req, max_context_len, nhead, head_dim]
+                1,  # kv_cat: [total_tokens, nhead, head_dim]
+                block_dim=48,
+                src_ptr=self.dev_ptr_list[layer_idx],
+            )
+
+        k_nope, k_pe = kv_cat.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+        return k_nope.contiguous(), k_pe.contiguous()
+
+    def materialize_selected_kv(
+        self,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        topk_indices: torch.Tensor,
+        selected_kv_buffer: torch.Tensor,
+        stream: torch.npu.Stream,
+    ) -> None:
+        """Materialize top-k KV entries and refresh the device cache metadata."""
+        layer_idx = layer.layer_id - self.start_layer
+        stream = stream if stream is not None else torch.npu.current_stream()
+
+        with torch.npu.stream(stream):
+            # Route invalid requests to sentinel rows without changing graph shape.
+            # slot_map_row_indices: invalid -> self.size (reserved slot-map row)
+            # device_cache_row_indices: invalid -> 0 (masked by valid_topk_mask)
+            req_pool_indices = forward_batch.req_pool_indices
+            req_pool_indices = req_pool_indices.to(torch.long).contiguous()
+            valid_req_mask = (req_pool_indices >= 0) & (req_pool_indices < self.size)
+            slot_map_row_indices = torch.where(
+                valid_req_mask,
+                req_pool_indices,
+                torch.full_like(req_pool_indices, self.size),
+            )
+            device_cache_row_indices = torch.where(
+                valid_req_mask,
+                req_pool_indices,
+                torch.zeros_like(req_pool_indices),
+            )
+
+            # Normalize top-k indices and mask invalid requests and token IDs.
+            topk_indices = normalize_batch_topk_indices(topk_indices)
+            batch_size, topk_len = topk_indices.shape
+            if topk_len > self.sparse_context_len:
+                raise RuntimeError(
+                    "DSA top-k length exceeds sparse KV device cache capacity: "
+                    f"topk_len={topk_len}, sparse_context_len={self.sparse_context_len}."
+                )
+            valid_topk_mask = (
+                (topk_indices >= 0)
+                & (topk_indices < self.max_context_len)
+                & valid_req_mask.unsqueeze(1)
+            )
+
+            # Query the slot map for device-cache hits and their slot positions.
+            slot_lookup_req_indices = slot_map_row_indices.to(
+                dtype=torch.int32
+            ).contiguous()
+            slot_lookup_topk_indices = topk_indices.to(dtype=torch.int32).contiguous()
+            token_on_device, device_token_pos = slot_map_lookup(
+                self.device_slot_map[layer_idx],
+                slot_lookup_req_indices,
+                slot_lookup_topk_indices,
+            )
+            token_on_device = token_on_device.to(torch.bool) & valid_topk_mask
+
+            # Build copy indices on the main stream, then protect their use on
+            # the hit and miss streams with copy_ready.
+            hit_src_index, hit_dst_index, hit_valid_mask = _build_hit_src_dst_index(
+                token_on_device,
+                device_token_pos,
+                device_cache_row_indices,
+                self.sparse_context_len,
+            )
+
+            host_miss_mask = (~token_on_device) & valid_topk_mask
+            miss_src_index, miss_dst_index, miss_valid_mask = _build_miss_src_dst_index(
+                host_miss_mask,
+                topk_indices,
+                device_cache_row_indices,
+                self.max_context_len,
+            )
+
+            cache_slot_ids = self._device_cache_slot_ids[:topk_len]
+            request_cache_offsets = (
+                device_cache_row_indices.unsqueeze(1) * self.sparse_context_len
+            )
+            refill_src_index = torch.arange(
+                batch_size * topk_len,
+                dtype=torch.long,
+                device=topk_indices.device,
+            )
+            refill_dst_index = (
+                (request_cache_offsets + cache_slot_ids).reshape(-1).contiguous()
+            )
+            refill_valid_mask = valid_topk_mask.reshape(-1).contiguous()
+
+            copy_ready = torch.npu.Event()
+            _record_stream_event(stream, copy_ready)
+
+        # Copy device-cache hits into the selected KV buffer.
+        with torch.npu.stream(self._materialize_d2d_hit_stream):
+            _wait_stream_event(self._materialize_d2d_hit_stream, copy_ready)
+            unidex_copy_inplace(
+                self.device_kv_buffer[layer_idx],
+                selected_kv_buffer,
+                hit_src_index,
+                hit_dst_index,
+                hit_valid_mask,
+                2,
+                2,  #
+                block_dim=24,
+            )
+            _record_stream_event(self._materialize_d2d_hit_stream, self.hit_done)
+
+        # Copy host shared-memory misses into the selected KV buffer.
+        with torch.npu.stream(self._materialize_h2d_miss_stream):
+            _wait_stream_event(self._materialize_h2d_miss_stream, copy_ready)
+            unidex_copy_inplace(
+                self.host_kv_buffer[layer_idx],
+                selected_kv_buffer,
+                miss_src_index,
+                miss_dst_index,
+                miss_valid_mask,
+                2,
+                2,
+                block_dim=24,
+                src_ptr=self.dev_ptr_list[layer_idx],
+            )
+            _record_stream_event(self._materialize_h2d_miss_stream, self.miss_done)
+
+        # Refill the device cache with the current top-k after hit and miss
+        # copies complete, so the next step can reuse these entries.
+        with torch.npu.stream(self._materialize_refill_stream):
+            _wait_stream_event(self._materialize_refill_stream, self.hit_done)
+            _wait_stream_event(self._materialize_refill_stream, self.miss_done)
+            unidex_copy_inplace(
+                selected_kv_buffer,
+                self.device_kv_buffer[layer_idx],
+                refill_src_index,
+                refill_dst_index,
+                refill_valid_mask,
+                2,
+                2,
+                block_dim=24,
+            )
+            _record_stream_event(self._materialize_refill_stream, self.refill_done)
+
+        # Replace the slot-map row with the current top-k mapping. Invalid
+        # entries use the max_context_len sentinel column to preserve shape.
+        with torch.npu.stream(self._materialize_slot_map_stream):
+            _wait_stream_event(self._materialize_slot_map_stream, copy_ready)
+            self.device_slot_map[layer_idx].copy_(self._device_slot_map_minus_one)
+
+            slot_map_token_indices = torch.where(
+                valid_topk_mask,
+                topk_indices.to(torch.long),
+                torch.full_like(topk_indices, self.max_context_len, dtype=torch.long),
+            )
+            slot_map_slot_values = torch.where(
+                valid_topk_mask,
+                cache_slot_ids.to(torch.int32),
+                torch.full_like(cache_slot_ids, -1, dtype=torch.int32),
+            )
+            slot_map_flat_indices = (
+                slot_map_row_indices.unsqueeze(1) * self._slot_map_width
+                + slot_map_token_indices
+            ).reshape(-1)
+            unidex_copy_inplace(
+                slot_map_slot_values.reshape(-1, 1).contiguous(),
+                self.device_slot_map[layer_idx].view(-1, 1),
+                torch.arange(
+                    slot_map_slot_values.numel(),
+                    dtype=torch.long,
+                    device=topk_indices.device,
+                ),
+                slot_map_flat_indices,
+                valid_topk_mask.reshape(-1),
+                1,
+                1,
+                block_dim=48,
+            )
+            _record_stream_event(self._materialize_slot_map_stream, self.slot_map_done)
+
+
+_global_sparse_kv_manager: Optional[SparseKVCacheManager] = None
+
+
+def register_sparse_kv_manager(manager: SparseKVCacheManager) -> None:
+    global _global_sparse_kv_manager
+    _global_sparse_kv_manager = manager
+
+
+def get_sparse_kv_manager() -> Optional[SparseKVCacheManager]:
+    return _global_sparse_kv_manager
+
+
+def _build_hit_src_dst_index(
+    token_on_device: torch.Tensor,
+    device_token_pos: torch.Tensor,
+    current_req_indices: torch.Tensor,
+    sparse_context_len: int,
+):
+    """
+    token_on_device: [bs, topk], bool
+    device_token_pos: [bs, topk], int64 or int32
+    current_req_indices: [bs], int64
+
+    Return:
+        src_index_full: [bs * topk], int64
+        dst_index_full: [bs * topk], int64
+        valid_mask: [bs * topk], bool
+
+    Flattening rule:
+        src row = req_id * sparse_context_len + device_token_pos
+        dst row = batch_id * topk + topk_pos
+    """
+    if token_on_device.dim() != 2 or device_token_pos.dim() != 2:
+        raise RuntimeError(
+            f"token_on_device and device_token_pos must be 2-D, got "
+            f"{token_on_device.dim()} and {device_token_pos.dim()}"
+        )
+    if token_on_device.shape != device_token_pos.shape:
+        raise RuntimeError(
+            f"token_on_device and device_token_pos must have the same shape, got "
+            f"{tuple(token_on_device.shape)} and {tuple(device_token_pos.shape)}"
+        )
+    if current_req_indices.dim() != 1:
+        raise RuntimeError(
+            f"current_req_indices must be 1-D, got {current_req_indices.dim()}"
+        )
+
+    bs, topk = token_on_device.shape
+    if current_req_indices.numel() != bs:
+        raise RuntimeError(
+            f"current_req_indices length mismatch: "
+            f"{current_req_indices.numel()} vs batch {bs}"
+        )
+    if sparse_context_len <= 0:
+        raise RuntimeError(
+            f"sparse_context_len must be positive, got {sparse_context_len}"
+        )
+
+    device = token_on_device.device
+
+    valid_mask = token_on_device.reshape(-1).contiguous()
+
+    flat_dst_index_all = torch.arange(
+        bs * topk,
+        device=device,
+        dtype=torch.int64,
+    )
+
+    req_offsets = current_req_indices.to(torch.int64).unsqueeze(1) * sparse_context_len
+    src_index_2d = req_offsets + device_token_pos.to(torch.int64)
+    flat_src_index_all = src_index_2d.reshape(-1).contiguous()
+
+    return flat_src_index_all, flat_dst_index_all, valid_mask
+
+
+def _build_miss_src_dst_index(
+    token_from_host: torch.Tensor,
+    topk_indices: torch.Tensor,
+    current_req_indices: torch.Tensor,
+    max_context_len: int,
+):
+    if token_from_host.dim() != 2 or topk_indices.dim() != 2:
+        raise RuntimeError(
+            f"token_from_host and topk_indices must be 2-D, got "
+            f"{token_from_host.dim()} and {topk_indices.dim()}"
+        )
+    if token_from_host.shape != topk_indices.shape:
+        raise RuntimeError(
+            f"token_from_host and topk_indices must have the same shape, got "
+            f"{tuple(token_from_host.shape)} and {tuple(topk_indices.shape)}"
+        )
+    if current_req_indices.dim() != 1:
+        raise RuntimeError(
+            f"current_req_indices must be 1-D, got {current_req_indices.dim()}"
+        )
+    if current_req_indices.numel() != token_from_host.shape[0]:
+        raise RuntimeError(
+            f"current_req_indices length mismatch: "
+            f"{current_req_indices.numel()} vs batch {token_from_host.shape[0]}"
+        )
+
+    bs, topk = token_from_host.shape
+    device = token_from_host.device
+
+    valid_2d = token_from_host & (topk_indices >= 0) & (topk_indices < max_context_len)
+    valid_mask = valid_2d.reshape(-1).contiguous()
+
+    flat_dst_index_all = torch.arange(
+        bs * topk,
+        device=device,
+        dtype=torch.int64,
+    )
+
+    req_offsets = current_req_indices.to(torch.int64).unsqueeze(1) * max_context_len
+    src_index_2d = req_offsets + topk_indices.to(torch.int64)
+    flat_src_index_all = src_index_2d.reshape(-1).contiguous()
+
+    return flat_src_index_all, flat_dst_index_all, valid_mask
@@ -310,6 +310,21 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
        dcp_size = get_parallel().attn_dcp_size

        if kvc.use_mla_backend:
+            if envs.SGLANG_NPU_ENABLE_SPARSE_KV_OFFLOAD.get():
+                # NPU sparse KV offload uses an index-only device pool.
+                from sglang.srt.hardware_backend.npu.sparsity_driven_kv_offload.config import (
+                    get_sparsity_driven_kv_offload_cell_size,
+                )
+
+                offload_cell_size = get_sparsity_driven_kv_offload_cell_size(
+                    model_config=model_config,
+                    use_mla_backend=kvc.use_mla_backend,
+                    num_layers=num_layers,
+                    element_size=kv_size,
+                )
+                if offload_cell_size is not None:
+                    return offload_cell_size
+
            from sglang.srt.mem_cache.kv_cache_configurator import (
                calculate_mla_kv_cache_dim,
            )