[NPU] Support DeepSeek V4 Flash MTP on Ascend (#28980)

Co-authored-by: Kurkur <mccllm@qq.com> Co-authored-by: gjsheu <gjsheu@163.com> Co-authored-by: root <root@localhost.localdomain> Co-authored-by: khalil2ji3mp6 <khalilzhk@gmail.com>
2026-06-30 16:22:13 +08:00
parent 4b4b4af583
commit 89620b9169
13 changed files with 852 additions and 86 deletions
@@ -68,6 +68,7 @@ class ForwardMetadata:
    seq_lens_list_cumsum: Optional[List[int]] = None
    seq_lens: Optional[torch.Tensor] = None
    actual_seq_lengths_q: Optional[torch.Tensor] = None
+    actual_seq_lengths_q_pa: Optional[torch.Tensor] = None
    actual_seq_lengths_kv: Optional[torch.Tensor] = None

    # swa attention mask for graph mode decode
@@ -684,7 +685,7 @@ class AscendAttnBackend(AttentionBackend):
            metadata.swa_mask[:bs, 0, :].copy_(mask)
            metadata.swa_mask[bs:, :, :].fill_(True)
        metadata.block_tables[:bs, :max_seq_pages].copy_(
-            self.req_to_token[req_pool_indices[:bs], :max_len][:, :: self.page_size]
+            self.req_to_token[req_pool_indices[:bs], 0 : max_len : self.page_size]
            // self.page_size
        )

@@ -2417,6 +2418,7 @@ class AscendAttnBackend(AttentionBackend):
        topk_indices: Optional[torch.Tensor] = None,
        sinks: Optional[torch.Tensor] = None,
        slopes: Optional[torch.Tensor] = None,
+        **kwargs,
    ):
        if is_mla_preprocess_enabled() and self.use_mla:
            # MLAPO does saving kv_cache
@@ -11,11 +11,12 @@ from sglang.srt.hardware_backend.npu.attention.ascend_backend import AscendAttnB
 from sglang.srt.layers.attention.dsv4.compressor import CompressorBackendMixin
 from sglang.srt.layers.attention.dsv4.indexer import C4IndexerBackendMixin
 from sglang.srt.layers.dp_attention import get_attention_tp_size
+from sglang.srt.model_executor.forward_batch_info import DSV4OutCacheLoc, ForwardMode
 from sglang.srt.model_executor.forward_context import get_attn_backend

 if TYPE_CHECKING:
    from sglang.srt.layers.radix_attention import RadixAttention
-    from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
    from sglang.srt.model_executor.model_runner import ModelRunner

 logger = logging.getLogger(__name__)
@@ -67,11 +68,16 @@ class CompressorAscendBackendMixin(CompressorBackendMixin):
    def _build_npu_compress_metadata(self, forward_batch: ForwardBatch) -> None:
        fm = self.forward_metadata
        is_decode = forward_batch.forward_mode.is_decode()
+        is_verify = forward_batch.forward_mode.is_target_verify()
+        _verify_compress = is_verify and bool(self._dsv4_compress_ratios)
+        _seq_lens = forward_batch.seq_lens.to(torch.int32)
+        if _verify_compress:
+            _seq_lens = _seq_lens + self.speculative_num_draft_tokens
        result = self._compute_compress_locs(
            pool=self.token_to_kv_pool,
            req_to_token=self.req_to_token,
            req_pool_indices=forward_batch.req_pool_indices,
-            seq_lens=forward_batch.seq_lens.to(torch.int32),
+            seq_lens=_seq_lens,
            out_cache_loc=forward_batch.out_cache_loc,
            is_decode=is_decode,
            bs=forward_batch.batch_size,
@@ -89,6 +95,9 @@ class CompressorAscendBackendMixin(CompressorBackendMixin):
                    if f"c{ratio}_loc" not in result:
                        setattr(fm, f"c{ratio}_loc", None)

+        if _verify_compress:
+            self._build_npu_compress_metadata_verify(forward_batch)
+
    def _build_npu_compress_metadata_prefill(self, forward_batch: ForwardBatch) -> None:
        # eager-only: prefill is never graph-captured, host reads (cu_cpu) are safe here
        fm = self.forward_metadata
@@ -99,7 +108,9 @@ class CompressorAscendBackendMixin(CompressorBackendMixin):
        cu = fm.actual_seq_lengths_q_pa

        cu_cpu = cu.cpu().tolist()
-        ratio_lists: dict = {r: [] for r in self._dsv4_compress_ratios if r in (4, 128)}
+        ratio_lists: dict = {
+            r: [] for r in self._dsv4_unique_compress_ratios if r in (4, 128)
+        }
        for idx in range(bs):
            start = int(cu_cpu[idx])
            end = int(cu_cpu[idx + 1])
@@ -178,14 +189,18 @@ class CompressorAscendBackendMixin(CompressorBackendMixin):
        req_to_token_pool,
        out_cache_loc_dsv4,
        is_graph: bool = False,
+        seq_lens_max_override: Optional[int] = None,
    ) -> dict:
        result: dict = {}
        req_pool = req_pool_indices

+        if seq_lens_max_override is not None:
+            seq_lens_max = int(seq_lens_max_override)
+        else:
            seq_lens_max = int(seq_lens.max().item()) if bs > 0 else 0
        n_pages = max(1, (seq_lens_max + self.page_size - 1) // self.page_size)

-        for ratio in self._dsv4_compress_ratios:
+        for ratio in self._dsv4_unique_compress_ratios:
            if ratio not in (4, 128):
                continue
            # state table holds one slot per RAW token; block 0 is the skip sentinel reserved by NPUCompressStatePool
@@ -258,7 +273,7 @@ class CompressorAscendBackendMixin(CompressorBackendMixin):
        if is_decode:
            valid = seq_lens > 0
            positions_last = torch.clamp(seq_lens - 1, min=0)
-            for ratio in self._dsv4_compress_ratios:
+            for ratio in self._dsv4_unique_compress_ratios:
                if ratio not in (4, 128):
                    continue
                padding_size = min(bs, bs // ratio + bs)
@@ -291,7 +306,10 @@ class CompressorAscendBackendMixin(CompressorBackendMixin):
        x: torch.Tensor,
        forward_batch: ForwardBatch,
    ) -> None:
-        if not forward_batch.forward_mode.is_decode():
+        if (
+            forward_batch.forward_mode.is_prefill()
+            and not forward_batch.forward_mode.is_target_verify()
+        ):
            return self._forward_compress_native(compressor, x, forward_batch)

        from sglang.srt.layers.deepseek_v4_rope import (
@@ -350,8 +368,19 @@ class CompressorAscendBackendMixin(CompressorBackendMixin):

        # prefill output may be padded; trim to loc length
        loc = getattr(fm, f"c{ratio}_loc", None)
-        if loc is not None and loc.numel() < cmp_kv.shape[0]:
+        is_prefill = (
+            forward_batch.forward_mode.is_prefill()
+            and not forward_batch.forward_mode.is_target_verify()
+        )
+        if loc is not None:
+            if is_prefill and loc.numel() < cmp_kv.shape[0]:
                cmp_kv = cmp_kv[: loc.numel()]
+            elif loc.numel() != cmp_kv.shape[0]:
+                raise RuntimeError(
+                    "DSV4 NPU fused compressor loc/kv length mismatch before "
+                    f"epilog: mode={forward_batch.forward_mode}, ratio={ratio}, "
+                    f"loc={loc.numel()}, kv={cmp_kv.shape[0]}"
+                )

        if self.graph_mode or cmp_kv.shape[0] > 0:
            if compressor.rotate:
@@ -702,6 +731,30 @@ class CompressorAscendBackendMixin(CompressorBackendMixin):
        else:
            backend_fm = self.forward_metadata
            loc = backend_fm.c4_loc if compressor.ratio == 4 else backend_fm.c128_loc
+        if loc is not None:
+            if loc.numel() != kv.shape[0]:
+                raise RuntimeError(
+                    "DSV4 NPU fused compressor epilog loc/kv length mismatch: "
+                    f"mode={forward_batch.forward_mode}, "
+                    f"ratio={compressor.ratio}, loc={loc.numel()}, kv={kv.shape[0]}"
+                )
+            if forward_batch.forward_mode.is_target_verify():
+                valid = loc != 0
+                if self.graph_mode:
+                    kv_mask = valid.to(kv.dtype).view(
+                        valid.shape[0], *([1] * (kv.dim() - 1))
+                    )
+                    kv = kv * kv_mask
+                    if kv_scale is not None:
+                        scale_mask = valid.to(kv_scale.dtype).view(
+                            valid.shape[0], *([1] * (kv_scale.dim() - 1))
+                        )
+                        kv_scale = kv_scale * scale_mask
+                else:
+                    loc = loc[valid]
+                    kv = kv[valid]
+                    if kv_scale is not None:
+                        kv_scale = kv_scale[valid]
        self.token_to_kv_pool.set_compress_buffer(
            compressor.layer_id,
            loc,
@@ -916,6 +969,7 @@ class DeepseekV4AscendAttnBackend(
        speculative_step_id: int = 0,
    ):
        super().__init__(model_runner, speculative_step_id=speculative_step_id)
+        self.use_graph_swa_mask = False
        cfg = model_runner.model_config
        self._dsv4_config = cfg
        tp_size = get_attention_tp_size()
@@ -927,11 +981,16 @@ class DeepseekV4AscendAttnBackend(
        self._dsv4_index_n_heads = hf.index_n_heads
        self._dsv4_index_head_dim = hf.index_head_dim
        self._dsv4_compress_ratios = hf.compress_ratios
+        if getattr(model_runner, "is_draft_worker", False):
+            self._dsv4_compress_ratios = type(hf.compress_ratios)()
        self._dsv4_has_c4 = 4 in self._dsv4_compress_ratios
        self._dsv4_has_c128 = 128 in self._dsv4_compress_ratios
        self._dsv4_sliding_window_size = (
            cfg.sliding_window_size if cfg.sliding_window_size is not None else 128
        )
+        self._dsv4_unique_compress_ratios = list(
+            dict.fromkeys(self._dsv4_compress_ratios)
+        )

    def _init_dsv4_graph_buffers(self, *, max_bs: int, max_num_tokens: int) -> None:
        device = self.device
@@ -1024,14 +1083,14 @@ class DeepseekV4AscendAttnBackend(
        ]

        n_tok = bs * tokens_per_bs
+        c4_pad = min(n_tok, n_tok // 4 + bs)
+        c128_pad = min(n_tok, n_tok // 128 + bs)
        metadata.swa_loc = torch.zeros(n_tok, dtype=torch.int64, device=device)
-        metadata.c4_loc = torch.zeros(n_tok, dtype=torch.int64, device=device)
-        metadata.c128_loc = torch.zeros(n_tok, dtype=torch.int64, device=device)
+        metadata.c4_loc = torch.zeros(c4_pad, dtype=torch.int64, device=device)
+        metadata.c128_loc = torch.zeros(c128_pad, dtype=torch.int64, device=device)
        metadata.c4_state_loc = torch.zeros(n_tok, dtype=torch.int64, device=device)
        metadata.c128_state_loc = torch.zeros(n_tok, dtype=torch.int64, device=device)

-        c4_pad = min(n_tok, n_tok // 4 + bs)
-        c128_pad = min(n_tok, n_tok // 128 + bs)
        metadata.positions_cmp_padding_c4 = torch.zeros(
            c4_pad, dtype=torch.int64, device=device
        )
@@ -1055,7 +1114,13 @@ class DeepseekV4AscendAttnBackend(

    def _apply_dsv4_graph_metadata(self, forward_batch: ForwardBatch) -> None:
        fm = self.forward_metadata
-        forward_mode = forward_batch.forward_mode
+        forward_mode = (
+            getattr(forward_batch, "global_forward_mode", None)
+            or forward_batch.forward_mode
+        )
+        actual_forward_mode = getattr(forward_batch, "actual_forward_mode", None)
+        if actual_forward_mode is None:
+            actual_forward_mode = forward_batch.forward_mode
        bs = forward_batch.batch_size
        seq_lens = forward_batch.seq_lens
        req_pool_indices = forward_batch.req_pool_indices
@@ -1067,29 +1132,57 @@ class DeepseekV4AscendAttnBackend(
            tokens_per_bs = 1

        seq_lens_cpu = forward_batch.seq_lens_cpu
-        assert seq_lens_cpu is not None, (
-            "V4 graph replay requires seq_lens_cpu - buffers.seq_lens is stale on "
-            "NPU (Graph.update only refreshes fm.actual_seq_lengths_kv inside the "
-            "captured graph, not the device-side buffers.seq_lens)."
-        )
+        assert seq_lens_cpu is not None, "V4 graph replay requires seq_lens_cpu."
+        if forward_mode.is_target_verify():
+            # In graph replay, buffers.seq_lens already contains the attention KV
+            # length (live length + draft tokens). Padded rows therefore show up as
+            # tokens_per_bs instead of 0. Use the CPU live lengths as the source of
+            # truth so padded rows stay masked out.
            live_seq_lens = seq_lens_cpu[:bs].to(device=device, dtype=torch.int32)
-        fm.actual_seq_lengths_kv.copy_(live_seq_lens.clamp(min=1))
+        elif seq_lens is not None and seq_lens.device.type != "cpu":
+            live_seq_lens = seq_lens[:bs].to(dtype=torch.int32)
+        else:
+            live_seq_lens = seq_lens_cpu[:bs].to(device=device, dtype=torch.int32)
+        attn_seq_lens = live_seq_lens
+        if forward_mode.is_target_verify():
+            valid_verify_rows = live_seq_lens > 0
+            attn_seq_lens = live_seq_lens + int(tokens_per_bs)
+            attn_seq_lens = torch.where(valid_verify_rows, attn_seq_lens, live_seq_lens)
+            fm.seq_lens_cpu_int = (seq_lens_cpu[:bs] + int(tokens_per_bs)).int()
+            fm.seq_lens_cpu_int = torch.where(
+                seq_lens_cpu[:bs] > 0,
+                fm.seq_lens_cpu_int,
+                seq_lens_cpu[:bs].int(),
+            )
+        fm.actual_seq_lengths_kv.copy_(attn_seq_lens.clamp(min=1))

        pool = self.token_to_kv_pool
        out_cache_loc = forward_batch.out_cache_loc

+        _verify_compress = (
+            forward_mode.is_target_verify()
+            and actual_forward_mode.is_target_verify()
+            and bool(self._dsv4_compress_ratios)
+        )
+        _compress_seq_lens = live_seq_lens
+        _compress_seq_lens_max = int(seq_lens_cpu[:bs].max()) if bs > 0 else 0
+        if _verify_compress:
+            _compress_seq_lens = live_seq_lens + int(tokens_per_bs)
+            _compress_seq_lens_max += int(tokens_per_bs)
+
        result = self._compute_compress_locs(
            pool=pool,
            req_to_token=self.req_to_token,
            req_pool_indices=req_pool_indices[:bs],
-            seq_lens=live_seq_lens,
+            seq_lens=_compress_seq_lens,
            out_cache_loc=out_cache_loc,
            is_decode=forward_mode.is_decode(),
            bs=bs,
            device=device,
            req_to_token_pool=self.req_to_token_pool,
-            out_cache_loc_dsv4=forward_batch.out_cache_loc_dsv4,
+            out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None),
            is_graph=True,
+            seq_lens_max_override=_compress_seq_lens_max,
        )

        def _copy_2d(dst: torch.Tensor, src: torch.Tensor, val: int) -> None:
@@ -1098,6 +1191,10 @@ class DeepseekV4AscendAttnBackend(

        def _copy_1d(dst: torch.Tensor, src: torch.Tensor) -> None:
            dst.fill_(0)
+            assert src.shape[0] <= dst.shape[0], (
+                f"graph replay 1D metadata overflow: src={src.shape[0]} > "
+                f"dst={dst.shape[0]}"
+            )
            dst[: src.shape[0]].copy_(src)

        for key in (
@@ -1121,6 +1218,71 @@ class DeepseekV4AscendAttnBackend(
            if key in result and hasattr(fm, key) and getattr(fm, key) is not None:
                _copy_1d(getattr(fm, key), result[key])

+        if _verify_compress:
+            verify_seq_lens_cpu = seq_lens_cpu[:bs] + int(tokens_per_bs)
+            verify_seq_lens_cpu = torch.where(
+                seq_lens_cpu[:bs] > 0,
+                verify_seq_lens_cpu,
+                seq_lens_cpu[:bs],
+            )
+            self._fill_verify_positions_cmp_padding_one(
+                forward_batch.positions,
+                fm.positions_cmp_padding_c4,
+                4,
+                verify_seq_lens_cpu,
+                n_draft=tokens_per_bs,
+            )
+            self._fill_verify_positions_cmp_padding_one(
+                forward_batch.positions,
+                fm.positions_cmp_padding_c128,
+                128,
+                verify_seq_lens_cpu,
+                n_draft=tokens_per_bs,
+            )
+            fm.start_pos.copy_(live_seq_lens.to(torch.int32))
+            valid = live_seq_lens[:bs] > 0
+            fm.seqused.copy_(
+                (valid.to(torch.int32) * int(tokens_per_bs)).to(device=device)
+            )
+            _bundle = getattr(forward_batch, "out_cache_loc_dsv4", None)
+            if _bundle is not None:
+                for ratio in self._dsv4_unique_compress_ratios:
+                    if ratio not in (4, 128):
+                        continue
+                    bl = _bundle.out_c4_loc if ratio == 4 else _bundle.out_c128_loc
+                    if bl is not None:
+                        dst_loc = getattr(fm, f"c{ratio}_loc", None)
+                        if dst_loc is not None:
+                            dst_loc.zero_()
+                            bl32 = bl.to(torch.int32)
+                            assert bl32.numel() <= dst_loc.numel(), (
+                                f"replay verify c{ratio}_loc overflow: "
+                                f"{bl32.numel()} > {dst_loc.numel()}"
+                            )
+                            dst_loc[: bl32.numel()].copy_(bl32)
+
+        elif (
+            forward_mode.is_target_verify()
+            # The graph may replay a target-verify capture for an idle/padded
+            # DP rank. There is no real DSV4 allocation bundle in that case;
+            # zero the compressor metadata so captured writes land in the
+            # reserved dummy slot instead of reusing stale locs.
+            and not actual_forward_mode.is_target_verify()
+            and bool(self._dsv4_compress_ratios)
+        ):
+            for tensor in (
+                fm.positions_cmp_padding_c4,
+                fm.positions_cmp_padding_c128,
+                fm.c4_loc,
+                fm.c128_loc,
+                fm.c4_state_loc,
+                fm.c128_state_loc,
+            ):
+                if tensor is not None:
+                    tensor.zero_()
+            fm.start_pos.zero_()
+            fm.seqused.zero_()
+
        swa_loc = pool.translate_loc_from_full_to_swa(out_cache_loc).to(torch.int64)
        _copy_1d(fm.swa_loc, swa_loc)

@@ -1172,7 +1334,11 @@ class DeepseekV4AscendAttnBackend(

        device = forward_batch.seq_lens.device
        # cu_seqlens_q must hold per-request QUERY token counts, not KV lengths.
-        if forward_batch.forward_mode.is_extend():
+        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()
+        ):
            seq_lens_cpu = forward_batch.extend_seq_lens_cpu
            if isinstance(seq_lens_cpu, list):
                seq_lens_cpu = torch.tensor(seq_lens_cpu, dtype=torch.int32)
@@ -1197,7 +1363,7 @@ class DeepseekV4AscendAttnBackend(
            or forward_batch.forward_mode.is_draft_extend_v2()
        ):
            B = forward_batch.batch_size
-            from sglang.srt.utils.common import get_global_server_args
+            from sglang.srt.server_args import get_global_server_args

            n_draft = get_global_server_args().speculative_num_draft_tokens or 1
            actual_q = torch.arange(
@@ -1244,7 +1410,7 @@ class DeepseekV4AscendAttnBackend(
            forward_batch.forward_mode.is_target_verify()
            or forward_batch.forward_mode.is_draft_extend_v2()
        ):
-            from sglang.srt.utils.common import get_global_server_args
+            from sglang.srt.server_args import get_global_server_args

            max_seqlen_q = get_global_server_args().speculative_num_draft_tokens or 1
        else:
@@ -1469,6 +1635,182 @@ class DeepseekV4AscendAttnBackend(
            cache=swa_k,
        )

+    def _build_npu_compress_metadata_verify(self, forward_batch: ForwardBatch) -> None:
+        fm = self.forward_metadata
+        device = forward_batch.seq_lens.device
+        positions = forward_batch.positions
+        t = positions.shape[0]
+        bs = forward_batch.batch_size
+        n_draft = int(
+            getattr(
+                getattr(forward_batch, "spec_info", None),
+                "draft_token_num",
+                self.speculative_num_draft_tokens,
+            )
+        )
+        verify_seq_lens_cpu = forward_batch.seq_lens_cpu[:bs] + int(n_draft)
+        padding_sizes = {}
+        for ratio in (4, 128):
+            if ratio not in self._dsv4_compress_ratios:
+                continue
+            padding_size = max(1, min(t, t // ratio + bs))
+            padding_sizes[ratio] = padding_size
+            padding = torch.zeros(padding_size, dtype=torch.int64, device=device)
+            self._fill_verify_positions_cmp_padding_one(
+                positions, padding, ratio, verify_seq_lens_cpu, n_draft=n_draft
+            )
+            setattr(fm, f"positions_cmp_padding_c{ratio}", padding)
+        fm.start_pos = forward_batch.seq_lens.to(torch.int32)
+        valid = forward_batch.seq_lens[:bs] > 0
+        fm.seqused = valid.to(torch.int32) * int(n_draft)
+        _bundle = getattr(forward_batch, "out_cache_loc_dsv4", None)
+        if _bundle is not None:
+            for ratio in self._dsv4_unique_compress_ratios:
+                if ratio not in (4, 128):
+                    continue
+                bl = _bundle.out_c4_loc if ratio == 4 else _bundle.out_c128_loc
+                if bl is None:
+                    loc = None
+                else:
+                    padding_size = padding_sizes[ratio]
+                    loc = torch.zeros(padding_size, dtype=torch.int32, device=device)
+                    if bl.numel() > 0:
+                        assert bl.numel() <= padding_size, (
+                            f"verify c{ratio}_loc overflow: "
+                            f"{bl.numel()} > {padding_size}"
+                        )
+                        loc[: bl.numel()].copy_(bl.to(torch.int32))
+                setattr(fm, f"c{ratio}_loc", loc)
+
+    def _fill_verify_positions_cmp_padding(
+        self,
+        positions: torch.Tensor,
+        c4_positions: torch.Tensor,
+        c128_positions: torch.Tensor,
+        seq_lens_cpu: Optional[torch.Tensor] = None,
+    ) -> None:
+        c4_positions.fill_(0)
+        c128_positions.fill_(0)
+        if positions.numel() == 0:
+            return
+
+        n_draft = self.speculative_num_draft_tokens
+        request_num = positions.shape[0] // n_draft
+        if request_num == 0:
+            return
+
+        fm = self.forward_metadata
+        if seq_lens_cpu is None:
+            seq_lens_cpu = getattr(fm, "seq_lens_cpu", None)
+        if seq_lens_cpu is None:
+            seq_lens_cpu = getattr(fm, "seq_lens_cpu_int", None)
+        if seq_lens_cpu is None:
+            raise RuntimeError(
+                "DSV4 verify buffer refresh requires seq_lens_cpu or "
+                "seq_lens_cpu_int on forward metadata."
+            )
+        seq_lens_cpu = seq_lens_cpu[:request_num]
+        if seq_lens_cpu.device.type != "cpu":
+            seq_lens_cpu = seq_lens_cpu.cpu()
+
+        start_positions = seq_lens_cpu - n_draft + 1
+        abs_positions = start_positions.view(-1, 1) + torch.arange(
+            n_draft, dtype=start_positions.dtype
+        ).view(1, -1)
+        mask_c4 = (abs_positions % 4) != 0
+        mask_c128 = (abs_positions % 128) != 0
+
+        gather_shape_c4 = min(
+            positions.shape[0], mask_c4.numel(), c4_positions.shape[0]
+        )
+        gather_shape_c128 = min(
+            positions.shape[0], mask_c128.numel(), c128_positions.shape[0]
+        )
+        sorted_indices_c4 = (
+            torch.argsort(mask_c4.flatten(), dim=0, stable=True)[:gather_shape_c4]
+            .pin_memory()
+            .to(device=positions.device, non_blocking=True)
+        )
+        sorted_indices_c128 = (
+            torch.argsort(mask_c128.flatten(), dim=0, stable=True)[:gather_shape_c128]
+            .pin_memory()
+            .to(device=positions.device, non_blocking=True)
+        )
+
+        c4_positions[:gather_shape_c4].copy_(
+            torch.gather(positions, 0, sorted_indices_c4)
+        )
+        c128_positions[:gather_shape_c128].copy_(
+            torch.gather(positions, 0, sorted_indices_c128)
+        )
+
+    def _fill_verify_positions_cmp_padding_one(
+        self,
+        positions: torch.Tensor,
+        dst: torch.Tensor,
+        ratio: int,
+        seq_lens_cpu: torch.Tensor,
+        n_draft: Optional[int] = None,
+    ) -> None:
+        dst.zero_()
+        if ratio not in self._dsv4_compress_ratios or positions.numel() == 0:
+            return
+
+        if n_draft is None:
+            n_draft = self.speculative_num_draft_tokens
+        n_draft = int(n_draft)
+        request_num = positions.shape[0] // n_draft
+        if request_num == 0:
+            return
+        seq_lens_cpu = seq_lens_cpu[:request_num]
+        if seq_lens_cpu.device.type != "cpu":
+            seq_lens_cpu = seq_lens_cpu.cpu()
+
+        start_positions = seq_lens_cpu - n_draft + 1
+        abs_positions = start_positions.view(-1, 1) + torch.arange(
+            n_draft, dtype=start_positions.dtype
+        ).view(1, -1)
+        boundary_mask = abs_positions % ratio == 0
+        indices = torch.nonzero(boundary_mask.flatten(), as_tuple=False).flatten()
+
+        if indices.numel() == 0:
+            return
+        # This tiny H2D copy runs on the verify metadata path. Keep it blocking:
+        # on NPU, a non-blocking copy from a short-lived pinned CPU tensor can
+        # surface later as an unrelated CopyKernel stream failure.
+        indices = indices[: dst.numel()].to(device=positions.device)
+        dst[: indices.numel()].copy_(torch.gather(positions, 0, indices))
+
+    def update_verify_buffers_to_fill_after_draft(
+        self, spec_info, cuda_graph_bs: Optional[int]
+    ):
+        fm = self.forward_metadata
+        positions = spec_info.positions
+        c4_positions = getattr(fm, "positions_cmp_padding_c4", None)
+        c128_positions = getattr(fm, "positions_cmp_padding_c128", None)
+        if c4_positions is None or c128_positions is None:
+            return
+
+        n_draft = int(
+            getattr(spec_info, "draft_token_num", self.speculative_num_draft_tokens)
+        )
+        seq_lens_cpu = getattr(fm, "seq_lens_cpu_int", None)
+        if seq_lens_cpu is None:
+            seq_lens_cpu = getattr(spec_info, "seq_lens_cpu", None)
+            if seq_lens_cpu is None:
+                raise RuntimeError(
+                    "DSV4 verify buffer refresh requires seq_lens_cpu_int on "
+                    "forward metadata or seq_lens_cpu on spec_info."
+                )
+            seq_lens_cpu = seq_lens_cpu + n_draft
+
+        self._fill_verify_positions_cmp_padding_one(
+            positions, c4_positions, 4, seq_lens_cpu, n_draft=n_draft
+        )
+        self._fill_verify_positions_cmp_padding_one(
+            positions, c128_positions, 128, seq_lens_cpu, n_draft=n_draft
+        )
+

 def _get_kv_indices(
    forward_batch: ForwardBatch,
@@ -1486,3 +1828,224 @@ def _get_kv_indices(
    block_id = logic_pos // page_size
    offset_in_block = logic_pos % page_size
    return page_table[req_idx, block_id] * page_size + offset_in_block
+
+
+class DeepseekV4AscendMultiStepDraftBackend:
+
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        topk: int,
+        speculative_num_steps: int,
+    ):
+        self.topk = topk
+        self.speculative_num_steps = speculative_num_steps
+        self.attn_backends = [
+            DeepseekV4AscendAttnBackend(model_runner, speculative_step_id=step_id)
+            for step_id in range(speculative_num_steps)
+        ]
+
+    def common_template(self, forward_batch: ForwardBatch, call_fn):
+        assert forward_batch.spec_info is not None
+
+        for i in range(self.speculative_num_steps - 1):
+            call_fn(i, forward_batch)
+
+    def _step_out_cache_loc(self, forward_batch: ForwardBatch, step_id: int):
+        out_cache_loc = forward_batch.out_cache_loc
+        if out_cache_loc is None:
+            return None
+
+        single_step_width = forward_batch.batch_size * self.topk
+        if out_cache_loc.numel() <= single_step_width:
+            return out_cache_loc
+
+        step_layout_width = self.topk * self.speculative_num_steps
+        if step_layout_width == 0 or out_cache_loc.numel() % step_layout_width != 0:
+            return out_cache_loc
+
+        from sglang.srt.speculative.eagle_utils import per_step_draft_out_cache_loc
+
+        batch_size = out_cache_loc.numel() // step_layout_width
+        return per_step_draft_out_cache_loc(
+            out_cache_loc,
+            batch_size,
+            self.topk,
+            self.speculative_num_steps,
+        )[step_id]
+
+    def _step_out_cache_loc_dsv4(self, forward_batch: ForwardBatch, step_id: int):
+        bundle = forward_batch.out_cache_loc_dsv4
+        if bundle is None or forward_batch.out_cache_loc is None:
+            return None
+
+        step_width = forward_batch.batch_size * self.topk
+        total_width = step_width * self.speculative_num_steps
+        raw_total_width = bundle.out_full_loc.numel()
+        if (
+            raw_total_width < total_width
+            and raw_total_width % self.speculative_num_steps == 0
+            and (raw_total_width // self.speculative_num_steps) % self.topk == 0
+        ):
+            step_width = raw_total_width // self.speculative_num_steps
+            total_width = raw_total_width
+        if step_width == 0 or bundle.out_full_loc.numel() < total_width:
+            return bundle
+
+        full_steps = bundle.out_full_loc[:total_width].reshape(
+            step_width // self.topk, self.topk, self.speculative_num_steps
+        )
+        full_steps = full_steps.permute((2, 0, 1)).reshape(
+            self.speculative_num_steps, -1
+        )
+        swa_steps = bundle.out_swa_loc[:total_width].reshape(
+            step_width // self.topk, self.topk, self.speculative_num_steps
+        )
+        swa_steps = swa_steps.permute((2, 0, 1)).reshape(self.speculative_num_steps, -1)
+
+        def step_state(loc):
+            if loc is None or loc.numel() < total_width:
+                return loc
+            steps = loc[:total_width].reshape(
+                step_width // self.topk, self.topk, self.speculative_num_steps
+            )
+            return steps.permute((2, 0, 1)).reshape(self.speculative_num_steps, -1)[
+                step_id
+            ]
+
+        def step_compress(loc, ratio: int):
+            if loc is None or loc.numel() == 0:
+                return loc
+            raw_bs = step_width // self.topk
+            seq_lens = forward_batch.seq_lens[:raw_bs].to(torch.int64)
+            positions = seq_lens[:, None, None] + torch.arange(
+                self.speculative_num_steps,
+                device=seq_lens.device,
+                dtype=seq_lens.dtype,
+            )
+            positions = positions.expand(-1, self.topk, -1)
+            should_compress = ((positions + 1) % ratio) == 0
+            counts = should_compress.reshape(-1).to(torch.int64)
+            offsets = torch.cumsum(counts, dim=0) - counts
+            step_mask = should_compress[:, :, step_id].reshape(-1)
+            step_offsets = offsets.reshape(
+                raw_bs, self.topk, self.speculative_num_steps
+            )[:, :, step_id].reshape(-1)
+            return loc[step_offsets[step_mask].to(torch.int64)]
+
+        return DSV4OutCacheLoc(
+            out_full_loc=full_steps[step_id],
+            out_swa_loc=swa_steps[step_id],
+            out_c4_loc=step_compress(bundle.out_c4_loc, 4),
+            out_c128_loc=step_compress(bundle.out_c128_loc, 128),
+            out_c4_state_loc=step_state(bundle.out_c4_state_loc),
+            out_c128_state_loc=step_state(bundle.out_c128_state_loc),
+        )
+
+    def _with_step_cache_locs(self, forward_batch: ForwardBatch, step_id: int, call_fn):
+        old_out_cache_loc = forward_batch.out_cache_loc
+        old_out_cache_loc_dsv4 = forward_batch.out_cache_loc_dsv4
+        step_out_cache_loc = self._step_out_cache_loc(forward_batch, step_id)
+        if step_out_cache_loc is not None:
+            forward_batch.out_cache_loc = step_out_cache_loc
+        forward_batch.out_cache_loc_dsv4 = self._step_out_cache_loc_dsv4(
+            forward_batch, step_id
+        )
+        try:
+            return call_fn()
+        finally:
+            forward_batch.out_cache_loc = old_out_cache_loc
+            forward_batch.out_cache_loc_dsv4 = old_out_cache_loc_dsv4
+
+    def _build_step_forward_batch(
+        self, forward_batch: ForwardBatch, step_id: int
+    ) -> ForwardBatch:
+        from sglang.srt.model_executor.forward_batch_info import build_inner_fb_view
+
+        step_fb = build_inner_fb_view(
+            forward_batch,
+            bs=forward_batch.batch_size,
+            forward_mode=ForwardMode.DECODE,
+        )
+        old_bundle = forward_batch.out_cache_loc_dsv4
+        step_out_cache_loc = self._step_out_cache_loc(forward_batch, step_id)
+        step_bundle = self._step_out_cache_loc_dsv4(forward_batch, step_id)
+        step_fb.out_cache_loc_dsv4 = step_bundle
+        step_fb.global_forward_mode = getattr(
+            forward_batch, "global_forward_mode", None
+        )
+        if (
+            step_bundle is not None
+            and step_bundle is not old_bundle
+            and step_bundle.out_full_loc is not None
+        ):
+            step_fb.out_cache_loc = step_bundle.out_full_loc
+        elif step_out_cache_loc is not None:
+            step_fb.out_cache_loc = step_out_cache_loc
+        return step_fb
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        def call_fn(i, forward_batch):
+            self._with_step_cache_locs(
+                forward_batch,
+                i,
+                lambda: self.attn_backends[i].init_forward_metadata(forward_batch),
+            )
+
+        self.common_template(forward_batch, call_fn)
+
+    def init_cuda_graph_state(self, max_bs, max_num_tokens):
+        for i in range(self.speculative_num_steps):
+            self.attn_backends[i].init_cuda_graph_state(max_bs, max_num_tokens)
+
+    def init_forward_metadata_out_graph(
+        self,
+        forward_batch: ForwardBatch,
+        in_capture: bool = False,
+    ):
+        def call_fn(i, forward_batch):
+            self.attn_backends[i].init_forward_metadata_out_graph(
+                self._build_step_forward_batch(forward_batch, i),
+                in_capture=in_capture,
+            )
+
+        self.common_template(forward_batch, call_fn)
+
+    def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch) -> None:
+        def call_fn(i, forward_batch):
+            self.attn_backends[i].init_forward_metadata_in_graph(forward_batch)
+
+        self.common_template(forward_batch, call_fn)
+
+    def init_forward_metadata_replay_cuda_graph(
+        self, forward_batch: ForwardBatch, bs: int
+    ):
+        def call_fn(i, forward_batch):
+            old_oc = forward_batch.out_cache_loc
+            old_bundle = forward_batch.out_cache_loc_dsv4
+            step_bundle = self._step_out_cache_loc_dsv4(forward_batch, i)
+            forward_batch.out_cache_loc_dsv4 = step_bundle
+            if (
+                step_bundle is not None
+                and step_bundle is not old_bundle
+                and step_bundle.out_full_loc is not None
+            ):
+                forward_batch.out_cache_loc = step_bundle.out_full_loc
+            self.attn_backends[i]._replay_forward_batch = forward_batch
+            try:
+                self.attn_backends[i].init_forward_metadata_replay_cuda_graph(
+                    bs,
+                    forward_batch.req_pool_indices,
+                    forward_batch.seq_lens,
+                    seq_lens_sum=-1,
+                    encoder_lens=None,
+                    forward_mode=ForwardMode.DECODE,
+                    spec_info=forward_batch.spec_info,
+                    seq_lens_cpu=forward_batch.seq_lens_cpu,
+                )
+            finally:
+                self.attn_backends[i]._replay_forward_batch = None
+                forward_batch.out_cache_loc = old_oc
+                forward_batch.out_cache_loc_dsv4 = old_bundle
+
+        self.common_template(forward_batch, call_fn)
@@ -28,8 +28,13 @@ 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:
@@ -60,6 +65,60 @@ def get_last_loc(
    )


+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."""

@@ -289,30 +348,7 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
            "DSV4NPUTokenToKVPoolAllocator requires req_pool_indices "
            "(forwarded from batch.req_pool_indices)."
        )
-        assert dsv4_state_lens is not None, (
-            "DSV4NPUTokenToKVPoolAllocator requires dsv4_state_lens "
-            "(ScheduleBatch._compute_dsv4_state_lens_*)."
-        )
-        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,
-        )
+        if dsv4_state_lens is not None:
            out_c4_state_loc = self._alloc_state_extend(
                self.c4_state_attn_allocator,
                prefix_lens,
@@ -337,6 +373,29 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
                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,
@@ -442,6 +501,38 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
            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],
@@ -572,7 +663,7 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator):

        if req is None or req_to_token_pool is None:
            return
-        kv_len = req.kv_committed_len
+        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
@@ -610,6 +701,39 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
                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;
@@ -24,7 +24,7 @@ into the allocator itself.

 from __future__ import annotations

-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Sequence

 import torch

@@ -37,6 +37,9 @@ def maybe_write_dsv4_extend(
    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.

@@ -87,15 +90,24 @@ def maybe_write_dsv4_extend(
    )

    # c4_state / c128_state writes: tail-only. Bundle length is
-    # sum(c{N}_state_alloc_len_i), NOT total raw extend tokens; each req's slots
-    # go at raw positions [req.c{N}_state_alloc_offset, seq_len).
+    # 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,
-            [getattr(r, "c4_state_alloc_offset", 0) for r in batch.reqs],
+            c4_state_alloc_offsets,
            seq_lens_cpu,
            bundle.out_c4_state_loc,
        )
@@ -105,7 +117,7 @@ def maybe_write_dsv4_extend(
        _write_state_tail_per_req(
            req_to_token_pool.write_c128_state,
            req_pool_indices_cpu,
-            [getattr(r, "c128_state_alloc_offset", 0) for r in batch.reqs],
+            c128_state_alloc_offsets,
            seq_lens_cpu,
            bundle.out_c128_state_loc,
        )
@@ -187,6 +199,42 @@ def maybe_write_dsv4_decode(
        )


+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,
@@ -1313,15 +1313,8 @@ def _mask_topk_ids_padded_region(
    # TODO: let the kernel support other dtypes
    if _is_cuda and topk_ids.dtype == torch.int32 and fill_value == -1:
        mask_topk_ids(topk_ids, num_token_non_padded)
-    elif _can_fuse_padded_region(topk_ids):
-        _fill_padded_rows(topk_ids, num_token_non_padded, fill_value)
    elif _is_npu:
-        # On NPU, bool-indexed scatter `topk_ids[bool_mask, :] = -1` lowers
-        # to aclnnNonzeroV2 and can trigger an aicore timeout under long
-        # workloads; `torch.where` avoids that nonzero scan.
-        indices = torch.arange(0, topk_ids.shape[0], device=topk_ids.device)
-        mask = (indices >= num_token_non_padded).unsqueeze(-1)
-        topk_ids = torch.where(mask, torch.full_like(topk_ids, -1), topk_ids)
+        return
    elif _can_fuse_padded_region(topk_ids):
        _fill_padded_rows(topk_ids, num_token_non_padded, fill_value)
    else:
@@ -1475,6 +1475,7 @@ def build_inner_fb_view(
        seq_lens_cpu=forward_batch.seq_lens_cpu,
        encoder_lens=encoder_lens,
        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,
    )

@@ -987,7 +987,7 @@ class ModelRunnerKVCacheMixin:
                    "logical_attn_allocator",
                    self.token_to_kv_pool_allocator,
                )
-                assert swa_allocator.__class__ == SWATokenToKVPoolAllocator
+                assert isinstance(swa_allocator, SWATokenToKVPoolAllocator)
                self.token_to_kv_pool.full_to_swa_index_mapping = (
                    swa_allocator.full_to_swa_index_mapping
                )
@@ -129,7 +129,7 @@ def build_replay_fb_view(
    fields like spec_info, out_cache_loc, and the runtime
    actual_forward_mode) with the padded capture-time buffers from
    buffers (for req_pool_indices, seq_lens, seq_lens_cpu,
-    encoder_lens).
+    positions, encoder_lens).

    forward_mode is the capture-time mode (used by backends for
    bucket / dispatch decisions); actual_forward_mode is the
@@ -144,6 +144,7 @@ def build_replay_fb_view(
        forward_mode=capture_forward_mode,
        actual_forward_mode=forward_batch.forward_mode,
        input_ids=buffers.input_ids[:num_tokens],
+        positions=buffers.positions[:num_tokens],
        req_pool_indices=buffers.req_pool_indices[:bs],
        seq_lens=buffers.seq_lens[:bs],
        seq_lens_sum=(
@@ -2145,16 +2145,16 @@ class DeepseekV4ForCausalLM(nn.Module):
        name = name.replace(".attn_norm.", ".input_layernorm.")
        name = name.replace(".ffn_norm.", ".post_attention_layernorm.")

-        if "self_attn" in name:
-            name = name.replace(".scale", ".weight_scale_inv")
+        if "self_attn" in name and name.endswith(".scale"):
+            name = name.removesuffix(".scale") + ".weight_scale_inv"

        name = name.replace(".gate.tid2eid", ".topk.tid2eid")
        name = name.replace(".gate.bias", ".gate.e_score_correction_bias")
        name = name.replace(".w1.", ".gate_proj.")
        name = name.replace(".w2.", ".down_proj.")
        name = name.replace(".w3.", ".up_proj.")
-        if "mlp" in name:
-            name = name.replace(".scale", ".weight_scale_inv")
+        if "mlp" in name and name.endswith(".scale"):
+            name = name.removesuffix(".scale") + ".weight_scale_inv"

        return name

@@ -23,6 +23,7 @@ from sglang.srt.layers.linear import ReplicatedLinear
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.moe.utils import get_moe_a2a_backend
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig
 from sglang.srt.layers.utils.cp_utils import (
    cp_all_gather_rerange_output,
    cp_round_robin_input_ids,
@@ -91,15 +92,17 @@ class DeepseekV4ModelNextN(nn.Module):
            quant_config=quant_config,
            prefix=add_prefix("h_proj", prefix),
        )
-
-        layer_name = "decoder"
+        if isinstance(quant_config, ModelSlimConfig):
+            prefix = "mtp.0"
+        else:
+            prefix = add_prefix("decoder", prefix)

        self.decoder = DeepseekV4DecoderLayer(
            config,
            layer_id=0,
            quant_config=quant_config,
            is_nextn=True,
-            prefix=add_prefix(layer_name, prefix),
+            prefix=prefix,
            alt_streams=None,
            compress_ratio_override=COMPRESS_RATIO_NEXTN_LAYER,
        )
@@ -236,10 +236,15 @@ class DraftBackendFactory:
        )

    def _create_dsv4_decode_backend(self):
-        # On NPU the "dsv4" backend resolves to the Ascend V4 subclass; its
-        # draft path reuses the Ascend multi-step draft backend.
+        # Decode here is the EAGLE multi-step draft decode path.
        if is_npu():
-            return self._create_ascend_decode_backend()
+            from sglang.srt.hardware_backend.npu.attention.ascend_dsv4_backend import (
+                DeepseekV4AscendMultiStepDraftBackend,
+            )
+
+            return DeepseekV4AscendMultiStepDraftBackend(
+                self.draft_model_runner, self.topk, self.speculative_num_steps
+            )
        elif is_hip():
            from sglang.srt.layers.attention.deepseek_v4_backend_hip_radix import (
                DeepseekV4MultiStepBackend,
@@ -338,9 +343,13 @@ class DraftBackendFactory:

    def _create_dsv4_prefill_backend(self):
        # On NPU the "dsv4" backend resolves to the Ascend V4 subclass; its
-        # draft-extend path reuses the Ascend prefill draft backend.
+        # draft-extend path uses the registered DSV4 prefill backend.
        if is_npu():
-            return self._create_ascend_prefill_backend()
+            from sglang.srt.layers.attention.attention_registry import (
+                ATTENTION_BACKENDS,
+            )
+
+            return ATTENTION_BACKENDS["dsv4"](self.draft_model_runner)
        elif is_hip():
            from sglang.srt.layers.attention.deepseek_v4_backend_hip_radix import (
                DeepseekV4HipRadixBackend,
@@ -556,7 +556,8 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
            seq_lens_sum=seq_lens_sum,
            seq_lens_cpu=buffers.seq_lens_cpu,
            encoder_lens=None,
-            out_cache_loc=forward_batch.out_cache_loc,
+            out_cache_loc=buffers.out_cache_loc[:num_tokens],
+            out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None),
            spec_info=forward_batch.spec_info,
        )
        self.draft_extend_attn_backend.init_forward_metadata_out_graph(fb_view)
@@ -1,11 +1,18 @@
 from __future__ import annotations

 import math
+from collections import defaultdict
 from enum import IntEnum
 from typing import TYPE_CHECKING, List, Optional

 import torch

+from sglang.srt.hardware_backend.npu.dsv4.dsv4_allocator import (
+    alloc_paged_token_slots_extend_npu,
+)
+from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import (
+    maybe_build_dsv4_verify_bundle,
+)
 from sglang.srt.mem_cache.common import (
    alloc_paged_token_slots_extend,
    alloc_token_slots,
@@ -34,6 +41,14 @@ if _is_cuda or _is_hip or _is_musa:
    )


+ALLOC_EXTEND_FUNCS = defaultdict(
+    lambda: alloc_paged_token_slots_extend,
+    {
+        "npu": alloc_paged_token_slots_extend_npu,
+    },
+)
+
+
 def per_step_draft_out_cache_loc(
    out_cache_loc: torch.Tensor,
    batch_size: int,
@@ -355,6 +370,10 @@ def eagle_prepare_for_verify(
            device=device,
        )

+        batch.out_cache_loc_dsv4 = maybe_build_dsv4_verify_bundle(
+            batch, verify_input.draft_token_num
+        )
+
        prepare_mamba_track_for_verify(batch)

        # TBO's split_spec_info reads these; no-verify-sync leaves both None.
@@ -663,7 +682,8 @@ def eagle_prepare_for_decode(batch: ScheduleBatch):
            batch.req_pool_indices,
            cur_kv_lens_device,
        )
-        out_cache_loc = alloc_paged_token_slots_extend(
+        device_type = getattr(batch.device, "type", str(batch.device).split(":", 1)[0])
+        out_cache_loc = ALLOC_EXTEND_FUNCS[device_type](
            batch.tree_cache,
            cur_kv_lens_device,
            cur_kv_lens_cpu,
@@ -671,8 +691,9 @@ def eagle_prepare_for_decode(batch: ScheduleBatch):
            nxt_kv_lens_cpu,
            last_loc,
            num_needed_tokens,
+            req_pool_indices=batch.req_pool_indices,
+            batch=batch,
        )
-
    assign_req_to_token_pool_func(
        batch.req_pool_indices,
        batch.req_to_token_pool.req_to_token,