[NPU] Add Ascend NPU support for DeepSeek-V4 (#25144)

Co-authored-by: khalil2ji3mp6 <khalilzhk@gmail.com>
Co-authored-by: randgun <kelonlu@163.com>
Co-authored-by: t00937989 <tanlei33@huawei.com>

python/sglang/srt/model_executor/forward_batch_info.py

@@ -217,6 +217,65 @@ def compute_local_num_token_non_padded(
     )
 
 
+@dataclass
+class DSV4OutCacheLoc:
+    """Per-forward-pass KV cache allocation for DeepSeek-V4 on NPU.
+
+    Bundles slot indices for full/SWA pools, the two compressed-KV pools
+    (c4/c128), and the two compressed-state pools (c4_state/c128_state).
+    Populated by the NPU V4 allocator (DSV4NPUTokenToKVPoolAllocator) when
+    the model is DeepSeek-V4 on NPU; left as ``None`` on ForwardBatch
+    otherwise. CUDA's DSV4 path doesn't construct this bundle (state is
+    derived via translate_kv_loc_to_compress_state_loc there).
+
+    All fields are token-level slot ids in their respective pools (NOT page
+    ids). Attention backends convert to page ids via ``// page_size`` when
+    constructing PA_ND block tables.
+
+    State fields default to ``None`` so the bundle is constructible from
+    paths that allocate KV but not state (or vice versa); the NPU allocator
+    fills all six on real alloc, CUDA paths leave state ones None and use
+    the ring-hash translation instead.
+    """
+
+    out_full_loc: torch.Tensor
+    out_swa_loc: torch.Tensor
+    out_c4_loc: torch.Tensor
+    out_c128_loc: torch.Tensor
+    out_c4_state_loc: Optional[torch.Tensor] = None
+    out_c128_state_loc: Optional[torch.Tensor] = None
+
+
+@dataclass
+class DSV4StateLens:
+    """Per-extend/decode c4/c128 compress-state pool allocation lens (DSV4-NPU).
+
+    Built by ``ScheduleBatch._compute_dsv4_state_lens_{extend,decode}`` and
+    threaded through ``mem_cache/common.py`` to
+    ``DSV4NPUTokenToKVPoolAllocator.alloc_{extend,decode}``, which consumes:
+
+      * ``c{4,128}_prefix_lens`` / ``..._cpu`` — per-req prev cumulative
+        state-slot count (the paged allocator's ``prefix`` contract).
+      * ``c{4,128}_seq_lens`` / ``..._cpu`` — per-req new cumulative count.
+      * ``c{4,128}_extend_num_tokens`` — total new state slots this step.
+
+    Replaces the 10 loose ``c{4,128}_state_*`` kwargs the allocator used to
+    take: scheduler only produces this object, common only forwards it, the
+    allocator only consumes it.
+    """
+
+    c4_prefix_lens: torch.Tensor
+    c4_prefix_lens_cpu: torch.Tensor
+    c4_seq_lens: torch.Tensor
+    c4_seq_lens_cpu: torch.Tensor
+    c4_extend_num_tokens: int
+    c128_prefix_lens: torch.Tensor
+    c128_prefix_lens_cpu: torch.Tensor
+    c128_seq_lens: torch.Tensor
+    c128_seq_lens_cpu: torch.Tensor
+    c128_extend_num_tokens: int
+
+
 @dataclass
 class NgramEmbeddingInfo:
     """Ngram embedding state for LongCat models."""
@@ -286,6 +345,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
     # The original sequence length without being chunked. Qwen-1M related.
     orig_seq_lens: Optional[torch.Tensor] = None
 
+    # DSV4-NPU only: per-pool slot bundle from DSV4NPUTokenToKVPoolAllocator,
+    # consumed by the Ascend backend for PA_ND block tables. None elsewhere.
+    out_cache_loc_dsv4: Optional[DSV4OutCacheLoc] = None
     # The indices to track mamba state with
     mamba_track_indices: Optional[torch.Tensor] = None  # shape: [b], int64
     # The mask to track mamba state if needed
@@ -615,6 +677,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             # Inputs aliased by reference from ScheduleBatch
             seq_lens_cpu=seq_lens_cpu,
             orig_seq_lens=batch.orig_seq_lens,
+            out_cache_loc_dsv4=batch.out_cache_loc_dsv4,
             mamba_track_indices=batch.mamba_track_indices,
             mamba_track_mask=batch.mamba_track_mask,
             mamba_track_seqlens=batch.mamba_track_seqlens,

python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py

@@ -391,7 +391,17 @@ class ModelRunnerKVCacheMixin:
                     start_layer=self.start_layer,
                 )
             else:
-                self.req_to_token_pool = ReqToTokenPool(
+                # DSV4 on NPU needs an extended ReqToTokenPool holding per-req
+                # swa/c4/c128/c{4,128}_state tables; others stay on the stock one.
+                req_to_token_pool_cls = ReqToTokenPool
+                if _is_npu and is_deepseek_v4(self.model_config.hf_config):
+                    from sglang.srt.hardware_backend.npu.dsv4.dsv4_req_to_token_pool import (
+                        DSV4NPUReqToTokenPool,
+                    )
+
+                    req_to_token_pool_cls = DSV4NPUReqToTokenPool
+
+                self.req_to_token_pool = req_to_token_pool_cls(
                     size=max_num_reqs,
                     max_context_len=self.model_config.context_len
                     + extra_max_context_len,
@@ -412,7 +422,8 @@ class ModelRunnerKVCacheMixin:
 
         if is_dsv4_model:
             swa_page_size = self.page_size
-            assert swa_page_size == 256, "In paged swa mode, page_size must be 256."
+            if not _is_npu:
+                assert swa_page_size == 256, "In paged swa mode, page_size must be 256."
 
             if self.is_draft_worker:
                 from sglang.srt.models.deepseek_v4_nextn import (
@@ -424,7 +435,40 @@ class ModelRunnerKVCacheMixin:
                 ] * self.num_effective_layers
             else:
                 compression_ratios = self.model_config.compress_ratios
-            self.token_to_kv_pool = DeepSeekV4TokenToKVPool(
+
+            # NPU + DSV4 → paged-state subclass: the fused compressor kernel
+            # needs cache_mode=1 (paged); Atlas A3 rejects cache_mode=2 (ring),
+            # so the CUDA ring-buffer state path can't be shared. CUDA keeps
+            # DeepSeekV4TokenToKVPool unchanged; NPU recomputes state sizes below.
+            if _is_npu:
+                from sglang.srt.hardware_backend.npu.dsv4.dsv4_memory_pool import (
+                    DSV4NPUTokenToKVPool,
+                    npu_state_pool_size,
+                )
+
+                pool_cls = DSV4NPUTokenToKVPool
+                # Recompute state pool sizes for the NPU paged formula (CUDA's
+                # ring sizes are dropped here). Tail-only allocation keeps the
+                # per-req-budget formula sufficient at any prefill length: long
+                # prompts allocate only ``tail+128`` (c4) / ``tail`` (c128)
+                # slots (tail = seq_len % 128), and decode is drained by
+                # sliding eviction in ``ScheduleBatch._evict_swa``.
+                c4_state_pool_size = npu_state_pool_size(
+                    ratio=4,
+                    page_size=self.page_size,
+                    max_num_reqs=self.max_running_requests,
+                )
+                c128_state_pool_size = npu_state_pool_size(
+                    ratio=128,
+                    page_size=self.page_size,
+                    max_num_reqs=self.max_running_requests,
+                )
+            else:
+                pool_cls = DeepSeekV4TokenToKVPool
+                c4_state_pool_size = self.c4_state_pool_size
+                c128_state_pool_size = self.c128_state_pool_size
+
+            self.token_to_kv_pool = pool_cls(
                 max_num_reqs=self.max_running_requests,
                 # SWA ring is indexed by req_pool_idx; PD decode inflates req_to_token
                 # past max_running_requests (pre-alloc), so size to the real capacity.
@@ -432,8 +476,8 @@ class ModelRunnerKVCacheMixin:
                 swa_size=self.swa_max_total_num_tokens,
                 c4_size=self.c4_max_total_num_tokens,
                 c128_size=self.c128_max_total_num_tokens,
-                c4_state_pool_size=self.c4_state_pool_size,
-                c128_state_pool_size=self.c128_state_pool_size,
+                c4_state_pool_size=c4_state_pool_size,
+                c128_state_pool_size=c128_state_pool_size,
                 page_size=self.page_size,
                 swa_page_size=swa_page_size,
                 sliding_window=self.model_config.window_size,
@@ -763,10 +807,21 @@ class ModelRunnerKVCacheMixin:
                 )
             elif _is_npu and (
                 self.server_args.attention_backend == "ascend"
+                or is_dsv4_model
                 or self.hybrid_gdn_config is not None
             ):
                 if self.is_hybrid_swa:
-                    self.token_to_kv_pool_allocator = SWATokenToKVPoolAllocator(
+                    # DSV4 on NPU: SWA allocator subclass that also drives the
+                    # c4/c128 allocators, producing a DSV4OutCacheLoc per alloc.
+                    if is_dsv4_model:
+                        from sglang.srt.hardware_backend.npu.dsv4.dsv4_allocator import (
+                            DSV4NPUTokenToKVPoolAllocator,
+                        )
+
+                        swa_allocator_cls = DSV4NPUTokenToKVPoolAllocator
+                    else:
+                        swa_allocator_cls = SWATokenToKVPoolAllocator
+                    self.token_to_kv_pool_allocator = swa_allocator_cls(
                         self.full_max_total_num_tokens,
                         self.swa_max_total_num_tokens,
                         page_size=self.page_size,
@@ -843,6 +898,13 @@ class ModelRunnerKVCacheMixin:
                     )
                 )
 
+            # DSV4-NPU: wire allocator back-ref into req_to_token_pool so its
+            # free(req) can release c4/c128 pool pages alongside the slot.
+            if hasattr(self.req_to_token_pool, "register_dsv4_allocator"):
+                self.req_to_token_pool.register_dsv4_allocator(
+                    self.token_to_kv_pool_allocator
+                )
+
         else:
             assert self.is_draft_worker
             if self.is_hybrid_swa:

python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py

@@ -156,6 +156,7 @@ def build_replay_fb_view(
         seq_lens_cpu=buffers.seq_lens_cpu[:bs],
         encoder_lens=buffers.encoder_lens[:bs] if is_encoder_decoder else None,
         out_cache_loc=getattr(forward_batch, "out_cache_loc", None),
+        out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None),
         spec_info=forward_batch.spec_info,
     )