diff --git a/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py b/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py index 38a9e94b0..0de315543 100644 --- a/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py +++ b/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py @@ -9,7 +9,11 @@ from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_f from sglang.kernels.ops.speculative.dspark.dispatch import inputs_on_cuda from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout -from sglang.srt.utils import is_npu +from sglang.srt.utils import ( + is_npu, +) + +_is_npu = is_npu() class RaggedVerifyWindow(msgspec.Struct, frozen=True): @@ -798,7 +802,7 @@ def build_commit_inject_layout_triton( class BuildOutTokens: @classmethod def execute(cls, *args, **kwargs) -> torch.Tensor: - if not is_npu() and inputs_on_cuda(*args, **kwargs): + if inputs_on_cuda(*args, **kwargs) and not _is_npu: return cls.triton(*args, **kwargs) return cls.torch(*args, **kwargs) diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index ad01b9c7e..c10348325 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -277,9 +277,9 @@ def _target_checkpoint_bundles_dspark_draft(server_args: ServerArgs) -> bool: def _handle_dspark(server_args: ServerArgs) -> None: _is_npu = server_args.device.startswith("npu") - if not server_args.device.startswith("cuda") and not _is_npu: + if not server_args.device.startswith(("cuda", "npu")): raise ValueError( - "DSpark speculative decoding only supports CUDA and NPU devices." + "DSpark speculative decoding only supports CUDA or NPU device." ) # dp_size==1 with dp_attention is a degenerate flag under DSV4 CP; skip DP-only checks. diff --git a/python/sglang/srt/disaggregation/ascend/conn.py b/python/sglang/srt/disaggregation/ascend/conn.py index 2ad9da407..8ab6b06f7 100644 --- a/python/sglang/srt/disaggregation/ascend/conn.py +++ b/python/sglang/srt/disaggregation/ascend/conn.py @@ -21,15 +21,9 @@ logger = logging.getLogger(__name__) class AscendStateType(str, enum.Enum): - """DSV4-on-NPU per-pool PD components, kept out of the cross-hardware - StateType enum. Sent via the same page-indexed path as SWA.""" + """DSV4-on-NPU PD components without a cross-hardware equivalent.""" - DSV4_SWA = "dsv4_swa" - DSV4_C4 = "dsv4_c4" DSV4_C128 = "dsv4_c128" - DSV4_INDEXER = "dsv4_indexer" - DSV4_C4_STATE = "dsv4_c4_state" - DSV4_C128_STATE = "dsv4_c128_state" _DSV4_KVCACHE_STATE_TYPES = tuple(AscendStateType) @@ -71,6 +65,32 @@ class AscendKVManager(MooncakeKVManager): def get_mla_kv_ptrs_with_pp( self, src_kv_ptrs: List[int], dst_kv_ptrs: List[int], state_type=None ) -> Tuple[List[int], List[int], int]: + mla_ratios = getattr(self.kv_args, "mla_compression_ratios", None) + if mla_ratios: + if len(src_kv_ptrs) == len(dst_kv_ptrs): + return src_kv_ptrs, dst_kv_ptrs, len(src_kv_ptrs) + + start_layer = self.kv_args.prefill_start_layer + end_layer = self.kv_args.prefill_end_layer + c4_full = sum(ratio == 4 for ratio in mla_ratios) + c4_start = sum(ratio == 4 for ratio in mla_ratios[:start_layer]) + c4_end = sum(ratio == 4 for ratio in mla_ratios[:end_layer]) + c128_start = sum(ratio == 128 for ratio in mla_ratios[:start_layer]) + c128_end = sum(ratio == 128 for ratio in mla_ratios[:end_layer]) + + if state_type == AscendStateType.DSV4_C128: + dst = dst_kv_ptrs[c128_start:c128_end] + return src_kv_ptrs, dst, len(src_kv_ptrs) + + # NPU main KV layout: [C4 KV, index K, index scale]. + if state_type is None and len(dst_kv_ptrs) == 3 * c4_full: + dst = [] + for offset in (0, c4_full, 2 * c4_full): + dst.extend(dst_kv_ptrs[offset + c4_start : offset + c4_end]) + return src_kv_ptrs, dst, len(src_kv_ptrs) + + return super().get_mla_kv_ptrs_with_pp(src_kv_ptrs, dst_kv_ptrs, state_type) + # src_kv_ptrs: k_data, v_data, index_k_data(optional) # dst_kv_ptrs: k_data, v_data, index_k_data(optional) # state_type is accepted for parity with the common disaggregation path; diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index c2c009b81..c26de286c 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1102,17 +1102,6 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): self.tree_cache.dec_lock_ref(decode_req.req.last_node) break - if total_prefix_len != 0 and hasattr( - self.token_to_kv_pool_allocator, "c4_attn_allocator" - ): - if prefix_len > 0: - self.tree_cache.dec_lock_ref(decode_req.req.last_node) - raise RuntimeError( - "DSV4 NPU PD disaggregation does not support decode-side " - "prefix cache yet; disable disaggregation decode radix/HiCache " - "for PD + chunked prefill." - ) - dst_kv_indices = self._pre_alloc( decode_req.req, prefix_indices, @@ -1176,7 +1165,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): def _swa_payload(): window_size = self.scheduler.sliding_window_size - window_start = max(0, seq_len - window_size) + window_start = max(total_prefix_len, seq_len - window_size) window_start = page_align_floor(window_start, page_size) window_kv_indices_full = self.req_to_token_pool.req_to_token[ decode_req.req.req_pool_idx, window_start:seq_len @@ -1234,15 +1223,6 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): StateType.SWA_RING: _swa_ring_payload, StateType.C128_STATE: _c128_state_payload, } - if hasattr(self.req_to_token_pool, "req_to_token_c4"): - # DSV4 on NPU: per-pool dst page indices, produced by the same - # shared builder prefill uses so src/dst line up positionally. - if total_prefix_len != 0: - raise RuntimeError( - "DSV4 NPU PD disaggregation does not support decode-side " - "prefix cache yet; disable disaggregation decode radix/HiCache " - "for PD + chunked prefill." - ) if _is_npu and isinstance(self.token_to_kv_pool, DeepSeekV4TokenToKVPool): from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( dsv4_state_payloads, @@ -1254,7 +1234,6 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): decode_req.req.req_pool_idx, seq_len, self.token_to_kv_pool_allocator.page_size, - self.scheduler.sliding_window_size, prefix_len=total_prefix_len, ) ) @@ -1743,7 +1722,7 @@ def alloc_for_decode_prealloc( ) extra_kwargs = {} dsv4_unwrap_prealloc = None - if hasattr(allocator, "c4_attn_allocator"): + if hasattr(allocator, "c128_attn_allocator"): assert req_to_token_pool is not None from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( dsv4_prealloc_kwargs, diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 2a15fa758..5c44bb1eb 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -1200,7 +1200,7 @@ class SchedulerDisaggregationPrefillMixin: def _swa_payload(): window_size = self.sliding_window_size - window_start = max(0, seq_len - window_size) + window_start = max(req.disagg_decode_prefix_len, seq_len - window_size) window_start = (window_start // page_size) * page_size window_kv_indices_full = self.req_to_token_pool.req_to_token[ req.req_pool_idx, window_start:seq_len @@ -1274,8 +1274,7 @@ class SchedulerDisaggregationPrefillMixin: req.req_pool_idx, seq_len, page_size, - self.sliding_window_size, - prefix_len=0, + prefix_len=req.disagg_decode_prefix_len, ) ) state_indices = [ diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 16e303377..96eaa1e5a 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -1064,17 +1064,7 @@ def setup_state_kv_args( kv_args.is_hybrid_mla_backend = False kv_args.state_conv_shard_groups = [] - if is_npu() and isinstance(token_to_kv_pool, DSV4NPUTokenToKVPool): - # Pool ships each sub-pool as its own page-indexed component (fixed order - # so prefill and decode register identically); skips get_state_buf_infos. - for ( - st, - comp_ptrs, - comp_lens, - comp_item_lens, - ) in token_to_kv_pool.get_pd_state_components(): - append_state_component(kv_args, st, comp_ptrs, comp_lens, comp_item_lens) - elif isinstance(token_to_kv_pool, MiniMaxSparseKVPool): + if isinstance(token_to_kv_pool, MiniMaxSparseKVPool): if token_to_kv_pool.index_kv_pool is not None: raise NotImplementedError( "PD disaggregation for MiniMax sparse layers with index value " @@ -1175,15 +1165,25 @@ def setup_state_kv_args( kv_args, StateType.DSA, data_ptrs, data_lens, item_lens ) + if is_npu() and isinstance(token_to_kv_pool, DSV4NPUTokenToKVPool): + from sglang.srt.disaggregation.ascend.conn import AscendStateType + + c128_ptrs, c128_lens, c128_item_lens = token_to_kv_pool.get_c128_kv_buf_infos() + if c128_ptrs: + append_state_component( + kv_args, + AscendStateType.DSV4_C128, + c128_ptrs, + c128_lens, + c128_item_lens, + ) + # DSV4 NextN shares the target allocator, so target and draft use the same # local SWA indices. Keep draft buffers in a separate positional component # to avoid mixing them into the target's heterogeneous state layout, while - # reusing the existing SWA transport dispatch. NPU has a different paged - # state layout and is intentionally left unchanged. - if ( - not is_npu() - and isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool) - and isinstance(draft_token_to_kv_pool, DeepSeekV4TokenToKVPool) + # reusing the existing SWA transport dispatch on both GPU and NPU. + if isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool) and isinstance( + draft_token_to_kv_pool, DeepSeekV4TokenToKVPool ): if not draft_token_to_kv_pool.compression_ratios or not all( ratio == 0 for ratio in draft_token_to_kv_pool.compression_ratios 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 5d6ae2a74..4cc25b584 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -443,15 +443,14 @@ class AscendAttnBackend(AttentionBackend): def init_forward_metadata(self, forward_batch: ForwardBatch): """Init the metadata for a forward pass.""" self.forward_metadata = ForwardMetadata() + seq_lens_max = forward_batch.seq_lens.max() if forward_batch.forward_mode.is_target_verify(): + spec_tokens_per_req = int(forward_batch.spec_info.draft_token_num) # Overlap scheduling can publish the CPU sequence length one step # ahead of the device tensor. FIA consumes seq_lens_cpu below, so # derive the block-table width from the same source. Otherwise a # page-aligned request can expose KV_S=N while asking FIA for N+1. - seq_lens_max = ( - forward_batch.seq_lens_cpu.max().item() - + self.speculative_num_draft_tokens - ) + seq_lens_max = forward_batch.seq_lens_cpu.max().item() + spec_tokens_per_req elif ( forward_batch.forward_mode.is_decode_or_idle() and forward_batch.spec_info is not None @@ -499,10 +498,10 @@ class AscendAttnBackend(AttentionBackend): seq_lens_list_cumsum = np.cumsum(forward_batch.extend_seq_lens_cpu) self.forward_metadata.seq_lens_list_cumsum = seq_lens_list_cumsum - if forward_batch.forward_mode.is_target_verify() and not _is_dflash_verify( - forward_batch.spec_info - ): - self.forward_metadata.seq_lens_cpu_int += self.speculative_num_draft_tokens + if forward_batch.forward_mode.is_target_verify(): + spec_algorithm = forward_batch.spec_algorithm + if spec_algorithm is None or not spec_algorithm.is_dspark(): + self.forward_metadata.seq_lens_cpu_int += spec_tokens_per_req elif ( forward_batch.forward_mode.is_decode_or_idle() and forward_batch.spec_info is not None @@ -519,11 +518,16 @@ class AscendAttnBackend(AttentionBackend): forward_batch.forward_mode.is_target_verify() or forward_batch.forward_mode.is_draft_extend_v2() ): + spec_tokens_per_req = ( + int(forward_batch.spec_info.draft_token_num) + if forward_batch.forward_mode.is_target_verify() + else self.speculative_num_draft_tokens + ) self.forward_metadata.actual_seq_lengths_q = torch.arange( - self.speculative_num_draft_tokens, - self.speculative_num_draft_tokens - + forward_batch.seq_lens.shape[0] * self.speculative_num_draft_tokens, - self.speculative_num_draft_tokens, + spec_tokens_per_req, + spec_tokens_per_req + + forward_batch.seq_lens.shape[0] * spec_tokens_per_req, + spec_tokens_per_req, dtype=torch.int32, device=self.device, ) @@ -720,12 +724,15 @@ class AscendAttnBackend(AttentionBackend): max_seq_pages = (max_len + self.page_size - 1) // self.page_size if self.is_hybrid_swa: - metadata.block_tables_swa[:bs, :max_seq_pages].copy_( - self.full_to_swa_index_mapping[ - self.req_to_token[req_pool_indices[:bs], :max_len] - ][:, :: self.page_size] - // self.page_size + full_page_locs = self.req_to_token[ + req_pool_indices[:bs], + 0 : max_len : self.page_size, + ] + swa_page_table = ( + self.full_to_swa_index_mapping[full_page_locs] // self.page_size ) + + metadata.block_tables_swa[:bs, :max_seq_pages].copy_(swa_page_table) metadata.block_tables_swa[:bs, max_seq_pages:].fill_(0) metadata.block_tables_swa[bs:, :].fill_(0) 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 5c147a2af..bc9c816f7 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 @@ -7,13 +7,19 @@ from typing import TYPE_CHECKING, Optional import torch import torch.nn.functional as F +import torch_npu +from sglang.kernels.ops.speculative.dspark.dspark_attn_metadata import ( + BuildBlockSeqLensCausal, + BuildDsparkSwaPageIndices, + ComputeDsparkWindowGather, +) +from sglang.srt.environ import envs from sglang.srt.hardware_backend.npu.attention.ascend_backend import AscendAttnBackend -from sglang.srt.layers.attention.dsv4.compressor import CompressorBackendMixin -from sglang.srt.layers.attention.dsv4.indexer import C4IndexerBackendMixin +from sglang.srt.hardware_backend.npu.dsv4.dsv4_rope import Dsv4NpuRoPE from sglang.srt.model_executor.forward_batch_info import DSV4OutCacheLoc, ForwardMode from sglang.srt.model_executor.forward_context import get_attn_backend -from sglang.srt.runtime_context import get_parallel, get_spec +from sglang.srt.runtime_context import get_parallel if TYPE_CHECKING: from sglang.srt.layers.radix_attention import RadixAttention @@ -50,21 +56,48 @@ def _apply_hadamard(inp: torch.Tensor, hadamard_matrix: torch.Tensor) -> torch.T return flat.matmul(hadamard_matrix).view(init_shape).to(torch.bfloat16) -def _overlap_transform( - tensor: torch.Tensor, value: float, head_dim: int +def _build_explicit_state_block_table( + *, + compress_ratio: int, + coff: int, + state_pool, + token_to_kv_pool, + req_to_token: torch.Tensor, + req_pool_indices: torch.Tensor, + start_pos: torch.Tensor, + cu_seqlens: torch.Tensor, + seqused: torch.Tensor, + max_input_capacity: int, ) -> torch.Tensor: - # Build (n_chunks, 2*ratio, d) from (n_chunks, ratio, coff*d): first ratio rows - # = current chunk left half (:d), last ratio rows = previous chunk right half (d:); - # first chunk's right half filled with `value`. - n_chunks, r, _ = tensor.shape - d = head_dim - out = tensor.new_full((n_chunks, 2 * r, d), value) - out[:, r:] = tensor[..., d:] - out[1:, :r] = tensor[:-1, :, :d] - return out + """Adapt GPU-style state locations to the A3 cache_mode=2 table ABI.""" + req_pool_indices = req_pool_indices.to(torch.int64) + capacities = cu_seqlens[1:] - cu_seqlens[:-1] + history_size = coff * compress_ratio + width = history_size + max_input_capacity + columns = torch.arange(width, dtype=torch.int64, device=req_to_token.device) + positions = start_pos[:, None] - history_size + columns + within_capacity = columns[None, :] < history_size + capacities[:, None] + valid = (seqused[:, None] > 0) & within_capacity & (positions >= 0) + + if compress_ratio == 4: + # Masked history/ragged columns are still indexed before torch.where. + safe_positions = positions.clamp(0, req_to_token.shape[1] - 1) + full_locs = req_to_token[req_pool_indices[:, None], safe_positions] + swa_locs = token_to_kv_pool.translate_loc_from_full_to_swa(full_locs) + state_locs = state_pool.translate_from_swa_loc_to_state_loc(swa_locs) + else: + state_locs = state_pool.translate_from_req_position_to_state_loc( + req_pool_indices[:, None], positions + ) + + return torch.where( + valid, + state_locs.to(torch.int32), + state_pool.dummy_state_loc, + ).contiguous() -class CompressorAscendBackendMixin(CompressorBackendMixin): +class CompressorAscendBackendMixin: @staticmethod def _to_cpu_int_list(values) -> Optional[list[int]]: @@ -98,10 +131,13 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): fm = self.forward_metadata is_decode = forward_batch.forward_mode.is_decode() is_verify = forward_batch.forward_mode.is_target_verify() + fm.dsv4_explicit_state_block_tables = {} + fm.dsv4_max_input_capacity = 1 if is_decode else None _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 + n_draft = int(forward_batch.spec_info.draft_token_num) + _seq_lens = _seq_lens + n_draft result = self._compute_compress_locs( pool=self.token_to_kv_pool, req_to_token=self.req_to_token, @@ -119,8 +155,6 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): if not is_decode: for ratio in self._dsv4_compress_ratios: if ratio in (4, 128): - if f"c{ratio}_state_loc" not in result: - setattr(fm, f"c{ratio}_state_loc", None) if f"c{ratio}_loc" not in result: setattr(fm, f"c{ratio}_loc", None) # _compute_compress_locs builds positions_cmp_padding / start_pos / @@ -144,6 +178,13 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): cu = fm.actual_seq_lengths_q_pa cu_cpu = cu.cpu().tolist() + fm.dsv4_max_input_capacity = max( + 1, + max( + (int(cu_cpu[idx + 1]) - int(cu_cpu[idx]) for idx in range(bs)), + default=0, + ), + ) prefix_cpu = self._extend_prefix_lens_cpu(forward_batch) ratio_lists: dict = { r: [] for r in self._dsv4_unique_compress_ratios if r in (4, 128) @@ -182,10 +223,8 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): padding[: cat.shape[0]].copy_(cat) setattr(fm, f"positions_cmp_padding_c{ratio}", padding) - # start_pos = each req's GLOBAL start (= extend_prefix_lens) so the fused op - # (cache_mode=1) reads the prior-chunk partial-block state and aligns blocks - # to the global grid; prefix==0 -> 0 (non-chunked, unchanged). Only the fused - # chunked path reads it. seqused=None -> op derives chunk len from cu_seqlens. + # start_pos is each request's global chunk start. cache_mode=2 uses it + # together with the explicit table to align history/current columns. if forward_batch.extend_prefix_lens is not None: fm.start_pos = forward_batch.extend_prefix_lens.to( device=device, dtype=torch.int32 @@ -198,7 +237,7 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): ) else: fm.start_pos = torch.zeros(bs, dtype=torch.int32, device=device) - fm.seqused = None + fm.seqused = (cu[1:] - cu[:-1]).to(torch.int32) # bundle out_c*_loc = the NEW c-pool slots allocated this extend (incremental), # densely packed in batch order to match cmp_kv. Valid under chunked prefill: @@ -218,27 +257,6 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): bundle_loc.to(torch.int32) if bundle_loc is not None else None, ) - # req_to_token_c*_state is not re-zeroed on slot reuse; zero pre-tail page cols so the kernel block-0 skip masks stale entries - page_size = self.page_size - for ratio in (4, 128): - spt = getattr(fm, f"c{ratio}_state_page_table", None) - if spt is None: - continue - for idx in range(bs): - chunk_len = int(cu_cpu[idx + 1] - cu_cpu[idx]) - if chunk_len == 0: - continue - seqlen = int(prefix_cpu[idx]) + chunk_len - tail = seqlen % 128 - if ratio == 4: - c_alloc_len = tail + 128 if (tail <= 3 and seqlen >= 128) else tail - else: - c_alloc_len = tail - c_alloc_offset = seqlen - c_alloc_len - first_tail_page = c_alloc_offset // page_size - if first_tail_page > 0: - spt[idx, :first_tail_page] = 0 - def _compute_compress_locs( self, *, @@ -263,32 +281,12 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): 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_unique_compress_ratios: if ratio not in (4, 128): continue - state_loc_src = None bundle_loc = None - # state table holds one slot per RAW token; block 0 is the skip sentinel reserved by NPUCompressStatePool - state_table = ( - req_to_token_pool.req_to_token_c4_state - if ratio == 4 - else req_to_token_pool.req_to_token_c128_state - ) - state_slots_2d = state_table[req_pool_64, : n_pages * self.page_size] - state_page_2d = (state_slots_2d[:, :: self.page_size] // self.page_size).to( - torch.int32 - ) if is_decode: - if out_cache_loc_dsv4 is not None: - state_loc_src = ( - out_cache_loc_dsv4.out_c4_state_loc - if ratio == 4 - else out_cache_loc_dsv4.out_c128_state_loc - ) - # bundle_loc and cmp_kv are both densely packed in batch order, so # write them densely; indexing by batch slot would misalign them. if out_cache_loc_dsv4 is not None: @@ -298,16 +296,7 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): else out_cache_loc_dsv4.out_c128_loc ) - result[f"c{ratio}_state_page_table"] = state_page_2d if is_decode: - if state_loc_src is None: - state_loc_decode = torch.zeros( - bs, - dtype=torch.int32, - device=device, - ) - else: - state_loc_decode = state_loc_src.to(torch.int32) compress_out_loc = torch.zeros( bs, dtype=torch.int32, @@ -317,23 +306,24 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): n_compress = bundle_loc.numel() if n_compress > 0: compress_out_loc[:n_compress] = bundle_loc.to(torch.int32) - result[f"c{ratio}_state_loc"] = state_loc_decode result[f"c{ratio}_loc"] = compress_out_loc - c_table = ( - req_to_token_pool.req_to_token_c4 - if ratio == 4 - else req_to_token_pool.req_to_token_c128 - ) # graph: keep shape aligned with the preallocated buffer; eager: clamp >=1 so kernels see a column if is_graph: n_c_tokens = seq_lens_max // ratio else: n_c_tokens = max(1, seq_lens_max // ratio) - slots = c_table[req_pool_64, :n_c_tokens] - c_page_table = (slots[:, :: self.page_size] // self.page_size).to( - torch.int32 - ) + if ratio == 4: + slots = req_to_token[req_pool_64, : n_c_tokens * ratio] + c_page_table = (slots[:, :: self.page_size] // self.page_size).to( + torch.int32 + ) + else: + c128_page_size = req_to_token_pool.c128_page_size + n_groups = (n_c_tokens + c128_page_size - 1) // c128_page_size + c_page_table = req_to_token_pool.req_to_c128_sidecar[ + req_pool_64, :n_groups + ].to(torch.int32) result[f"c{ratio}_page_table"] = c_page_table if is_decode: @@ -372,7 +362,6 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): x: torch.Tensor, forward_batch: ForwardBatch, ) -> None: - from sglang.srt.hardware_backend.npu.dsv4.dsv4_rope import Dsv4NpuRoPE ratio = compressor.ratio coff = 1 + int(compressor.overlap) @@ -381,28 +370,29 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): self._ensure_fused_caches(compressor) fm = self.forward_metadata - positions_cmp = getattr(fm, f"positions_cmp_padding_c{ratio}", None) - page_table = getattr(fm, f"c{ratio}_state_page_table", None) - start_pos = getattr(fm, "start_pos", None) - seqused = getattr(fm, "seqused", None) - cu_seqlens = getattr(fm, "actual_seq_lengths_q_pa", None) - assert positions_cmp is not None and page_table is not None, ( - "fused compressor needs backend metadata " - "(positions_cmp_padding / c*_state_page_table) — make sure " - "_build_npu_compress_metadata ran before this forward." - ) - assert start_pos is not None, "fused compressor needs start_pos" - assert cu_seqlens is not None, "fused compressor needs cu_seqlens" - pool = self.token_to_kv_pool - state_cache = pool.get_state_cache( - compressor.layer_id, compressor.is_in_indexer - ) + state_pool = pool._get_state_pool(compressor.layer_id, compressor.is_in_indexer) + state_cache = state_pool.state_cache_3d + table_cache = fm.dsv4_explicit_state_block_tables + if ratio not in table_cache: + table_cache[ratio] = _build_explicit_state_block_table( + compress_ratio=ratio, + coff=coff, + state_pool=state_pool, + token_to_kv_pool=pool, + req_to_token=self.req_to_token, + req_pool_indices=forward_batch.req_pool_indices, + start_pos=fm.start_pos, + cu_seqlens=fm.actual_seq_lengths_q_pa, + seqused=fm.seqused, + max_input_capacity=fm.dsv4_max_input_capacity, + ) + state_block_table = table_cache[ratio] cos, sin = Dsv4NpuRoPE.for_freqs( compressor.freqs_cis, getattr(compressor, "rotary_emb", None) ).get_cos_sin( - positions_cmp, + getattr(fm, f"positions_cmp_padding_c{ratio}"), torch.float32, view_4d=False, allow_build=False, @@ -419,14 +409,14 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): rope_cos=cos, rope_head_dim=compressor.rope_head_dim, cmp_ratio=ratio, - state_block_table=page_table, - cu_seqlens=cu_seqlens, - seqused=seqused, - start_pos=start_pos, + state_block_table=state_block_table, + cu_seqlens=fm.actual_seq_lengths_q_pa, + seqused=fm.seqused, + start_pos=fm.start_pos, coff=coff, norm_eps=compressor.norm.variance_epsilon, rotary_mode=2, - cache_mode=1, + cache_mode=2, ) # prefill output may be padded; trim to loc length @@ -450,317 +440,6 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): cmp_kv = _apply_hadamard(cmp_kv, compressor.hadamard_matrix) self._compressor_epilog_npu(compressor, cmp_kv, forward_batch) - def _forward_compress_native( - self, - compressor, - x: torch.Tensor, - forward_batch: ForwardBatch, - ) -> None: - """Reference per-request unfused compress path for precision ablations. - - Production dispatch no longer calls this path: ordinary prefill, every - chunked-prefill chunk, verify, and decode all use the fused compressor. - - * Prefill: split seq into ``cutoff = seqlen - seqlen % ratio`` to compress - + ``remainder`` stashed as state (overlap/ratio=4 also stashes the last - ``ratio`` of the cutoff). State writes via ``set_state_buffer``; cutoff gets - ape-weighted softmax over ratio, sum, norm+rope+(opt) hadamard, then write. - * Non-prefill (one token/req): append (kv, score) to the state ring; if it - completes a ratio-aligned chunk, gather the chunk (overlap: 2*ratio, else - ratio), ape-weighted softmax + sum, and write via ``set_compress_buffer``. - """ - import torch_npu # local: NPU-only, used for npu_rotary_mul below - - positions = forward_batch.positions - ratio, overlap, d = compressor.ratio, compressor.overlap, compressor.head_dim - device = x.device - self._ensure_compressor_hadamard(compressor, device) - dtype = x.dtype - x_f32 = x.float() - # wkv + wgate are fused into one wkv_gate.weight [2*coff*head_dim, hidden_size] - # (kv concatenated before wgate); split along the output dim to recover each. - coff = 1 + int(overlap) - W = compressor.wkv_gate.weight.float() - kv_full = F.linear(x_f32, W[: coff * d]) # [T, coff*d] - score_full = F.linear(x_f32, W[coff * d :]) # [T, coff*d] - - seq_lens_cpu = forward_batch.seq_lens_cpu - extend_prefix_lens_cpu = self._extend_prefix_lens_cpu(forward_batch) - is_prefill = forward_batch.forward_mode.is_prefill() - token_to_kv_pool = self.token_to_kv_pool - backend_fm = self.forward_metadata - if ratio == 4: - page_table = backend_fm.c4_state_page_table - else: - page_table = backend_fm.c128_state_page_table - - kv_out_list: list[torch.Tensor] = [] - kv_state_to_be_cached: list[torch.Tensor] = [] - score_state_to_be_cached: list[torch.Tensor] = [] - state_loc_list: list[torch.Tensor] = [] - kv_out_positions: list[torch.Tensor] = [] - # Per-token write loc: record (req_idx_in_batch, compressed_seq_pos_in_req) - # to derive the c{N}_kv_pool slot from the slab allocator, not out_cache_loc - # // ratio (correct only when raw kv allocation aligns to ratio). - write_req_indices: list[torch.Tensor] = [] - write_pos_in_req: list[torch.Tensor] = [] - seqlen_offset = 0 - # Running offset into the tail-only state bundle, flat layout - # ``[req0_alloc_len_slots, ...]`` where ``alloc_len_i = seqlen_i - - # c{ratio}_state_alloc_offset_i`` (NOT raw seqlen; see - # ScheduleBatch._compute_dsv4_state_lens_extend). - state_bundle_offset = 0 - - for idx, seqlen in enumerate(seq_lens_cpu): - seqlen = int(seqlen) - if seqlen == 0: - continue - if is_prefill: - # Chunked follow-up (prefix_len>0) is routed to the fused compressor - # by the forward_compress dispatch (main compressor + c4 indexer), so - # the native path only ever sees non-chunked / first-chunk prefill. - prefix_len = ( - int(extend_prefix_lens_cpu[idx]) - if extend_prefix_lens_cpu is not None - else 0 - ) - assert prefix_len == 0, ( - "native compress prefill reached with prefix_len=" - f"{prefix_len}; chunked prefill must route to the fused op" - ) - pos_req = positions[seqlen_offset : seqlen_offset + seqlen] - - # Per-req tail-only state alloc range; same formula as - # ScheduleBatch._compute_dsv4_state_lens_extend (recomputed to - # avoid threading another tensor through forward_batch). - tail_128 = seqlen % 128 - if ratio == 4: - c_alloc_len = ( - tail_128 + 128 - if (tail_128 <= 3 and seqlen >= 128) - else tail_128 - ) - else: # ratio == 128 - c_alloc_len = tail_128 - c_alloc_offset = seqlen - c_alloc_len - - # Bundle slice for this req. The NPU paged state pool emits real - # slot ids (no ring-hash); slice by ``state_bundle_offset`` (cumulative - # alloc_len), NOT ``seqlen_offset`` (cumulative raw seqlen). - bundle = forward_batch.out_cache_loc_dsv4 - assert bundle is not None, ( - "unfused compress prefill on NPU needs the DSV4 " - "alloc bundle; expected maybe_write_dsv4_extend to have " - "populated batch.out_cache_loc_dsv4 before forward." - ) - bundle_state_loc = ( - bundle.out_c4_state_loc if ratio == 4 else bundle.out_c128_state_loc - ) - if c_alloc_len > 0: - # Require a populated bundle only when this req allocates - # slots. A 128-aligned ratio==128 prefill has c_alloc_len==0 - # (no partial tail), so an all-128-aligned batch legitimately - # yields an empty bundle. Empty while c_alloc_len > 0 means - # c{ratio}_state_attn_allocator was never initialized. - assert ( - bundle_state_loc is not None and bundle_state_loc.numel() > 0 - ), ( - f"unfused compress prefill: bundle.out_c{ratio}_state_loc " - f"is empty/None — DSV4NPUTokenToKVPoolAllocator's " - f"c{ratio}_state_attn_allocator was not initialized (check " - f"pool_configurator's NPU branch + npu_state_pool_size)." - ) - out_cache_loc = bundle_state_loc[ - state_bundle_offset : state_bundle_offset + c_alloc_len - ] - state_bundle_offset += c_alloc_len - else: - # No tail to cache: empty slot view, never indexed below. - # Only reached for c128 (c4's c_alloc_len is always > 0). - out_cache_loc = torch.empty((0,), dtype=torch.int64, device=device) - remainder = seqlen % ratio - cutoff = seqlen - remainder - # ``cutoff`` is raw coords; subtract ``c_alloc_offset`` for - # slice-relative indexing into the per-req bundle slice. - cutoff_in_slice = cutoff - c_alloc_offset - should_compress = cutoff >= ratio - # ratio-strided positions for the cutoff chunks (one rope pos per token). - pos_compressed = pos_req[:cutoff:ratio] - kv = kv_full[seqlen_offset : seqlen_offset + seqlen] - score = score_full[seqlen_offset : seqlen_offset + seqlen] - - if overlap and should_compress: - # Stash the trailing ratio tokens of the cutoff so the next - # decode step can do overlap compression across the boundary - # (for ratio=4 this window is inside the state alloc range). - kv_state_to_be_cached.append(kv[cutoff - ratio : cutoff]) - score_state_to_be_cached.append( - score[cutoff - ratio : cutoff] + compressor.ape - ) - state_loc_list.append( - out_cache_loc[cutoff_in_slice - ratio : cutoff_in_slice] - ) - if remainder > 0: - kv_cut, kv_rem = kv.split([cutoff, remainder], dim=0) - score_cut, score_rem = score.split([cutoff, remainder], dim=0) - kv_state_to_be_cached.append(kv_rem) - score_state_to_be_cached.append( - score_rem + compressor.ape[:remainder] - ) - state_loc_list.append(out_cache_loc[-remainder:]) - kv = kv_cut - score = score_cut - - if should_compress: - kv = kv.unflatten(0, (-1, ratio)) # [n_chunks, ratio, coff*d] - score = score.unflatten(0, (-1, ratio)) + compressor.ape - if overlap: - kv = _overlap_transform(kv, value=0.0, head_dim=d) - score = _overlap_transform( - score, value=float("-inf"), head_dim=d - ) - kv_compressed = (kv * score.softmax(dim=1)).sum( - dim=1 - ) # [n_chunks, d] - n_compressed_this_req = kv_compressed.shape[0] - kv_out_list.append(kv_compressed) - kv_out_positions.append(pos_compressed) - write_req_indices.append( - torch.full( - (n_compressed_this_req,), - idx, - dtype=torch.int64, - device=device, - ) - ) - write_pos_in_req.append( - torch.arange( - n_compressed_this_req, - dtype=torch.int64, - device=device, - ) - ) - seqlen_offset += seqlen - else: - # Decode: append (kv, score+ape[pos%r]) to the state ring at - # c{4,128}_state_loc[idx]; if this completes a ratio-aligned - # chunk, gather it and produce one compressed kv via ape-softmax-sum. - start_pos = seqlen - 1 - should_compress = (start_pos + 1) % ratio == 0 - pos_req = positions[idx : idx + 1] + (1 - ratio) - kv = kv_full[idx : idx + 1] - score = score_full[idx : idx + 1] + compressor.ape[start_pos % ratio] - if ratio == 4: - state_loc_decode = backend_fm.c4_state_loc - else: - state_loc_decode = backend_fm.c128_state_loc - token_to_kv_pool.set_state_buffer( - compressor.layer_id, - state_loc_decode[idx : idx + 1], - kv.view(1, 1, -1), - score.view(1, 1, -1), - compressor.is_in_indexer, - ) - if should_compress: - if overlap: - kv_indices = _get_kv_indices( - forward_batch, 2 * ratio, page_table, idx, seqlen - ) - kv_state, score_state = token_to_kv_pool.get_state_buffer( - compressor.layer_id, compressor.is_in_indexer, kv_indices - ) - # kv_state / score_state: [2*r, 1, coff*d] → [2*r, d] - kv_state = kv_state.squeeze(1) - score_state = score_state.squeeze(1) - kv_state = torch.cat( - [kv_state[:ratio, :d], kv_state[ratio:, d:]], dim=0 - ) - score_state = torch.cat( - [score_state[:ratio, :d], score_state[ratio:, d:]], - dim=0, - ) - kv_compressed = (kv_state * score_state.softmax(dim=0)).sum( - dim=0, keepdim=True - ) - else: - kv_indices = _get_kv_indices( - forward_batch, ratio, page_table, idx, seqlen - ) - kv_state, score_state = token_to_kv_pool.get_state_buffer( - compressor.layer_id, compressor.is_in_indexer, kv_indices - ) - kv_compressed = ( - kv_state[:, 0] * score_state[:, 0].softmax(dim=0) - ).sum(dim=0, keepdim=True) - kv_out_list.append(kv_compressed) - kv_out_positions.append(pos_req) - # Decode: 1 compressed token at compressed_seq_pos = seqlen//ratio - 1 - decode_pos = seqlen // ratio - 1 - write_req_indices.append( - torch.tensor([idx], dtype=torch.int64, device=device) - ) - write_pos_in_req.append( - torch.tensor([decode_pos], dtype=torch.int64, device=device) - ) - - # Flush the prefill state stash to the pool in one shot. - if kv_state_to_be_cached: - kv_state_cat = torch.cat(kv_state_to_be_cached, dim=0).unsqueeze(1) - score_state_cat = torch.cat(score_state_to_be_cached, dim=0).unsqueeze(1) - state_loc_cat = torch.cat(state_loc_list, dim=0) - token_to_kv_pool.set_state_buffer( - compressor.layer_id, - state_loc_cat, - kv_state_cat, - score_state_cat, - compressor.is_in_indexer, - ) - - # Norm + rope + optional hadamard on the freshly compressed tokens, - # then write via _compressor_epilog_npu with explicit slab-derived locs. - if kv_out_list: - kv_out = torch.cat(kv_out_list, dim=0).to(dtype) - pos_out = torch.cat(kv_out_positions, dim=0) - kv_out = compressor.norm(kv_out) - rope_dim = compressor.rope_head_dim - from sglang.srt.hardware_backend.npu.dsv4.dsv4_rope import Dsv4NpuRoPE - - cos, sin = Dsv4NpuRoPE.for_freqs( - compressor.freqs_cis, getattr(compressor, "rotary_emb", None) - ).get_cos_sin( - pos_out, - kv_out.dtype, - view_4d=True, - allow_build=False, - cache_dtype=torch.float32, - ) - rope_slice = kv_out[..., -rope_dim:] - rope_view = rope_slice.unsqueeze(-2).unsqueeze(1) # (T, 1, 1, rope_dim) - rope_rot = torch_npu.npu_rotary_mul( - rope_view, cos, sin, rotary_mode="interleave" - ) - rope_slice.copy_(rope_rot.view_as(rope_slice)) - if compressor.rotate: - kv_out = _apply_hadamard(kv_out, compressor.hadamard_matrix) - # c{N}_kv_pool slot per compressed token. DSV4NPUReqToTokenPool's - # token-level slot id table is indexed directly by compressed-seq - # position (elements already are c-pool slot ids; no page indirection). - req_indices_flat = torch.cat(write_req_indices, dim=0) - pos_in_req_flat = torch.cat(write_pos_in_req, dim=0) - req_pool_flat = forward_batch.req_pool_indices[req_indices_flat] - c_table = ( - self.req_to_token_pool.req_to_token_c4 - if ratio == 4 - else self.req_to_token_pool.req_to_token_c128 - ) - write_locs = c_table[ - req_pool_flat.to(torch.int64), pos_in_req_flat.to(torch.int64) - ].to(torch.int32) - self._compressor_epilog_npu( - compressor, kv_out, forward_batch, override_loc=write_locs - ) - return None - def _ensure_compressor_hadamard(self, compressor, device: torch.device) -> None: if getattr(compressor, "hadamard_matrix", None) is None: H = _walsh_hadamard_matrix(compressor.head_dim, torch.float32, device) @@ -789,8 +468,6 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): kv_scale: Optional[torch.Tensor] = None li_kv_dtype = getattr(compressor, "li_kv_dtype", "bf16") if li_kv_dtype == "int8" and compressor.is_in_indexer: - import torch_npu - kv, kv_scale = torch_npu.npu_dynamic_quant(kv) kv_scale = kv_scale.to(torch.float16) @@ -832,41 +509,91 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): ) -class C4IndexerAscendBackendMixin(C4IndexerBackendMixin): +class C4IndexerAscendBackendMixin: def init_forward_metadata_indexer(self, core_attn_metadata): # li_quant_metadata is built in _compute_kernel_metadata; None satisfies the mixin contract return None - def forward_c4_indexer_npu( + def _forward_prepare( self, c4_indexer, x: torch.Tensor, q_lora: torch.Tensor, forward_batch: ForwardBatch, - skip_compressor: bool = False, - ) -> torch.Tensor: - assert ( - not skip_compressor - ), "skip_compressor=True is not supported by forward_c4_indexer_npu" + ) -> tuple[torch.Tensor, torch.Tensor]: + q = self._compute_q_npu(c4_indexer, q_lora, forward_batch.positions) + weights, _ = c4_indexer.weights_proj(x) + weights = weights * (c4_indexer.softmax_scale * c4_indexer.n_heads**-0.5) + c4_indexer.compressor(x, forward_batch) + return q, weights + def _can_use_indexer_multi_stream(self) -> bool: + return envs.SGLANG_NPU_USE_MULTI_STREAM.get() + + def _get_npu_indexer_q_stream(self): + s = getattr(self, "_npu_indexer_q_stream_obj", None) + if s is None: + s = torch.npu.Stream() + self._npu_indexer_q_stream_obj = s + return s + + def _forward_prepare_multi_stream( + self, + c4_indexer, + x: torch.Tensor, + q_lora: torch.Tensor, + forward_batch: ForwardBatch, + q_lora_ready, + ) -> tuple[torch.Tensor, torch.Tensor]: + from sglang.srt.hardware_backend.npu.utils import ( + get_indexer_weight_stream, + ) + + cur = torch.npu.current_stream() + stream_q = self._get_npu_indexer_q_stream() + stream_w = get_indexer_weight_stream() + + # q_lora/x are produced on cur; workers wait for them. + stream_q.wait_stream(cur) + stream_w.wait_stream(cur) + + # route-KV write on cur; ordered before the topk read by cur's program order. + c4_indexer.compressor(x, forward_batch) + + # weights_proj + scale on stream_w. + with torch.npu.stream(stream_w): + weights = c4_indexer.weights_proj(x)[0] + weights = weights * (c4_indexer.softmax_scale * c4_indexer.n_heads**-0.5) + weights.record_stream(stream_w) + + # q (wq_b + rope + hadamard) on stream_q. + with torch.npu.stream(stream_q): + if q_lora_ready is not None: + stream_q.wait_event(q_lora_ready) + q = self._compute_q_npu(c4_indexer, q_lora, forward_batch.positions) + q.record_stream(stream_q) + + cur.wait_stream(stream_w) + cur.wait_stream(stream_q) + return q, weights + + def _forward_indexer( + self, + c4_indexer, + x: torch.Tensor, + q: torch.Tensor, + weights: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: ratio = c4_indexer.compressor.ratio device = x.device - self._ensure_npu_c4_indexer(c4_indexer, device) bs = x.shape[0] is_prefill = ( forward_batch.forward_mode.is_extend() and not forward_batch.forward_mode.is_target_verify() ) - q = self._compute_q_npu(c4_indexer, q_lora, forward_batch.positions) - - weights, _ = c4_indexer.weights_proj(x) - weights = weights * (c4_indexer.softmax_scale * c4_indexer.n_heads**-0.5) - - if not skip_compressor: - c4_indexer.compressor(x, forward_batch) - li_kv_dtype = getattr(c4_indexer.compressor, "li_kv_dtype", "bf16") if li_kv_dtype == "int8": # Empty/idle rank (T=0) must skip the indexer kernel; test is_idle @@ -897,7 +624,12 @@ class C4IndexerAscendBackendMixin(C4IndexerBackendMixin): for i, _end_token in enumerate(end_pos): seq_i = int(seqlens_cpu[i]) kv_indices = _get_kv_indices( - forward_batch, seq_i // ratio, page_table, i, seq_i // ratio + forward_batch, + seq_i // ratio, + page_table, + i, + seq_i // ratio, + page_size=self.page_size // ratio, ) kv_cache_value = self.token_to_kv_pool.get_compress_buffer( c4_indexer.layer_id, True, kv_indices @@ -958,7 +690,6 @@ class C4IndexerAscendBackendMixin(C4IndexerBackendMixin): def _compute_q_npu( self, c4_indexer, q_lora: torch.Tensor, positions: torch.Tensor ) -> torch.Tensor: - from sglang.srt.hardware_backend.npu.dsv4.dsv4_rope import Dsv4NpuRoPE bs = q_lora.shape[0] q, _ = c4_indexer.wq_b(q_lora) @@ -994,8 +725,6 @@ class C4IndexerAscendBackendMixin(C4IndexerBackendMixin): weights: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: - import torch_npu - q_int8, q_scale = torch_npu.npu_dynamic_quant(q) fm = self.forward_metadata li_quant_metadata = fm.kernel_metadata["li_quant_metadata"] @@ -1034,9 +763,17 @@ class C4IndexerAscendBackendMixin(C4IndexerBackendMixin): ) -> None: if forward_batch.forward_mode.is_idle(): return - topk_idxs = self.forward_c4_indexer_npu( - c4_indexer, x, q_lora, forward_batch, skip_compressor=skip_compressor - ) + assert ( + not skip_compressor + ), "skip_compressor=True is not supported on the NPU indexer path" + self._ensure_npu_c4_indexer(c4_indexer, x.device) + if self._can_use_indexer_multi_stream(): + q, weights = self._forward_prepare_multi_stream( + c4_indexer, x, q_lora, forward_batch, q_lora_ready + ) + else: + q, weights = self._forward_prepare(c4_indexer, x, q_lora, forward_batch) + topk_idxs = self._forward_indexer(c4_indexer, x, q, weights, forward_batch) self.forward_metadata.c4_topk_indices = topk_idxs @@ -1072,6 +809,103 @@ class DeepseekV4AscendAttnBackend( self._dsv4_unique_compress_ratios = list( dict.fromkeys(self._dsv4_compress_ratios) ) + self._is_dspark_algorithm = bool( + model_runner.spec_algorithm is not None + and model_runner.spec_algorithm.is_dspark() + ) + self._is_dspark_draft_worker = bool( + getattr(model_runner, "is_draft_worker", False) + and self._is_dspark_algorithm + ) + self._dsv4_graph_tokens_per_req = int(model_runner.decode_num_tokens_per_req()) + self._dsv4_state_pools_by_ratio = { + pool.ratio: pool + for pool in self.token_to_kv_pool.compress_state_pools + if pool is not None + } + + def _is_dspark_draft_block(self, forward_batch: ForwardBatch) -> bool: + spec_algorithm = forward_batch.spec_algorithm + return ( + self._is_dspark_draft_worker + and forward_batch.forward_mode.is_target_verify() + and spec_algorithm is not None + and spec_algorithm.is_dspark() + ) + + def _init_dspark_sparse_metadata(self, forward_batch: ForwardBatch) -> None: + """Build block-noncausal SWA slot ids for a DSpark draft forward. + + Every token in a DSpark draft block attends to the trailing SWA + context and to the whole current draft block. The Ascend + sparse-attention operator consumes physical SWA slot ids with shape + [T, N_kv, K], where K must be 128-aligned. + """ + fm = self.forward_metadata + fm.ori_sparse_indices = None + fm.ori_win_left = self._dsv4_sliding_window_size - 1 + fm.ori_win_right = 0 + + if not self._is_dspark_draft_block(forward_batch): + return + + block_size = int(forward_batch.spec_info.draft_token_num) + out_cache_loc = forward_batch.out_cache_loc + + ori_sparse_indices = self._build_dspark_sparse_indices( + seq_lens=forward_batch.seq_lens, + req_pool_indices=forward_batch.req_pool_indices, + out_cache_loc=out_cache_loc, + block_size=block_size, + ) + ori_sparse_indices = ori_sparse_indices.unsqueeze(1).contiguous() + + fm.ori_sparse_indices = ori_sparse_indices + fm.ori_win_left = self._dsv4_sliding_window_size + block_size - 1 + fm.ori_win_right = 0 + + def _build_dspark_sparse_indices( + self, + *, + seq_lens: torch.Tensor, + req_pool_indices: torch.Tensor, + out_cache_loc: torch.Tensor, + block_size: int, + ) -> torch.Tensor: + """Return [bs * block_size, K] physical SWA slots for one draft block. + + This helper is deliberately allocation-producing. Eager forwards use + the returned tensor directly; graph replay copies it into the stable + graph-owned ``ori_sparse_indices`` storage. + """ + bs = int(seq_lens.shape[0]) + expected_tokens = bs * block_size + + seq_lens_causal = BuildBlockSeqLensCausal.execute( + seq_lens=seq_lens, + block_size=block_size, + device=seq_lens.device, + ) + req_pool_indices_repeated = req_pool_indices.repeat_interleave(block_size) + gather = ComputeDsparkWindowGather.execute( + seq_lens_casual=seq_lens_causal, + req_pool_indices_repeated=req_pool_indices_repeated, + block_size=block_size, + swa_window=self._dsv4_sliding_window_size, + ) + ori_sparse_indices, _ = BuildDsparkSwaPageIndices.execute( + req_to_token=self.req_to_token, + full_to_swa_mapping=self.token_to_kv_pool.full_to_swa_index_mapping, + req_pool_indices_per_request=gather.req_pool_indices_per_request, + offsets=gather.offsets, + invalid=gather.invalid, + out_loc=out_cache_loc[:expected_tokens], + context_lens=gather.context_lens, + block_size=block_size, + swa_window=self._dsv4_sliding_window_size, + page_index_aligned_size=128, + ) + return ori_sparse_indices def _init_dsv4_graph_buffers(self, *, max_bs: int, max_num_tokens: int) -> None: device = self.device @@ -1090,12 +924,6 @@ class DeepseekV4AscendAttnBackend( self.graph_metadata["c128_page_table"] = torch.full( (max_bs, max_pages), -1, dtype=torch.int32, device=device ) - self.graph_metadata["c4_state_page_table"] = torch.zeros( - (max_bs, max_pages), dtype=torch.int32, device=device - ) - self.graph_metadata["c128_state_page_table"] = torch.zeros( - (max_bs, max_pages), dtype=torch.int32, device=device - ) # 1024 int32 per kernel-metadata buffer (fixed op metadata size) for key in ( @@ -1115,6 +943,22 @@ class DeepseekV4AscendAttnBackend( device=device, ) + if self._is_dspark_draft_worker: + block_size = self._dsv4_graph_tokens_per_req + sparse_width = ( + (self._dsv4_sliding_window_size + block_size + 127) // 128 * 128 + ) + self.graph_metadata["ori_sparse_indices"] = torch.full( + ( + max_bs * block_size, + self._dsv4_kv_head_num, + sparse_width, + ), + -1, + dtype=torch.int32, + device=device, + ) + def init_forward_metadata_out_graph( self, forward_batch: ForwardBatch, @@ -1134,7 +978,7 @@ class DeepseekV4AscendAttnBackend( device = self.device if forward_mode.is_target_verify() or forward_mode.is_draft_extend_v2(): - tokens_per_req = self.speculative_num_draft_tokens + tokens_per_req = self._dsv4_graph_tokens_per_req else: tokens_per_req = 1 @@ -1156,12 +1000,6 @@ class DeepseekV4AscendAttnBackend( metadata.swa_page_table = self.graph_metadata["swa_page_table"][:bs, :] metadata.c4_page_table = self.graph_metadata["c4_page_table"][:bs, :] metadata.c128_page_table = self.graph_metadata["c128_page_table"][:bs, :] - metadata.c4_state_page_table = self.graph_metadata["c4_state_page_table"][ - :bs, : - ] - metadata.c128_state_page_table = self.graph_metadata["c128_state_page_table"][ - :bs, : - ] n_tok = bs * tokens_per_req c4_pad = min(n_tok, n_tok // 4 + bs) @@ -1169,8 +1007,19 @@ class DeepseekV4AscendAttnBackend( metadata.swa_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) + metadata.dsv4_max_input_capacity = tokens_per_req + metadata.dsv4_explicit_state_block_tables = { + ratio: torch.full( + ( + bs, + (2 if ratio == 4 else 1) * ratio + tokens_per_req, + ), + state_pool.dummy_state_loc, + dtype=torch.int32, + device=device, + ) + for ratio, state_pool in self._dsv4_state_pools_by_ratio.items() + } metadata.positions_cmp_padding_c4 = torch.zeros( c4_pad, dtype=torch.int64, device=device @@ -1191,6 +1040,15 @@ class DeepseekV4AscendAttnBackend( T = bs * tokens_per_req metadata.c4_topk_indices = self.graph_metadata["c4_topk_indices"][:T, :] + metadata.ori_sparse_indices = None + metadata.ori_win_left = self._dsv4_sliding_window_size - 1 + metadata.ori_win_right = 0 + if self._is_dspark_draft_worker and forward_mode.is_target_verify(): + metadata.ori_sparse_indices = self.graph_metadata["ori_sparse_indices"][:T] + metadata.ori_sparse_indices.fill_(-1) + metadata.ori_sparse_indices[:, :, 0] = 0 + metadata.ori_win_left = self._dsv4_sliding_window_size + tokens_per_req - 1 + self.forward_metadata = metadata @staticmethod @@ -1222,33 +1080,77 @@ class DeepseekV4AscendAttnBackend( graph_mode = forward_batch.forward_mode runtime_mode = getattr(forward_batch, "actual_forward_mode", None) or graph_mode bs = forward_batch.batch_size + num_padding = int(getattr(forward_batch, "num_padding", 0) or 0) + + raw_bs = bs - num_padding seq_lens_cpu = forward_batch.seq_lens_cpu assert seq_lens_cpu is not None, "V4 graph replay requires seq_lens_cpu." device = forward_batch.seq_lens.device tokens_per_bs = ( - self.speculative_num_draft_tokens + self._dsv4_graph_tokens_per_req if graph_mode.is_target_verify() or graph_mode.is_draft_extend_v2() else 1 ) - seq_lens = forward_batch.seq_lens - if graph_mode.is_target_verify(): - 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) + raw_seq_lens_cpu = seq_lens_cpu[:bs] is_idle_replay = runtime_mode.is_idle() + if graph_mode.is_target_verify(): + explicit_live_cpu = getattr( + getattr(forward_batch, "spec_info", None), + "live_seq_lens_cpu", + None, + ) + if is_idle_replay: + live_seq_lens_cpu = torch.zeros_like(raw_seq_lens_cpu) + final_seq_lens_cpu = live_seq_lens_cpu + elif self._is_dspark_algorithm or explicit_live_cpu is not None: + # DSpark/DFLASH temporarily expand seq_lens_cpu to the final + # target-verify length and carry the committed/live prefix + # separately. Graph replay batches are lightweight namespace + # views and do not necessarily retain spec_algorithm. + final_seq_lens_cpu = raw_seq_lens_cpu + if explicit_live_cpu is None: + live_seq_lens_cpu = torch.clamp( + final_seq_lens_cpu - int(tokens_per_bs), min=0 + ) + else: + explicit_live_cpu = torch.as_tensor( + explicit_live_cpu, + dtype=final_seq_lens_cpu.dtype, + device=final_seq_lens_cpu.device, + ).flatten() + live_seq_lens_cpu = torch.zeros_like(final_seq_lens_cpu) + num_live_rows = min(bs, explicit_live_cpu.numel()) + if num_live_rows > 0: + live_seq_lens_cpu[:num_live_rows].copy_( + explicit_live_cpu[:num_live_rows] + ) + else: + # EAGLE and the other uniform verify callers keep + # seq_lens_cpu at the committed/live prefix length. + live_seq_lens_cpu = raw_seq_lens_cpu + final_seq_lens_cpu = live_seq_lens_cpu + int(tokens_per_bs) + live_seq_lens = live_seq_lens_cpu.to(device=device, dtype=torch.int32) + elif ( + forward_batch.seq_lens is not None + and forward_batch.seq_lens.device.type != "cpu" + ): + live_seq_lens_cpu = seq_lens_cpu[:bs] + final_seq_lens_cpu = live_seq_lens_cpu + live_seq_lens = forward_batch.seq_lens[:bs].to(dtype=torch.int32) + else: + live_seq_lens_cpu = seq_lens_cpu[:bs] + final_seq_lens_cpu = live_seq_lens_cpu + live_seq_lens = live_seq_lens_cpu.to(device=device, dtype=torch.int32) has_compress = self._dsv4_has_c4 or self._dsv4_has_c128 active_target_verify = ( graph_mode.is_target_verify() and not is_idle_replay and has_compress ) compress_seq_lens = live_seq_lens - compress_seq_lens_max = int(seq_lens_cpu[:bs].max()) if bs > 0 else 0 + compress_seq_lens_max = int(final_seq_lens_cpu.max()) if bs > 0 else 0 if active_target_verify: - compress_seq_lens = live_seq_lens + int(tokens_per_bs) - compress_seq_lens_max += int(tokens_per_bs) + compress_seq_lens = final_seq_lens_cpu.to(device=device, dtype=torch.int32) return SimpleNamespace( forward_batch=forward_batch, @@ -1259,9 +1161,12 @@ class DeepseekV4AscendAttnBackend( has_compress=has_compress, active_target_verify=active_target_verify, bs=bs, + raw_bs=raw_bs, tokens_per_bs=tokens_per_bs, device=device, seq_lens_cpu=seq_lens_cpu, + final_seq_lens_cpu=final_seq_lens_cpu, + live_seq_lens_cpu=live_seq_lens_cpu, live_seq_lens=live_seq_lens, compress_seq_lens=compress_seq_lens, compress_seq_lens_max=compress_seq_lens_max, @@ -1272,17 +1177,17 @@ class DeepseekV4AscendAttnBackend( attn_seq_lens = ctx.live_seq_lens if ctx.graph_mode.is_target_verify(): valid_verify_rows = ctx.live_seq_lens > 0 - attn_seq_lens = ctx.live_seq_lens + int(ctx.tokens_per_bs) - attn_seq_lens = torch.where( - valid_verify_rows, attn_seq_lens, ctx.live_seq_lens + final_seq_lens = ctx.final_seq_lens_cpu.to( + device=ctx.device, dtype=torch.int32 ) - fm.seq_lens_cpu_int = ( - ctx.seq_lens_cpu[: ctx.bs] + int(ctx.tokens_per_bs) - ).int() + attn_seq_lens = torch.where( + valid_verify_rows, final_seq_lens, ctx.live_seq_lens + ) + fm.seq_lens_cpu_int = ctx.final_seq_lens_cpu.int() fm.seq_lens_cpu_int = torch.where( - ctx.seq_lens_cpu[: ctx.bs] > 0, + ctx.live_seq_lens_cpu > 0, fm.seq_lens_cpu_int, - ctx.seq_lens_cpu[: ctx.bs].int(), + ctx.live_seq_lens_cpu.int(), ) fm.actual_seq_lengths_kv.copy_(attn_seq_lens.clamp(min=1)) @@ -1301,16 +1206,9 @@ class DeepseekV4AscendAttnBackend( is_graph=True, seq_lens_max_override=ctx.compress_seq_lens_max, ) - for key in ( - "c4_page_table", - "c128_page_table", - "c4_state_page_table", - "c128_state_page_table", - ): + for key in ("c4_page_table", "c128_page_table"): if key in result: - self._copy_2d_with_tail( - getattr(ctx.fm, key), result[key], 0 if "state" in key else -1 - ) + self._copy_2d_with_tail(getattr(ctx.fm, key), result[key], -1) def _refresh_graph_decode_compress_1d_direct(self, ctx) -> None: fm = ctx.fm @@ -1318,14 +1216,9 @@ class DeepseekV4AscendAttnBackend( for ratio in self._dsv4_unique_compress_ratios: if ratio not in (4, 128): continue - state_loc = None loc = None if bundle is not None: - state_loc = ( - bundle.out_c4_state_loc if ratio == 4 else bundle.out_c128_state_loc - ) loc = bundle.out_c4_loc if ratio == 4 else bundle.out_c128_loc - self._copy_1d_with_zero_tail(getattr(fm, f"c{ratio}_state_loc"), state_loc) self._copy_1d_with_zero_tail(getattr(fm, f"c{ratio}_loc"), loc) valid = ctx.live_seq_lens > 0 @@ -1343,11 +1236,11 @@ class DeepseekV4AscendAttnBackend( def _refresh_graph_target_verify_compress_1d_direct(self, ctx) -> None: fm = ctx.fm - verify_seq_lens_cpu = ctx.seq_lens_cpu[: ctx.bs] + int(ctx.tokens_per_bs) + verify_seq_lens_cpu = ctx.final_seq_lens_cpu verify_seq_lens_cpu = torch.where( - ctx.seq_lens_cpu[: ctx.bs] > 0, + ctx.live_seq_lens_cpu > 0, verify_seq_lens_cpu, - ctx.seq_lens_cpu[: ctx.bs], + ctx.live_seq_lens_cpu, ) self._fill_verify_positions_cmp_padding_one( ctx.forward_batch.positions, @@ -1385,14 +1278,30 @@ class DeepseekV4AscendAttnBackend( 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_() + def _refresh_graph_explicit_state_block_tables(self, ctx) -> None: + fm = ctx.fm + for ratio, fixed_table in fm.dsv4_explicit_state_block_tables.items(): + fixed_table.copy_( + _build_explicit_state_block_table( + compress_ratio=ratio, + coff=2 if ratio == 4 else 1, + state_pool=self._dsv4_state_pools_by_ratio[ratio], + token_to_kv_pool=self.token_to_kv_pool, + req_to_token=self.req_to_token, + req_pool_indices=ctx.forward_batch.req_pool_indices[: ctx.bs], + start_pos=fm.start_pos, + cu_seqlens=fm.actual_seq_lengths_q_pa, + seqused=fm.seqused, + max_input_capacity=fm.dsv4_max_input_capacity, + ) + ) + def _refresh_graph_swa_metadata_direct(self, ctx) -> None: fm = ctx.fm swa_loc = self.token_to_kv_pool.translate_loc_from_full_to_swa( @@ -1404,13 +1313,53 @@ class DeepseekV4AscendAttnBackend( fm.block_tables_swa if fm.block_tables_swa is not None else fm.block_tables ) if ctx.bs > 0: - spec = int(getattr(self, "speculative_num_draft_tokens", 0) or 0) - max_len = int(ctx.seq_lens_cpu[: ctx.bs].max()) + spec + max_len = int(ctx.final_seq_lens_cpu.max()) max_seq_pages = (max_len + self.page_size - 1) // self.page_size if 0 < max_seq_pages < swa_src.shape[1]: swa_src = swa_src[:, :max_seq_pages] self._copy_2d_with_tail(fm.swa_page_table, swa_src, -1) + def _refresh_graph_dspark_sparse_metadata(self, ctx) -> None: + if not (self._is_dspark_draft_worker and ctx.graph_mode.is_target_verify()): + return + + dst = getattr(ctx.fm, "ori_sparse_indices", None) + if dst is None: + raise RuntimeError( + "DSpark NPU graph replay is missing its captured " + "ori_sparse_indices buffer." + ) + + if ctx.is_idle_replay: + dst.fill_(-1) + dst[:, :, 0] = 0 + return + + src = self._build_dspark_sparse_indices( + # ``out_cache_loc`` is deliberately left unpadded by the graph + # replay view. Build sparse indices for real requests only, then + # populate the remaining rows of the captured buffer below. + seq_lens=ctx.live_seq_lens[: ctx.raw_bs], + req_pool_indices=ctx.forward_batch.req_pool_indices[: ctx.raw_bs], + out_cache_loc=ctx.forward_batch.out_cache_loc, + block_size=ctx.tokens_per_bs, + ).unsqueeze(1) + if ( + src.ndim != dst.ndim + or src.shape[0] > dst.shape[0] + or src.shape[1:] != dst.shape[1:] + ): + raise RuntimeError( + "DSpark NPU graph sparse-index shape mismatch: " + f"runtime={tuple(src.shape)}, captured={tuple(dst.shape)}." + ) + # A graph bucket can contain padded request rows. Keep those rows + # valid for capture/replay while replacing only the live prefix. + dst.fill_(-1) + dst[: src.shape[0]].copy_(src) + if src.shape[0] < dst.shape[0]: + dst[src.shape[0] :, :, 0] = 0 + def _refresh_graph_kernel_metadata(self, ctx) -> None: fm = ctx.fm kernel_metadata_new = self._kernel_metadata_from_parts( @@ -1445,7 +1394,10 @@ class DeepseekV4AscendAttnBackend( elif ctx.active_target_verify: self._refresh_graph_target_verify_compress_1d_direct(ctx) + self._refresh_graph_explicit_state_block_tables(ctx) + self._refresh_graph_swa_metadata_direct(ctx) + self._refresh_graph_dspark_sparse_metadata(ctx) self._refresh_graph_kernel_metadata(ctx) self.forward_metadata = ctx.fm @@ -1493,8 +1445,11 @@ class DeepseekV4AscendAttnBackend( or forward_batch.forward_mode.is_draft_extend_v2() ): B = forward_batch.batch_size - - n_draft = get_spec().speculative_num_draft_tokens or 1 + n_draft = ( + int(forward_batch.spec_info.draft_token_num) + if forward_batch.forward_mode.is_target_verify() + else self.speculative_num_draft_tokens + ) actual_q = torch.arange( n_draft, B * n_draft + 1, n_draft, dtype=torch.int32, device=device ) @@ -1503,15 +1458,6 @@ class DeepseekV4AscendAttnBackend( [torch.zeros(1, dtype=torch.int32, device=device), actual_q], dim=0, ) - elif forward_batch.forward_mode.is_idle(): - B = forward_batch.batch_size - fm.actual_seq_lengths_q = torch.arange( - 1, B + 1, dtype=torch.int32, device=device - ) - fm.actual_seq_lengths_q_pa = torch.arange( - 0, B + 1, dtype=torch.int32, device=device - ) - fm.actual_seq_lengths_kv = torch.ones(B, dtype=torch.int32, device=device) else: fm.actual_seq_lengths_q = None fm.actual_seq_lengths_q_pa = None @@ -1528,6 +1474,7 @@ class DeepseekV4AscendAttnBackend( else: fm.actual_seq_lengths_kv = forward_batch.seq_lens.to(torch.int32) + self._init_dspark_sparse_metadata(forward_batch) fm.kernel_metadata = self._compute_kernel_metadata(forward_batch) if self._dsv4_compress_ratios: @@ -1539,8 +1486,11 @@ class DeepseekV4AscendAttnBackend( forward_batch.forward_mode.is_target_verify() or forward_batch.forward_mode.is_draft_extend_v2() ): - - max_seqlen_q = get_spec().speculative_num_draft_tokens or 1 + max_seqlen_q = ( + int(forward_batch.spec_info.draft_token_num) + if forward_batch.forward_mode.is_target_verify() + else self.speculative_num_draft_tokens + ) else: max_seqlen_q = 1 return self._kernel_metadata_from_parts( @@ -1562,14 +1512,17 @@ class DeepseekV4AscendAttnBackend( max_seqlen_q: int, is_nextn: bool, ) -> dict: + fm = self.forward_metadata common = { "cu_seqlens_q": actual_seq_lengths_q_pa, "seqused_kv": actual_seq_lengths_kv, "cmp_ratio": 1, "ori_mask_mode": 4, "cmp_mask_mode": 3, - "ori_win_left": self._dsv4_sliding_window_size - 1, - "ori_win_right": 0, + "ori_win_left": getattr( + fm, "ori_win_left", self._dsv4_sliding_window_size - 1 + ), + "ori_win_right": getattr(fm, "ori_win_right", 0), "layout_q": "TND", "layout_kv": "PA_ND", } @@ -1582,11 +1535,27 @@ class DeepseekV4AscendAttnBackend( "has_cmp_kv": False, } c1a_kwargs = base_kwargs | common - kernel_metadata = { - "c1a_metadata": torch.ops.custom.npu_sparse_attn_sharedkv_metadata( - **c1a_kwargs + if self._is_dspark_draft_worker: + seq_lens_cpu = getattr(fm, "seq_lens_cpu_int", None) + max_seqlen_kv = ( + int(seq_lens_cpu[:bs].max().item()) + if seq_lens_cpu is not None and bs > 0 + else int(actual_seq_lengths_kv[:bs].max().item()) ) - } + c1a_kwargs.update( + cu_seqlens_ori_kv=actual_seq_lengths_q_pa, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + ) + c1a_metadata = torch.ops._C_ascend.npu_sparse_attn_sharedkv_metadata( + device=str(actual_seq_lengths_kv.device), + **c1a_kwargs, + ) + else: + c1a_metadata = torch.ops.custom.npu_sparse_attn_sharedkv_metadata( + **c1a_kwargs, + ) + kernel_metadata = {"c1a_metadata": c1a_metadata} if self._dsv4_has_c4: c4a_overrides = { @@ -1656,12 +1625,12 @@ class DeepseekV4AscendAttnBackend( layer_id=layer.layer_id, swa_k=k, forward_batch=forward_batch ) if compress_ratio == 0: - return self._forward_dense(q, layer, forward_batch, attn_sink) + return self._forward_swa(q, layer, forward_batch, attn_sink) return self._forward_compressed( q, layer, forward_batch, attn_sink, compress_ratio ) - def _forward_dense( + def _forward_swa( self, q: torch.Tensor, layer: RadixAttention, @@ -1676,8 +1645,10 @@ class DeepseekV4AscendAttnBackend( cu_seqlens_q=fm.actual_seq_lengths_q_pa, seqused_kv=fm.actual_seq_lengths_kv, ori_mask_mode=4, - ori_win_left=self._dsv4_sliding_window_size - 1, - ori_win_right=0, + ori_win_left=getattr( + fm, "ori_win_left", self._dsv4_sliding_window_size - 1 + ), + ori_win_right=getattr(fm, "ori_win_right", 0), layout_q="TND", layout_kv="PA_ND", q=q, @@ -1686,8 +1657,18 @@ class DeepseekV4AscendAttnBackend( sinks=attn_sink, metadata=fm.kernel_metadata["c1a_metadata"], softmax_scale=layer.scaling, + cmp_ratio=1, ) - out, _ = torch.ops.custom.npu_sparse_attn_sharedkv(**attn_kwargs) + if self._is_dspark_draft_worker: + attn_kwargs["cu_seqlens_ori_kv"] = fm.actual_seq_lengths_q_pa + ori_sparse_indices = getattr(fm, "ori_sparse_indices", None) + if ori_sparse_indices is not None: + attn_kwargs["ori_sparse_indices"] = ori_sparse_indices + q_arg = attn_kwargs.pop("q") + if self._is_dspark_draft_worker: + out, _ = torch.ops._C_ascend.npu_sparse_attn_sharedkv(q_arg, **attn_kwargs) + else: + out, _ = torch.ops.custom.npu_sparse_attn_sharedkv(q_arg, **attn_kwargs) return out def _forward_compressed( @@ -1720,9 +1701,15 @@ class DeepseekV4AscendAttnBackend( ori_page_size = ori_kv.shape[1] cmp_native_page_size = cmp_kv.shape[1] cmp_block_table = getattr(fm, f"c{compress_ratio}_page_table") - assert cmp_native_page_size == ori_page_size, ( - f"cmp page_size={cmp_native_page_size} != ori page_size={ori_page_size}; " - "c{N}_kv_pool must be allocated with the global page_size on NPU " + expected_cmp_page_size = ( + ori_page_size // 4 + if compress_ratio == 4 + else pool.c128_kv_pool.kernel_page_size + ) + assert cmp_native_page_size == expected_cmp_page_size, ( + f"c{compress_ratio} page_size={cmp_native_page_size} != " + f"expected={expected_cmp_page_size} for ori page_size={ori_page_size}; " + "C4 and C128 must use their configured native page layouts " "(see NPUDeepSeekV4SingleKVPool.kernel_page_size)" ) @@ -1751,12 +1738,34 @@ class DeepseekV4AscendAttnBackend( attn_kwargs["cmp_sparse_indices"] = topk.view(-1, 1, topk.shape[-1]) else: attn_kwargs["cmp_sparse_indices"] = None - out, _ = torch.ops.custom.npu_sparse_attn_sharedkv(**attn_kwargs) + q_arg = attn_kwargs.pop("q") + out, _ = torch.ops.custom.npu_sparse_attn_sharedkv(q_arg, **attn_kwargs) return out + def get_swa_out_cache_loc(self, forward_batch: ForwardBatch) -> torch.Tensor: + """Return the SWA KV write locations used by DeepSeek-V4 draft layers. + + During NPU graph capture/replay, ``swa_loc`` is stable graph storage + whose contents are refreshed by ``_apply_dsv4_graph_metadata``. Eager + forwards do not necessarily build that metadata, so translate the + current full-pool locations on demand as a fallback. + """ + out_cache_loc = forward_batch.out_cache_loc + metadata = self.forward_metadata + cached = getattr(metadata, "swa_loc", None) + if ( + cached is not None + and not forward_batch.forward_mode.is_idle() + and cached.shape[0] == out_cache_loc.shape[0] + ): + return cached + return self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc).to( + torch.int64 + ) + def store_cache(self, *, layer_id: int, swa_k: torch.Tensor, forward_batch): pool = self.token_to_kv_pool - swa_loc = pool.translate_loc_from_full_to_swa(forward_batch.out_cache_loc) + swa_loc = self.get_swa_out_cache_loc(forward_batch) pool.set_swa_buffer( layer_id=layer_id, loc=swa_loc, @@ -1769,14 +1778,11 @@ class DeepseekV4AscendAttnBackend( 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) + n_draft = int(forward_batch.spec_info.draft_token_num) + # The parent backend normalizes this to final KV lengths for every + # algorithm: it adds n_draft for EAGLE/NGRAM, while DSpark/DFLASH + # already pass expanded lengths and are not incremented again. + verify_seq_lens_cpu = fm.seq_lens_cpu_int[:bs] padding_sizes = {} for ratio in (4, 128): if ratio not in self._dsv4_compress_ratios: @@ -1791,6 +1797,7 @@ class DeepseekV4AscendAttnBackend( 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) + fm.dsv4_max_input_capacity = max(1, n_draft) _bundle = getattr(forward_batch, "out_cache_loc_dsv4", None) if _bundle is not None: for ratio in self._dsv4_unique_compress_ratios: @@ -1810,82 +1817,18 @@ class DeepseekV4AscendAttnBackend( 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, + n_draft: int, ) -> 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: @@ -1919,9 +1862,7 @@ class DeepseekV4AscendAttnBackend( 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) - ) + n_draft = int(spec_info.draft_token_num) 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) @@ -1946,10 +1887,12 @@ def _get_kv_indices( page_table: torch.Tensor, req_idx: int, seqlen: int, + page_size: Optional[int] = None, ) -> torch.Tensor: logic_start = max(0, seqlen - kv_len) logic_end = seqlen - page_size = get_attn_backend().page_size + if page_size is None: + page_size = get_attn_backend().page_size if page_size == 1: return page_table[req_idx, logic_start:logic_end] logic_pos = torch.arange(logic_start, logic_end, device=page_table.device) @@ -2031,16 +1974,6 @@ class DeepseekV4AscendMultiStepDraftBackend: ) 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 @@ -2066,8 +1999,6 @@ class DeepseekV4AscendMultiStepDraftBackend: 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): diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/c128_sidecar_component.py b/python/sglang/srt/hardware_backend/npu/dsv4/c128_sidecar_component.py new file mode 100644 index 000000000..b4fda7839 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/dsv4/c128_sidecar_component.py @@ -0,0 +1,294 @@ +"""C128 sidecar ownership for the DSV4 NPU Unified Radix Cache. + +The component deliberately exposes only complete physical C128 pages to the +radix tree; partial tail pages remain request-owned. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Callable, Optional + +import torch + +from sglang.srt.mem_cache.base_prefix_cache import ( + InsertParams, + InsertResult, + MatchPrefixParams, + MatchResult, +) +from sglang.srt.mem_cache.unified_cache.cache_action import ( + FreeComponentDeviceSlot, + SWARebuild, +) +from sglang.srt.mem_cache.unified_cache.components import ( + BASE_COMPONENT_TYPE, + ComponentType, + EvictLayer, + TreeComponent, +) + +if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import Req + from sglang.srt.mem_cache.unified_cache.cache_action import ( + CacheAction, + ComponentAction, + ) + from sglang.srt.mem_cache.unified_radix_cache import ( + UnifiedTreeNode, + ) + + +class C128SidecarComponent(TreeComponent): + component_type = ComponentType.C128 + + @property + def allocator(self): + return self.cache.token_to_kv_pool_allocator + + def _adjust_session_path( + self, + leaf: UnifiedTreeNode, + stop: UnifiedTreeNode, + delta: int, + ) -> None: + """Adjust session protection for every C128 boundary on a radix path.""" + node = leaf + while node is not stop and node is not self.tree_core.root_node: + cd = node.component_data[self.component_type] + if delta < 0: + assert cd.session_ref > 0 + prev_ref = cd.session_ref + cd.session_ref += delta + if (prev_ref == 0) != (cd.session_ref == 0): + self._refresh_session_partition(node) + node = node.parent + + def _dec_session_coverage(self, session_id: str, leaf: UnifiedTreeNode) -> None: + self._adjust_session_path(leaf, self.tree_core.root_node, -1) + + def _advance_session_coverage( + self, + session_id: str, + leaf: UnifiedTreeNode, + old_ancestor: Optional[UnifiedTreeNode], + ) -> None: + stop = old_ancestor or self.tree_core.root_node + self._adjust_session_path(leaf, stop, 1) + + def _recede_session_coverage( + self, + session_id: str, + leaf: UnifiedTreeNode, + fallback: Optional[UnifiedTreeNode], + ) -> None: + stop = fallback or self.tree_core.root_node + self._adjust_session_path(leaf, stop, -1) + + def _attach(self, node: UnifiedTreeNode, pages: torch.Tensor) -> None: + if pages.numel() == 0: + return + ct = self.component_type + cd = node.component_data[ct] + assert cd.value is None + value = pages.clone() + self.tree_core.set_component_device_value(node.id, ct, value) + self.allocator.retain_c128_pages(value) + + def create_match_validator( + self, match_device_only: bool = False + ) -> Callable[[UnifiedTreeNode], bool]: + # A page is attached only to the node ending its full physical group. + return lambda node: node.component_data[self.component_type].value is not None + + def finalize_match_result_in_cache( + self, params: MatchPrefixParams, result: MatchResult + ) -> MatchResult: + req = params.req + if req is None: + return result + + chunks = [] + node = self.tree_core.node_by_id(result.best_match_node) + root = self.tree_core.root_node + while node is not root: + value = node.component_data[self.component_type].value + if value is not None: + chunks.append(value) + node = node.parent + chunks.reverse() + pages = ( + torch.cat(chunks) + if chunks + else self.allocator.c128_attn_allocator.free_pages.new_empty((0,)) + ) + group_tokens = 128 * self.allocator.c128_attn_allocator.page_size + assert pages.numel() == len(result.device_indices) // group_tokens + self.cache.req_to_token_pool.set_c128_prefix_pages(req, pages) + return result + + def recover_after_unevict( + self, + node: UnifiedTreeNode, + prefix_len: int, + total_prefix_len: int, + params: InsertParams, + cache_actions: list[CacheAction | ComponentAction], + ) -> None: + pages = params.c128_value + assert pages is not None + group_tokens = 128 * self.allocator.c128_attn_allocator.page_size + start = total_prefix_len // group_tokens + end = (total_prefix_len + prefix_len) // group_tokens + self._attach(node, pages[start:end]) + + @staticmethod + def _node_depth(node: UnifiedTreeNode) -> int: + depth = 0 + while node.parent is not None: + depth += len(node.key) + node = node.parent + return depth + + @staticmethod + def _split_pending_swa_rebuild( + new_parent: UnifiedTreeNode, + child: UnifiedTreeNode, + cache_actions: list[CacheAction | ComponentAction], + ) -> None: + """Keep a deferred SWA rebuild aligned when C128 splits its source node.""" + for i, pending in enumerate(cache_actions): + if isinstance(pending, SWARebuild) and pending.node_id == child.id: + cache_actions[i : i + 1] = [ + SWARebuild( + new_parent.id, + new_parent.component_data[BASE_COMPONENT_TYPE].value, + ), + SWARebuild( + child.id, + child.component_data[BASE_COMPONENT_TYPE].value, + ), + ] + return + + def _ensure_boundary_node( + self, + tail: UnifiedTreeNode, + boundary: int, + cache_actions: list[CacheAction | ComponentAction], + ) -> UnifiedTreeNode: + node = tail + while node.parent is not None and self._node_depth(node.parent) >= boundary: + node = node.parent + + node_end = self._node_depth(node) + if node_end == boundary: + return node + + node_start = node_end - len(node.key) + assert node_start < boundary < node_end + new_parent, action = self.tree_core._split_node( + node.key, node, boundary - node_start + ) + if action is not None: + cache_actions.append(action) + self._split_pending_swa_rebuild(new_parent, node, cache_actions) + return new_parent + + def commit_insert_component_data( + self, + node: UnifiedTreeNode, + is_new_leaf: bool, + params: InsertParams, + result: InsertResult, + cache_actions: list[CacheAction | ComponentAction], + ) -> None: + if not is_new_leaf: + return + assert params.key is not None + assert params.c128_value is not None + + # Full/SWA may initially represent the new suffix as one long leaf. + # Materialize every complete C128 group boundary so a later branch can + # always match the nearest full C128-page prefix instead of falling back to + # the previous, potentially much shorter, Radix node. + group_tokens = 128 * self.allocator.c128_attn_allocator.page_size + first_boundary = (result.prefix_len // group_tokens + 1) * group_tokens + for boundary in range(first_boundary, len(params.key) + 1, group_tokens): + boundary_node = self._ensure_boundary_node(node, boundary, cache_actions) + page_index = boundary // group_tokens - 1 + self._attach(boundary_node, params.c128_value[page_index : page_index + 1]) + + def redistribute_on_node_split( + self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode + ) -> None: + # Every stored value belongs to the old child's end boundary. Splitting + # inside that group leaves the page on the child. + ct = self.component_type + new_parent.component_data[ct].session_ref = child.component_data[ct].session_ref + assert new_parent.component_data[ct].session_ids is None + + def evict_component( + self, + node: UnifiedTreeNode, + device_frees: dict[ComponentType, list[torch.Tensor]], + host_frees: dict[ComponentType, list[torch.Tensor]], + target: EvictLayer = EvictLayer.DEVICE, + ) -> tuple[int, int]: + cd = node.component_data[self.component_type] + if EvictLayer.DEVICE in target and cd.value is not None: + device_frees[self.component_type].append(cd.value) + self.tree_core.component_evictable_size_[self.component_type] -= len( + cd.value + ) + cd.value = None + # C128 pages are auxiliary to Full tokens and must not inflate the + # public token-eviction count. + return 0, 0 + + def prepare_for_caching_req( + self, + req: Req, + insert_params: InsertParams, + token_ids_len: int, + is_finished: bool, + ) -> int: + logical_len = token_ids_len + if self.tree_core.is_eagle and logical_len > 0: + logical_len -= 1 + group_tokens = 128 * self.allocator.c128_attn_allocator.page_size + cache_len = logical_len // group_tokens * group_tokens + num_pages = cache_len // group_tokens + insert_params.c128_value = self.cache.req_to_token_pool.req_to_c128_sidecar[ + int(req.req_pool_idx), :num_pages + ].clone() + return cache_len + 1 if self.tree_core.is_eagle and cache_len > 0 else cache_len + + def apply_component_action(self, action: ComponentAction) -> None: + if isinstance(action, FreeComponentDeviceSlot): + for page_ids in action.indices: + self.allocator.release_c128_pages(page_ids) + return + raise AssertionError( + f"C128SidecarComponent: unhandled action {type(action).__name__}" + ) + + def eviction_priority(self, is_leaf: bool) -> int: + return 0 if is_leaf else 2 + + def _evict_device_start(self, request_cnt): + pass + + def _evict_device_next_node(self, tracker, device_frees, host_frees): + return None + + def _evict_device_end(self) -> None: + pass + + def acquire_component_lock(self, node, result, lock_host=False): + return result + + def release_component_lock(self, node, params, lock_host=False) -> None: + pass + + def free_host_values(self, host_values) -> None: + pass 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 a6e0dcbbd..bcfb5baba 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py @@ -1,30 +1,26 @@ -"""DSV4-NPU SWA + c4/c128 paged allocator. +"""DSV4-NPU SWA + c128 paged KV allocator. Subclasses :class:`SWATokenToKVPoolAllocator` and adds paged allocation for the -c4/c128 compressed-KV pools and their tail-only compress-state pools, alongside -the parent's full + SWA pools. +C128 compressed-KV pool alongside the parent's full + SWA pools. C4 KV slots +are derived from full slots. Compressor state is fixed ring storage owned by +the KV pool and never enters this token allocator. Per ``alloc_extend`` / ``alloc_decode``: 1. super() allocates the full + SWA slots (``out_full_loc``). - 2. Allocate c4/c128 KV slots — one compressed token per ``ratio`` raw tokens - (``seq_len // ratio - prefix_len // ratio``) — via the standard - :class:`NPUPagedTokenToKVPoolAllocator` over the pool's c4/c128 KV buffers. - 3. Allocate the c4/c128 compress-state slots the same way, tail-only per req, - using the per-req lens the scheduler packed into ``DSV4StateLens``. - 4. Return a :class:`DSV4OutCacheLoc` bundling all five slot families. + 2. Derive c4 KV slots from full and allocate c128 KV slots — one compressed + token per ``ratio`` raw tokens + (``seq_len // ratio - prefix_len // ratio``). + 3. Return a :class:`DSV4OutCacheLoc` containing only KV slot families. -State slots are paged because the NPU fused compressor runs ``cache_mode=1``; the -base class' ``translate_kv_loc_to_compress_state_loc`` ring-hash is the CUDA-only -path and is unused on NPU. The bundle is the explicit return value: +The bundle is the explicit return value: mem_cache/common.py unpacks ``out_full_loc`` and stashes the bundle on ``batch.out_cache_loc_dsv4``; ``DSV4NPUReqToTokenPool`` writes the per-req -``req_to_token_c{4,128}[_state]`` tables that :meth:`free` and the last_loc -lookups read back. +``req_to_c128_sidecar`` table that :meth:`free` and the last_loc lookup read back. """ from __future__ import annotations -from typing import TYPE_CHECKING, List, Optional +from typing import Optional import torch @@ -35,34 +31,29 @@ from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( ) from sglang.srt.mem_cache.allocation import alloc_paged_token_slots_extend from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator -from sglang.srt.model_executor.forward_batch_info import DSV4OutCacheLoc, DSV4StateLens - -if TYPE_CHECKING: - from sglang.srt.managers.schedule_batch import Req +from sglang.srt.model_executor.forward_batch_info import DSV4OutCacheLoc def get_last_loc( - req_to_token: torch.Tensor, + req_to_c128_sidecar: torch.Tensor, req_pool_indices: torch.Tensor, prefix_lens: torch.Tensor, + page_size: int, ) -> torch.Tensor: """Slot id of each req's last already-allocated token, or -1 when ``prefix_lens[i] == 0`` (fresh req). - Looks up ``req_to_token[req, prefix_lens - 1]`` to anchor the paged - allocator's ``alloc_extend`` on the real previous tail slot, preserving the - intra-page slot continuity the kernel's ``cmp_block_table`` relies on (the - allocator debug-asserts ``(last_loc + 1) % page_size == prefix_lens % - page_size``). Result dtype matches ``prefix_lens``. + Looks up the C128 sidecar page to anchor the paged allocator's + ``alloc_extend`` on the real previous tail slot, preserving intra-page slot + continuity. Result dtype matches ``prefix_lens``. """ req_pool_indices = req_pool_indices.to(torch.int64) - safe_idx = (prefix_lens.to(torch.int64) - 1).clamp(min=0) - looked_up = req_to_token[req_pool_indices, safe_idx].to(prefix_lens.dtype) - return torch.where( - prefix_lens > 0, - looked_up, - torch.full_like(prefix_lens, -1), + last_pos = (prefix_lens.to(torch.int64) - 1).clamp(min=0) + page_ids = req_to_c128_sidecar[req_pool_indices, last_pos // page_size].to( + prefix_lens.dtype ) + last_loc = page_ids * page_size + last_pos.to(prefix_lens.dtype) % page_size + return torch.where(prefix_lens > 0, last_loc, torch.full_like(prefix_lens, -1)) def alloc_paged_token_slots_extend_npu(*args, batch=None, **kwargs): @@ -81,20 +72,9 @@ def alloc_paged_token_slots_reserve_extend( 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 - ) - + """Allocate reserved draft KV slots and update DSV4 KV tables.""" out_cache_loc = alloc_paged_token_slots_extend( tree_cache, prefix_lens, @@ -104,7 +84,6 @@ def alloc_paged_token_slots_reserve_extend( last_loc, extend_num_tokens, req_pool_indices=req_pool_indices, - dsv4_state_lens=dsv4_state_lens, batch=batch, ) if batch is not None: @@ -113,14 +92,12 @@ def alloc_paged_token_slots_reserve_extend( 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.""" + """SWA allocator + C128 KV allocator and full-derived C4 locations.""" def __init__( self, @@ -143,54 +120,29 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): ) def mk(pool_size, pool): - # c4/c128 KV and state sub-pools implement KVCache, so they drop into - # the standard paged allocator. pool_size is in compressed-token units. + # C128 KV sub-pool implements KVCache, so it drops into the standard + # paged allocator. pool_size is in compressed-token units. return NPUPagedTokenToKVPoolAllocator( pool_size, - page_size=page_size, + page_size=pool.kernel_page_size, dtype=dtype, device=device, kvcache=pool, need_sort=need_sort, ) - self.c4_attn_allocator = mk(kvcache.c4_size, kvcache.c4_kv_pool) self.c128_attn_allocator = mk(kvcache.c128_size, kvcache.c128_kv_pool) - - # State allocators (paged, NPU-only). Any layer's pool works as KVCache - # pointer (slot alloc is layer-agnostic); None when no c{ratio} layers or - # zero budget. - self.c4_state_attn_allocator: Optional[NPUPagedTokenToKVPoolAllocator] = None - self.c128_state_attn_allocator: Optional[NPUPagedTokenToKVPoolAllocator] = None - state_pools = getattr(kvcache, "compress_state_pools", None) - if state_pools: - - def first_state_pool(want_ratio): - return next( - ( - p - for r, p in zip(kvcache.compression_ratios, state_pools) - if r == want_ratio and p is not None - ), - None, - ) - - c4_state_pool = first_state_pool(4) - c128_state_pool = first_state_pool(128) - if c4_state_pool is not None and kvcache.c4_state_pool_size > 0: - self.c4_state_attn_allocator = mk( - kvcache.c4_state_pool_size, c4_state_pool - ) - if c128_state_pool is not None and kvcache.c128_state_pool_size > 0: - self.c128_state_attn_allocator = mk( - kvcache.c128_state_pool_size, c128_state_pool - ) + self.c128_page_refcount = torch.zeros( + self.c128_attn_allocator.num_pages + 1, + dtype=torch.int32, + device=device, + ) # Returned by the c-pool helpers when a step adds no compressed tokens. self._empty_loc = torch.empty((0,), dtype=torch.int64, device=device) # Per-call handle to the DSV4NPUReqToTokenPool, stashed by alloc_extend/ - # alloc_decode for last_loc lookups; avoids a permanent allocator->pool ref. + # alloc_decode for the C128 KV last_loc lookup. self._cur_req_to_token_pool = None @staticmethod @@ -206,6 +158,49 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): diff = ((seq_lens_cpu // ratio) - (prefix_lens_cpu // ratio)).clamp(min=0) return int(diff.sum().item()) + @staticmethod + def _derive_c4_loc_from_full(out_full_loc: torch.Tensor) -> torch.Tensor: + """Map full slots closing a 4-token group to their C4 slots.""" + completed_group = (out_full_loc >= 0) & ((out_full_loc % 4) == 3) + return out_full_loc[completed_group] // 4 + + def retain_c128_pages(self, page_ids: torch.Tensor) -> None: + page_ids = page_ids.to(torch.int64).view(-1) + if page_ids.numel() == 0: + return + self.c128_page_refcount.index_add_( + 0, + page_ids, + torch.ones_like(page_ids, dtype=self.c128_page_refcount.dtype), + ) + + def release_c128_pages(self, page_ids: torch.Tensor) -> None: + page_ids = torch.unique(page_ids.to(torch.int64).view(-1)) + page_ids = page_ids[page_ids > 0] + if page_ids.numel() == 0: + return + self.c128_page_refcount.index_add_( + 0, + page_ids, + -torch.ones_like(page_ids, dtype=self.c128_page_refcount.dtype), + ) + free_pages = page_ids[self.c128_page_refcount[page_ids] == 0] + if free_pages.numel() > 0: + self.c128_attn_allocator.free( + free_pages * self.c128_attn_allocator.page_size + ) + + def replace_req_c128_prefix( + self, req_pool_idx: int, page_ids: torch.Tensor, req_to_token_pool + ) -> None: + table = req_to_token_pool.req_to_c128_sidecar + page_ids = page_ids.to(device=table.device, dtype=table.dtype).view(-1) + old = table[req_pool_idx, : page_ids.numel()].clone() + changed = old != page_ids + self.release_c128_pages(old[changed]) + self.retain_c128_pages(page_ids[changed]) + table[req_pool_idx, : page_ids.numel()] = page_ids + @staticmethod def _pool_exhausted( ratio: int, kind: str, need: int, available: int @@ -218,60 +213,6 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): f"on req finish." ) - def _alloc_state_extend( - self, - allocator: Optional[NPUPagedTokenToKVPoolAllocator], - raw_prefix_lens: torch.Tensor, - state_prefix_lens: torch.Tensor, - state_prefix_lens_cpu: torch.Tensor, - state_seq_lens: torch.Tensor, - state_seq_lens_cpu: torch.Tensor, - req_pool_indices: torch.Tensor, - last_loc_dtype: torch.dtype, - state_extend_num_tokens: int, - ratio: int, - ) -> torch.Tensor: - """Allocate tail-only state-pool slots for an extend at ``ratio``. - - The state pool is a separate paged slot space; each req allocates only - its trailing window (cumulative lens precomputed by - ``ScheduleBatch._compute_dsv4_state_lens_*`` and passed via - ``DSV4StateLens``). ``state_last_loc`` is looked up from - ``req_to_token_c{ratio}_state`` at the RAW position - ``raw_prefix_lens - 1`` (the last position the previous extend/decode - populated). Returns ``_empty_loc`` when the allocator is absent (no - c{ratio} layers) or there is nothing to add. - """ - if allocator is None or state_extend_num_tokens == 0: - return self._empty_loc - - assert self._cur_req_to_token_pool is not None, ( - "alloc_extend/alloc_decode must be called with req_to_token_pool= " - "for the state-pool last_loc lookup." - ) - state_table = ( - self._cur_req_to_token_pool.req_to_token_c4_state - if ratio == 4 - else self._cur_req_to_token_pool.req_to_token_c128_state - ) - state_last_loc = get_last_loc( - state_table, req_pool_indices, raw_prefix_lens - ).to(last_loc_dtype) - - result = allocator.alloc_extend( - state_prefix_lens, - state_prefix_lens_cpu, - state_seq_lens, - state_seq_lens_cpu, - state_last_loc, - state_extend_num_tokens, - ) - if result is None: - raise self._pool_exhausted( - ratio, "state", state_extend_num_tokens, allocator.available_size() - ) - return result - def _alloc_c_extend( self, allocator: NPUPagedTokenToKVPoolAllocator, @@ -286,7 +227,7 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): """Allocate compressed-KV slots for an extend at ``ratio``. Prefix/seq lens are translated to compressed units (``// ratio``); the - c-pool last_loc comes from ``req_to_token_c{ratio}`` via + c-pool last_loc comes from ``req_to_c128_sidecar`` via :func:`get_last_loc` so the paged allocator continues in-page (or opens a fresh page at a ratio boundary), keeping the intra-page continuity the ``cmp_block_table`` reader relies on. Returns ``_empty_loc`` when this @@ -300,16 +241,12 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): "alloc_extend/alloc_decode must be called with req_to_token_pool= " "for the c-pool last_loc lookup." ) - c_table = ( - self._cur_req_to_token_pool.req_to_token_c4 - if ratio == 4 - else self._cur_req_to_token_pool.req_to_token_c128 - ) + c_table = self._cur_req_to_token_pool.req_to_c128_sidecar c_prefix = (prefix_lens // ratio).to(prefix_lens.dtype) c_seq = (seq_lens // ratio).to(seq_lens.dtype) - c_last_loc = get_last_loc(c_table, req_pool_indices, c_prefix).to( - last_loc_dtype - ) + c_last_loc = get_last_loc( + c_table, req_pool_indices, c_prefix, allocator.page_size + ).to(last_loc_dtype) result = allocator.alloc_extend( c_prefix, @@ -325,7 +262,17 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): ) return result - def _alloc_c_and_state( + def _has_c128_sidecar_capacity( + self, prefix_lens_cpu: torch.Tensor, seq_lens_cpu: torch.Tensor + ) -> bool: + ratio = 128 + page_size = self.c128_attn_allocator.page_size + prefix_groups = (prefix_lens_cpu // ratio + page_size - 1) // page_size + seq_groups = (seq_lens_cpu // ratio + page_size - 1) // page_size + need = int((seq_groups - prefix_groups).clamp(min=0).sum().item()) + return need <= self.c128_attn_allocator.available_size() // page_size + + def _alloc_compressed_kv( self, out_full_loc: torch.Tensor, out_swa_loc: torch.Tensor, @@ -335,57 +282,13 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): seq_lens_cpu: torch.Tensor, last_loc_dtype: torch.dtype, req_pool_indices: Optional[torch.Tensor], - dsv4_state_lens: Optional[DSV4StateLens], ) -> DSV4OutCacheLoc: - """Allocate c4/c128 KV + state slots and bundle them with full/swa loc. - - Shared by alloc_extend / alloc_decode (which differ only in how - prefix_lens is derived). State lens are tail-only, precomputed by - ScheduleBatch._compute_dsv4_state_lens_*; raw prefix_lens drives the - state last_loc lookup. - """ + """Allocate C128 KV and derive C4 KV, then bundle all KV locations.""" assert req_pool_indices is not None, ( "DSV4NPUTokenToKVPoolAllocator requires req_pool_indices " "(forwarded from batch.req_pool_indices)." ) - if dsv4_state_lens is not None: - out_c4_state_loc = self._alloc_state_extend( - self.c4_state_attn_allocator, - prefix_lens, - dsv4_state_lens.c4_prefix_lens, - dsv4_state_lens.c4_prefix_lens_cpu, - dsv4_state_lens.c4_seq_lens, - dsv4_state_lens.c4_seq_lens_cpu, - req_pool_indices, - last_loc_dtype, - dsv4_state_lens.c4_extend_num_tokens, - ratio=4, - ) - out_c128_state_loc = self._alloc_state_extend( - self.c128_state_attn_allocator, - prefix_lens, - dsv4_state_lens.c128_prefix_lens, - dsv4_state_lens.c128_prefix_lens_cpu, - dsv4_state_lens.c128_seq_lens, - dsv4_state_lens.c128_seq_lens_cpu, - req_pool_indices, - last_loc_dtype, - dsv4_state_lens.c128_extend_num_tokens, - ratio=128, - ) - else: - out_c4_state_loc = self._empty_loc - out_c128_state_loc = self._empty_loc - out_c4_loc = self._alloc_c_extend( - self.c4_attn_allocator, - prefix_lens, - prefix_lens_cpu, - seq_lens, - seq_lens_cpu, - req_pool_indices, - last_loc_dtype, - ratio=4, - ) + out_c4_loc = self._derive_c4_loc_from_full(out_full_loc) out_c128_loc = self._alloc_c_extend( self.c128_attn_allocator, prefix_lens, @@ -401,173 +304,6 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): out_swa_loc=out_swa_loc, out_c4_loc=out_c4_loc, out_c128_loc=out_c128_loc, - out_c4_state_loc=out_c4_state_loc, - out_c128_state_loc=out_c128_state_loc, - ) - - def compute_dsv4_state_lens_extend( - self, reqs: List[Req], seq_lens: List[int], prefix_lens: List[int] - ) -> Optional[DSV4StateLens]: - """Per-req c{4,128}_state pool alloc lens for extend (tail-only). - - State pool stores only the trailing portion of each sequence (the c{N} - compressor's read/write window); the tail length depends on raw - seq_len's alignment to the SWA page boundary (128):: - - c4_alloc_len = tail + 128 if (tail <= 3 and seq_len >= 128) else tail - c128_alloc_len = tail where tail = seq_len % 128 - - Long prefills allocate only the trailing partial window, not slots for - already-compressed positions, so the small paged state pool (~256 - slots/req) stays sufficient even for 28k-token prompts. - - Mutates per-req cumulative state via getattr/setattr so the community - ``Req`` needs no DSV4 field declarations: - * ``req.c{4,128}_state_kv_len`` — cumulative slot count (prefix for - the paged allocator; never decreases on eviction). - * ``req.c{4,128}_state_alloc_offset`` — low-water raw-position mark - for eviction (see ``dsv4_common_hooks.maybe_evict_dsv4_state``). - - Returns None when this model has no paged state pools (CUDA / non-V4 / - zero budget) — callers pass that straight through as ``dsv4_state_lens``. - """ - if self.c4_state_attn_allocator is None: - return None - c4_prefix: List[int] = [] - c4_seq: List[int] = [] - c128_prefix: List[int] = [] - c128_seq: List[int] = [] - for req, seq_len, prefix_len in zip(reqs, seq_lens, prefix_lens): - tail = seq_len % 128 - c4_alloc_len = tail + 128 if (tail <= 3 and seq_len >= 128) else tail - c128_alloc_len = tail - chunk_len = seq_len - prefix_len - - if prefix_len > 0: - c4_count = min(c4_alloc_len, chunk_len) - c128_count = min(c128_alloc_len, chunk_len) - else: - c4_count = c4_alloc_len - c128_count = c128_alloc_len - req.c4_state_alloc_offset = seq_len - c4_alloc_len - req.c128_state_alloc_offset = seq_len - c128_alloc_len - - prev_c4 = getattr(req, "c4_state_kv_len", 0) - prev_c128 = getattr(req, "c128_state_kv_len", 0) - new_c4 = prev_c4 + c4_count - new_c128 = prev_c128 + c128_count - - c4_prefix.append(prev_c4) - c4_seq.append(new_c4) - c128_prefix.append(prev_c128) - c128_seq.append(new_c128) - - req.c4_state_kv_len = new_c4 - req.c128_state_kv_len = new_c128 - req.c4_state_write_offset = seq_len - c4_count - req.c128_state_write_offset = seq_len - c128_count - - return self._pack_state_lens( - c4_prefix, - c4_seq, - c128_prefix, - c128_seq, - c4_extend_num_tokens=int(sum(s - p for s, p in zip(c4_seq, c4_prefix))), - c128_extend_num_tokens=int( - sum(s - p for s, p in zip(c128_seq, c128_prefix)) - ), - ) - - def compute_dsv4_state_lens_decode( - self, reqs: List[Req] - ) -> Optional[DSV4StateLens]: - """Per-req c{4,128}_state pool alloc lens for decode: exactly 1 new - state slot per req per pool. ``c{N}_state_alloc_offset`` does NOT - advance here (only eviction advances it). Returns None when there are - no paged state pools.""" - if self.c4_state_attn_allocator is None: - return None - c4_prefix: List[int] = [] - c4_seq: List[int] = [] - c128_prefix: List[int] = [] - c128_seq: List[int] = [] - for req in reqs: - prev_c4 = getattr(req, "c4_state_kv_len", 0) - prev_c128 = getattr(req, "c128_state_kv_len", 0) - c4_prefix.append(prev_c4) - c4_seq.append(prev_c4 + 1) - c128_prefix.append(prev_c128) - c128_seq.append(prev_c128 + 1) - req.c4_state_kv_len = prev_c4 + 1 - req.c128_state_kv_len = prev_c128 + 1 - - bs = len(reqs) - return self._pack_state_lens( - c4_prefix, - c4_seq, - c128_prefix, - c128_seq, - c4_extend_num_tokens=bs, - c128_extend_num_tokens=bs, - ) - - def compute_dsv4_state_lens_reserve( - self, reqs: List[Req], prefix_lens: List[int], seq_lens: List[int] - ) -> Optional[DSV4StateLens]: - """Allocate state slots for a speculative pre-reserved raw interval.""" - if self.c4_state_attn_allocator is None: - return None - - c4_prefix: List[int] = [] - c4_seq: List[int] = [] - c128_prefix: List[int] = [] - c128_seq: List[int] = [] - for req, prefix_len, seq_len in zip(reqs, prefix_lens, seq_lens): - reserve = max(0, int(seq_len) - int(prefix_len)) - prev_c4 = getattr(req, "c4_state_kv_len", 0) - prev_c128 = getattr(req, "c128_state_kv_len", 0) - c4_prefix.append(prev_c4) - c4_seq.append(prev_c4 + reserve) - c128_prefix.append(prev_c128) - c128_seq.append(prev_c128 + reserve) - req.c4_state_kv_len = prev_c4 + reserve - req.c128_state_kv_len = prev_c128 + reserve - - total = sum(max(0, int(s) - int(p)) for p, s in zip(prefix_lens, seq_lens)) - return self._pack_state_lens( - c4_prefix, - c4_seq, - c128_prefix, - c128_seq, - c4_extend_num_tokens=total, - c128_extend_num_tokens=total, - ) - - def _pack_state_lens( - self, - c4_prefix: List[int], - c4_seq: List[int], - c128_prefix: List[int], - c128_seq: List[int], - *, - c4_extend_num_tokens: int, - c128_extend_num_tokens: int, - ) -> DSV4StateLens: - c4_prefix_cpu = torch.tensor(c4_prefix, dtype=torch.int64) - c4_seq_cpu = torch.tensor(c4_seq, dtype=torch.int64) - c128_prefix_cpu = torch.tensor(c128_prefix, dtype=torch.int64) - c128_seq_cpu = torch.tensor(c128_seq, dtype=torch.int64) - return DSV4StateLens( - c4_prefix_lens=c4_prefix_cpu.to(self.device, non_blocking=True), - c4_prefix_lens_cpu=c4_prefix_cpu, - c4_seq_lens=c4_seq_cpu.to(self.device, non_blocking=True), - c4_seq_lens_cpu=c4_seq_cpu, - c4_extend_num_tokens=c4_extend_num_tokens, - c128_prefix_lens=c128_prefix_cpu.to(self.device, non_blocking=True), - c128_prefix_lens_cpu=c128_prefix_cpu, - c128_seq_lens=c128_seq_cpu.to(self.device, non_blocking=True), - c128_seq_lens_cpu=c128_seq_cpu, - c128_extend_num_tokens=c128_extend_num_tokens, ) def alloc_extend( @@ -580,12 +316,13 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): extend_num_tokens: int, *, req_pool_indices: Optional[torch.Tensor] = None, - dsv4_state_lens: Optional[DSV4StateLens] = None, req_to_token_pool=None, ) -> Optional[DSV4OutCacheLoc]: # Stash per-req tables for this call's last_loc lookups (read by # _alloc_c_extend / _alloc_state_extend); no permanent allocator->pool ref. self._cur_req_to_token_pool = req_to_token_pool + if not self._has_c128_sidecar_capacity(prefix_lens_cpu, seq_lens_cpu): + return None out_full_loc = super().alloc_extend( prefix_lens, prefix_lens_cpu, @@ -602,7 +339,6 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): seq_lens_cpu, last_loc.dtype, req_pool_indices, - dsv4_state_lens, ) def _wrap_full_alloc( @@ -614,10 +350,9 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): seq_lens_cpu, loc_dtype, req_pool_indices, - dsv4_state_lens, ) -> Optional[DSV4OutCacheLoc]: # Shared tail of alloc_extend / alloc_extend_swa_tail: translate the full - # loc to swa, then add the c4/c128(+state) pools into a DSV4OutCacheLoc. + # loc to swa, then add the c4/c128 KV pools into a DSV4OutCacheLoc. if out_full_loc is None: return None out_swa_loc = self.translate_loc_from_full_to_swa(out_full_loc) @@ -625,7 +360,7 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): "translate_loc_from_full_to_swa returned None — " "full_to_swa_index_mapping not initialized?" ) - return self._alloc_c_and_state( + return self._alloc_compressed_kv( out_full_loc, out_swa_loc, prefix_lens, @@ -634,7 +369,6 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): seq_lens_cpu, loc_dtype, req_pool_indices, - dsv4_state_lens, ) def alloc_decode( @@ -644,10 +378,13 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): last_loc: torch.Tensor, *, req_pool_indices: Optional[torch.Tensor] = None, - dsv4_state_lens: Optional[DSV4StateLens] = None, req_to_token_pool=None, ) -> Optional[DSV4OutCacheLoc]: self._cur_req_to_token_pool = req_to_token_pool + if not self._has_c128_sidecar_capacity( + (seq_lens_cpu - 1).clamp(min=0), seq_lens_cpu + ): + return None out_full_loc = super().alloc_decode(seq_lens, seq_lens_cpu, last_loc) if out_full_loc is None: return None @@ -657,7 +394,7 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): # seq_len//ratio so _alloc_c_extend anchors on the real c-pool last_loc. prefix_lens = (seq_lens - 1).clamp(min=0) prefix_lens_cpu = (seq_lens_cpu - 1).clamp(min=0) - return self._alloc_c_and_state( + return self._alloc_compressed_kv( out_full_loc, out_swa_loc, prefix_lens, @@ -666,7 +403,6 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): seq_lens_cpu, last_loc.dtype, req_pool_indices, - dsv4_state_lens, ) def alloc_extend_swa_tail( @@ -680,13 +416,14 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): swa_tail_len: int, *, req_pool_indices: Optional[torch.Tensor] = None, - dsv4_state_lens: Optional[DSV4StateLens] = None, req_to_token_pool=None, ) -> Optional[DSV4OutCacheLoc]: """Disagg-decode prealloc variant of :meth:`alloc_extend`: super() does - full+swa-tail, then _alloc_c_and_state adds c4/c128(+state) → DSV4OutCacheLoc. + full+swa-tail, then _alloc_compressed_kv adds c4/c128 KV → DSV4OutCacheLoc. """ self._cur_req_to_token_pool = req_to_token_pool + if not self._has_c128_sidecar_capacity(prefix_lens_cpu, seq_lens_cpu): + return None out_full_loc = super().alloc_extend_swa_tail( prefix_lens, prefix_lens_cpu, @@ -704,7 +441,6 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): seq_lens_cpu, last_loc.dtype, req_pool_indices, - dsv4_state_lens, ) def free( @@ -714,20 +450,15 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): req=None, req_to_token_pool=None, ): - """Unified free for full/swa/c4/c128 pools. Two forms (may co-fire): + """Unified free for full/SWA/C4/C128 KV and C128 request state. + + Two forms may co-fire: * ``free(free_index)`` — full + SWA only (tail/radix eviction; no req identity, so c-pool free can't run). * ``free(req=, req_to_token_pool=)`` — from DSV4NPUReqToTokenPool.free - on req finish: reads the per-req slot lists from - ``req_to_token_c{4,128}[_state]`` and returns them to the c-pools - (the paged allocator dedupes by page). - - KV pools free ``[0, kv_len // ratio)``. State pools are 1-per-raw-token - and free only the tail ``[c{N}_state_alloc_offset, kv_len)`` — the prefix - was already returned by ScheduleBatch._evict_swa (state rides SWA - eviction); freeing it again would double-free (caught by the paged - allocator's debug_mode assert, corrupts the free list otherwise). + on request finish: returns C128 KV pages and clears that request's + fixed C128 state bank before the req_pool_idx can be reused. """ if free_index is not None: super().free(free_index) @@ -739,53 +470,31 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): if kv_len <= 0 or req_pool_idx is None: return - # KV pools: free the leading [0, kv_len // ratio) compressed slots. - for ratio, allocator, table_attr in ( - (4, self.c4_attn_allocator, "req_to_token_c4"), - (128, self.c128_attn_allocator, "req_to_token_c128"), - ): - n = kv_len // ratio - if n > 0 and hasattr(req_to_token_pool, table_attr): - slots = getattr(req_to_token_pool, table_attr)[req_pool_idx, :n] - slots = slots[slots > 0] - # to int64 — paged allocator's free does cpu()//page_size on it. - if slots.numel() > 0: - allocator.free(slots.to(torch.int64)) + row = req_to_token_pool.req_to_c128_sidecar[int(req_pool_idx)] + self.release_c128_pages(row[row > 0]) + row.zero_() + self.get_kvcache().clear_c128_req_state(int(req_pool_idx)) - # State pools: free only the tail [c{N}_state_alloc_offset, kv_len). - for ratio, allocator, table_attr, off_attr in ( - ( - 4, - self.c4_state_attn_allocator, - "req_to_token_c4_state", - "c4_state_alloc_offset", - ), - ( - 128, - self.c128_state_attn_allocator, - "req_to_token_c128_state", - "c128_state_alloc_offset", - ), - ): - if allocator is None or not hasattr(req_to_token_pool, table_attr): - continue - off = getattr(req, off_attr, 0) - if kv_len > off: - slots = getattr(req_to_token_pool, table_attr)[req_pool_idx, off:kv_len] - slots = slots[slots > 0] - if slots.numel() > 0: - allocator.free(slots.to(torch.int64)) + def available_size(self): + return min( + super().available_size(), + self.c128_attn_allocator.available_size() * 128, + ) + + def resize(self, config) -> None: + self.c128_attn_allocator.size = int(config.c128_max_total_num_tokens) + self.c128_attn_allocator.num_pages = ( + self.c128_attn_allocator.size // self.c128_attn_allocator.page_size + ) + super().resize(config) def clear(self): super().clear() - # super().__init__ calls clear() before our sub-allocators exist; - # getattr(..., None) tolerates that and the always-None state allocators. - for attr in ( - "c4_attn_allocator", - "c128_attn_allocator", - "c4_state_attn_allocator", - "c128_state_attn_allocator", - ): + # super().__init__ calls clear() before our C128 allocator exists. + for attr in ("c128_attn_allocator",): allocator = getattr(self, attr, None) if allocator is not None: allocator.clear() + refcount = getattr(self, "c128_page_refcount", None) + if refcount is not None: + refcount.zero_() 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 393fe0744..5cc5761aa 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 @@ -1,4 +1,4 @@ -"""Helpers used by mem_cache/common.py to wire DSV4-NPU per-req tables. +"""Helpers used by mem_cache/common.py to wire DSV4-NPU KV tables. mem_cache/common.py runs platform-agnostic alloc flow. When the model is DSV4 on NPU, ``alloc_paged_token_slots_{extend,decode}`` already stashed the @@ -7,20 +7,23 @@ DSV4 on NPU, ``alloc_paged_token_slots_{extend,decode}`` already stashed the these hooks then: 1. Read the bundle from ``batch.out_cache_loc_dsv4``. - 2. Write the per-pool slot ids into the per-req tables on the - :class:`DSV4NPUReqToTokenPool`. + 2. Write newly allocated C128 page ids into the per-request sidecar. + +Compressor state is fixed ring storage and does not participate in this +allocation/write path. PD reuses the public SWA/C128-state payloads and only +builds an NPU-specific payload for the independently addressed C128 KV pool. Non-DSV4 paths leave ``batch.out_cache_loc_dsv4`` None, so this module is a no-op for them. The disagg per-req prealloc path does not build a ``ScheduleBatch`` and so -bypasses the batch hook; it writes the same tables via +bypasses the batch hook; it writes the same sidecar via ``write_dsv4_prealloc_tables`` (driven by ``dsv4_unwrap_prealloc``). """ from __future__ import annotations -from typing import TYPE_CHECKING, Sequence +from typing import TYPE_CHECKING import torch @@ -33,19 +36,13 @@ 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. - For each compressed pool (c4 / c128), spreads the flat - ``out_c{4,128}_loc`` tensor across requests using per-req extend - counts (``seq_lens[i] // ratio - prefix_lens[i] // ratio``) and writes - the resulting slot ids into ``req_to_token_c{4,128}[req, prefix:seq]``. + Spreads the flat ``out_c128_loc`` tensor across requests and writes newly + allocated page ids into ``req_to_c128_sidecar``. C4 locations are derived + from the full-token table. - Also writes ``req_to_token_swa[req, prefix:seq]`` with the swa slots - derived from out_full_loc via the SWA index mapping. """ # Bundle stashed on batch.out_cache_loc_dsv4 by mem_cache/common.py; # None on CUDA / non-V4 paths → no-op. @@ -54,33 +51,15 @@ def maybe_write_dsv4_extend( return req_to_token_pool = batch.req_to_token_pool - if not hasattr(req_to_token_pool, "write_c4"): + if not hasattr(req_to_token_pool, "write_c128"): return # non-DSV4 pool; skip defensively (shouldn't happen) - # c4_state / c128_state writes: tail-only. Bundle length is - # sum(c{N}_state_alloc_len_i), NOT total raw extend tokens. Normal extend - # uses the per-Req low-water marks; reserve callers can pass explicit raw - # offsets for the pre-reserved interval. - if c4_state_alloc_offsets is None: - c4_state_alloc_offsets = [ - getattr(r, "c4_state_write_offset", 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_write_offset", getattr(r, "c128_state_alloc_offset", 0) - ) - for r in batch.reqs - ] _write_dsv4_tables( req_to_token_pool, req_pool_indices_cpu, prefix_lens_cpu, seq_lens_cpu, bundle, - c4_state_offsets=c4_state_alloc_offsets, - c128_state_offsets=c128_state_alloc_offsets, ) @@ -89,18 +68,10 @@ def dsv4_state_payloads( req_pool_idx: int, seq_len: int, page_size: int, - window_size: int, *, prefix_len: int = 0, ): - """Per-StateType PD-payload builders for DSV4-on-NPU. - - For chunked prefill, intermediate chunks can leave old C4/C128 state pages in - the req table. PD only needs the final active tail state; scanning the whole - prompt span would transfer stale state pages and can perturb decode accuracy. - """ - if not hasattr(req_to_token_pool, "req_to_token_c4"): - return {} + """Build the only NPU-specific DSV4 PD payload: C128 KV pages.""" import numpy as np @@ -109,103 +80,38 @@ def dsv4_state_payloads( seq_len = max(0, int(seq_len)) prefix_len = max(0, min(int(prefix_len), seq_len)) - def empty_pages(): - return np.empty((0,), dtype=np.int32) - - def pages(table, lo: int, hi: int, *, drop_zero_pages: bool = False): + def c128_kv_pages(): + c128_page_size = req_to_token_pool.c128_page_size + lo = prefix_len // (128 * c128_page_size) + hi = (seq_len // 128 + c128_page_size - 1) // c128_page_size if hi <= lo: - return empty_pages() + return np.empty((0,), dtype=np.int32) + pages = ( + req_to_token_pool.req_to_c128_sidecar[req_pool_idx, lo:hi] + .cpu() + .numpy() + .astype(np.int32) + ) + return pages[pages > 0] - lo = max(0, int(lo)) - hi = max(lo, int(hi)) - page_lo = (lo // page_size) * page_size - page_hi = ((hi + page_size - 1) // page_size) * page_size - if page_hi <= page_lo: - return empty_pages() - - slots = table[req_pool_idx, page_lo:page_hi:page_size].cpu().numpy() - if slots.size == 0: - return empty_pages() - - page_indices = (slots // page_size).astype(np.int32) - if drop_zero_pages: - page_indices = page_indices[page_indices > 0] - return page_indices - - def state_tail_range(compress_ratio: int): - tail_len = seq_len % 128 - if compress_ratio == 4: - state_len = tail_len + 128 if tail_len <= 3 and seq_len >= 128 else tail_len - elif compress_ratio == 128: - state_len = tail_len - else: - raise ValueError(f"Unsupported DSV4 state compress ratio: {compress_ratio}") - - if state_len == 0: - return None - - start = max(prefix_len, seq_len - state_len) - if start >= seq_len: - return None - return start, seq_len - - def state_pages(table, compress_ratio: int): - state_range = state_tail_range(compress_ratio) - if state_range is None: - return empty_pages() - lo, hi = state_range - return pages(table, lo, hi, drop_zero_pages=True) - - if window_size is None or window_size <= 0: - window_start = prefix_len - else: - window_start = max(prefix_len, seq_len - window_size) - window_start = (window_start // page_size) * page_size - - # DSV4_INDEXER shares the c4 slot space (written at the c4 loc). - return { - AscendStateType.DSV4_SWA: lambda: pages( - req_to_token_pool.req_to_token_swa, - window_start, - seq_len, - drop_zero_pages=True, - ), - AscendStateType.DSV4_C4: lambda: pages( - req_to_token_pool.req_to_token_c4, prefix_len // 4, seq_len // 4 - ), - AscendStateType.DSV4_C128: lambda: pages( - req_to_token_pool.req_to_token_c128, prefix_len // 128, seq_len // 128 - ), - AscendStateType.DSV4_INDEXER: lambda: pages( - req_to_token_pool.req_to_token_c4, prefix_len // 4, seq_len // 4 - ), - AscendStateType.DSV4_C4_STATE: lambda: state_pages( - req_to_token_pool.req_to_token_c4_state, 4 - ), - AscendStateType.DSV4_C128_STATE: lambda: state_pages( - req_to_token_pool.req_to_token_c128_state, 128 - ), - } + return {AscendStateType.DSV4_C128: c128_kv_pages} def dsv4_prealloc_kwargs(allocator, req, fill_len, req_to_token_pool, *, device): """Extra ``alloc_extend(_swa_tail)`` kwargs for the DSV4 allocator; ``{}`` for non-DSV4 so callers can splat it unconditionally.""" - if not hasattr(allocator, "c4_attn_allocator"): + if not hasattr(allocator, "c128_attn_allocator"): return {} return dict( req_pool_indices=torch.tensor( [req.req_pool_idx], dtype=torch.int64, device=device ), - dsv4_state_lens=allocator.compute_dsv4_state_lens_extend( - [req], [fill_len], [0] - ), req_to_token_pool=req_to_token_pool, ) def dsv4_unwrap_prealloc(kv_loc, req_to_token_pool, req, prefix_len, fill_len): - """Unwrap a DSV4OutCacheLoc bundle to its full-pool loc and write the five + """Unwrap a DSV4OutCacheLoc bundle to its full-pool loc and write the per-req tables; a plain tensor (non-DSV4) passes through unchanged.""" if kv_loc is None or not hasattr(kv_loc, "out_full_loc"): return kv_loc @@ -220,9 +126,9 @@ def write_dsv4_prealloc_tables( fill_len: int, bundle, ) -> None: - """Write the five DSV4 per-req tables for one request on the disagg-decode + """Write the DSV4 per-req tables for one request on the disagg-decode prealloc path (no ScheduleBatch); no-op without bundle / DSV4 tables.""" - if bundle is None or not hasattr(req_to_token_pool, "write_c4"): + if bundle is None or not hasattr(req_to_token_pool, "write_c128"): return rp = torch.tensor([req.req_pool_idx]) pl = torch.tensor([prefix_len]) @@ -234,8 +140,6 @@ def write_dsv4_prealloc_tables( pl, sl, bundle, - c4_state_offsets=[getattr(req, "c4_state_alloc_offset", 0)], - c128_state_offsets=[getattr(req, "c128_state_alloc_offset", 0)], ) @@ -245,27 +149,8 @@ def _write_dsv4_tables( prefix_lens_cpu: torch.Tensor, seq_lens_cpu: torch.Tensor, bundle, - *, - c4_state_offsets: Sequence[int] | torch.Tensor, - c128_state_offsets: Sequence[int] | torch.Tensor, ) -> None: - """Write DSV4 SWA, compressed-KV, and compression-state tables.""" - _write_per_req_slice( - req_to_token_pool.write_swa, - req_pool_indices_cpu, - prefix_lens_cpu, - seq_lens_cpu, - bundle.out_swa_loc, - ratio=1, - ) - _write_per_req_slice( - req_to_token_pool.write_c4, - req_pool_indices_cpu, - prefix_lens_cpu, - seq_lens_cpu, - bundle.out_c4_loc, - ratio=4, - ) + """Write newly allocated C128 page ids into the request sidecar.""" _write_per_req_slice( req_to_token_pool.write_c128, req_pool_indices_cpu, @@ -275,36 +160,14 @@ def _write_dsv4_tables( ratio=128, ) - if bundle.out_c4_state_loc is not None and hasattr( - req_to_token_pool, "write_c4_state" - ): - _write_state_tail_per_req( - req_to_token_pool.write_c4_state, - req_pool_indices_cpu, - c4_state_offsets, - seq_lens_cpu, - bundle.out_c4_state_loc, - ) - if bundle.out_c128_state_loc is not None and hasattr( - req_to_token_pool, "write_c128_state" - ): - _write_state_tail_per_req( - req_to_token_pool.write_c128_state, - req_pool_indices_cpu, - c128_state_offsets, - seq_lens_cpu, - bundle.out_c128_state_loc, - ) - def maybe_write_dsv4_decode( batch: ScheduleBatch, seq_lens_cpu: torch.Tensor, token_per_req: int, ) -> None: - """Post-alloc_decode hook for DSV4. Spreads the new token slot ids - (one per req for swa, gated by ratio boundary for c4/c128) into the - per-req tables on DSV4NPUReqToTokenPool. + """Post-alloc_decode hook for DSV4. Spreads new C128 KV slot ids into + the per-req sidecar on DSV4NPUReqToTokenPool. ``seq_lens_cpu`` is the POST-decode seq len (already incremented by ``token_per_req``); the new compressed tokens go at positions @@ -317,28 +180,12 @@ def maybe_write_dsv4_decode( return req_to_token_pool = batch.req_to_token_pool - if not hasattr(req_to_token_pool, "write_c4"): + if not hasattr(req_to_token_pool, "write_c128"): return prefix_lens_cpu = (seq_lens_cpu - token_per_req).clamp(min=0) req_pool_indices_cpu = batch.req_pool_indices.cpu() - _write_per_req_slice( - req_to_token_pool.write_swa, - req_pool_indices_cpu, - prefix_lens_cpu, - seq_lens_cpu, - bundle.out_swa_loc, - ratio=1, - ) - _write_per_req_slice( - req_to_token_pool.write_c4, - req_pool_indices_cpu, - prefix_lens_cpu, - seq_lens_cpu, - bundle.out_c4_loc, - ratio=4, - ) _write_per_req_slice( req_to_token_pool.write_c128, req_pool_indices_cpu, @@ -348,32 +195,13 @@ def maybe_write_dsv4_decode( ratio=128, ) - # State table decode writes: one slot per raw decode token (ratio=1). - if bundle.out_c4_state_loc is not None and hasattr( - req_to_token_pool, "write_c4_state" - ): - _write_per_req_slice( - req_to_token_pool.write_c4_state, - req_pool_indices_cpu, - prefix_lens_cpu, - seq_lens_cpu, - bundle.out_c4_state_loc, - ratio=1, - ) - if bundle.out_c128_state_loc is not None and hasattr( - req_to_token_pool, "write_c128_state" - ): - _write_per_req_slice( - req_to_token_pool.write_c128_state, - req_pool_indices_cpu, - prefix_lens_cpu, - seq_lens_cpu, - bundle.out_c128_state_loc, - ratio=1, - ) - -def maybe_build_dsv4_verify_bundle(batch: ScheduleBatch, draft_token_num: int): +def maybe_build_dsv4_verify_bundle( + batch: ScheduleBatch, + draft_token_num: int, + *, + live_seq_lens_cpu: torch.Tensor | None = None, +): """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 @@ -381,31 +209,45 @@ def maybe_build_dsv4_verify_bundle(batch: ScheduleBatch, draft_token_num: int): the larger allocation bundle produced during decode preparation. """ pool = batch.req_to_token_pool - if not hasattr(pool, "req_to_token_c4"): + if not hasattr(pool, "req_to_c128_sidecar"): 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() + + if live_seq_lens_cpu is None: + live_seq_lens_cpu = batch.seq_lens_cpu + if live_seq_lens_cpu is None: + live_seq_lens_cpu = batch.seq_lens[: len(req_indices)].cpu() + live_seq_lens = live_seq_lens_cpu[: len(req_indices)].tolist() + + verify_lens = [int(draft_token_num)] * len(req_indices) def flatten_interval(table: torch.Tensor, ratio: int) -> torch.Tensor: + page_size = pool.c128_page_size 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 + for req_idx, live_seq_len, verify_len in zip( + req_indices, live_seq_lens, verify_lens + ): + start = int(live_seq_len) // ratio + end = (int(live_seq_len) + int(verify_len)) // ratio if end > start: - chunks.append(table[int(req_idx), start:end]) + positions = torch.arange(start, end, device=table.device) + pages = table[int(req_idx), positions // page_size] + chunks.append(pages * page_size + positions % page_size) return torch.cat(chunks) if chunks else table.new_empty((0,)) + out_full_loc = batch.out_cache_loc + out_c4_loc = out_full_loc[(out_full_loc >= 0) & ((out_full_loc % 4) == 3)] // 4 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), + out_full_loc=out_full_loc, + out_swa_loc=batch.token_to_kv_pool_allocator.translate_loc_from_full_to_swa( + out_full_loc + ), + out_c4_loc=out_c4_loc, + out_c128_loc=flatten_interval(pool.req_to_c128_sidecar, 128), ) @@ -437,23 +279,6 @@ def _write_per_req( pt += alloc_len -def _write_state_tail_per_req( - write_fn, - req_pool_indices_cpu: torch.Tensor, - state_alloc_offsets: list, - seq_lens_cpu: torch.Tensor, - flat_loc: torch.Tensor, -) -> None: - """Tail-only state write: req i's slots go at ``[state_alloc_offsets[i], - seq_lens[i])`` in ``req_to_token_c{N}_state``.""" - _write_per_req( - write_fn, - req_pool_indices_cpu, - flat_loc, - lambda i: (int(state_alloc_offsets[i]), int(seq_lens_cpu[i].item())), - ) - - def _write_per_req_slice( write_fn, req_pool_indices_cpu: torch.Tensor, @@ -473,108 +298,3 @@ def _write_per_req_slice( int(seq_lens_cpu[i].item()) // ratio, ), ) - - -def maybe_evict_dsv4_state(batch: ScheduleBatch, req: Req, pre_len: int) -> None: - """Per-decode evict for the DSV4-NPU compress-state pools, independent of - SWA evict cadence. Called every decode step from ``ScheduleBatch``. - - The state pool is small (~2 pages c4 / ~3 pages c128 of raw positions per - req) — with a large sliding_window (SWA evict fires every - ``eviction_interval`` and needs ``pre_len > sliding_window + page_size`` to - free anything) the pool exhausts before the first SWA frontier advance, so - we drain it here on its own cadence. - - Retention windows (kernel read window + decode lookahead margin): - c4 = 8 + 16, c128 = 128 + 64 raw positions — intentionally smaller than one - SWA page so the first eviction fires before the small pool fills. Watermarks - are page-aligned so freed slots are whole pages reclaimable by the paged - allocator. ``req.c{4,128}_state_alloc_offset`` (read/written via getattr/ - setattr) is the low-water mark. No-op on non-DSV4-NPU paths. - """ - allocator = batch.token_to_kv_pool_allocator - pool = batch.req_to_token_pool - if not hasattr(allocator, "c4_state_attn_allocator") or ( - allocator.c4_state_attn_allocator is None - and allocator.c128_state_attn_allocator is None - ): - return - - page_size = batch.tree_cache.page_size - c4_watermark = ((max(0, pre_len - (8 + 16))) // page_size) * page_size - c128_watermark = ((max(0, pre_len - (128 + 64))) // page_size) * page_size - - _free_state_range( - allocator.c4_state_attn_allocator, - pool, - "req_to_token_c4_state", - req, - "c4_state_alloc_offset", - c4_watermark, - ) - _free_state_range( - allocator.c128_state_attn_allocator, - pool, - "req_to_token_c128_state", - req, - "c128_state_alloc_offset", - c128_watermark, - ) - - -def maybe_evict_dsv4_state_on_swa( - allocator, pool, req: Req, new_swa_evicted_seqlen: int -) -> None: - """Free compress-state slots that ride along with SWA eviction. - - State at raw positions < ``swa_evicted_seqlen`` is no longer readable (the - compressor only reads the trailing ``2*ratio`` window) and is returned to - its paged allocator to keep the small state pool from exhausting on long - generations. No-op when the DSV4-NPU state allocators are absent. - - This path is needed for small-sliding-window models where - ``sliding_window < retention`` (e.g. c128 retention 192 > window 128): - in that case the watermark-based eviction alone may not free slots - fast enough, and the SWA-ride eviction is the primary reclaim mechanism. - For typical large-window models (DS-V4 with window >> 192), the - watermark eviction always runs first, making this path a no-op. - """ - if not hasattr(allocator, "c4_state_attn_allocator"): - return - _free_state_range( - allocator.c4_state_attn_allocator, - pool, - "req_to_token_c4_state", - req, - "c4_state_alloc_offset", - new_swa_evicted_seqlen, - ) - _free_state_range( - allocator.c128_state_attn_allocator, - pool, - "req_to_token_c128_state", - req, - "c128_state_alloc_offset", - new_swa_evicted_seqlen, - ) - - -def _free_state_range( - state_allocator, - pool, - table_attr: str, - req: Req, - offset_attr: str, - watermark: int, -) -> None: - """Free ``[alloc_offset, watermark)`` raw-position state slots for ``req`` - and advance its low-water mark. No-op when the allocator/table is absent or - the watermark hasn't advanced past the current offset.""" - offset = getattr(req, offset_attr, 0) - if state_allocator is None or not hasattr(pool, table_attr) or watermark <= offset: - return - free_slots = getattr(pool, table_attr)[req.req_pool_idx, offset:watermark] - free_slots = free_slots[free_slots > 0] - if free_slots.numel() > 0: - state_allocator.free(free_slots.to(torch.int64)) - setattr(req, offset_attr, watermark) diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py index c8b1172a5..228e3cd9a 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py @@ -1,34 +1,19 @@ """NPU-only KV pool variant for DeepSeek-V4. -Subclasses :class:`DeepSeekV4TokenToKVPool` to swap the ring-buffered -:class:`CompressStatePool` for the paged :class:`NPUCompressStatePool` that -the on-NPU fused compressor kernel (``torch.ops.custom.compressor`` with -``cache_mode=1``) requires. Atlas A3 rejects ``cache_mode=2`` (ring) entirely, -so this is the only valid layout on that hardware. +The full/SWA/C4/C128 KV buffers keep their Ascend-specific PA_ND layout. The +Compressor state buffers, however, use the same ownership and flat ``state_loc`` +rules as the GPU implementation: -Selected at pool construction time by -:meth:`ModelRunnerKVCacheMixin._init_pools` when the model is DSV4 AND the -device is NPU. CUDA continues to use the unchanged base class. +* C4A/C4Li state follows SWA physical pages. +* C128A state follows ``req_pool_idx`` and absolute position. -The subclass overrides only: - - * ``_make_attn_state_pool`` / ``_make_indexer_state_pool`` — the per-ratio - state-pool factories the base ``_init_paged_compress_states`` loop calls. - Both return :class:`NPUCompressStatePool` (paged, ``cache_mode=1``) - instead of the base's ring-buffered :class:`CompressStatePool`. - * ``translate_kv_loc_to_compress_state_loc`` — raise loudly. The ring - hash this method implements is meaningless on the paged kernel; callers - must consume ``out_cache_loc_dsv4.out_c{4,128}_state_loc`` from the - allocator bundle instead. Currently the only NPU caller that still - invokes translate is the unfused Python compressor decode path - (``layers/attention/dsv4/compressor.py``); with USE_FUSED_COMPRESSOR=1 - that path is dead. If someone disables the fused compressor, they hit - the raise with a clear message. +``NPUCompressStatePool`` only adds the contiguous 3-D view and positive dummy +location required by the Atlas A3 ``cache_mode=2`` operator. There is no paged +state allocator or ``cache_mode=1`` compatibility storage. """ from __future__ import annotations -import math from typing import List, Optional, Tuple import torch @@ -42,6 +27,7 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( DeepSeekV4SingleKVPool, DeepSeekV4TokenToKVPool, ) +from sglang.srt.runtime_context import get_schedule class NPUDeepSeekV4SingleKVPool(DeepSeekV4SingleKVPool): @@ -49,10 +35,9 @@ class NPUDeepSeekV4SingleKVPool(DeepSeekV4SingleKVPool): ``npu_sparse_attn_sharedkv`` reads KV in PA_ND layout ``(num_pages, kernel_page_size, num_kv_heads=1, dim)`` with ``dim`` packing - K_nope + K_rope as bf16, and requires ``cmp_kv.shape[1] == ori_kv.shape[1]``. - So the c4/c128 pools (whose token-level page_size is ``page_size // ratio``) - are allocated at the GLOBAL ``kernel_page_size`` rather than their own - per-ratio page_size; the SWA pool uses ``kernel_page_size == page_size``. + K_nope + K_rope as bf16. C4 uses its native page so its physical page id can + be shared with the corresponding full page. C128 uses its independently + configured physical page size; Full/SWA use the global page size. The CUDA fp8-packed-bytes layout (the base ``create_buffer``) is untouched. """ @@ -68,8 +53,8 @@ class NPUDeepSeekV4SingleKVPool(DeepSeekV4SingleKVPool): return super().create_buffer(num_pages=num_pages) kv_dim = self.qk_nope_head_dim + self.qk_rope_head_dim self.kv_cache_total_dim = kv_dim - # GLOBAL kernel_page_size keeps cmp_kv.shape[1] == ori_kv.shape[1]; writes - # are flat-indexed by loc, so page granularity affects shape not location. + # Writes are flat-indexed by loc; kernel_page_size controls the physical + # page layout exposed to the NPU operators. npu_num_pages = (self.size + self.kernel_page_size + 1) // self.kernel_page_size return torch.zeros( npu_num_pages, @@ -81,59 +66,17 @@ class NPUDeepSeekV4SingleKVPool(DeepSeekV4SingleKVPool): ) -def npu_state_pool_size( - *, - ratio: int, - page_size: int, - max_num_reqs: int, -) -> int: - """Per-pool state slot count for the NPU paged state pool's - :class:`NPUPagedTokenToKVPoolAllocator`. - - Sizing formula:: - - max(2, ceil(1.8 * ratio / page_size) + 1) * max_num_reqs * page_size - - Sized for steady-state during decode: each req keeps roughly the trailing - ``sliding_window_size`` worth of state slots live at any one time (SWA - eviction in :meth:`ScheduleBatch._evict_swa` frees state slots as it - advances), and the 1.8x factor adds headroom for the tail-only allocation - pattern across page boundaries. - - Prefill no longer drives sizing because allocation is tail-only — long - prompts only allocate ``c{ratio}_alloc_len`` slots (``≤ tail + 128`` for - c4, ``≤ tail`` for c128, where ``tail = seq_len % 128``), not the full raw - seqlen. See :meth:`ScheduleBatch._compute_dsv4_state_lens_extend` for the - per-req formula. - - Result is in TOKEN units (matches the SGLang allocator - ``PagedTokenToKVPoolAllocator(size, ...)`` convention where - ``num_pages = size // page_size`` is the count of USABLE pages handed out - by ``free_pages = arange(1, num_pages+1)``). The BUFFER allocates one extra - page (see :class:`NPUCompressStatePool`, sized ``(num_pages + 1) * - page_size`` — page 0 is the kernel's skip-sentinel). - """ - blocks_per_req = max(2, math.ceil(1.8 * ratio / page_size) + 1) - num_usable_pages = blocks_per_req * max_num_reqs - return num_usable_pages * page_size - - class NPUCompressStatePool(CompressStatePool): - """Paged compress-state pool for the NPU fused compressor kernel. + """Thin A3 adapter over the shared GPU-style ring state pool. - ``torch.ops.custom.compressor`` (cache_mode=1) reads/writes the compress - state via ``state_cache`` shape ``(block_num, page_size, 2*coff*head_dim)`` - indexed by a paged ``state_block_table`` (block ids from 1; value 0 means - "skip this slot"). The CUDA :class:`CompressStatePool` sizes itself - ring-style, which misaddresses slots under cache_mode=1 (ring is also - unsupported on Atlas A3). This subclass keeps the parent's buffer layout - (``(self._size, 2*coff*head_dim)`` flat; ``state_cache_3d`` reshapes to - ``(num_blocks, page_size, 2*coff*head_dim)``) but replaces the size formula - with a paged one derived from ``max_num_reqs``. Block 0 is reserved as the - kernel's skip-sentinel (zero kv / -inf score) so any ``state_block_table`` - entry defaulting to 0 lands in a deterministic, attention-neutral place. + Allocation, sizing, ring ownership and address translation are inherited + from :class:`CompressStatePool`. NPU only requests a contiguous 3-D view, + enforces the A3 FP32 contract and replaces invalid locations with a cleared + positive dummy row. - NPU-only; CUDA keeps using the unchanged :class:`CompressStatePool`. + Location 0 is valid in explicit mode. Invalid/history-padding locations map + to the final cleared row instead of ``-1`` because the A3 kernel consumes + unsigned offsets. """ def __init__( @@ -146,56 +89,66 @@ class NPUCompressStatePool(CompressStatePool): device: str, enable_memory_saver: bool, ratio: int, - page_size: int, + ring_size: int, + swa_page_size: int, ): - # Bypass parent __init__ — its ring-based sizing is incompatible with the - # kernel's paged block-id contract. We redo buffer alloc and set the same - # fields so the parent API (state_cache_3d, kv_score_buffer) stays intact. assert ratio in ( 4, 128, ), f"NPUCompressStatePool only supports ratio in (4, 128); got {ratio}" - assert page_size > 1, ( - "NPUCompressStatePool requires page_size>1 (kernel's " - "state_cache_3d view is (block_num, page_size, slot_dim)). " - "Got page_size=%d." % page_size + assert dtype == torch.float32, ( + "Atlas A3 custom.compressor requires FP32 state_cache, " + f"but NPUCompressStatePool got {dtype}." + ) + assert ring_size > 0, f"ring_size must be positive, got {ring_size}" + super().__init__( + size=size, + ring_size=ring_size, + overlap=overlap, + head_dim=head_dim, + dtype=dtype, + device=device, + enable_memory_saver=enable_memory_saver, + ratio=ratio, + online=False, + swa_page_size=swa_page_size, + state_cache_page_size=ring_size, + ) + self.dummy_state_loc = self._size - 1 + + # The shared pool initializes its dummy row. A cold C128 request bank + # additionally needs every row initialized before its first partial use. + if ratio == 128: + self.kv_score_buffer.clear() + + def _replace_invalid_with_dummy(self, state_loc: torch.Tensor) -> torch.Tensor: + return torch.where( + state_loc < 0, + torch.full_like(state_loc, self.dummy_state_loc), + state_loc, ) - # ``size`` is the ALLOCATOR's size (npu_state_pool_size output). Buffer - # needs one EXTRA page so the free list arange(1, num_pages+1) indexes it - # without OOB (page 0 = skip sentinel; pages 1..num_pages handed out). - num_usable_pages = (size + page_size - 1) // page_size - num_buffer_pages = num_usable_pages + 1 - self._size = num_buffer_pages * page_size - self.ratio = ratio - self.page_size = page_size - # ring_size=0 marks "not ring-buffered" (paged allocator replaces the - # parent's ring hashing); kept so downstream hasattr probes don't break. - self.ring_size = 0 - # online compress is a CUDA-only opt with no NPU fused-compressor support; - # force off so layout matches kernel expectations. - self.online = False - - # Slot dim = 2 * coff * head_dim = [kv | score]; coff = 1 (no overlap) or - # 2 (overlap). Matches CompressStatePool non-online layout. - self.last_dim = 2 * (1 + int(overlap)) * head_dim - - # Reuse parent's buffer-alloc helper; only self._size differs from the - # ring-based parent path. - self._alloc_kv_score_buffer( - dtype=dtype, device=device, enable_memory_saver=enable_memory_saver + def translate_from_swa_loc_to_state_loc( + self, swa_loc: torch.Tensor + ) -> torch.Tensor: + return self._replace_invalid_with_dummy( + super().translate_from_swa_loc_to_state_loc(swa_loc) ) - # Block 0 = kernel skip-sentinel: kv zeroed, score -inf (softmax → 0). - # The free list excludes it; only stale state_block_table entries land here. - self.kv_score_buffer.kv[:page_size].zero_() - self.kv_score_buffer.score[:page_size].fill_(float("-inf")) + def translate_from_req_position_to_state_loc( + self, req_pool_indices: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: + return self._replace_invalid_with_dummy( + super().translate_from_req_position_to_state_loc( + req_pool_indices, positions + ) + ) class NPUDeepSeekV4IndexerPool(DeepSeekV4IndexerPool): """NPU c4-indexer pool. Keeps the base packed CUDA buffer (read by get_contiguous_buf_infos / NSA) and ADDS dedicated int8 K + float16 scale - buffers in PA_ND layout at the global ``kernel_page_size``, written by + buffers in PA_ND layout at the native C4 ``kernel_page_size``, written by ``torch_npu.npu_scatter_nd_update_`` and read by ``torch.ops.custom.npu_quant_lightning_indexer``. """ @@ -270,12 +223,12 @@ class NPUDeepSeekV4IndexerPool(DeepSeekV4IndexerPool): class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): - """NPU-only DSV4 KV pool with paged compress-state buffers. + """NPU-only DSV4 KV pool with explicit-location ring state buffers. The full / SWA / c4 / c128 KV pools use the NPU bf16 PA_ND layout - (:class:`NPUDeepSeekV4SingleKVPool`); the compress-state pool is paged - (:class:`NPUCompressStatePool`) rather than ring-buffered; and the indexer - pool adds dedicated int8 K + fp16 scale buffers + (:class:`NPUDeepSeekV4SingleKVPool`); :class:`NPUCompressStatePool` + exposes the explicit-location ring view required by A3; + and the indexer pool adds dedicated int8 K + fp16 scale buffers (:class:`NPUDeepSeekV4IndexerPool`). The generic-accessor / port-hook methods at the bottom of this class are the NPU equivalents of the CUDA DSV4 store-cache chain — kept here, not in the community base, which raises @@ -283,6 +236,16 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): accessors instead). """ + def __init__(self, *args, **kwargs): + c128_page_size = get_schedule().c128_page_size + if c128_page_size <= 0 or c128_page_size % 16 != 0: + raise ValueError( + "c128_page_size must be a positive multiple of 16 for the NPU " + f"sparse-attention operator, got {c128_page_size}" + ) + self.c128_page_size = c128_page_size + super().__init__(*args, **kwargs) + def _make_kv_pool( self, *, @@ -301,6 +264,16 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): "enable_hisparse is not supported on the NPU DSV4 KV pool " f"(got c4 pool class {cls.__name__})." ) + # Full/SWA use the global page size, C4 uses its native compressed page, + # and C128 has an independent physical page size. + is_c4_pool = page_size * 4 == global_page_size + is_c128_pool = page_size * 128 == global_page_size + if is_c4_pool: + kernel_page_size = page_size + elif is_c128_pool: + kernel_page_size = self.c128_page_size + else: + kernel_page_size = global_page_size return NPUDeepSeekV4SingleKVPool( size, page_size, @@ -310,7 +283,7 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): layer_num, device, enable_memory_saver, - kernel_page_size=global_page_size, + kernel_page_size=kernel_page_size, ) def _get_state_pool(self, layer_id: int, from_indexer: bool) -> CompressStatePool: @@ -332,13 +305,14 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): ) return NPUCompressStatePool( size=self._state_pool_size(ratio), + ring_size=self.get_ring_size(ratio), overlap=ratio == 4, head_dim=self.qk_nope_head_dim + self.qk_rope_head_dim, dtype=self.c4_state_dtype if ratio == 4 else self.c128_state_dtype, device=self.device, enable_memory_saver=enable_memory_saver, ratio=ratio, - page_size=self.swa_page_size, + swa_page_size=self.swa_page_size, ) def _make_indexer_state_pool( @@ -348,24 +322,16 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): # slot_dim (indexer_head_dim vs attention head_dim). return NPUCompressStatePool( size=self.c4_state_pool_size, + ring_size=self.get_ring_size(ratio), overlap=ratio == 4, head_dim=self.indexer_head_dim, device=self.device, dtype=self.c4_state_dtype, enable_memory_saver=enable_memory_saver, ratio=ratio, - page_size=self.swa_page_size, + swa_page_size=self.swa_page_size, ) - def clear_unaccepted_c128_draft_states( - self, - req_pool_indices: torch.Tensor, - seq_lens: torch.Tensor, - accept_lens: torch.Tensor, - num_draft_tokens: int, - ) -> None: - pass - def _make_indexer_pool( self, size: int, @@ -376,8 +342,7 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): device: str, enable_memory_saver: bool, ) -> NPUDeepSeekV4IndexerPool: - # NPU dedicated int8 K + fp16 scale buffers use the GLOBAL page_size - # (= self.page_size) as kernel_page_size, matching ori_kv for the kernel. + # Indexer shares C4 addresses and therefore uses the same native page. return NPUDeepSeekV4IndexerPool( size, page_size, @@ -386,86 +351,58 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): layer_num, device, enable_memory_saver, - kernel_page_size=self.page_size, + kernel_page_size=page_size, ) def get_contiguous_buf_infos(self) -> Tuple[List[int], List[int], List[int]]: - # No full-token contiguous space on NPU; everything ships per-pool via - # get_pd_state_components(), so the contiguous path is empty. - return [], [], [] - - def get_pd_state_components( - self, - ) -> List[Tuple[str, List[int], List[int], List[int]]]: - """Ordered ``(AscendStateType, data_ptrs, data_lens, item_lens)`` per pool, in a - fixed order so prefill and decode register identically (empty pools skipped).""" - from sglang.srt.disaggregation.ascend.conn import AscendStateType - - components: List[Tuple[str, List[int], List[int], List[int]]] = [] - - def kv_entry(bufs): - return ( - [b.data_ptr() for b in bufs], - [b.nbytes for b in bufs], - [b[0].nbytes for b in bufs], - ) - - def state_entry(want_ratio: int, include_indexer: bool): - ptrs: List[int] = [] - lens: List[int] = [] - ilens: List[int] = [] - - def add(pool): - t = pool.kv_score_buffer.kv_score - ptrs.append(t.data_ptr()) - lens.append(t.nbytes) - ilens.append(t[0].nbytes * pool.page_size) - - for ratio, pool in zip(self.compression_ratios, self.compress_state_pools): - if pool is not None and ratio == want_ratio: - add(pool) - if include_indexer: - # indexer compress-state pools are all ratio 4 and share the - # c4_state slot space. - for pool in self.indexer_compress_state_pools: - if pool is not None: - add(pool) - return ptrs, lens, ilens - - # KV pools (4D PA_ND). - if self.swa_kv_pool is not None: - components.append( - (AscendStateType.DSV4_SWA, *kv_entry(self.swa_kv_pool.kv_buffer)) - ) - if self.c4_kv_pool is not None: - components.append( - (AscendStateType.DSV4_C4, *kv_entry(self.c4_kv_pool.kv_buffer)) - ) - if self.c128_kv_pool is not None: - components.append( - (AscendStateType.DSV4_C128, *kv_entry(self.c128_kv_pool.kv_buffer)) - ) - if self.c4_indexer_kv_pool is not None: - idx_bufs = list(self.c4_indexer_kv_pool.index_k_buffer) + list( - self.c4_indexer_kv_pool.index_scale_buffer - ) - components.append((AscendStateType.DSV4_INDEXER, *kv_entry(idx_bufs))) - - # Compress-state pools (paged, flat 2D). c4_state bundles attn-c4-state + - # indexer-c4-state (same req_to_token_c4_state slot space). - components.append( - (AscendStateType.DSV4_C4_STATE, *state_entry(4, include_indexer=True)) + """Main PD buffers addressed by the full KV page id.""" + buffers = ( + self.c4_kv_pool.kv_buffer + + self.c4_indexer_kv_pool.index_k_buffer + + self.c4_indexer_kv_pool.index_scale_buffer ) - components.append( - (AscendStateType.DSV4_C128_STATE, *state_entry(128, include_indexer=False)) + return ( + [buf.data_ptr() for buf in buffers], + [buf.nbytes for buf in buffers], + [buf[0].nbytes for buf in buffers], ) - # Drop empty components (e.g. a ratio with no layers) so every shipped - # component has non-zero item_lens; the set is identical on both sides. - return [c for c in components if c[1]] + def get_state_buf_infos(self) -> Tuple[List[int], List[int], List[int]]: + """GPU-compatible ``StateType.SWA`` component. + + SWA KV, C4 attention state and C4 indexer state retain separate buffers + but share the same SWA page/state index. + """ + data_ptrs: List[int] = [] + data_lens: List[int] = [] + item_lens: List[int] = [] + + for buf in self.swa_kv_pool.kv_buffer: + data_ptrs.append(buf.data_ptr()) + data_lens.append(buf.nbytes) + item_lens.append(buf[0].nbytes) + + for pools in (self.compress_state_pools, self.indexer_compress_state_pools): + for pool in pools: + if pool is None or pool.ratio != 4: + continue + state = pool.kv_score_buffer.kv_score + data_ptrs.append(state.data_ptr()) + data_lens.append(state.nbytes) + item_lens.append(state[0].nbytes * pool.ring_size) + + return data_ptrs, data_lens, item_lens + + def get_c128_kv_buf_infos(self) -> Tuple[List[int], List[int], List[int]]: + buffers = self.c128_kv_pool.kv_buffer + return ( + [buf.data_ptr() for buf in buffers], + [buf.nbytes for buf in buffers], + [buf[0].nbytes for buf in buffers], + ) def get_state_cache(self, layer_id: int, from_indexer: bool) -> torch.Tensor: - """fp32 ``[block_num, page_size, 2*coff*D]`` view of this layer's + """FP32 ``[block_num, ring_size, 2*coff*D]`` view of this layer's kv+score buffer — the fused compressor op (``torch.ops.custom.compressor``)'s ``state_cache`` argument.""" return self._get_state_pool(layer_id, from_indexer).state_cache_3d @@ -559,42 +496,43 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): cache = cache.unsqueeze(1) buf_flat[loc] = cache.to(buf_flat.dtype) - # ------------------------------------------------------------------ - # NPU port hooks — used by dsv4/{compressor,indexer}.py forward_npu. - # CompressStatePool stores a fused [kv | score] tensor; split is a last-dim slice. - # ------------------------------------------------------------------ - - def set_state_buffer( + def set_swa_key_buffer_radix_fused_norm_rope( self, layer_id: int, - loc: torch.Tensor, + swa_loc: torch.Tensor, kv: torch.Tensor, - score: torch.Tensor, - from_indexer: bool, + kv_weight: torch.Tensor, + eps: float, + freqs_cis: torch.Tensor, + positions: torch.Tensor, ) -> None: - # KVAndScore.kv_score is [..., 2*coff*head_dim] = [kv | score]. - kv_score = self._get_state_pool(layer_id, from_indexer).kv_score_buffer.kv_score - last_dim = kv_score.shape[-1] - half = last_dim // 2 - kv_view = kv.reshape(-1, half).to(kv_score.dtype) - score_view = score.reshape(-1, half).to(kv_score.dtype) - kv_score[loc, :half] = kv_view - kv_score[loc, half:] = score_view + kv_out = torch_npu.npu_rms_norm(kv, kv_weight, eps)[0] - def get_state_buffer( - self, - layer_id: int, - from_indexer: bool, - kv_indices: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: - kv_score = self._get_state_pool(layer_id, from_indexer).kv_score_buffer.kv_score - if kv_indices is not None: - kv_score = kv_score[kv_indices] - last_dim = kv_score.shape[-1] - half = last_dim // 2 - kv = kv_score[..., :half].unsqueeze(-2) # add num_kv_heads=1 axis - score = kv_score[..., half:].unsqueeze(-2) - return kv, score + rope_dim = freqs_cis.shape[-1] * 2 + + from sglang.srt.hardware_backend.npu.dsv4.dsv4_rope import Dsv4NpuRoPE + + cos, sin = Dsv4NpuRoPE.for_freqs(freqs_cis).get_cos_sin( + positions, + kv_out.dtype, + view_4d=True, + allow_build=True, + cache_dtype=torch.float32, + ) + Dsv4NpuRoPE.apply_rotary_mul_inplace( + kv_out.reshape(kv_out.shape[0], -1, kv_out.shape[-1]), + None, + cos, + sin, + qk_nope_dim=kv_out.shape[-1] - rope_dim, + ) + + safe_swa_loc = swa_loc.clamp_min(0).to(torch.int64) + self.set_swa_buffer( + layer_id, + safe_swa_loc, + kv_out, + ) def set_compress_buffer( self, @@ -648,22 +586,3 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): assert from_indexer, "only indexer compress pool has dequant scale" compress_layer_id = self.layer_mapping[layer_id].compress_layer_id return self.c4_indexer_kv_pool.get_index_scale(compress_layer_id) - - def translate_kv_loc_to_compress_state_loc( - self, - kv_loc: torch.Tensor, - compress_ratio: int, - ) -> torch.Tensor: - # Parent's ring-buffer hash is meaningless under the paged cache_mode=1 - # contract; returning a stale value would silently corrupt state. Fail loud. - raise RuntimeError( - "DSV4NPUTokenToKVPool.translate_kv_loc_to_compress_state_loc was " - "called, but the NPU fused compressor kernel uses a paged state " - "pool (cache_mode=1) and does not support ring-buffer state " - "addressing (cache_mode=2 is explicitly unsupported on Atlas A3). " - "Callers must consume out_cache_loc_dsv4.out_c{4,128}_state_loc " - "from the allocator bundle (set during alloc_extend/alloc_decode) " - "and read state_page_table from req_to_token_c{4,128}_state on " - "the DSV4NPUReqToTokenPool instead. See " - "hardware_backend/npu/dsv4_memory_pool.py for the rationale." - ) diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_req_to_token_pool.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_req_to_token_pool.py index a5efc5448..35369cc17 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_req_to_token_pool.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_req_to_token_pool.py @@ -1,33 +1,13 @@ """DSV4-NPU per-request mapping pool. -Subclass of ``ReqToTokenPool`` that adds five auxiliary per-request tables +Subclass of ``ReqToTokenPool`` that adds the one auxiliary per-request table needed by the DSV4 attention backend: - * ``req_to_token_swa`` — slot ids in the SWA full-pool view - * ``req_to_token_c4`` — slot ids in the c4 compressed-KV pool - * ``req_to_token_c128`` — slot ids in the c128 compressed-KV pool - * ``req_to_token_c4_state`` — c4 state-pool slot ids, 1 per raw token - * ``req_to_token_c128_state`` — c128 state-pool slot ids, 1 per raw token + * ``req_to_c128_sidecar`` — one page id per C128 physical page -Compressed KV pools store 1 slot per ``ratio`` raw tokens, so their per-req -table column count is ``max_context_len // ratio``. swa mirrors the raw -token count. Elements are token-level slot ids; the attention backend -converts to page ids via ``// page_size`` when constructing PA_ND block -tables. - -The c4/c128 STATE pools also have per-req tables here: the NPU fused -compressor uses a paged state pool (``cache_mode=1``), so each raw token's -state slot id is recorded (1 column per raw token) and the backend builds -``state_block_table = req_to_token_c{N}_state[req, ::page_size] // page_size`` -to feed the kernel. (The base class' ``translate_kv_loc_to_compress_state_loc`` -ring-hash is the CUDA-only path; it is disabled on NPU.) - -Memory cost example (size=64, max_context_len=32K): swa 8MB + c4 2MB + -c128 64KB ≈ 10MB extra on top of the base req_to_token (8MB). - -The tables are populated by the ``dsv4_common_hooks`` writers (driven from -``mem_cache/common.py``) immediately after a successful alloc_extend / -alloc_decode, using the per-pool slot indices returned in ``DSV4OutCacheLoc``. +C4 locations are derived from the base full-token table, and SWA locations use +the existing full-to-SWA mapping. The sidecar is populated by the existing +``dsv4_common_hooks`` flow from the slot indices in ``DSV4OutCacheLoc``. """ from __future__ import annotations @@ -37,6 +17,7 @@ import torch from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.disaggregation.decode import DecodeReqToTokenPool from sglang.srt.mem_cache.memory_pool import ReqToTokenPool +from sglang.srt.runtime_context import get_schedule from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter @@ -51,7 +32,11 @@ class DSV4ReqToTokenTablesMixin: """ def _init_dsv4_tables( - self, max_context_len: int, device: str, enable_memory_saver: bool + self, + max_context_len: int, + device: str, + enable_memory_saver: bool, + c128_page_size: int, ) -> None: memory_saver_adapter = TorchMemorySaverAdapter.create( enable=enable_memory_saver @@ -59,70 +44,83 @@ class DSV4ReqToTokenTablesMixin: # Back-ref to DSV4NPUTokenToKVPoolAllocator, wired via # register_dsv4_allocator after both exist, so free(req) can release - # c4/c128 pages. None at construction so base clear() runs safely. + # c128 pages. None at construction so base clear() runs safely. self._dsv4_allocator = None + self.c128_page_size = c128_page_size - # (name, columns). swa + state tables: 1 slot per raw token; c4/c128: - # 1 slot per `ratio` raw tokens. Init zero so unallocated columns map to - # block 0 (kernel skip sentinel cleared by NPUCompressStatePool). + group_tokens = 128 * c128_page_size with memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): - for name, cols in ( - ("req_to_token_swa", max_context_len), - ("req_to_token_c4", max(1, max_context_len // 4)), - ("req_to_token_c128", max(1, max_context_len // 128)), - ("req_to_token_c4_state", max_context_len), - ("req_to_token_c128_state", max_context_len), - ): - setattr( - self, - name, - torch.zeros( - (self._alloc_size, cols), - dtype=torch.int32, - device=device, - ), - ) - - # Per-pool write helpers, called by mem_cache/common.py after alloc, using - # slot indices from DSV4OutCacheLoc. Args: (req_pool_idx, token_offset), slot. - def write_swa(self, indices, values: torch.Tensor) -> None: - self.req_to_token_swa[indices] = values - - def write_c4(self, indices, values: torch.Tensor) -> None: - self.req_to_token_c4[indices] = values + self.req_to_c128_sidecar = torch.zeros( + ( + self._alloc_size, + max(1, (max_context_len + group_tokens - 1) // group_tokens), + ), + dtype=torch.int32, + device=device, + ) def write_c128(self, indices, values: torch.Tensor) -> None: - self.req_to_token_c128[indices] = values - - def write_c4_state(self, indices, values: torch.Tensor) -> None: - self.req_to_token_c4_state[indices] = values - - def write_c128_state(self, indices, values: torch.Tensor) -> None: - self.req_to_token_c128_state[indices] = values + req_pool_idx, token_slice = indices + page_size = self.c128_page_size + first_group = (token_slice.start + page_size - 1) // page_size + end_group = (token_slice.stop + page_size - 1) // page_size + if first_group == end_group: + return + groups = torch.arange(first_group, end_group, device=values.device) + pages = values[groups * page_size - token_slice.start] // page_size + prefix_pages = self.req_to_c128_sidecar[req_pool_idx, :end_group].clone() + prefix_pages[groups] = pages + self._dsv4_allocator.replace_req_c128_prefix(req_pool_idx, prefix_pages, self) def register_dsv4_allocator(self, allocator) -> None: """Wire the DSV4NPUTokenToKVPoolAllocator ref so ``free(req)`` can - release c4/c128 pool pages alongside the req_pool_idx slot.""" + release C128 KV pages.""" self._dsv4_allocator = allocator + def set_c128_prefix_pages(self, req, page_ids: torch.Tensor) -> None: + """Install pages returned by a Radix match. + + Prefix matching can happen before a request slot is allocated, so the + page ids are temporarily carried by ``Req`` and installed by ``alloc``. + """ + if req.req_pool_idx is None: + req.c128_prefix_page_ids = page_ids + return + self._dsv4_allocator.replace_req_c128_prefix( + int(req.req_pool_idx), page_ids, self + ) + + def alloc(self, reqs): + fresh = [req.req_pool_idx is None for req in reqs] + indices = super().alloc(reqs) + if indices is None: + return None + for is_fresh, req, req_pool_idx in zip(fresh, reqs, indices): + if is_fresh: + self.req_to_c128_sidecar[int(req_pool_idx)].zero_() + pages = getattr(req, "c128_prefix_page_ids", None) + if pages is not None: + self._dsv4_allocator.replace_req_c128_prefix( + int(req_pool_idx), pages, self + ) + req.c128_prefix_page_ids = None + return indices + def _dsv4_free(self, req) -> None: - # Trigger c4/c128 free via the allocator's unified free path. May be None + # Trigger C128 KV free/state clear via the allocator's unified path. May be None # between __init__ and register_dsv4_allocator — defensive None check. if self._dsv4_allocator is not None: self._dsv4_allocator.free(req=req, req_to_token_pool=self) class DSV4NPUReqToTokenPool(DSV4ReqToTokenTablesMixin, ReqToTokenPool): - """ReqToTokenPool extended with DSV4 SWA + c4/c128 per-req tables. + """ReqToTokenPool extended with the DSV4 C128 group sidecar mapping. Drop-in replacement for ReqToTokenPool when the model is DeepSeek-V4 on NPU. Selected by ``model_runner_kv_cache_mixin`` based on model arch + device. Non-DSV4 and non-NPU paths continue to use the base class. - The auxiliary tables are intentionally NOT zeroed on ``clear()``: they are - indexed only by active rows (via req_pool_idx) and only each row's - ``[:seq_len]`` prefix is read, so stale entries past kv_committed_len are - unreachable by the attention metadata builder. + Each freshly allocated request row is cleared before use. """ def __init__( @@ -133,7 +131,12 @@ class DSV4NPUReqToTokenPool(DSV4ReqToTokenTablesMixin, ReqToTokenPool): enable_memory_saver: bool, ): super().__init__(size, max_context_len, device, enable_memory_saver) - self._init_dsv4_tables(max_context_len, device, enable_memory_saver) + self._init_dsv4_tables( + max_context_len, + device, + enable_memory_saver, + get_schedule().c128_page_size, + ) def free(self, req): self._dsv4_free(req) @@ -141,7 +144,7 @@ class DSV4NPUReqToTokenPool(DSV4ReqToTokenTablesMixin, ReqToTokenPool): class DSV4NPUDecodeReqToTokenPool(DSV4ReqToTokenTablesMixin, DecodeReqToTokenPool): - """DecodeReqToTokenPool with the DSV4 swa/c4/c128(+state) per-req tables. + """DecodeReqToTokenPool with the C128 group sidecar mapping. The disagg-decode counterpart of DSV4NPUReqToTokenPool; DecodeReqToTokenPool pre-allocates extra req slots for in-flight prefill transfers. @@ -162,7 +165,12 @@ class DSV4NPUDecodeReqToTokenPool(DSV4ReqToTokenTablesMixin, DecodeReqToTokenPoo enable_memory_saver=enable_memory_saver, pre_alloc_size=pre_alloc_size, ) - self._init_dsv4_tables(max_context_len, device, enable_memory_saver) + self._init_dsv4_tables( + max_context_len, + device, + enable_memory_saver, + get_schedule().c128_page_size, + ) def free(self, req): self._dsv4_free(req) diff --git a/python/sglang/srt/hardware_backend/npu/extra_ops_loader.py b/python/sglang/srt/hardware_backend/npu/extra_ops_loader.py new file mode 100644 index 000000000..b7ac5e528 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/extra_ops_loader.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import logging +import os +import re +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + +import torch + +logger = logging.getLogger(__name__) + + +@dataclass +class OpLibSpec: + """Configuration for a standalone operator library loaded into ``torch.ops``.""" + + name: str # human-readable id, used in logs/errors + so_env: str # env var that points to the standalone .so path + namespace: str # torch.ops. the operators register into + required_ops: tuple[str, ...] + pre_load_imports: tuple[str, ...] = () # modules to import before loading + + +class TorchOpLoader: + """Load a standalone .so into ``torch.ops`` and validate its operators.""" + + def __init__(self, spec: OpLibSpec) -> None: + self._spec = spec + self._loaded_library: Optional[Path] = None + + def _missing_ops(self) -> list[str]: + namespace = getattr(torch.ops, self._spec.namespace, None) + if namespace is None: + return list(self._spec.required_ops) + return [op for op in self._spec.required_ops if not hasattr(namespace, op)] + + def registered(self) -> bool: + """Return whether the required operators are already registered.""" + return not self._missing_ops() + + def _resolve_so_path(self) -> Path: + explicit = os.environ.get(self._spec.so_env) + if not explicit: + raise RuntimeError( + f"The {self._spec.name} operators are not registered. Set " + f"{self._spec.so_env} to the standalone .so library path." + ) + path = Path(explicit).expanduser().resolve() + if not path.is_file(): + raise RuntimeError(f"{self._spec.so_env} points to a missing file: {path}") + return path + + def _validate_python_abi(self, library_path: Path) -> None: + abi_match = re.search(r"\.cpython-(\d+)-", library_path.name) + current_abi = f"{sys.version_info.major}{sys.version_info.minor}" + if abi_match is not None and abi_match.group(1) != current_abi: + raise RuntimeError( + f"{library_path} was built for CPython {abi_match.group(1)}, " + f"but SGLang is running CPython {current_abi}. Rebuild the " + "extension with the SGLang Python/Torch/torch-npu environment." + ) + + def initialize(self) -> Optional[Path]: + """Register the operators before backend execution. + + Idempotent: returns ``None`` if the operators are already registered + (e.g. by another package). Otherwise loads the standalone .so pointed + to by ``so_env`` into ``torch.ops`` and validates the required operators. + + Returns the loaded library path when this call loaded it, else ``None``. + """ + if self.registered(): + return None + if self._loaded_library is not None: + missing = self._missing_ops() + raise RuntimeError( + f"Loaded {self._loaded_library}, but required " + f"{self._spec.namespace} operators are missing: {missing}." + ) + + for module in self._spec.pre_load_imports: + __import__(module) # noqa: F401 side-effect imports (e.g. torch_npu) + + library_path = self._resolve_so_path() + self._validate_python_abi(library_path) + try: + torch.ops.load_library(str(library_path)) + except Exception as exc: + raise RuntimeError( + f"Failed to load the {self._spec.name} operator library " + f"{library_path}. Ensure its dependent CANN/custom-op libraries " + "are visible through LD_LIBRARY_PATH and the Ascend OPP setup." + ) from exc + + missing = self._missing_ops() + if missing: + raise RuntimeError( + f"Loaded {library_path}, but required " + f"{self._spec.namespace} operators are missing: {missing}." + ) + + self._loaded_library = library_path + logger.info("Registered %s operators from %s", self._spec.name, library_path) + return library_path + + +def initialize_dspark_sparse_attn_ops() -> Optional[Path]: + """Register the DSpark sparse-attention ops before backend execution.""" + spec = OpLibSpec( + name="DSpark sparse-attention", + so_env="SGLANG_DSPARK_EXTRA_OPS_SO", + namespace="_C_ascend", + required_ops=( + "npu_sparse_attn_sharedkv_metadata", + "npu_sparse_attn_sharedkv", + ), + pre_load_imports=("torch_npu",), + ) + return TorchOpLoader(spec).initialize() diff --git a/python/sglang/srt/layers/quantization/modelslim/modelslim.py b/python/sglang/srt/layers/quantization/modelslim/modelslim.py index 65f74a306..d90b3d7fb 100644 --- a/python/sglang/srt/layers/quantization/modelslim/modelslim.py +++ b/python/sglang/srt/layers/quantization/modelslim/modelslim.py @@ -138,6 +138,66 @@ class ModelSlimConfig(QuantizationConfig): "forward_npu", [npu_wrapper_rmsnorm_forward], ) + # DSpark checkpoint weights use mtp..*, while the runtime draft + # model constructs canonical modules under stages..*. Keep + # this transformation in sync with + # DeepseekV4ForCausalLMDSpark._remap_dspark_weight_name. Merely + # replacing ``mtp`` with ``stages`` is insufficient: it silently + # misses ModelSlim lookups such as stages.0.mlp.experts and + # stages.0.self_attn. + dspark_quant_aliases = {} + for name, scheme in quant_config.items(): + if not isinstance(name, str) or not name.startswith("mtp."): + continue + + parts = name.split(".", 2) + if len(parts) != 3: + continue + + stage_id, rest = parts[1], parts[2] + if not stage_id.isdigit(): + continue + + # The draft model attaches the target model's shared embedding and + # LM head; mtp-local copies are not runtime draft modules. + if rest.startswith(("embed.", "embed_tokens.", "head.", "lm_head.")): + continue + if rest.startswith("markov_head."): + alias = f"markov_head.{rest[len('markov_head.'):]}" + elif rest.startswith("confidence_head."): + alias = f"confidence_head.{rest[len('confidence_head.'):]}" + else: + mapped_rest = rest + if mapped_rest.startswith("attn."): + mapped_rest = "self_attn." + mapped_rest.removeprefix("attn.") + elif mapped_rest.startswith("ffn."): + mapped_rest = "mlp." + mapped_rest.removeprefix("ffn.") + elif mapped_rest.startswith("attn_norm."): + mapped_rest = "input_layernorm." + mapped_rest.removeprefix( + "attn_norm." + ) + elif mapped_rest.startswith("ffn_norm."): + mapped_rest = ( + "post_attention_layernorm." + + mapped_rest.removeprefix("ffn_norm.") + ) + mapped_rest = mapped_rest.replace(".w1.", ".gate_proj.") + mapped_rest = mapped_rest.replace(".w2.", ".down_proj.") + mapped_rest = mapped_rest.replace(".w3.", ".up_proj.") + mapped_rest = mapped_rest.replace(".gate.tid2eid", ".topk.tid2eid") + mapped_rest = mapped_rest.replace( + ".gate.bias", ".gate.e_score_correction_bias" + ) + alias = f"stages.{stage_id}.{mapped_rest}" + + dspark_quant_aliases[alias] = scheme + + quant_config = { + **dspark_quant_aliases, + **quant_config, + } + + self.quant_description = quant_config def update_packed_modules_mapping(self, mapping: Dict[str, List[str]]) -> None: self.packed_modules_mapping.update(mapping) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index ee5445d5e..e927daeff 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -81,9 +81,6 @@ from sglang.srt.disaggregation.decode_schedule_batch_mixin import ( from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST, DisaggregationMode from sglang.srt.dllm.mixin.req import ReqDllmMixin from sglang.srt.environ import envs -from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( - maybe_evict_dsv4_state, -) from sglang.srt.managers.embed_types import PositionalEmbeds from sglang.srt.managers.scheduler_components.new_token_ratio_tracker import ( NewTokenRatioTracker, @@ -1164,6 +1161,7 @@ class Req(ReqDllmMixin): # kv_send(req.input_ids[req.start_send_idx:req.extend_range.end]) # start_send_idx = req.extend_range.end self.start_send_idx: int = 0 + self.disagg_decode_prefix_len: int = 0 # For overlap schedule, we delay the kv transfer until `process_batch_result_disagg_prefill` rather than `process_prefill_chunk` in non-overlap # This is because kv is not ready in `process_prefill_chunk`. @@ -2073,8 +2071,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # The output locations of the KV cache out_cache_loc: torch.Tensor = None # shape: [b], int64 - # DSV4-NPU: per-pool slot bundle from DSV4NPUTokenToKVPoolAllocator (None - # elsewhere); c4/c128 state lens ride on ``batch.dsv4_state_lens``. + # DSV4-NPU: KV-only per-pool slot bundle from + # DSV4NPUTokenToKVPoolAllocator (None elsewhere). out_cache_loc_dsv4: Optional[Any] = None # For hybrid GDN prefix cache @@ -3330,10 +3328,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): ): self._evict_swa(req, req.seqlen - 1) - # DSV4-NPU only (no-op elsewhere): the small paged compress-state - # pool must drain every decode step, independent of SWA cadence. - maybe_evict_dsv4_state(self, req, req.seqlen - 1) - # Once the decode position has moved past the sliding window, # the SWA portion of the prefill-time tree lock is no longer # needed by this request. Convert it from protected to diff --git a/python/sglang/srt/mem_cache/allocation.py b/python/sglang/srt/mem_cache/allocation.py index f29b74c84..fc5428a34 100644 --- a/python/sglang/srt/mem_cache/allocation.py +++ b/python/sglang/srt/mem_cache/allocation.py @@ -47,7 +47,6 @@ if _is_cpu: if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req, ScheduleBatch - from sglang.srt.model_executor.forward_batch_info import DSV4StateLens logger = logging.getLogger(__name__) @@ -172,30 +171,6 @@ def alloc_token_slots( return out_cache_loc -def _compute_dsv4_state_lens(batch, *, is_decode: bool): - """Per-req c{4,128}_state pool alloc lens (``DSV4StateLens``) for this step. - None on CUDA / non-V4 paths (allocator has no ``compute_dsv4_state_lens_*``). - """ - allocator = batch.token_to_kv_pool_allocator - if not hasattr(allocator, "compute_dsv4_state_lens_extend"): - return None - from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( - maybe_evict_dsv4_state, - ) - - if is_decode: - for req in batch.reqs: - maybe_evict_dsv4_state(batch, req, req.seqlen - 1) - return allocator.compute_dsv4_state_lens_decode(batch.reqs) - prefix_lens = batch.prefix_lens - for req, prefix_len in zip(batch.reqs, prefix_lens): - if prefix_len > 0: - maybe_evict_dsv4_state(batch, req, prefix_len) - return allocator.compute_dsv4_state_lens_extend( - batch.reqs, batch.seq_lens_cpu.tolist(), prefix_lens - ) - - def alloc_paged_token_slots_extend( tree_cache: BasePrefixCache, prefix_lens: torch.Tensor, @@ -205,7 +180,6 @@ def alloc_paged_token_slots_extend( last_loc: torch.Tensor, extend_num_tokens: int, req_pool_indices: Optional[torch.Tensor] = None, - dsv4_state_lens: Optional[DSV4StateLens] = None, batch=None, ): # Over estimate the number of tokens: assume each request needs a new page. @@ -213,15 +187,13 @@ def alloc_paged_token_slots_extend( num_tokens = extend_num_tokens + len(seq_lens_cpu) * allocator.page_size evict_from_tree_cache(tree_cache, num_tokens) - is_dsv4 = req_pool_indices is not None and hasattr(allocator, "c4_attn_allocator") + is_dsv4 = req_pool_indices is not None and hasattr(allocator, "c128_attn_allocator") extra_alloc_kwargs = {} if is_dsv4: extra_alloc_kwargs["req_pool_indices"] = req_pool_indices - # Per-call per-req tables for the c-pool / state last_loc lookup. + # Per-call per-req table for the C128 KV last_loc lookup. if batch is not None: extra_alloc_kwargs["req_to_token_pool"] = batch.req_to_token_pool - if dsv4_state_lens is not None: - extra_alloc_kwargs["dsv4_state_lens"] = dsv4_state_lens out = allocator.alloc_extend( prefix_lens, @@ -370,7 +342,6 @@ def alloc_for_extend( last_loc=torch.cat(last_loc), extend_num_tokens=batch.extend_num_tokens, req_pool_indices=req_pool_indices_device, - dsv4_state_lens=_compute_dsv4_state_lens(batch, is_decode=False), batch=batch, ) @@ -466,7 +437,6 @@ def _alloc_extend_loc_with_kv_reuse( last_loc=torch.cat(last_loc), extend_num_tokens=alloc_extend_num_tokens, req_pool_indices=req_pool_indices_device, - dsv4_state_lens=_compute_dsv4_state_lens(batch, is_decode=False), batch=batch, ) @@ -497,7 +467,6 @@ def alloc_paged_token_slots_decode( last_loc: torch.Tensor, token_per_req: int = 1, req_pool_indices: Optional[torch.Tensor] = None, - dsv4_state_lens: Optional[DSV4StateLens] = None, batch=None, ) -> torch.Tensor: """Allocate paged KV cache for decode batch.""" @@ -506,17 +475,15 @@ def alloc_paged_token_slots_decode( num_tokens = len(seq_lens) * allocator.page_size evict_from_tree_cache(tree_cache, num_tokens) - # DSV4-NPU allocator also needs req_pool_indices + per-req state lens and + # DSV4-NPU allocator also needs req_pool_indices for C128 KV allocation and # returns a DSV4OutCacheLoc bundle; hasattr-gated so others stay unchanged. - is_dsv4 = req_pool_indices is not None and hasattr(allocator, "c4_attn_allocator") + is_dsv4 = req_pool_indices is not None and hasattr(allocator, "c128_attn_allocator") extra_alloc_kwargs = {} if is_dsv4: extra_alloc_kwargs["req_pool_indices"] = req_pool_indices - # Per-call per-req tables for the last_loc lookup. + # Per-call per-req C128 table for the last_loc lookup. if batch is not None: extra_alloc_kwargs["req_to_token_pool"] = batch.req_to_token_pool - if dsv4_state_lens is not None: - extra_alloc_kwargs["dsv4_state_lens"] = dsv4_state_lens out = allocator.alloc_decode(seq_lens, seq_lens_cpu, last_loc, **extra_alloc_kwargs) @@ -571,7 +538,6 @@ def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor: last_loc=last_loc, token_per_req=token_per_req, req_pool_indices=batch.req_pool_indices, - dsv4_state_lens=_compute_dsv4_state_lens(batch, is_decode=True), batch=batch, ) diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 5ef8e5ddf..73341b7f3 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -67,6 +67,9 @@ class InsertParams: # Mamba specific mamba_value: Optional[torch.Tensor] = None + # DSV4 NPU C128 sidecar pages, one page id per physical C128 page group. + c128_value: Optional[torch.Tensor] = None + # SWA specific prev_prefix_len: int = 0 swa_evicted_seqlen: int = 0 diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 19e1c3e12..88f6a8e30 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -10,9 +10,6 @@ from sglang.kernels.ops.memory.common import ( _get_last_loc_safe_kernel as _get_last_loc_safe_kernel, ) from sglang.kernels.ops.memory.common import get_last_loc_kernel as get_last_loc_kernel -from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( - maybe_evict_dsv4_state_on_swa, -) from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool @@ -96,9 +93,6 @@ def free_swa_out_of_window_slots( req.req_pool_idx, req.kv.swa_evicted_seqlen : new_swa_evicted_seqlen ] token_to_kv_pool_allocator.free_swa(free_slots) - maybe_evict_dsv4_state_on_swa( - token_to_kv_pool_allocator, req_to_token_pool, req, new_swa_evicted_seqlen - ) req.kv.swa_evicted_seqlen = new_swa_evicted_seqlen diff --git a/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py b/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py index bb0f703e3..4c5dd567b 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py @@ -8,11 +8,10 @@ import torch from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.mem_cache.utils import maybe_init_custom_mem_pool -from sglang.srt.utils import is_hip, is_npu +from sglang.srt.utils import is_hip from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter _is_hip = is_hip() -_is_npu = is_npu() def _lcm(a: int, b: int) -> int: @@ -95,11 +94,14 @@ class CompressStatePool: online: bool = False, swa_page_size: int = 0, online_mtp_max_draft_tokens: int = 0, + state_cache_page_size: int = 1, ): self.ratio = ratio self.ring_size = ring_size self.swa_page_size = swa_page_size + self.page_size = state_cache_page_size self.enable_memory_saver = enable_memory_saver + self.online = online self.online_mtp_state_slot_offset = 0 self.online_mtp_max_draft_tokens = 0 @@ -115,11 +117,12 @@ class CompressStatePool: last_dim = 3 * head_dim else: self._size = size + self.ring_size + 1 - # Pad to lcm(ratio, page_size) so the flat buffer reshapes cleanly into - # [block_num, page_size, last_dim] for the fused compressor op; page_size=1 falls back to ratio-only padding. - pad_to = ( - _lcm(ratio, swa_page_size) if (swa_page_size > 1 and _is_npu) else ratio - ) + # The common GPU pool is flat by default. A backend that also needs + # a physical 3-D cache view can request its second-axis page size; + # allocation and ring ownership still stay in this shared class. + pad_to = ratio + if state_cache_page_size > 1: + pad_to = _lcm(pad_to, state_cache_page_size) self._size = (self._size + pad_to - 1) // pad_to * pad_to self._logical_size = self._size last_dim = 2 * (1 + overlap) * head_dim @@ -148,10 +151,9 @@ class CompressStatePool: :class:`KVAndScore`. Sets ``self.memory_saver_adapter``, ``self.custom_mem_pool`` and ``self.kv_score_buffer``. - Subclasses (e.g. :class:`NPUCompressStatePool`) that compute a - different ``self._size`` reuse this instead of duplicating the - allocation boilerplate. Requires ``self._size`` and ``self.last_dim`` - to be set already. + The shared constructor computes ``self._size`` and ``self.last_dim`` + before entering this helper. Backend subclasses should normally call + that constructor instead of duplicating this allocation path. """ self.memory_saver_adapter = TorchMemorySaverAdapter.create( enable=enable_memory_saver diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 1b4dfd78c..18dedb76d 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -1082,37 +1082,18 @@ class KVCacheConfigurator: else: compression_ratios = self.model_config.compress_ratios - # NPU + DSV4 → paged-state subclass: the fused compressor kernel - # needs cache_mode=1 (paged); Atlas A3 rejects cache_mode=2 (ring), - # so the CUDA ring-buffer state path can't be shared. CUDA keeps - # DeepSeekV4TokenToKVPool unchanged; NPU recomputes state sizes below. + # NPU keeps its PA_ND KV-pool subclass, while Compressor state sizing + # follows the same fixed ring ownership as GPU. Do not replace the + # configurator's C4-SWA/C128-request budgets with a paged allocator + # estimate: Atlas A3 cache_mode=2 consumes explicit flat state_locs. if _is_npu: from sglang.srt.hardware_backend.npu.dsv4.dsv4_memory_pool import ( DSV4NPUTokenToKVPool, - npu_state_pool_size, ) pool_cls = DSV4NPUTokenToKVPool - # Recompute state pool sizes for the NPU paged formula (CUDA's - # ring sizes are dropped here). Tail-only allocation keeps the - # per-req-budget formula sufficient at any prefill length: long - # prompts allocate only ``tail+128`` (c4) / ``tail`` (c128) - # slots (tail = seq_len % 128), and decode is drained by - # sliding eviction in ``ScheduleBatch._evict_swa``. - c4_state_pool_size = npu_state_pool_size( - ratio=4, - page_size=get_schedule().page_size, - max_num_reqs=max_running_requests, - ) - c128_state_pool_size = npu_state_pool_size( - ratio=128, - page_size=get_schedule().page_size, - max_num_reqs=max_running_requests, - ) else: pool_cls = DeepSeekV4TokenToKVPool - c4_state_pool_size = c4_state_pool_size - c128_state_pool_size = c128_state_pool_size token_to_kv_pool = pool_cls( max_num_reqs=max_running_requests, diff --git a/python/sglang/srt/mem_cache/registry.py b/python/sglang/srt/mem_cache/registry.py index 43386f2c7..e979e08ad 100644 --- a/python/sglang/srt/mem_cache/registry.py +++ b/python/sglang/srt/mem_cache/registry.py @@ -176,6 +176,17 @@ def _create_unified_radix_cache( if ctx.is_hybrid_ssm: tree_components.append(ComponentType.MAMBA) + if hasattr(params.req_to_token_pool, "req_to_c128_sidecar"): + from sglang.srt.hardware_backend.npu.dsv4.c128_sidecar_component import ( + C128SidecarComponent, + ) + + tree_components.append(ComponentType.C128) + params.component_registry_override = { + **(params.component_registry_override or {}), + ComponentType.C128: C128SidecarComponent, + } + params.tree_components = tuple(tree_components) if use_mlx() and ctx.is_hybrid_ssm: from sglang.srt.hardware_backend.mlx.kv_cache.auxiliary_state import ( diff --git a/python/sglang/srt/mem_cache/unified_cache/component_type.py b/python/sglang/srt/mem_cache/unified_cache/component_type.py index 625e4d782..f09a63f0c 100644 --- a/python/sglang/srt/mem_cache/unified_cache/component_type.py +++ b/python/sglang/srt/mem_cache/unified_cache/component_type.py @@ -9,6 +9,7 @@ class ComponentType(int, Enum): FULL = 0 SWA = 1 MAMBA = 2 + C128 = 3 def __str__(self) -> str: # keep human-readable logging return self.name.lower() diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index b78f906b0..1ca6e6583 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -454,6 +454,7 @@ class UnifiedRadixCache(BasePrefixCache): ComponentType.FULL: params.num_tokens, ComponentType.SWA: params.swa_num_tokens, ComponentType.MAMBA: params.mamba_num, + ComponentType.C128: 0, } self._evict_components(request_by_type, tracker) @@ -808,7 +809,7 @@ class UnifiedRadixCache(BasePrefixCache): result = self.insert(insert_params) # Match prefix - match_result = self.match_prefix(MatchPrefixParams(key=radix_key)) + match_result = self.match_prefix(MatchPrefixParams(key=radix_key, req=req)) new_indices = match_result.device_indices new_last_node = match_result.last_device_node new_prefix_len = result.prefix_len diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 82c88b01d..9bf587240 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -301,59 +301,23 @@ def compute_local_num_token_non_padded_cpu( class DSV4OutCacheLoc: """Per-forward-pass KV cache allocation for DeepSeek-V4 on NPU. - Bundles slot indices for full/SWA pools, the two compressed-KV pools - (c4/c128), and the two compressed-state pools (c4_state/c128_state). + Bundles slot indices for full/SWA pools and the two compressed-KV pools + (C4/C128). Compressor state uses fixed ring storage and explicit + ``state_loc`` metadata, so it is not part of the token-allocation bundle. Populated by the NPU V4 allocator (DSV4NPUTokenToKVPoolAllocator) when the model is DeepSeek-V4 on NPU; left as ``None`` on ForwardBatch - otherwise. CUDA's DSV4 path doesn't construct this bundle (state is - derived via translate_kv_loc_to_compress_state_loc there). + otherwise. All fields are token-level slot ids in their respective pools (NOT page ids). Attention backends convert to page ids via ``// page_size`` when constructing PA_ND block tables. - State fields default to ``None`` so the bundle is constructible from - paths that allocate KV but not state (or vice versa); the NPU allocator - fills all six on real alloc, CUDA paths leave state ones None and use - the ring-hash translation instead. """ out_full_loc: torch.Tensor out_swa_loc: torch.Tensor out_c4_loc: torch.Tensor out_c128_loc: torch.Tensor - out_c4_state_loc: Optional[torch.Tensor] = None - out_c128_state_loc: Optional[torch.Tensor] = None - - -@dataclass -class DSV4StateLens: - """Per-extend/decode c4/c128 compress-state pool allocation lens (DSV4-NPU). - - Built by ``ScheduleBatch._compute_dsv4_state_lens_{extend,decode}`` and - threaded through ``mem_cache/common.py`` to - ``DSV4NPUTokenToKVPoolAllocator.alloc_{extend,decode}``, which consumes: - - * ``c{4,128}_prefix_lens`` / ``..._cpu`` — per-req prev cumulative - state-slot count (the paged allocator's ``prefix`` contract). - * ``c{4,128}_seq_lens`` / ``..._cpu`` — per-req new cumulative count. - * ``c{4,128}_extend_num_tokens`` — total new state slots this step. - - Replaces the 10 loose ``c{4,128}_state_*`` kwargs the allocator used to - take: scheduler only produces this object, common only forwards it, the - allocator only consumes it. - """ - - c4_prefix_lens: torch.Tensor - c4_prefix_lens_cpu: torch.Tensor - c4_seq_lens: torch.Tensor - c4_seq_lens_cpu: torch.Tensor - c4_extend_num_tokens: int - c128_prefix_lens: torch.Tensor - c128_prefix_lens_cpu: torch.Tensor - c128_seq_lens: torch.Tensor - c128_seq_lens_cpu: torch.Tensor - c128_extend_num_tokens: int @dataclass diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index a557671b3..0c48e90dc 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -734,7 +734,10 @@ class MQALayer(MqaAttentionBase): self.register_buffer("cos_cache", cos_cache, persistent=False) self.register_buffer("sin_cache", sin_cache, persistent=False) - if envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get() and alt_streams is not None: + if alt_streams is not None and ( + (_is_cuda and envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get()) + or (_is_npu and envs.SGLANG_NPU_USE_MULTI_STREAM.get()) + ): self.alt_streams = alt_streams[:3] self.alt_streams_indexer = alt_streams[-2:] else: @@ -957,6 +960,108 @@ class MQALayer(MqaAttentionBase): return q + def _forward_prepare_multi_stream_npu( + self, + x: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + attn_backend, + q_out: Optional[torch.Tensor] = None, + x_quant=None, + ) -> torch.Tensor: + # NPU multi-stream: KV on stream_kv, Q on stream_q, overlapped with + # indexer/compressor on current. rope is split; the kv-only call passes + # kv.unsqueeze(1) as q_rope so the op sees [T,1,1,head_dim] like the + # fused path. + assert self.alt_streams is not None + current_stream = torch.npu.current_stream() + stream_kv = self.alt_streams[0] + stream_q = self.alt_streams[1] + stream_kv.wait_stream(current_stream) + stream_q.wait_stream(current_stream) + + x_linear = x_quant if x_quant is not None else x + qkv_a: Optional[torch.Tensor] = None + qkv_a_ready = None + if self.fuse_wqa_wkv: + qkv_a, _ = self.wqkv_a(x_linear) + qkv_a_ready = current_stream.record_event() + if qkv_a is not None: + q_lora = qkv_a[..., : self.q_lora_rank] + else: + q_lora, _ = self.wq_a(x_linear) + q_lora = self.q_norm(q_lora) + q_lora_ready = current_stream.record_event() + + # KV block on stream_kv. + with torch.npu.stream(stream_kv): + if qkv_a_ready is not None: + stream_kv.wait_event(qkv_a_ready) + if qkv_a is not None: + kv = qkv_a[..., self.q_lora_rank :] + else: + kv, _ = self.wkv(x) + kv = self.kv_norm(kv) + cos4_k, sin4_k = self._get_npu_rope_position_cache( + positions, kv.dtype, inverse=False + ) + Dsv4NpuRoPE.apply_rotary_mul_inplace( + kv.unsqueeze(1), + None, + cos4_k, + sin4_k, + qk_nope_dim=self.qk_nope_head_dim, + ) + attn_backend.store_cache( + layer_id=self.layer_id, + swa_k=kv, + forward_batch=forward_batch, + ) + + # Q block on stream_q (needs only q_lora). + with torch.npu.stream(stream_q): + stream_q.wait_event(q_lora_ready) + q, _ = self.wq_b(q_lora) + q = q.view(-1, self.n_local_heads, self.head_dim) + _dummy = q.new_ones(q.shape[-1]) + q = torch_npu.npu_rms_norm(q, _dummy, self.eps)[0] + cos4_q, sin4_q = self._get_npu_rope_position_cache( + positions, q.dtype, inverse=False + ) + Dsv4NpuRoPE.apply_rotary_mul_inplace( + q, + None, + cos4_q, + sin4_q, + qk_nope_dim=self.qk_nope_head_dim, + ) + if q_out is not None: + q_out.copy_(q) + q.record_stream(stream_q) + + del qkv_a + + # Indexer + compressor: serial on current. + if self.indexer is not None: + self.indexer( + x=x, + q_lora=q_lora, + forward_batch=forward_batch, + attn_backend=attn_backend, + ) + if self.compressor is not None: + attn_backend.forward_core_compressor( + x, + forward_batch, + self.layer_id, + self.compressor, + ) + + # Join stream_kv + stream_q before downstream attention. + current_stream.wait_stream(stream_kv) + current_stream.wait_stream(stream_q) + return q + def _forward_prepare_multi_stream_hip( self, x: torch.Tensor, @@ -1306,6 +1411,12 @@ class MQALayer(MqaAttentionBase): ) and not (self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch)) and not (_is_hip and self.compressor is None) + ) or ( + _is_npu + and envs.SGLANG_NPU_USE_MULTI_STREAM.get() + and self.alt_streams is not None + and x.shape[0] <= self._multi_stream_bs_limit + and not forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed() ) tp_slice, q_padded, q_out = slice(None), None, None @@ -1339,6 +1450,15 @@ class MQALayer(MqaAttentionBase): q_out, x_quant=x_quant, ) + elif _is_npu: + q = self._forward_prepare_multi_stream_npu( + x, + positions, + forward_batch, + attn_backend, + q_out, + x_quant=x_quant, + ) else: q = self._forward_prepare_multi_stream( x, @@ -1506,7 +1626,7 @@ class DeepseekV4DecoderLayer(nn.Module): layer_id=layer_id, quant_config=quant_config, prefix=add_prefix("self_attn", prefix), - alt_streams=None if _is_npu else alt_streams, + alt_streams=alt_streams, compress_ratio_override=compress_ratio_override, ) moe_alt_stream = ( @@ -2353,7 +2473,7 @@ class DeepseekV4Model(nn.Module): or (_is_npu and envs.SGLANG_NPU_USE_MULTI_STREAM.get()) ) device_module = torch.get_device_module() - num_alt_streams = 5 if _is_cuda else 2 + num_alt_streams = 5 if (_is_cuda or _is_npu) else 2 self.alt_streams = ( [device_module.Stream() for _ in range(num_alt_streams)] if use_stream_pool diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py index c93c84fad..f57c069c6 100644 --- a/python/sglang/srt/models/deepseek_v4_dspark.py +++ b/python/sglang/srt/models/deepseek_v4_dspark.py @@ -18,12 +18,16 @@ from sglang.kernels.ops.speculative.dspark.dspark_draft_model import ( ) from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config from sglang.srt.environ import envs +from sglang.srt.layers.dp_attention import is_dp_attention_enabled from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.moe.utils import is_shared_experts_fusion_disabled from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.radix_attention import RadixAttention -from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_context import get_token_to_kv_pool @@ -53,7 +57,7 @@ from sglang.srt.speculative.ragged_verify import ( RaggedVerifyMode, read_ragged_verify_mode, ) -from sglang.srt.utils import add_prefix, is_blackwell_supported +from sglang.srt.utils import add_prefix, is_blackwell_supported, is_npu from sglang.srt.utils.invariants import Bucket, InClosedRange, Invariant, expect logger = logging.getLogger(__name__) @@ -64,6 +68,7 @@ _PAD_NUM_HEADS = 64 _CONFIDENCE = Invariant( "dspark.model.confidence", Bucket.GUARD, InClosedRange(0.0, 1.0) ) +_is_npu = is_npu() def apply_rotary_emb( @@ -126,6 +131,12 @@ class DSparkAttention(MqaAttentionBase): self._use_fast_kernel = envs.SGLANG_DSPARK_FAST_KERNEL.get() self.alt_streams = alt_streams self._multi_stream_bs_limit = 128 if is_blackwell_supported() else 64 + if _is_npu: + self.register_buffer( + "_q_post_norm_weight", + torch.ones(self.head_dim), + persistent=False, + ) def kv_proj_only(self, x: torch.Tensor) -> torch.Tensor: kv, _ = self.wkv(x) @@ -180,10 +191,36 @@ class DSparkAttention(MqaAttentionBase): fused_q_norm_rope(q, q_out, self.eps, self.freqs_cis, positions) return q_out else: - q = q * torch.rsqrt( - q.float().square().mean(-1, keepdim=True) + self.eps - ).to(q.dtype) - apply_rotary_emb(q[..., -self.rope_head_dim :], self.freqs_cis[positions]) + if _is_npu: + import torch_npu + + from sglang.srt.hardware_backend.npu.dsv4.dsv4_rope import ( + Dsv4NpuRoPE, + ) + + q = torch_npu.npu_rms_norm(q, self._q_post_norm_weight, self.eps)[0] + cos4, sin4 = Dsv4NpuRoPE.for_freqs(self.freqs_cis).get_cos_sin( + positions, + q.dtype, + view_4d=True, + inverse=False, + allow_build=True, + cache_dtype=torch.float32, + ) + Dsv4NpuRoPE.apply_rotary_mul_inplace( + q, + None, + cos4, + sin4, + qk_nope_dim=q.shape[-1] - self.rope_head_dim, + ) + else: + q = q * torch.rsqrt( + q.float().square().mean(-1, keepdim=True) + self.eps + ).to(q.dtype) + apply_rotary_emb( + q[..., -self.rope_head_dim :], self.freqs_cis[positions] + ) if q_out is not None: q_out.copy_(q) return q_out @@ -195,6 +232,10 @@ class DSparkAttention(MqaAttentionBase): hidden_states: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: + + if _is_npu and forward_batch.forward_mode.is_idle(): + return torch.zeros_like(hidden_states) + from sglang.srt.model_executor.forward_context import get_attn_backend pool = _resolve_dspark_pool() @@ -255,6 +296,7 @@ class DSparkAttention(MqaAttentionBase): attn_sink=attn_sink, save_kv_cache=False, ) + if o.shape[1] != self.n_local_heads: o = o[:, : self.n_local_heads, :] @@ -262,6 +304,24 @@ class DSparkAttention(MqaAttentionBase): fused_rope_inplace( o[..., -rd:], None, self.freqs_cis, positions=positions, inverse=True ) + elif _is_npu: + from sglang.srt.hardware_backend.npu.dsv4.dsv4_rope import Dsv4NpuRoPE + + cos4, sin4 = Dsv4NpuRoPE.for_freqs(self.freqs_cis).get_cos_sin( + positions, + o.dtype, + view_4d=True, + inverse=True, + allow_build=True, + cache_dtype=torch.float32, + ) + Dsv4NpuRoPE.apply_rotary_mul_inplace( + o, + None, + cos4, + sin4, + qk_nope_dim=o.shape[-1] - rd, + ) else: apply_rotary_emb(o[..., -rd:], self.freqs_cis[positions], inverse=True) @@ -581,6 +641,26 @@ class DSparkV4Stage(DeepseekV4DecoderLayer): class DeepseekV4ForCausalLMDSpark(nn.Module): + # ModelSlim NPU checkpoints carry QuaRot-aligned, MTP-local + # embedding/head weights. The native CUDA path keeps the original DSpark + # behavior and shares the target model's vocabulary modules. + uses_own_vocab_modules = _is_npu + + @classmethod + def shared_experts_fusion_disable_reason( + cls, + hf_config, + quant_config, + ): + if _is_npu: + return ( + "NPU DSpark ModelSlim weight loading does not support mapping " + "shared experts into fused expert slots." + ) + + return DeepseekV4ForCausalLM.shared_experts_fusion_disable_reason( + hf_config, quant_config + ) @classmethod def shared_experts_fusion_disable_reason(cls, hf_config, quant_config): @@ -662,10 +742,26 @@ class DeepseekV4ForCausalLMDSpark(nn.Module): self.norm_eps = float(config.rms_norm_eps) self.hc_eps = float(config.hc_eps) - self.embed_tokens: Optional[nn.Module] = None - self.lm_head: Optional[nn.Module] = None + if self.uses_own_vocab_modules: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=add_prefix("embed_tokens", prefix), + enable_tp=not is_dp_attention_enabled(), + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + prefix=add_prefix("lm_head", prefix), + use_attn_tp_group=get_parallel().enable_dp_lm_head, + ) + else: + self.embed_tokens: Optional[nn.Module] = None + self.lm_head: Optional[nn.Module] = None self._use_fp32_lm_head = envs.SGLANG_DSPARK_FP32_LM_HEAD.get() self._opt_markov_w2_tp_shard = envs.SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD.get() + if self.lm_head is not None: + self.markov_head.configure_tp_shard(lm_head=self.lm_head) @property def enable_confidence_head(self) -> bool: @@ -674,9 +770,10 @@ class DeepseekV4ForCausalLMDSpark(nn.Module): def attach_shared_modules( self, *, embed_tokens: nn.Module, lm_head: nn.Module ) -> None: - self.embed_tokens = embed_tokens - self.lm_head = lm_head - self.markov_head.configure_tp_shard(lm_head=lm_head) + if not self.uses_own_vocab_modules: + self.embed_tokens = embed_tokens + self.lm_head = lm_head + self.markov_head.configure_tp_shard(lm_head=self.lm_head) def project_target_hidden(self, main_hidden: torch.Tensor) -> torch.Tensor: stage0 = self.stages[0] @@ -820,7 +917,11 @@ class DeepseekV4ForCausalLMDSpark(nn.Module): ) for name, loaded_weight in weights: - mapped = self._remap_dspark_weight_name(name) + mapped = ( + self._remap_dspark_weight_name_npu(name) + if _is_npu + else self._remap_dspark_weight_name(name) + ) if mapped is None: continue if self.num_fused_shared_experts > 0 and ".mlp.shared_experts." in mapped: @@ -929,5 +1030,63 @@ class DeepseekV4ForCausalLMDSpark(nn.Module): mapped_rest = mapped_rest.replace(".scale", ".weight_scale_inv") return f"stages.{stage_id}.{mapped_rest}" + def _remap_dspark_weight_name_npu(self, name: str) -> Optional[str]: + if name.startswith(("embed.", "embed_tokens.", "head.", "lm_head.")): + return None + if "rotary_emb.inv_freq" in name: + return None + + if not name.startswith("mtp."): + return None + parts = name.split(".", 2) + if len(parts) < 3: + return None + stage_id, rest = parts[1], parts[2] + if not stage_id.isdigit(): + return None + stage_idx = int(stage_id) + if rest in ("embed.weight", "embed_tokens.weight"): + return ( + "embed_tokens.weight" + if self.uses_own_vocab_modules and stage_idx == 0 + else None + ) + if rest in ("head.weight", "lm_head.weight"): + return ( + "lm_head.weight" + if self.uses_own_vocab_modules and stage_idx == self.num_stages - 1 + else None + ) + if rest.startswith(("embed.", "embed_tokens.", "head.", "lm_head.")): + return None + + if rest.startswith("markov_head."): + return f"markov_head.{rest[len('markov_head.'):]}" + + if rest.startswith("confidence_head."): + if self.confidence_head is None: + return None + return f"confidence_head.{rest[len('confidence_head.'):]}" + + mapped_rest = rest + if mapped_rest.startswith("attn."): + mapped_rest = "self_attn." + mapped_rest.removeprefix("attn.") + elif mapped_rest.startswith("ffn."): + mapped_rest = "mlp." + mapped_rest.removeprefix("ffn.") + elif mapped_rest.startswith("attn_norm."): + mapped_rest = "input_layernorm." + mapped_rest.removeprefix("attn_norm.") + elif mapped_rest.startswith("ffn_norm."): + mapped_rest = "post_attention_layernorm." + mapped_rest.removeprefix( + "ffn_norm." + ) + mapped_rest = mapped_rest.replace(".w1.", ".gate_proj.") + mapped_rest = mapped_rest.replace(".w2.", ".down_proj.") + mapped_rest = mapped_rest.replace(".w3.", ".up_proj.") + mapped_rest = mapped_rest.replace(".gate.tid2eid", ".topk.tid2eid") + mapped_rest = mapped_rest.replace(".gate.bias", ".gate.e_score_correction_bias") + if mapped_rest.endswith(".scale"): + mapped_rest = mapped_rest.removesuffix(".scale") + ".weight_scale_inv" + return f"stages.{stage_id}.{mapped_rest}" + EntryClass = [DeepseekV4ForCausalLMDSpark] diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 589db0b1b..8596fd959 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -893,6 +893,11 @@ class ServerArgs: Arg(help="The number of tokens in a page.", resolvable=True), NS("schedule"), ] = None + c128_page_size: A[ + int, + "The physical page size of the NPU DSV4 C128 KV cache. Must be a positive multiple of 16.", + NS("schedule"), + ] = 16 swa_full_tokens_ratio: A[ float, Arg( diff --git a/python/sglang/srt/speculative/dflash_disaggregation.py b/python/sglang/srt/speculative/dflash_disaggregation.py new file mode 100644 index 000000000..1cb2be7b2 --- /dev/null +++ b/python/sglang/srt/speculative/dflash_disaggregation.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.srt.managers.overlap_utils import RelayPayload +from sglang.srt.speculative.draft_worker_common import make_draft_input_v2 + +if TYPE_CHECKING: + from sglang.srt.managers.overlap_utils import FutureMap + from sglang.srt.managers.schedule_batch import ScheduleBatch + from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2 + + +def build_dflash_family_disagg_draft_input( + batch: ScheduleBatch, + last_tokens_tensor: torch.Tensor, + future_map: FutureMap, +) -> DFlashDraftInputV2: + spec_info = make_draft_input_v2( + bonus_tokens=last_tokens_tensor, + new_seq_lens=batch.seq_lens, + ) + if batch.enable_overlap: + spec_info.future_indices = batch.req_pool_indices + future_map.publish(spec_info.future_indices, batch.seq_lens) + future_map.stash( + spec_info.future_indices, + RelayPayload(bonus_tokens=last_tokens_tensor), + ) + return spec_info diff --git a/python/sglang/srt/speculative/dflash_info.py b/python/sglang/srt/speculative/dflash_info.py index 4fd0d1a27..21baee79c 100644 --- a/python/sglang/srt/speculative/dflash_info.py +++ b/python/sglang/srt/speculative/dflash_info.py @@ -13,11 +13,14 @@ from sglang.srt.model_executor.forward_batch_info import ( ForwardMode, ) from sglang.srt.speculative.spec_info import SpecInput, SpecInputType +from sglang.srt.utils import is_npu if TYPE_CHECKING: from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout +_is_npu = is_npu() + @dataclass class DFlashVerifyInput(SpecInput): @@ -43,6 +46,9 @@ class DFlashVerifyInput(SpecInput): num_tokens_per_req: int = -1 ragged_verify_layout: Optional[RaggedVerifyLayout] = None + # Committed/live lengths before the verify caller temporarily expands + # batch.seq_lens_cpu to the target-attention KV lengths. + live_seq_lens_cpu: Optional[torch.Tensor] = None def __post_init__(self): super().__init__(spec_input_type=SpecInputType.DFLASH_VERIFY) @@ -58,14 +64,25 @@ class DFlashVerifyInput(SpecInput): """Prepare a DFLASH verify forward batch for overlap scheduling. The caller computes and stores `batch.out_cache_loc` before this - method is called. This helper only packages the verify forward and pre-initializes either CUDA-graph replay - metadata or eager attention metadata so the actual forward can run with - `skip_attn_backend_init=True`. + method is called. GPU keeps the original pre-planning path. NPU leaves + attention/graph metadata initialization to ModelRunner because DP/EP + padding can still change the compressor's runtime shapes. """ from sglang.srt.speculative.spec_utils import prepare_mamba_track_for_verify batch.input_ids = self.draft_token batch.spec_info = self + if _is_npu and not batch.forward_mode.is_idle(): + from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( + maybe_build_dsv4_verify_bundle, + ) + + batch.out_cache_loc_dsv4 = maybe_build_dsv4_verify_bundle( + batch, + self.draft_token_num, + live_seq_lens_cpu=self.live_seq_lens_cpu, + ) + batch.forward_mode = ( ForwardMode.IDLE if batch.forward_mode.is_idle() @@ -90,7 +107,13 @@ class DFlashVerifyInput(SpecInput): verify_forward_batch ) ) - if can_run_cuda_graph: + if _is_npu: + # Do not pre-plan target verify on NPU. DP/EP padding can change + # the compressor's logical batch without changing ForwardBatch's + # stale-plan shape fields. Let ModelRunner select graph/eager and + # initialize metadata after final batch preparation. + return verify_forward_batch, can_run_cuda_graph + elif can_run_cuda_graph: target_worker.model_runner.decode_cuda_graph_runner.load_batch( verify_forward_batch ) diff --git a/python/sglang/srt/speculative/dflash_info_v2.py b/python/sglang/srt/speculative/dflash_info_v2.py index a59034ee8..e30ed182d 100644 --- a/python/sglang/srt/speculative/dflash_info_v2.py +++ b/python/sglang/srt/speculative/dflash_info_v2.py @@ -124,6 +124,9 @@ class DFlashDraftInputV2(SpecInput): bs = batch.batch_size() if bs == 0: return + + batch.maybe_evict_swa() + self._ensure_prepare_length_buffers(bs, batch.device) assert self._prepare_batch_seq_lens_cpu_buf is not None assert self._prepare_cur_kv_lens_cpu_buf is not None @@ -205,7 +208,8 @@ class DFlashDraftInputV2(SpecInput): # plan-stream context, so forward work cannot observe partially # prepared req_to_token / KV allocation state. caller_stream.wait_stream(plan_stream) - + for req in batch.reqs: + req.decode_batch_idx += 1 # Seed committed; overlap's resolve overwrites it with the published value. batch.seq_lens_cpu = batch_seq_lens_cpu_t batch.seq_lens_sum = committed_seq_lens_sum diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft.py b/python/sglang/srt/speculative/dspark_components/dspark_draft.py index 87af5b7d8..2fd55f221 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft.py @@ -396,6 +396,10 @@ class DraftBlockProposer: forward_batch.can_run_dp_cuda_graph = batch.can_run_dp_cuda_graph if not self._dp_moe_sync or batch.global_num_tokens is None: return + # Graph bucket selection uses the raw per-rank request counts. Keep + # them separate from global_num_tokens_cpu below, which is scaled into + # draft-token units for DP/MoE synchronization. + forward_batch.original_global_num_tokens_cpu = batch.global_num_tokens gnt, gnt_logprob = spec_scale_global_num_tokens( self._draft_block_spec_info, batch.global_num_tokens, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py b/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py index a01c93a7a..4ecce20b6 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py @@ -146,7 +146,9 @@ def _resolve_folded_sampling(*, model, gamma, max_bs, device, tp_rank) -> bool: noise_bytes = max_bs * vocab * 4 logits_bytes = max_bs * gamma * vocab * model.lm_head.weight.dtype.itemsize need_gb = (noise_bytes + logits_bytes) / (1 << 30) - available_gb = get_available_gpu_memory(device, torch.cuda.current_device()) + available_gb = get_available_gpu_memory( + device, torch.get_device_module().current_device() + ) if available_gb - need_gb >= _CAPTURE_HEADROOM_GB: return True if tp_rank == 0: diff --git a/python/sglang/srt/speculative/dspark_components/dspark_verify.py b/python/sglang/srt/speculative/dspark_components/dspark_verify.py index 281937848..9939f3750 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_verify.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_verify.py @@ -45,8 +45,11 @@ from sglang.srt.speculative.spec_utils import ( SIMULATE_ACC_METHOD, sample_simulated_acc_len, ) +from sglang.srt.utils import is_npu from sglang.srt.utils.invariants import Bucket, Invariant, NotNaN, expect +_is_npu = is_npu() + # Draft proposal probs feeding rejection sampling; the data layer is the # in-kernel NaN-q guard in reject_sampling.py, so this is signal-only. _VERIFY_DRAFT_PROBS = Invariant("dspark.verify.draft_probs", Bucket.GUARD, NotNaN()) @@ -214,6 +217,7 @@ class TargetVerifyExecutor: batch.seq_lens_cpu = torch.ones((num_dummy_slots,), dtype=torch.int64) batch.seq_lens_sum = num_dummy_slots batch.forward_mode = ForwardMode.TARGET_VERIFY + verify_input.live_seq_lens_cpu = batch.seq_lens_cpu verify_forward_batch, _ = verify_input.prepare_for_verify( batch, self.target_worker ) @@ -221,7 +225,7 @@ class TargetVerifyExecutor: batch=None, forward_batch=verify_forward_batch, is_verify=True, - skip_attn_backend_init=True, + skip_attn_backend_init=True if not _is_npu else None, ) def run_non_compact( @@ -243,6 +247,7 @@ class TargetVerifyExecutor: draft_token_num=verify_w, custom_mask=None, capture_hidden_mode=CaptureHiddenMode.FULL, + live_seq_lens_cpu=batch.seq_lens_cpu, ) batch.out_cache_loc = verify_cache_loc seq_lens_cpu_backup = batch.seq_lens_cpu @@ -289,7 +294,7 @@ class TargetVerifyExecutor: batch=None, forward_batch=verify_forward_batch, is_verify=True, - skip_attn_backend_init=True, + skip_attn_backend_init=True if not _is_npu else None, ) return TargetVerifyResult( logits_output=target_out.logits_output, @@ -355,6 +360,7 @@ class TargetVerifyExecutor: custom_mask=None, capture_hidden_mode=CaptureHiddenMode.FULL, ragged_verify_layout=layout, + live_seq_lens_cpu=batch.seq_lens_cpu, ) batch.out_cache_loc = ragged_window.verify_cache_loc seq_lens_cpu_backup = batch.seq_lens_cpu diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index f5dc7c795..ee346be66 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -66,10 +66,12 @@ from sglang.srt.speculative.spec_utils import ( draft_tp_context, prepare_mamba_track_for_verify, ) -from sglang.srt.utils import get_available_gpu_memory, is_cuda +from sglang.srt.utils import get_available_gpu_memory, is_cuda, is_npu logger = logging.getLogger(__name__) +_is_npu = is_npu() + class DSparkWorkerV2(BaseSpecWorker): @@ -173,16 +175,22 @@ class DSparkWorkerV2(BaseSpecWorker): draft_token_num=int(self.gamma), device=self.device ) - target_model = self.target_worker.model_runner.model - lm_head = getattr(target_model, "lm_head", None) - if lm_head is None or not hasattr(lm_head, "weight"): - raise RuntimeError( - "DSpark requires the target model to expose `lm_head` with `weight`." + if getattr(self.draft_model, "uses_own_vocab_modules", False): + if self.ps.tp_rank == 0: + logger.info( + "DSpark draft uses its checkpoint-local embedding and LM head." + ) + else: + target_model = self.target_worker.model_runner.model + lm_head = getattr(target_model, "lm_head", None) + if lm_head is None or not hasattr(lm_head, "weight"): + raise RuntimeError( + "DSpark requires the target model to expose `lm_head` with `weight`." + ) + self.draft_model.attach_shared_modules( + embed_tokens=self._resolve_target_embed_tokens(target_model), + lm_head=lm_head, ) - self.draft_model.attach_shared_modules( - embed_tokens=self._resolve_target_embed_tokens(target_model), - lm_head=lm_head, - ) self._verify_planner = DSparkVerifyPlanner( draft_model=self.draft_model, @@ -329,6 +337,12 @@ class DSparkWorkerV2(BaseSpecWorker): def init_attention_backends(self): with self._draft_context(): + if _is_npu: + from sglang.srt.hardware_backend.npu.extra_ops_loader import ( + initialize_dspark_sparse_attn_ops, + ) + + initialize_dspark_sparse_attn_ops() self._draft_worker.init_attention_backends() self._need_mamba_verify_commit = mambaish_config( self.model_runner.model_config @@ -570,7 +584,6 @@ class DSparkWorkerV2(BaseSpecWorker): self._observers.begin_step() target_model = self.target_worker.model_runner.model - verify_window = alloc_verify_window( batch=batch, bs=bs, @@ -675,7 +688,6 @@ class DSparkWorkerV2(BaseSpecWorker): hidden_strided = None logits_output = target_verify.logits_output can_run_cuda_graph = target_verify.can_run_cuda_graph - if batch.has_grammar: # run_compact scatters its rows back to (bs * chain_len), so the mask # lines up with the logits on both verify paths. diff --git a/test/registered/unit/npu/attention/test_npu_ascend_dsv4_backend.py b/test/registered/unit/npu/attention/test_npu_ascend_dsv4_backend.py index c244c1d24..fc834b971 100644 --- a/test/registered/unit/npu/attention/test_npu_ascend_dsv4_backend.py +++ b/test/registered/unit/npu/attention/test_npu_ascend_dsv4_backend.py @@ -59,7 +59,6 @@ from sglang.srt.hardware_backend.npu.attention.ascend_dsv4_backend import ( DeepseekV4AscendMultiStepDraftBackend, _apply_hadamard, _get_kv_indices, - _overlap_transform, _walsh_hadamard_matrix, ) @@ -182,78 +181,6 @@ class TestApplyHadamard(unittest.TestCase): self.assertTrue(torch.equal(out, expected)) -class TestOverlapTransform(unittest.TestCase): - def test_shape(self): - # (n_chunks, ratio, 2*d) -> (n_chunks, 2*ratio, d) - n_chunks, r, d = 3, 2, 4 - tensor = torch.randn(n_chunks, r, 2 * d) - out = _overlap_transform(tensor, value=0.0, head_dim=d) - self.assertEqual(out.shape, (n_chunks, 2 * r, d)) - - def test_first_chunk_left_half_filled_with_value(self): - n_chunks, r, d = 3, 2, 4 - tensor = torch.randn(n_chunks, r, 2 * d) - fill = float("-inf") - out = _overlap_transform(tensor, value=fill, head_dim=d) - self.assertTrue(torch.equal(out[0, :r], torch.full((r, d), fill))) - - def test_first_chunk_left_half_filled_with_zero(self): - n_chunks, r, d = 2, 2, 4 - tensor = torch.randn(n_chunks, r, 2 * d) - out = _overlap_transform(tensor, value=0.0, head_dim=d) - self.assertTrue(torch.equal(out[0, :r], torch.zeros(r, d))) - - def test_right_half_mirrors_tensor_second_half(self): - n_chunks, r, d = 3, 2, 4 - tensor = torch.randn(n_chunks, r, 2 * d) - out = _overlap_transform(tensor, value=0.0, head_dim=d) - self.assertTrue(torch.equal(out[:, r:], tensor[..., d:])) - - def test_previous_chunk_left_half(self): - n_chunks, r, d = 3, 2, 4 - tensor = torch.randn(n_chunks, r, 2 * d) - out = _overlap_transform(tensor, value=0.0, head_dim=d) - self.assertTrue(torch.equal(out[1:, :r], tensor[:-1, :, :d])) - - def test_single_chunk(self): - n_chunks, r, d = 1, 2, 4 - tensor = torch.randn(n_chunks, r, 2 * d) - fill = 7.0 - out = _overlap_transform(tensor, value=fill, head_dim=d) - self.assertEqual(out.shape, (1, 2 * r, d)) - self.assertTrue(torch.equal(out[0, :r], torch.full((r, d), fill))) - self.assertTrue(torch.equal(out[0, r:], tensor[0, :, d:])) - - def test_full_element_mapping(self): - n_chunks, r, d = 2, 2, 3 - tensor = torch.arange(n_chunks * r * 2 * d, dtype=torch.float32).reshape( - n_chunks, r, 2 * d - ) - fill = -1.0 - out = _overlap_transform(tensor, value=fill, head_dim=d) - - for c in range(n_chunks): - for row in range(2 * r): - for col in range(d): - if c == 0 and row < r: - expected = fill - elif row >= r: - expected = tensor[c, row - r, d + col].item() - else: - expected = tensor[c - 1, row, col].item() - self.assertEqual( - out[c, row, col].item(), - expected, - f"mismatch at (c={c}, row={row}, col={col})", - ) - - def test_preserves_input_dtype(self): - n_chunks, r, d = 2, 2, 4 - tensor = torch.randn(n_chunks, r, 2 * d, dtype=torch.bfloat16) - out = _overlap_transform(tensor, value=0.0, head_dim=d) - self.assertEqual(out.dtype, torch.bfloat16) - - class TestGetKvIndices(unittest.TestCase): _PATCH_TARGET = ( "sglang.srt.hardware_backend.npu.attention.ascend_dsv4_backend.get_attn_backend" diff --git a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py index 90a52cea5..d60cf6b1b 100644 --- a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py +++ b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py @@ -41,6 +41,10 @@ _EVICT_METHOD = "maybe_evict_swa" # Any added/removed/recounted site fails until reviewed here. _SB = "managers/schedule_batch.py" _EAGLE_DECODE = ("speculative/eagle_utils.py", "eagle_prepare_for_decode") +_DFLASH_DECODE = ( + "speculative/dflash_info_v2.py", + "DFlashDraftInputV2.prepare_for_decode", +) _RESOLVE = ( "managers/scheduler_components/batch_result_processor.py", "SchedulerBatchResultProcessor._resolve_spec_v2_tokens", @@ -62,6 +66,11 @@ _OWNER_SITES = { # inside the owned-kv alloc_for_spec_decode function (op42). (*_EAGLE_DECODE, "decode_batch_idx"): 1, (*_EAGLE_DECODE, "evict"): 1, + # DFlash uses its stateful scheduler-side preparation instead of + # eagle_prepare_for_decode. spec_prepare_for_decode dispatches to exactly + # one of these two owners for each speculative decode iteration. + (*_DFLASH_DECODE, "decode_batch_idx"): 1, + (*_DFLASH_DECODE, "evict"): 1, ( "mem_cache/allocation.py", "alloc_for_spec_decode",