diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index 512eb4c77..a3211ad2c 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -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 diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py index f3bfca468..4cc1c26a0 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py @@ -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 - seq_lens_max = int(seq_lens.max().item()) if bs > 0 else 0 + 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]: - cmp_kv = cmp_kv[: loc.numel()] + 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)." - ) - 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)) + 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) + 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) diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py index e7f2e5962..20d39e7b2 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py @@ -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,10 +348,34 @@ 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_*)." - ) + if dsv4_state_lens is not None: + out_c4_state_loc = self._alloc_state_extend( + self.c4_state_attn_allocator, + prefix_lens, + dsv4_state_lens.c4_prefix_lens, + dsv4_state_lens.c4_prefix_lens_cpu, + dsv4_state_lens.c4_seq_lens, + dsv4_state_lens.c4_seq_lens_cpu, + req_pool_indices, + last_loc_dtype, + dsv4_state_lens.c4_extend_num_tokens, + ratio=4, + ) + out_c128_state_loc = self._alloc_state_extend( + self.c128_state_attn_allocator, + prefix_lens, + dsv4_state_lens.c128_prefix_lens, + dsv4_state_lens.c128_prefix_lens_cpu, + dsv4_state_lens.c128_seq_lens, + dsv4_state_lens.c128_seq_lens_cpu, + req_pool_indices, + last_loc_dtype, + dsv4_state_lens.c128_extend_num_tokens, + ratio=128, + ) + else: + out_c4_state_loc = self._empty_loc + out_c128_state_loc = self._empty_loc out_c4_loc = self._alloc_c_extend( self.c4_attn_allocator, prefix_lens, @@ -313,30 +396,6 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): last_loc_dtype, ratio=128, ) - out_c4_state_loc = self._alloc_state_extend( - self.c4_state_attn_allocator, - prefix_lens, - dsv4_state_lens.c4_prefix_lens, - dsv4_state_lens.c4_prefix_lens_cpu, - dsv4_state_lens.c4_seq_lens, - dsv4_state_lens.c4_seq_lens_cpu, - req_pool_indices, - last_loc_dtype, - dsv4_state_lens.c4_extend_num_tokens, - ratio=4, - ) - out_c128_state_loc = self._alloc_state_extend( - self.c128_state_attn_allocator, - prefix_lens, - dsv4_state_lens.c128_prefix_lens, - dsv4_state_lens.c128_prefix_lens_cpu, - dsv4_state_lens.c128_seq_lens, - dsv4_state_lens.c128_seq_lens_cpu, - req_pool_indices, - last_loc_dtype, - dsv4_state_lens.c128_extend_num_tokens, - ratio=128, - ) 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; diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_common_hooks.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_common_hooks.py index c7aa90e6e..4e5e7e6dc 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_common_hooks.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_common_hooks.py @@ -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, diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index e740cded7..4318218ee 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -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: diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 721b1e447..ef273f0fb 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -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, ) diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index 89932d921..63877b48f 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -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 ) diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index c5d727a6b..e34f6826d 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -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=( diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index bb4754ccd..d01d17836 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -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 diff --git a/python/sglang/srt/models/deepseek_v4_nextn.py b/python/sglang/srt/models/deepseek_v4_nextn.py index a11880b22..46ea9e52f 100644 --- a/python/sglang/srt/models/deepseek_v4_nextn.py +++ b/python/sglang/srt/models/deepseek_v4_nextn.py @@ -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, ) diff --git a/python/sglang/srt/speculative/draft_utils.py b/python/sglang/srt/speculative/draft_utils.py index fad4f5b9e..291620c28 100644 --- a/python/sglang/srt/speculative/draft_utils.py +++ b/python/sglang/srt/speculative/draft_utils.py @@ -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, diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index 1394f8150..d546b0cd4 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -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) diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index 14f8ebb31..b75bce92e 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -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,