diff --git a/python/sglang/kernels/ops/attention/deepseek_v4_rope.py b/python/sglang/kernels/ops/attention/deepseek_v4_rope.py index c3555e589..1dd647465 100644 --- a/python/sglang/kernels/ops/attention/deepseek_v4_rope.py +++ b/python/sglang/kernels/ops/attention/deepseek_v4_rope.py @@ -649,123 +649,3 @@ def fused_norm_rope_inplace_triton( HAS_WEIGHT=(weight is not None), USE_POS=(positions is not None), ) - - -# Cache contiguous real/imag halves of each freqs_cis (its .real/.imag are -# strided views, stride=2 on the interleaved layout), keyed by id. -_NPU_ROPE_CONTIG_CACHE: dict[int, tuple] = {} - - -def _get_contig_freqs_real_imag( - freqs_cis: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - """Return contiguous (real, imag) halves of ``freqs_cis``, cached by id. - - Used by NPU rope paths to avoid the per-call StridedSlice materialization - triggered by aclnnIndex over the strided ``.real`` / ``.imag`` views of - the complex ``freqs_cis`` buffer. First call per freqs_cis pays the - contiguous() once; later calls reuse the cached tensors. - - All callers within a single MQALayer (outer rope, indexer inner rope, - compressor epilog rope) get the same freqs_cis instance, so each layer - materializes at most one (real, imag) pair. - """ - cache_key = id(freqs_cis) - cached = _NPU_ROPE_CONTIG_CACHE.get(cache_key) - if cached is None: - cached = (freqs_cis.real.contiguous(), freqs_cis.imag.contiguous()) - _NPU_ROPE_CONTIG_CACHE[cache_key] = cached - return cached - - -def get_fused_compressor_rope_cos_sin( - freqs_cis: torch.Tensor, - positions_cmp: torch.Tensor, - dtype: torch.dtype, -) -> tuple[torch.Tensor, torch.Tensor]: - """Build (cos, sin) tensors shaped ``[T, rope_head_dim]`` for the fused - compressor op (``torch.ops.custom.compressor``). - - The op consumes ``rope_cos`` / ``rope_sin`` of shape - ``[min(T, T//cmp_ratio + B), rope_head_dim]`` in bf16/fp16. We index - the cached contig real/imag halves of the complex ``freqs_cis`` and - interleave-double the last dim to match the kernel's expected layout - (matches dsv4_release ``ComplexExpRotaryEmbedding.cos_cache``, which - is built as ``complex_cache.real.repeat_interleave(2, dim=-1)``). - - Safe to call from inside a captured aclgraph: both ``index_select`` and - ``repeat_interleave`` over a graph-input ``positions_cmp`` of fixed - capture-time shape produce static-shape outputs. Identical to what the - existing inplace_partial_rotary_mul fallback does at - :func:`v4_rope_inplace_npu`, just without the inverse / 4D-view step. - """ - real_contig, imag_contig = _get_contig_freqs_real_imag(freqs_cis) - cos_half = real_contig.index_select(0, positions_cmp) - sin_half = imag_contig.index_select(0, positions_cmp) - cos = cos_half.repeat_interleave(2, dim=-1).to(dtype) - sin = sin_half.repeat_interleave(2, dim=-1).to(dtype) - return cos, sin - - -def v4_rope_inplace_npu( - q_rope: torch.Tensor, - kv_rope: Optional[torch.Tensor], - freqs_cis: torch.Tensor, - positions: torch.Tensor, - inverse: bool = False, -) -> None: - """In-place interleaved RoPE for V4 — torch fallback used on NPU. - - Mirrors main's CUDA `fused_rope` kernel: consecutive (even, odd) pairs - of x form complex pairs, with `freqs_cis` a complex tensor where - `freqs_cis.real[t, k]` = cos(theta_{t,k}), `freqs_cis.imag` = sin(...) - indexed by frequency pair k in [0, rope_dim/2). - - NOTE on V4-Flash YARN `mscale`: when the model was trained with the - YARN magnitude-scale `mscale` ≠ 1.0, the cos/sin values stored in - `freqs_cis` MUST already be pre-multiplied by `mscale` at precompute - time — see `precompute_freqs_cis`. This function - just reads what's stored; it does NOT apply mscale here. - - Prefer the NPU-native `torch.ops.custom.inplace_partial_rotary_mul`: - the torch fallback differs by ~1 ULP per element vs the kernel because - torch does bf16*bf16 muls with bf16 accumulation while the NPU kernel - accumulates in fp32; 43 layers × (Q + K) = 86 rope calls compound that - drift enough to flip argmax on marginal prompts. - """ - # Build cos/sin caches in the kernel's expected (T, 1, 1, rope_dim) layout, - # each freq value repeated twice for the interleaved pairing convention. - freqs_real_contig, freqs_imag_contig = _get_contig_freqs_real_imag(freqs_cis) - cos_half = freqs_real_contig[positions] # (T, rope_dim/2) - sin_half = freqs_imag_contig[positions] - if inverse: - sin_half = -sin_half - cos_full = cos_half.repeat_interleave(2, dim=-1).to(q_rope.dtype) - sin_full = sin_half.repeat_interleave(2, dim=-1).to(q_rope.dtype) - rope_dim = cos_full.shape[-1] - # repeat_interleave produces a contiguous tensor, so the .view() - # below already returns a contiguous result — no .contiguous() needed. - cos4 = cos_full.view(-1, 1, 1, rope_dim) - sin4 = sin_full.view(-1, 1, 1, rope_dim) - # q_rope: (T, n_heads, rope_dim) → (T, 1, n_heads, rope_dim) view - # kv_rope: (T, 1, rope_dim) → (T, 1, 1, rope_dim) view - q_view = q_rope.unsqueeze(1) - torch.ops.custom.inplace_partial_rotary_mul( - q_view, - cos4, - sin4, - rotary_mode="interleave", - partial_slice=[0, rope_dim], - ) - if kv_rope is not None: - if kv_rope.dim() == 3: - kv_view = kv_rope.unsqueeze(1) - else: - kv_view = kv_rope.view(-1, 1, 1, rope_dim) - torch.ops.custom.inplace_partial_rotary_mul( - kv_view, - cos4, - sin4, - rotary_mode="interleave", - partial_slice=[0, rope_dim], - ) diff --git a/python/sglang/srt/disaggregation/ascend/conn.py b/python/sglang/srt/disaggregation/ascend/conn.py index 40251c587..2162a49bd 100644 --- a/python/sglang/srt/disaggregation/ascend/conn.py +++ b/python/sglang/srt/disaggregation/ascend/conn.py @@ -1,4 +1,5 @@ import concurrent.futures +import enum import logging from typing import List, Optional, Tuple @@ -6,6 +7,7 @@ import numpy as np import numpy.typing as npt from sglang.srt.disaggregation.ascend.transfer_engine import AscendTransferEngine +from sglang.srt.disaggregation.base.conn import StateType from sglang.srt.disaggregation.common.utils import group_concurrent_contiguous from sglang.srt.disaggregation.mooncake.conn import ( MooncakeKVBootstrapServer, @@ -18,7 +20,28 @@ from sglang.srt.utils.network import get_local_ip_auto 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_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) + + class AscendKVManager(MooncakeKVManager): + def _requires_exact_state_index_match(self, st: StateType) -> bool: + return ( + super()._requires_exact_state_index_match(st) + or st in _DSV4_KVCACHE_STATE_TYPES + ) + def init_engine(self): # TransferEngine initialized on ascend. local_ip = get_local_ip_auto() @@ -29,23 +52,29 @@ class AscendKVManager(MooncakeKVManager): ) def register_buffer_to_engine(self): - self.engine.batch_register(self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens) - # The Ascend backend optimize batch registration for small memory blocks. - self.engine.batch_register( - self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens - ) - # Batch register state/extra pool data buffers + # MemFabric aligns registered buffers to 2 MiB. Register everything in + # one batch so overlapping aligned ranges from small tensors are merged + # before they are published to the peer. + ptrs = list(self.kv_args.kv_data_ptrs) + lens = list(self.kv_args.kv_data_lens) + ptrs.extend(self.kv_args.aux_data_ptrs) + lens.extend(self.kv_args.aux_data_lens) for component_ptrs, component_lens in zip( self.kv_args.state_data_ptrs or [], self.kv_args.state_data_lens or [], ): - self.engine.batch_register(component_ptrs, component_lens) + ptrs.extend(component_ptrs) + lens.extend(component_lens) + if ptrs: + self.engine.batch_register(ptrs, lens) def get_mla_kv_ptrs_with_pp( - self, src_kv_ptrs: List[int], dst_kv_ptrs: List[int] + self, src_kv_ptrs: List[int], dst_kv_ptrs: List[int], state_type=None ) -> Tuple[List[int], List[int], int]: # 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; + # the NPU kv_buf_groups slicing below is state-type agnostic. start_layer = self.kv_args.prefill_start_layer kv_buf_groups = getattr(self.kv_args, "kv_buf_groups", 1) total_kv_layers = getattr(self.kv_args, "total_kv_layers", 0) @@ -179,6 +208,13 @@ class AscendKVManager(MooncakeKVManager): return 0 + def _is_generic_kvcache_state_type(self, st) -> bool: + # DSV4 per-pool components also use the page-indexed send path. + return ( + super()._is_generic_kvcache_state_type(st) + or st in _DSV4_KVCACHE_STATE_TYPES + ) + class AscendKVSender(MooncakeKVSender): pass diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 00cf04097..bbb30ef8b 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -87,13 +87,15 @@ from sglang.srt.observability.req_time_stats import ( set_time_batch, ) from sglang.srt.runtime_context import get_parallel -from sglang.srt.utils import get_num_new_pages +from sglang.srt.utils import get_num_new_pages, is_npu from sglang.srt.utils.network import NetworkAddress from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter logger = logging.getLogger(__name__) +_is_npu = is_npu() + if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.scheduler import Scheduler @@ -1047,6 +1049,17 @@ 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, @@ -1145,30 +1158,49 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): ) state_types = self.kv_manager.kv_args.state_types - state_indices: Optional[List] = [] if StateType.C128_STATE in state_types: clear_c128_state = getattr( self.token_to_kv_pool, "clear_c128_req_state", None ) if clear_c128_state is not None: clear_c128_state(int(decode_req.req.req_pool_idx)) - for st in state_types: - if st == StateType.MAMBA: - state_indices.append(_mamba_payload()) - elif st == StateType.SWA: - state_indices.append(_swa_payload()) - elif st == StateType.DSA: - state_indices.append(_dsa_payload()) - elif st == StateType.MINIMAX_INDEX_K: - # Index rows live at the same loc as main KV on the same - # page_size, so reuse the full-seq page-ids. - state_indices.append(_dsa_payload()) - elif st == StateType.SWA_RING: - state_indices.append(_swa_ring_payload()) - elif st == StateType.C128_STATE: - state_indices.append(_c128_state_payload()) - else: - state_indices.append(None) + # MINIMAX_INDEX_K reuses _dsa_payload: index rows live at the same loc + # as main KV on the same page_size. + payloads = { + StateType.MAMBA: _mamba_payload, + StateType.SWA: _swa_payload, + StateType.DSA: _dsa_payload, + StateType.MINIMAX_INDEX_K: _dsa_payload, + 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, + ) + + payloads.update( + dsv4_state_payloads( + self.req_to_token_pool, + 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, + ) + ) + state_indices: Optional[List] = [ + payloads[st]() if st in payloads else None for st in state_types + ] decode_req.metadata_buffer_index = ( self.req_to_metadata_buffer_idx_allocator.alloc() @@ -1503,6 +1535,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): prefix_indices=prefix_indices, uses_swa_tail=uses_swa_tail, swa_tail_len=swa_tail_len, + req_to_token_pool=self.req_to_token_pool, ) assert kv_loc is not None, ( f"KV cache is full! Bug in memory estimation. " @@ -1594,6 +1627,7 @@ def alloc_for_decode_prealloc( prefix_indices: Optional[torch.Tensor], uses_swa_tail: bool, swa_tail_len: int, + req_to_token_pool: Optional[ReqToTokenPool] = None, ) -> torch.Tensor: if req.kv is None: req.kv = ReqKvInfo(kv_allocated_len=fill_len, swa_evicted_seqlen=0) @@ -1608,6 +1642,22 @@ def alloc_for_decode_prealloc( if prefix_len > 0 else torch.tensor([-1], dtype=torch.int64, device=device) ) + extra_kwargs = {} + dsv4_unwrap_prealloc = None + if hasattr(allocator, "c4_attn_allocator"): + assert req_to_token_pool is not None + from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( + dsv4_prealloc_kwargs, + dsv4_unwrap_prealloc, + ) + + extra_kwargs = dsv4_prealloc_kwargs( + allocator, + req, + fill_len, + req_to_token_pool, + device=device, + ) if uses_swa_tail: # Tail-only SWA allocation: only valid when prefix_len == 0. # When prefix_len > 0 (radix cache hit), we fall back to @@ -1621,6 +1671,7 @@ def alloc_for_decode_prealloc( last_loc=last_loc, extend_num_tokens=fill_len, swa_tail_len=swa_tail_len, + **extra_kwargs, ) req.kv.swa_evicted_seqlen = fill_len - swa_tail_len else: @@ -1633,6 +1684,11 @@ def alloc_for_decode_prealloc( seq_lens_cpu=torch.tensor([fill_len], dtype=torch.int64), last_loc=last_loc, extend_num_tokens=delta_len, + **extra_kwargs, + ) + if dsv4_unwrap_prealloc is not None: + kv_loc = dsv4_unwrap_prealloc( + kv_loc, req_to_token_pool, req, total_prefix_len, fill_len ) return kv_loc diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index a1ccc4bbf..3a92fd2a5 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -995,6 +995,20 @@ class MooncakeKVManager(CommonKVManager): return skip_kv, skip_state + def _is_generic_kvcache_state_type(self, st: StateType) -> bool: + """State types sent via the page-indexed ``_send_kvcache_generic`` path + (not the mamba-state path); subclasses extend for hardware components.""" + return st in ( + StateType.SWA, + StateType.DSA, + StateType.SWA_RING, + StateType.C128_STATE, + ) + + def _requires_exact_state_index_match(self, st: StateType) -> bool: + """State types whose page lists are positional and must not be truncated.""" + return st in (StateType.SWA_RING, StateType.C128_STATE) + def maybe_send_extra( self, req: TransferInfo, @@ -1099,12 +1113,7 @@ class MooncakeKVManager(CommonKVManager): ) or rc ) - elif st in ( - StateType.SWA, - StateType.DSA, - StateType.SWA_RING, - StateType.C128_STATE, - ): + elif self._is_generic_kvcache_state_type(st): if ( target_rank_registration_info is not None and not self.is_mla_backend @@ -1127,7 +1136,7 @@ class MooncakeKVManager(CommonKVManager): # truncating silently misaligns rows and corrupts KV. # Paged SWA/DSA tolerate a 1-page drift -> keep the # lenient truncation below. - if st in (StateType.SWA_RING, StateType.C128_STATE): + if self._requires_exact_state_index_match(st): raise RuntimeError( f"{st.upper()} state index length mismatch: " f"prefill={len(src_indices)}, dst={len(dst_indices_local)}" @@ -2003,6 +2012,8 @@ class MooncakeKVReceiver(CommonKVReceiver): ) # Note(shangming): No need to add pp rank here since decode pp size should be equal to prefill pp size or 1 tp_rank = self.kv_mgr.kv_args.engine_rank + # Some pools have no full-token contiguous KV (kv_item_lens empty) + # and ship per-pool instead, so report 0. kv_item_len = ( self.kv_mgr.kv_args.kv_item_lens[0] if self.kv_mgr.kv_args.kv_item_lens diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index b17879278..3e95298e2 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -64,6 +64,7 @@ from sglang.srt.mem_cache.common import ( ) from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.observability.req_time_stats import set_schedule_time_batch +from sglang.srt.utils import is_npu from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method if TYPE_CHECKING: @@ -74,6 +75,8 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +_is_npu = is_npu() + def should_force_retry(req: Req) -> bool: """Test hook to force a request into optimistic prefill retry.""" @@ -1180,24 +1183,37 @@ class SchedulerDisaggregationPrefillMixin: state_types = ( self.disagg_prefill_bootstrap_queue.kv_manager.kv_args.state_types ) - state_indices = [] - for st in state_types: - if st == StateType.MAMBA: - state_indices.append(_mamba_payload()) - elif st == StateType.SWA: - state_indices.append(_swa_payload()) - elif st == StateType.DSA: - state_indices.append(_dsa_payload()) - elif st == StateType.MINIMAX_INDEX_K: - # Index rows live at the same loc as main KV on the same - # page_size, so reuse the full-seq page-ids. - state_indices.append(_dsa_payload()) - elif st == StateType.SWA_RING: - state_indices.append(_swa_ring_payload()) - elif st == StateType.C128_STATE: - state_indices.append(_c128_state_payload()) - else: - state_indices.append(None) + # MINIMAX_INDEX_K reuses _dsa_payload: index rows live at the same loc + # as main KV on the same page_size. + payloads = { + StateType.MAMBA: _mamba_payload, + StateType.SWA: _swa_payload, + StateType.DSA: _dsa_payload, + StateType.MINIMAX_INDEX_K: _dsa_payload, + StateType.SWA_RING: _swa_ring_payload, + StateType.C128_STATE: _c128_state_payload, + } + if _is_npu and isinstance( + self.token_to_kv_pool_allocator.get_kvcache(), + DeepSeekV4TokenToKVPool, + ): + from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( + dsv4_state_payloads, + ) + + payloads.update( + dsv4_state_payloads( + self.req_to_token_pool, + req.req_pool_idx, + seq_len, + page_size, + self.sliding_window_size, + prefix_len=0, + ) + ) + state_indices = [ + payloads[st]() if st in payloads else None for st in state_types + ] kv_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, start_idx:end_idx diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 2c47a97eb..cbb4abf75 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -27,6 +27,11 @@ if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req from sglang.srt.server_args import ServerArgs +if is_npu(): + from sglang.srt.hardware_backend.npu.dsv4.dsv4_memory_pool import ( + DSV4NPUTokenToKVPool, + ) + ######################### # Constants & Enums ######################### @@ -956,7 +961,17 @@ def setup_state_kv_args( kv_args.is_hybrid_mla_backend = False kv_args.state_conv_shard_groups = [] - if isinstance(token_to_kv_pool, MiniMaxSparseKVPool): + 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 token_to_kv_pool.index_kv_pool is not None: raise NotImplementedError( "PD disaggregation for MiniMax sparse layers with index value " 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 8b403c38b..6b9e3e50c 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -684,11 +684,11 @@ class AscendAttnBackend(AttentionBackend): metadata.block_tables_swa[bs:, :].fill_(0) # Update SWA mask: True = masked out (don't attend), False = attend - seq_lens_int = seq_lens_cpu[:bs].int() + seq_lens_int = seq_lens[:bs].int() starts = torch.clamp(seq_lens_int - self.sliding_window_size, min=0) indices = self.graph_metadata["swa_indices"] - start_exp = starts.unsqueeze(1).to(self.device) - seq_exp = seq_lens_int.unsqueeze(1).to(self.device) + start_exp = starts.unsqueeze(1) + seq_exp = seq_lens_int.unsqueeze(1) mask = (indices.unsqueeze(0) < start_exp) | ( indices.unsqueeze(0) >= seq_exp ) 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 d45ec4fda..eb5ee46ce 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 @@ -2,6 +2,7 @@ from __future__ import annotations import logging import math +from types import SimpleNamespace from typing import TYPE_CHECKING, Optional import torch @@ -65,6 +66,34 @@ def _overlap_transform( class CompressorAscendBackendMixin(CompressorBackendMixin): + @staticmethod + def _to_cpu_int_list(values) -> Optional[list[int]]: + if values is None: + return None + if isinstance(values, torch.Tensor): + values = values.cpu().tolist() + return [int(v) for v in values] + + def _extend_prefix_lens_cpu( + self, forward_batch: ForwardBatch + ) -> Optional[list[int]]: + prefix_lens = self._to_cpu_int_list( + getattr(forward_batch, "extend_prefix_lens_cpu", None) + ) + if prefix_lens is not None: + return prefix_lens + + seq_lens = self._to_cpu_int_list(getattr(forward_batch, "seq_lens_cpu", None)) + extend_lens = self._to_cpu_int_list( + getattr(forward_batch, "extend_seq_lens_cpu", None) + ) + if seq_lens is None or extend_lens is None or len(seq_lens) != len(extend_lens): + return None + return [ + max(0, seq_len - extend_len) + for seq_len, extend_len in zip(seq_lens, extend_lens) + ] + def _build_npu_compress_metadata(self, forward_batch: ForwardBatch) -> None: fm = self.forward_metadata is_decode = forward_batch.forward_mode.is_decode() @@ -94,6 +123,13 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): 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 / + # seqused only for decode. Every eager prefill uses the fused compressor, + # so build its global block positions, state metadata, and output locs + # here. Exclude speculative modes so the cu.cpu() host read never + # runs for target_verify / draft_extend (potentially graph-captured). + if forward_batch.forward_mode.is_extend_without_speculative(): + self._build_npu_compress_metadata_prefill(forward_batch) if _verify_compress: self._build_npu_compress_metadata_verify(forward_batch) @@ -108,6 +144,7 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): cu = fm.actual_seq_lengths_q_pa cu_cpu = cu.cpu().tolist() + 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) } @@ -116,12 +153,20 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): end = int(cu_cpu[idx + 1]) if end == start: continue - seq = end - start - req_positions = positions[start:end] + prefix = ( + int(prefix_cpu[idx]) + if prefix_cpu is not None and idx < len(prefix_cpu) + else 0 + ) + total = prefix + (end - start) for ratio in ratio_lists: - cutoff = seq - (seq % ratio) - if cutoff > 0: - ratio_lists[ratio].append(req_positions[:cutoff:ratio]) + first_k = prefix // ratio + last_k = total // ratio + if last_k > first_k: + ratio_lists[ratio].append( + torch.arange(first_k, last_k, device=device, dtype=torch.int64) + * ratio + ) for ratio in (4, 128): if ratio not in ratio_lists: @@ -137,18 +182,36 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): padding[: cat.shape[0]].copy_(cat) setattr(fm, f"positions_cmp_padding_c{ratio}", padding) - # start_pos=0: chunked prefill unsupported; seqused=None -> op derives lens from cu_seqlens - fm.start_pos = torch.zeros(bs, dtype=torch.int32, device=device) + # 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. + if forward_batch.extend_prefix_lens is not None: + fm.start_pos = forward_batch.extend_prefix_lens.to( + device=device, dtype=torch.int32 + ) + elif prefix_cpu is not None: + fm.start_pos = torch.tensor( + prefix_cpu[:bs] + [0] * max(0, bs - len(prefix_cpu)), + dtype=torch.int32, + device=device, + ) + else: + fm.start_pos = torch.zeros(bs, dtype=torch.int32, device=device) fm.seqused = None - # bundle out_c*_loc is densely packed in batch order (matches cmp_kv); invalid under chunked prefill + # 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: + # each chunk writes only the ratio-blocks it newly completed. bundle = forward_batch.out_cache_loc_dsv4 for ratio in (4, 128): if ratio not in ratio_lists: continue - bundle_loc = None - if bundle is not None: - bundle_loc = bundle.out_c4_loc if ratio == 4 else bundle.out_c128_loc + bundle_loc = ( + (bundle.out_c4_loc if ratio == 4 else bundle.out_c128_loc) + if bundle is not None + else None + ) setattr( fm, f"c{ratio}_loc", @@ -162,9 +225,10 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): if spt is None: continue for idx in range(bs): - seqlen = int(cu_cpu[idx + 1] - cu_cpu[idx]) - if seqlen == 0: + 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 @@ -193,6 +257,7 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): ) -> dict: result: dict = {} req_pool = req_pool_indices + req_pool_64 = req_pool.to(torch.int64) if seq_lens_max_override is not None: seq_lens_max = int(seq_lens_max_override) @@ -203,40 +268,27 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): 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.to(torch.int64), : n_pages * self.page_size - ] + 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: - state_loc_decode = None if out_cache_loc_dsv4 is not None: - state_loc_decode = ( + state_loc_src = ( out_cache_loc_dsv4.out_c4_state_loc if ratio == 4 else out_cache_loc_dsv4.out_c128_state_loc ) - if state_loc_decode is None: - state_loc_decode = torch.zeros( - bs, - dtype=torch.int32, - device=device, - ) - else: - state_loc_decode = state_loc_decode.to(torch.int32) - compress_out_loc = torch.zeros( - bs, - dtype=torch.int32, - device=device, - ) + # 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: @@ -245,12 +297,26 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): if ratio == 4 else out_cache_loc_dsv4.out_c128_loc ) - n_compress = bundle_loc.numel() - if n_compress > 0: - compress_out_loc[:n_compress] = bundle_loc.to(torch.int32) 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, + device=device, + ) + if bundle_loc is not None: + 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 @@ -264,7 +330,7 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): n_c_tokens = seq_lens_max // ratio else: n_c_tokens = max(1, seq_lens_max // ratio) - slots = c_table[req_pool.to(torch.int64), :n_c_tokens] + slots = c_table[req_pool_64, :n_c_tokens] c_page_table = (slots[:, :: self.page_size] // self.page_size).to( torch.int32 ) @@ -277,9 +343,9 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): if ratio not in (4, 128): continue padding_size = min(bs, bs // ratio + bs) - padding = torch.zeros(padding_size, dtype=torch.int64, device=device) should_compress = ((seq_lens % ratio) == 0) & valid pos_cmp = positions_last[should_compress].to(torch.int64) + (1 - ratio) + padding = torch.zeros(padding_size, dtype=torch.int64, device=device) if pos_cmp.numel() > 0: padding[: pos_cmp.shape[0]].copy_(pos_cmp) result[f"positions_cmp_padding_c{ratio}"] = padding @@ -306,15 +372,7 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): x: torch.Tensor, forward_batch: ForwardBatch, ) -> None: - if ( - forward_batch.forward_mode.is_prefill() - and not forward_batch.forward_mode.is_target_verify() - ): - return self._forward_compress_native(compressor, x, forward_batch) - - from sglang.kernels.ops.attention.deepseek_v4_rope import ( - get_fused_compressor_rope_cos_sin, - ) + from sglang.srt.hardware_backend.npu.dsv4.dsv4_rope import Dsv4NpuRoPE ratio = compressor.ratio coff = 1 + int(compressor.overlap) @@ -341,8 +399,13 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): compressor.layer_id, compressor.is_in_indexer ) - cos, sin = get_fused_compressor_rope_cos_sin( - compressor.freqs_cis, positions_cmp, dtype=torch.float32 + cos, sin = Dsv4NpuRoPE.for_freqs( + compressor.freqs_cis, getattr(compressor, "rotary_emb", None) + ).get_cos_sin( + positions_cmp, + torch.float32, + view_4d=False, + allow_build=False, ) cmp_kv = torch.ops.custom.compressor( @@ -393,7 +456,10 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): x: torch.Tensor, forward_batch: ForwardBatch, ) -> None: - """Per-request unfused compress path. + """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 @@ -419,6 +485,7 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): 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 @@ -449,6 +516,18 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): 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 @@ -510,7 +589,7 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): kv = kv_full[seqlen_offset : seqlen_offset + seqlen] score = score_full[seqlen_offset : seqlen_offset + seqlen] - if overlap and cutoff >= ratio: + 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). @@ -643,28 +722,17 @@ class CompressorAscendBackendMixin(CompressorBackendMixin): 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) - # npu_rotary_mul wants cos/sin in repeat_interleave(2) layout, reshaped - # to (T, 1, 1, rope_dim); cos=real, sin=imag of the complex freqs_cis. rope_dim = compressor.rope_head_dim - # Use the same contig cache as the outer rope path; .real/.imag on a - # complex tensor are strided views and aclnnIndex over them triggers - # StridedSlice (see _get_contig_freqs_real_imag in deepseek_v4_rope.py). - from sglang.kernels.ops.attention.deepseek_v4_rope import ( - _get_contig_freqs_real_imag, - ) + from sglang.srt.hardware_backend.npu.dsv4.dsv4_rope import Dsv4NpuRoPE - freqs_real, freqs_imag = _get_contig_freqs_real_imag(compressor.freqs_cis) - cos_half = freqs_real[pos_out].to(kv_out.dtype) - sin_half = freqs_imag[pos_out].to(kv_out.dtype) - cos = ( - cos_half.repeat_interleave(2, dim=-1) - .view(-1, 1, 1, rope_dim) - .contiguous() - ) - sin = ( - sin_half.repeat_interleave(2, dim=-1) - .view(-1, 1, 1, rope_dim) - .contiguous() + 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) @@ -890,16 +958,30 @@ class C4IndexerAscendBackendMixin(C4IndexerBackendMixin): def _compute_q_npu( self, c4_indexer, q_lora: torch.Tensor, positions: torch.Tensor ) -> torch.Tensor: - from sglang.kernels.ops.attention.deepseek_v4_rope import v4_rope_inplace_npu + from sglang.srt.hardware_backend.npu.dsv4.dsv4_rope import Dsv4NpuRoPE bs = q_lora.shape[0] q, _ = c4_indexer.wq_b(q_lora) q = q.view(bs, c4_indexer.n_local_heads, c4_indexer.head_dim) - v4_rope_inplace_npu( - q[..., -c4_indexer.rope_head_dim :], - None, - c4_indexer.freqs_cis, + qk_nope = c4_indexer.head_dim - c4_indexer.rope_head_dim + # Position-gathered RoPE values are forward-local. The rotary embedding + # object is shared, so retaining them there can leak target positions into + # NextN (or a previous graph replay) when the next batch has the same shape. + cos4, sin4 = Dsv4NpuRoPE.for_freqs( + c4_indexer.freqs_cis, getattr(c4_indexer, "rotary_emb", None) + ).get_cos_sin( positions, + q.dtype, + view_4d=True, + allow_build=False, + cache_dtype=torch.float32, + ) + Dsv4NpuRoPE.apply_rotary_mul_inplace( + q, + None, + cos4, + sin4, + qk_nope_dim=qk_nope, ) return _apply_hadamard(q, c4_indexer.hadamard_matrix) @@ -1111,91 +1193,114 @@ class DeepseekV4AscendAttnBackend( self.forward_metadata = metadata - def _apply_dsv4_graph_metadata(self, forward_batch: ForwardBatch) -> None: - fm = self.forward_metadata - forward_mode = ( - getattr(forward_batch, "global_forward_mode", None) - or forward_batch.forward_mode - ) - actual_forward_mode = getattr(forward_batch, "actual_forward_mode", None) - if actual_forward_mode is None: - actual_forward_mode = forward_batch.forward_mode + @staticmethod + def _copy_2d_with_tail(dst: torch.Tensor, src: torch.Tensor, val: int) -> None: + # Graph replay metadata buffers are sliced to the active bs; only the + # page-column tail needs the sentinel refresh. + r, c = src.shape + dst[:r, :c].copy_(src) + if c < dst.shape[1]: + dst[:, c:].fill_(val) + + @staticmethod + def _copy_1d_with_zero_tail(dst: torch.Tensor, src: Optional[torch.Tensor]) -> None: + if src is None: + dst.zero_() + return + n = src.numel() + assert ( + n <= dst.shape[0] + ), f"graph replay 1D metadata overflow: src={n} > dst={dst.shape[0]}" + if n > 0: + if src.dtype != dst.dtype: + src = src.to(dst.dtype) + dst[:n].copy_(src) + if n < dst.shape[0]: + dst[n:].fill_(0) + + def _build_dsv4_graph_replay_ctx(self, forward_batch: ForwardBatch): + graph_mode = forward_batch.forward_mode + runtime_mode = getattr(forward_batch, "actual_forward_mode", None) or graph_mode bs = forward_batch.batch_size - seq_lens = forward_batch.seq_lens - req_pool_indices = forward_batch.req_pool_indices - device = seq_lens.device - - if forward_mode.is_target_verify() or forward_mode.is_draft_extend_v2(): - tokens_per_req = self.speculative_num_draft_tokens - else: - tokens_per_req = 1 - seq_lens_cpu = forward_batch.seq_lens_cpu assert seq_lens_cpu is not None, "V4 graph replay requires seq_lens_cpu." - if forward_mode.is_target_verify(): - # In graph replay, buffers.seq_lens already contains the attention KV - # length (live length + draft tokens). Padded rows therefore show up as - # tokens_per_req instead of 0. Use the CPU live lengths as the source of - # truth so padded rows stay masked out. + + device = forward_batch.seq_lens.device + tokens_per_bs = ( + self.speculative_num_draft_tokens + 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) - attn_seq_lens = live_seq_lens - if forward_mode.is_target_verify(): - valid_verify_rows = live_seq_lens > 0 - attn_seq_lens = live_seq_lens + int(tokens_per_req) - attn_seq_lens = torch.where(valid_verify_rows, attn_seq_lens, live_seq_lens) - fm.seq_lens_cpu_int = (seq_lens_cpu[:bs] + int(tokens_per_req)).int() + is_idle_replay = runtime_mode.is_idle() + 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 + if active_target_verify: + compress_seq_lens = live_seq_lens + int(tokens_per_bs) + compress_seq_lens_max += int(tokens_per_bs) + + return SimpleNamespace( + forward_batch=forward_batch, + fm=self.forward_metadata, + graph_mode=graph_mode, + runtime_mode=runtime_mode, + is_idle_replay=is_idle_replay, + has_compress=has_compress, + active_target_verify=active_target_verify, + bs=bs, + tokens_per_bs=tokens_per_bs, + device=device, + seq_lens_cpu=seq_lens_cpu, + live_seq_lens=live_seq_lens, + compress_seq_lens=compress_seq_lens, + compress_seq_lens_max=compress_seq_lens_max, + ) + + def _refresh_graph_seq_metadata(self, ctx) -> None: + fm = ctx.fm + 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 + ) + fm.seq_lens_cpu_int = ( + ctx.seq_lens_cpu[: ctx.bs] + int(ctx.tokens_per_bs) + ).int() fm.seq_lens_cpu_int = torch.where( - seq_lens_cpu[:bs] > 0, + ctx.seq_lens_cpu[: ctx.bs] > 0, fm.seq_lens_cpu_int, - seq_lens_cpu[:bs].int(), + ctx.seq_lens_cpu[: ctx.bs].int(), ) fm.actual_seq_lengths_kv.copy_(attn_seq_lens.clamp(min=1)) - pool = self.token_to_kv_pool - out_cache_loc = forward_batch.out_cache_loc - - _verify_compress = ( - forward_mode.is_target_verify() - and actual_forward_mode.is_target_verify() - and bool(self._dsv4_compress_ratios) - ) - _compress_seq_lens = live_seq_lens - _compress_seq_lens_max = int(seq_lens_cpu[:bs].max()) if bs > 0 else 0 - if _verify_compress: - _compress_seq_lens = live_seq_lens + int(tokens_per_req) - _compress_seq_lens_max += int(tokens_per_req) - + def _refresh_graph_compress_page_tables_direct(self, ctx) -> None: result = self._compute_compress_locs( - pool=pool, + pool=self.token_to_kv_pool, req_to_token=self.req_to_token, - req_pool_indices=req_pool_indices[:bs], - seq_lens=_compress_seq_lens, - out_cache_loc=out_cache_loc, - is_decode=forward_mode.is_decode(), - bs=bs, - device=device, + req_pool_indices=ctx.forward_batch.req_pool_indices[: ctx.bs], + seq_lens=ctx.compress_seq_lens, + out_cache_loc=ctx.forward_batch.out_cache_loc, + is_decode=False, + bs=ctx.bs, + device=ctx.device, req_to_token_pool=self.req_to_token_pool, - out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None), + out_cache_loc_dsv4=getattr(ctx.forward_batch, "out_cache_loc_dsv4", None), is_graph=True, - seq_lens_max_override=_compress_seq_lens_max, + seq_lens_max_override=ctx.compress_seq_lens_max, ) - - def _copy_2d(dst: torch.Tensor, src: torch.Tensor, val: int) -> None: - dst.fill_(val) - dst[: src.shape[0], : src.shape[1]].copy_(src) - - def _copy_1d(dst: torch.Tensor, src: torch.Tensor) -> None: - dst.fill_(0) - assert src.shape[0] <= dst.shape[0], ( - f"graph replay 1D metadata overflow: src={src.shape[0]} > " - f"dst={dst.shape[0]}" - ) - dst[: src.shape[0]].copy_(src) - for key in ( "c4_page_table", "c128_page_table", @@ -1203,106 +1308,117 @@ class DeepseekV4AscendAttnBackend( "c128_state_page_table", ): if key in result: - _copy_2d(getattr(fm, key), result[key], 0 if "state" in key else -1) - for key in ("c4_loc", "c128_loc", "c4_state_loc", "c128_state_loc"): - if key in result: - _copy_1d(getattr(fm, key), result[key]) + self._copy_2d_with_tail( + getattr(ctx.fm, key), result[key], 0 if "state" in key else -1 + ) - for key in ( - "positions_cmp_padding_c4", - "positions_cmp_padding_c128", - "start_pos", - "seqused", + def _refresh_graph_decode_compress_1d_direct(self, ctx) -> None: + fm = ctx.fm + bundle = getattr(ctx.forward_batch, "out_cache_loc_dsv4", None) + 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 + positions_last = torch.clamp(ctx.live_seq_lens - 1, min=0) + for ratio in self._dsv4_unique_compress_ratios: + if ratio not in (4, 128): + continue + should_compress = ((ctx.live_seq_lens % ratio) == 0) & valid + pos_cmp = positions_last[should_compress].to(torch.int64) + (1 - ratio) + self._copy_1d_with_zero_tail( + getattr(fm, f"positions_cmp_padding_c{ratio}"), pos_cmp + ) + fm.start_pos.copy_(positions_last.to(torch.int32)) + fm.seqused.copy_(valid.to(torch.int32)) + + 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 = torch.where( + ctx.seq_lens_cpu[: ctx.bs] > 0, + verify_seq_lens_cpu, + ctx.seq_lens_cpu[: ctx.bs], + ) + self._fill_verify_positions_cmp_padding_one( + ctx.forward_batch.positions, + fm.positions_cmp_padding_c4, + 4, + verify_seq_lens_cpu, + n_draft=ctx.tokens_per_bs, + ) + self._fill_verify_positions_cmp_padding_one( + ctx.forward_batch.positions, + fm.positions_cmp_padding_c128, + 128, + verify_seq_lens_cpu, + n_draft=ctx.tokens_per_bs, + ) + fm.start_pos.copy_(ctx.live_seq_lens.to(torch.int32)) + valid = ctx.live_seq_lens[: ctx.bs] > 0 + fm.seqused.copy_( + (valid.to(torch.int32) * int(ctx.tokens_per_bs)).to(device=ctx.device) + ) + bundle = getattr(ctx.forward_batch, "out_cache_loc_dsv4", None) + if bundle is None: + return + for ratio in self._dsv4_unique_compress_ratios: + if ratio not in (4, 128): + continue + loc = bundle.out_c4_loc if ratio == 4 else bundle.out_c128_loc + self._copy_1d_with_zero_tail(getattr(fm, f"c{ratio}_loc"), loc) + + @staticmethod + def _clear_graph_target_verify_metadata(ctx) -> None: + fm = ctx.fm + for tensor in ( + fm.positions_cmp_padding_c4, + fm.positions_cmp_padding_c128, + fm.c4_loc, + fm.c128_loc, + fm.c4_state_loc, + fm.c128_state_loc, ): - if key in result and hasattr(fm, key) and getattr(fm, key) is not None: - _copy_1d(getattr(fm, key), result[key]) + if tensor is not None: + tensor.zero_() + fm.start_pos.zero_() + fm.seqused.zero_() - if _verify_compress: - verify_seq_lens_cpu = seq_lens_cpu[:bs] + int(tokens_per_req) - verify_seq_lens_cpu = torch.where( - seq_lens_cpu[:bs] > 0, - verify_seq_lens_cpu, - seq_lens_cpu[:bs], - ) - self._fill_verify_positions_cmp_padding_one( - forward_batch.positions, - fm.positions_cmp_padding_c4, - 4, - verify_seq_lens_cpu, - n_draft=tokens_per_req, - ) - self._fill_verify_positions_cmp_padding_one( - forward_batch.positions, - fm.positions_cmp_padding_c128, - 128, - verify_seq_lens_cpu, - n_draft=tokens_per_req, - ) - fm.start_pos.copy_(live_seq_lens.to(torch.int32)) - valid = live_seq_lens[:bs] > 0 - fm.seqused.copy_( - (valid.to(torch.int32) * int(tokens_per_req)).to(device=device) - ) - _bundle = getattr(forward_batch, "out_cache_loc_dsv4", None) - if _bundle is not None: - for ratio in self._dsv4_unique_compress_ratios: - if ratio not in (4, 128): - continue - bl = _bundle.out_c4_loc if ratio == 4 else _bundle.out_c128_loc - if bl is not None: - dst_loc = getattr(fm, f"c{ratio}_loc", None) - if dst_loc is not None: - dst_loc.zero_() - bl32 = bl.to(torch.int32) - assert bl32.numel() <= dst_loc.numel(), ( - f"replay verify c{ratio}_loc overflow: " - f"{bl32.numel()} > {dst_loc.numel()}" - ) - dst_loc[: bl32.numel()].copy_(bl32) - - elif ( - forward_mode.is_target_verify() - # The graph may replay a target-verify capture for an idle/padded - # DP rank. There is no real DSV4 allocation bundle in that case; - # zero the compressor metadata so captured writes land in the - # reserved dummy slot instead of reusing stale locs. - and not actual_forward_mode.is_target_verify() - and bool(self._dsv4_compress_ratios) - ): - for tensor in ( - fm.positions_cmp_padding_c4, - fm.positions_cmp_padding_c128, - fm.c4_loc, - fm.c128_loc, - fm.c4_state_loc, - fm.c128_state_loc, - ): - if tensor is not None: - tensor.zero_() - fm.start_pos.zero_() - fm.seqused.zero_() - - swa_loc = pool.translate_loc_from_full_to_swa(out_cache_loc).to(torch.int64) - _copy_1d(fm.swa_loc, swa_loc) + 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( + ctx.forward_batch.out_cache_loc + ).to(torch.int64) + self._copy_1d_with_zero_tail(fm.swa_loc, swa_loc) swa_src = ( fm.block_tables_swa if fm.block_tables_swa is not None else fm.block_tables ) - _copy_2d(fm.swa_page_table, swa_src, -1) - # base replay 0-pads the tail but page 0 is a real page; restore the -1 sentinel beyond valid pages - if bs > 0: - _spec = int(getattr(self, "speculative_num_draft_tokens", 0) or 0) - max_len = int(seq_lens_cpu[:bs].max()) + _spec + 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_seq_pages = (max_len + self.page_size - 1) // self.page_size - if 0 < max_seq_pages < fm.swa_page_table.shape[1]: - fm.swa_page_table[:, max_seq_pages:].fill_(-1) + 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_kernel_metadata(self, ctx) -> None: + fm = ctx.fm kernel_metadata_new = self._kernel_metadata_from_parts( - bs=bs, + bs=ctx.bs, actual_seq_lengths_q_pa=fm.actual_seq_lengths_q_pa, actual_seq_lengths_kv=fm.actual_seq_lengths_kv, block_tables=fm.block_tables, - max_seqlen_q=tokens_per_req, + max_seqlen_q=ctx.tokens_per_bs, is_nextn=False, ) for key in ( @@ -1313,11 +1429,26 @@ class DeepseekV4AscendAttnBackend( ): if key in kernel_metadata_new: fm.kernel_metadata[key].copy_(kernel_metadata_new[key]) - - # -1 sentinel; the indexer overwrites valid rows each step fm.c4_topk_indices.fill_(-1) - self.forward_metadata = fm + def _apply_dsv4_graph_metadata(self, forward_batch: ForwardBatch) -> None: + ctx = self._build_dsv4_graph_replay_ctx(forward_batch) + + self._refresh_graph_seq_metadata(ctx) + self._refresh_graph_compress_page_tables_direct(ctx) + + if ctx.graph_mode.is_decode(): + self._refresh_graph_decode_compress_1d_direct(ctx) + elif ctx.graph_mode.is_target_verify(): + if ctx.is_idle_replay and ctx.has_compress: + self._clear_graph_target_verify_metadata(ctx) + elif ctx.active_target_verify: + self._refresh_graph_target_verify_compress_1d_direct(ctx) + + self._refresh_graph_swa_metadata_direct(ctx) + self._refresh_graph_kernel_metadata(ctx) + + self.forward_metadata = ctx.fm def init_forward_metadata(self, forward_batch: ForwardBatch) -> None: super().init_forward_metadata(forward_batch) 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 7060f90ce..7b57e709f 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py @@ -406,7 +406,7 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): ) def compute_dsv4_state_lens_extend( - self, reqs: List[Req], seq_lens: List[int] + 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). @@ -437,15 +437,25 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): c4_seq: List[int] = [] c128_prefix: List[int] = [] c128_seq: List[int] = [] - for req, seq_len in zip(reqs, seq_lens): + 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_alloc_len - new_c128 = prev_c128 + c128_alloc_len + new_c4 = prev_c4 + c4_count + new_c128 = prev_c128 + c128_count c4_prefix.append(prev_c4) c4_seq.append(new_c4) @@ -454,8 +464,8 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): req.c4_state_kv_len = new_c4 req.c128_state_kv_len = new_c128 - req.c4_state_alloc_offset = seq_len - c4_alloc_len - req.c128_state_alloc_offset = seq_len - c128_alloc_len + 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, @@ -584,9 +594,32 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): last_loc, extend_num_tokens, ) + return self._wrap_full_alloc( + out_full_loc, + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + last_loc.dtype, + req_pool_indices, + dsv4_state_lens, + ) + + def _wrap_full_alloc( + self, + out_full_loc, + prefix_lens, + prefix_lens_cpu, + seq_lens, + 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. if out_full_loc is None: return None - out_swa_loc = self.translate_loc_from_full_to_swa(out_full_loc) assert out_swa_loc is not None, ( "translate_loc_from_full_to_swa returned None — " @@ -599,7 +632,7 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): prefix_lens_cpu, seq_lens, seq_lens_cpu, - last_loc.dtype, + loc_dtype, req_pool_indices, dsv4_state_lens, ) @@ -636,6 +669,44 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): dsv4_state_lens, ) + def alloc_extend_swa_tail( + self, + prefix_lens: torch.Tensor, + prefix_lens_cpu: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, + extend_num_tokens: int, + 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. + """ + self._cur_req_to_token_pool = req_to_token_pool + out_full_loc = super().alloc_extend_swa_tail( + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + last_loc, + extend_num_tokens, + swa_tail_len, + ) + return self._wrap_full_alloc( + out_full_loc, + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + last_loc.dtype, + req_pool_indices, + dsv4_state_lens, + ) + def free( self, free_index: Optional[torch.Tensor] = None, @@ -676,8 +747,10 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): 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. - allocator.free(slots.to(torch.int64)) + if slots.numel() > 0: + allocator.free(slots.to(torch.int64)) # State pools: free only the tail [c{N}_state_alloc_offset, kv_len). for ratio, allocator, table_attr, off_attr in ( @@ -699,7 +772,9 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): off = getattr(req, off_attr, 0) if kv_len > off: slots = getattr(req_to_token_pool, table_attr)[req_pool_idx, off:kv_len] - allocator.free(slots.to(torch.int64)) + slots = slots[slots > 0] + if slots.numel() > 0: + allocator.free(slots.to(torch.int64)) def backup_state(self): # EAGLE/NEXTN draft preprocess allocates speculative c{4,128} KV via 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 4e5e7e6dc..393fe0744 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 @@ -13,13 +13,9 @@ these hooks then: Non-DSV4 paths leave ``batch.out_cache_loc_dsv4`` None, so this module is a no-op for them. -TODO: the disagg DSV4 path bypasses these hooks — it calls -``allocator.alloc_extend`` directly then ``req_to_token_pool.write`` without -going through ``mem_cache/common.py`` (see ``disaggregation/decode.py``). The -DSV4OutCacheLoc bundle is still produced but never written into the per-req -tables, so disagg + DSV4 is unsupported here (c-pages leak). Fixing requires -calling these hooks from disagg's per-req alloc loop, or moving the write -into the allocator itself. +The disagg per-req prealloc path does not build a ``ScheduleBatch`` and so +bypasses the batch hook; it writes the same tables via +``write_dsv4_prealloc_tables`` (driven by ``dsv4_unwrap_prealloc``). """ from __future__ import annotations @@ -61,7 +57,199 @@ def maybe_write_dsv4_extend( if not hasattr(req_to_token_pool, "write_c4"): return # non-DSV4 pool; skip defensively (shouldn't happen) - # SWA writes: prefix..seq token positions, one slot per raw token. + # 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, + ) + + +def dsv4_state_payloads( + req_to_token_pool, + 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 {} + + import numpy as np + + from sglang.srt.disaggregation.ascend.conn import AscendStateType + + 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): + if hi <= lo: + return empty_pages() + + 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 + ), + } + + +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"): + 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 + 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 + write_dsv4_prealloc_tables(req_to_token_pool, req, prefix_len, fill_len, kv_loc) + return kv_loc.out_full_loc + + +def write_dsv4_prealloc_tables( + req_to_token_pool, + req: Req, + prefix_len: int, + fill_len: int, + bundle, +) -> None: + """Write the five 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"): + return + rp = torch.tensor([req.req_pool_idx]) + pl = torch.tensor([prefix_len]) + sl = torch.tensor([fill_len]) + + _write_dsv4_tables( + req_to_token_pool, + rp, + pl, + sl, + bundle, + c4_state_offsets=[getattr(req, "c4_state_alloc_offset", 0)], + c128_state_offsets=[getattr(req, "c128_state_alloc_offset", 0)], + ) + + +def _write_dsv4_tables( + req_to_token_pool, + req_pool_indices_cpu: torch.Tensor, + 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, @@ -70,8 +258,6 @@ def maybe_write_dsv4_extend( bundle.out_swa_loc, ratio=1, ) - - # c4 / c128 writes: prefix//ratio .. seq//ratio compressed positions. _write_per_req_slice( req_to_token_pool.write_c4, req_pool_indices_cpu, @@ -89,25 +275,13 @@ def maybe_write_dsv4_extend( ratio=128, ) - # 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_alloc_offset", 0) for r in batch.reqs - ] - if c128_state_alloc_offsets is None: - c128_state_alloc_offsets = [ - getattr(r, "c128_state_alloc_offset", 0) for r in batch.reqs - ] if bundle.out_c4_state_loc is not None and hasattr( req_to_token_pool, "write_c4_state" ): _write_state_tail_per_req( req_to_token_pool.write_c4_state, req_pool_indices_cpu, - c4_state_alloc_offsets, + c4_state_offsets, seq_lens_cpu, bundle.out_c4_state_loc, ) @@ -117,7 +291,7 @@ def maybe_write_dsv4_extend( _write_state_tail_per_req( req_to_token_pool.write_c128_state, req_pool_indices_cpu, - c128_state_alloc_offsets, + c128_state_offsets, seq_lens_cpu, bundle.out_c128_state_loc, ) @@ -400,5 +574,7 @@ def _free_state_range( 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] - state_allocator.free(free_slots.to(torch.int64)) + 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 a83771893..c8b1172a5 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 @@ -29,7 +29,7 @@ The subclass overrides only: from __future__ import annotations import math -from typing import Optional, Tuple +from typing import List, Optional, Tuple import torch import torch_npu @@ -167,6 +167,7 @@ class NPUCompressStatePool(CompressStatePool): 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. @@ -388,6 +389,81 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): kernel_page_size=self.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)) + ) + components.append( + (AscendStateType.DSV4_C128_STATE, *state_entry(128, include_indexer=False)) + ) + + # 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_cache(self, layer_id: int, from_indexer: bool) -> torch.Tensor: """fp32 ``[block_num, page_size, 2*coff*D]`` view of this layer's kv+score buffer — the fused compressor op 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 2fb68ea6d..a5efc5448 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 @@ -35,32 +35,24 @@ from __future__ import annotations 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.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter -class DSV4NPUReqToTokenPool(ReqToTokenPool): - """ReqToTokenPool extended with DSV4 SWA + c4/c128 per-req tables. +class DSV4ReqToTokenTablesMixin: + """Shared DSV4-NPU per-req table logic for the prefill/normal pool + (:class:`DSV4NPUReqToTokenPool`) and the disagg-decode pool + (:class:`DSV4NPUDecodeReqToTokenPool`), which differ only in their base. - 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. + Host class must call ``super().__init__(...)`` first (so ``_alloc_size`` + exists) then ``self._init_dsv4_tables(...)``; ``free`` should call + ``self._dsv4_free(req)`` before delegating to the base ``free``. """ - def __init__( - self, - size: int, - max_context_len: int, - device: str, - enable_memory_saver: bool, - ): - super().__init__(size, max_context_len, device, enable_memory_saver) - + def _init_dsv4_tables( + self, max_context_len: int, device: str, enable_memory_saver: bool + ) -> None: memory_saver_adapter = TorchMemorySaverAdapter.create( enable=enable_memory_saver ) @@ -91,11 +83,8 @@ class DSV4NPUReqToTokenPool(ReqToTokenPool): ), ) - # ------------------------------------------------------------------ # 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 @@ -113,16 +102,68 @@ class DSV4NPUReqToTokenPool(ReqToTokenPool): 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. This is a - one-way ref (pool -> allocator). The reverse direction (the allocator - reading these per-req tables for its c-pool / state last_loc lookup) is - no longer a stored back-ref: mem_cache/common.py passes this pool into - ``alloc_extend`` / ``alloc_decode`` per call instead.""" + release c4/c128 pool pages alongside the req_pool_idx slot.""" self._dsv4_allocator = allocator - def free(self, req): + def _dsv4_free(self, req) -> None: # Trigger c4/c128 free via the allocator's unified free 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. + + 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. + """ + + def __init__( + self, + size: int, + max_context_len: int, + device: str, + 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) + + def free(self, req): + self._dsv4_free(req) + super().free(req) + + +class DSV4NPUDecodeReqToTokenPool(DSV4ReqToTokenTablesMixin, DecodeReqToTokenPool): + """DecodeReqToTokenPool with the DSV4 swa/c4/c128(+state) per-req tables. + + The disagg-decode counterpart of DSV4NPUReqToTokenPool; DecodeReqToTokenPool + pre-allocates extra req slots for in-flight prefill transfers. + """ + + def __init__( + self, + size: int, + max_context_len: int, + device: str, + enable_memory_saver: bool, + pre_alloc_size: int, + ): + super().__init__( + size=size, + max_context_len=max_context_len, + device=device, + enable_memory_saver=enable_memory_saver, + pre_alloc_size=pre_alloc_size, + ) + self._init_dsv4_tables(max_context_len, device, enable_memory_saver) + + def free(self, req): + self._dsv4_free(req) super().free(req) diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_rope.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_rope.py new file mode 100644 index 000000000..aee422f3e --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_rope.py @@ -0,0 +1,181 @@ +"""NPU interleaved RoPE cos/sin cache for DeepSeek-V4 on Ascend. + +One Dsv4NpuRoPE per freqs_cis (singleton by id). Tables are built once at +init and registered as buffers on the shared rotary_emb, so model.to() moves +them and a captured aclgraph sees stable tensors; decode only does index_select. + +mscale: cos/sin stored in freqs_cis must already be pre-multiplied by the YARN +mscale at precompute time (see precompute_freqs_cis). We just read what's stored. +""" + +from typing import Optional + +import torch + + +class Dsv4NpuRoPE: + """Interleaved cos/sin tables, layout [c0,c0,c1,c1,...] / [s0,s0,s1,s1,...].""" + + # id(freqs_cis) -> instance. freqs_cis is a module buffer, lives with the model. + _instances: dict[int, "Dsv4NpuRoPE"] = {} + + def __init__( + self, freqs_cis: torch.Tensor, rotary_emb: Optional[object] = None + ) -> None: + self.freqs_cis = freqs_cis + # cos/sin registered as buffers on this module (None -> fall back to _tables). + self.rotary_emb = rotary_emb + # contiguous real/imag halves of complex freqs_cis [max_pos, rope_dim/2]; + # .real/.imag are strided views, materialize once to avoid per-call + # StridedSlice from aclnnIndex over the strided views. + self._real_imag: Optional[tuple[torch.Tensor, torch.Tensor]] = None + self._tables: dict[ + tuple[torch.dtype, torch.device], tuple[torch.Tensor, torch.Tensor] + ] = {} + + @classmethod + def for_freqs( + cls, freqs_cis: torch.Tensor, rotary_emb: Optional[object] = None + ) -> "Dsv4NpuRoPE": + # rotary_emb is only used at creation; callers sharing a warmed-up freqs_cis may omit it. + inst = cls._instances.get(id(freqs_cis)) + if inst is None or inst.freqs_cis is not freqs_cis: + inst = cls(freqs_cis, rotary_emb) + cls._instances[id(freqs_cis)] = inst + return inst + + def _contig_real_imag(self) -> tuple[torch.Tensor, torch.Tensor]: + if self._real_imag is None: + self._real_imag = ( + self.freqs_cis.real.contiguous(), + self.freqs_cis.imag.contiguous(), + ) + return self._real_imag + + @staticmethod + def _buffer_names(dtype: torch.dtype) -> tuple[str, str]: + suffix = str(dtype).replace("torch.", "").replace(".", "_") + return ( + f"_npu_interleaved_rope_cos_cache_{suffix}", + f"_npu_interleaved_rope_sin_cache_{suffix}", + ) + + def _register_or_set_buffer(self, name: str, tensor: torch.Tensor) -> None: + owner = self.rotary_emb + if hasattr(owner, "register_buffer"): + if name in getattr(owner, "_buffers", {}): + setattr(owner, name, tensor) + else: + owner.register_buffer(name, tensor, persistent=False) + else: + setattr(owner, name, tensor) + + def ensure_tables( + self, dtype: torch.dtype, *, allow_build: bool = True + ) -> tuple[torch.Tensor, torch.Tensor]: + # Returns [max_pos, rope_dim] tables. Call once at init (allow_build=True); + # decode uses allow_build=False (no repeat_interleave inside the captured graph). + expected_shape = (self.freqs_cis.shape[0], self.freqs_cis.shape[1] * 2) + + if self.rotary_emb is not None: + cos_name, sin_name = self._buffer_names(dtype) + cos = getattr(self.rotary_emb, cos_name, None) + sin = getattr(self.rotary_emb, sin_name, None) + if ( + cos is not None + and sin is not None + and tuple(cos.shape) == expected_shape + and tuple(sin.shape) == expected_shape + and cos.dtype == dtype + and sin.dtype == dtype + and cos.device == self.freqs_cis.device + and sin.device == self.freqs_cis.device + ): + return cos, sin + else: + cached = self._tables.get((dtype, self.freqs_cis.device)) + if cached is not None: + cos, sin = cached + if ( + tuple(cos.shape) == expected_shape + and tuple(sin.shape) == expected_shape + ): + return cached + + if not allow_build: + raise RuntimeError( + "NPU interleaved RoPE cache is missing in a no-build path. " + "Initialize it before forward to keep decode free of repeat_interleave." + ) + + real_contig, imag_contig = self._contig_real_imag() + cos = real_contig.repeat_interleave(2, dim=-1).to(dtype=dtype).contiguous() + sin = imag_contig.repeat_interleave(2, dim=-1).to(dtype=dtype).contiguous() + + if self.rotary_emb is not None: + cos_name, sin_name = self._buffer_names(dtype) + self._register_or_set_buffer(cos_name, cos) + self._register_or_set_buffer(sin_name, sin) + else: + self._tables[(dtype, self.freqs_cis.device)] = (cos, sin) + return cos, sin + + def get_cos_sin( + self, + positions: torch.Tensor, + dtype: torch.dtype, + *, + view_4d: bool = False, + inverse: bool = False, + allow_build: bool = True, + cache_dtype: Optional[torch.dtype] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # positions: [T]. Returns [T, rope_dim], or [T, 1, 1, rope_dim] if view_4d. + # Position-gathered tensors are forward-local; do not cache across forwards + # or MTP decode reuses the previous step's RoPE when only positions change. + cache_dtype = dtype if cache_dtype is None else cache_dtype + cos_cache, sin_cache = self.ensure_tables(cache_dtype, allow_build=allow_build) + cos = cos_cache.index_select(0, positions) + sin = sin_cache.index_select(0, positions) + if inverse: + sin = -sin + if cos.dtype != dtype: + cos = cos.to(dtype) + sin = sin.to(dtype) + if view_4d: + rope_dim = cos.shape[-1] + cos = cos.view(-1, 1, 1, rope_dim) + sin = sin.view(-1, 1, 1, rope_dim) + return cos, sin + + @staticmethod + def apply_rotary_mul_inplace( + q_rope: torch.Tensor, + kv_rope: Optional[torch.Tensor], + cos4: torch.Tensor, + sin4: torch.Tensor, + qk_nope_dim: int = 0, + ) -> None: + # q_rope: [T, n_heads, head_dim]; cos4/sin4: [T, 1, 1, rope_dim]; + # kv_rope: [T, 1, head_dim] or None. Prefer the NPU kernel: torch accumulates + # bf16 muls in bf16 while the kernel uses fp32; drift compounds and flips argmax. + rope_dim = cos4.shape[-1] + torch.ops.custom.inplace_partial_rotary_mul( + q_rope.unsqueeze(1), + cos4, + sin4, + rotary_mode="interleave", + partial_slice=[qk_nope_dim, qk_nope_dim + rope_dim], + ) + if kv_rope is not None: + if kv_rope.dim() == 3: + kv_view = kv_rope.unsqueeze(1) + else: + kv_view = kv_rope.view(-1, 1, 1, rope_dim) + torch.ops.custom.inplace_partial_rotary_mul( + kv_view, + cos4, + sin4, + rotary_mode="interleave", + partial_slice=[qk_nope_dim, qk_nope_dim + rope_dim], + ) diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py index 394ff813f..fd230be5b 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py @@ -18,7 +18,7 @@ from typing import TYPE_CHECKING import torch -from sglang.srt.configs.model_config import is_deepseek_dsa +from sglang.srt.configs.model_config import is_deepseek_dsa, is_deepseek_v4 from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import ( EAGLEDraftExtendCudaGraphRunner, ) @@ -35,7 +35,8 @@ class EAGLEDraftExtendNpuGraphRunner(EAGLEDraftExtendCudaGraphRunner): return torch.int32 def _replay_graph(self, shape_key, forward_batch): - if not is_deepseek_dsa(self.model_runner.model_config.hf_config): + hf_config = self.model_runner.model_config.hf_config + if not (is_deepseek_dsa(hf_config) and is_deepseek_v4(hf_config)): seq_lens = forward_batch.seq_lens_cpu.tolist() + [0] * ( self.bs - self.raw_bs ) diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py index 78c82d864..006f46ee6 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py @@ -18,7 +18,11 @@ from typing import TYPE_CHECKING, Dict, Union import torch -from sglang.srt.configs.model_config import AttentionArch, is_deepseek_dsa +from sglang.srt.configs.model_config import ( + AttentionArch, + is_deepseek_dsa, + is_deepseek_v4, +) from sglang.srt.speculative.eagle_draft_cuda_graph_runner import ( EAGLEDraftCudaGraphRunner, ) @@ -83,7 +87,8 @@ class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner): return bool(decision.item()) def _replay_graph(self, shape_key, forward_batch): - if not is_deepseek_dsa(self.model_runner.model_config.hf_config): + hf_config = self.model_runner.model_config.hf_config + if not (is_deepseek_dsa(hf_config) and is_deepseek_v4(hf_config)): seq_lens_for_each_draft_step = [] for speculative_step_id in range(self.speculative_num_steps - 1): seq_lens_cpu = ( diff --git a/python/sglang/srt/hardware_backend/npu/moe/topk.py b/python/sglang/srt/hardware_backend/npu/moe/topk.py index ffe772b04..aa4d85299 100644 --- a/python/sglang/srt/hardware_backend/npu/moe/topk.py +++ b/python/sglang/srt/hardware_backend/npu/moe/topk.py @@ -43,8 +43,31 @@ def fused_topk_npu( renormalize = topk_config.renormalize correction_bias = topk_config.correction_bias + # sqrtsoftplus (DSV4 noaux_tc): top-k over (scores + bias); weights from + # un-biased scores. The custom op fuses softplus/sqrt/topk/gather/norm/cast. + if topk_config.scoring_func == "sqrtsoftplus": + routed_scaling_factor = ( + topk_config.routed_scaling_factor + if topk_config.apply_routed_scaling_factor_on_output + else 1.0 + ) + topk_weights, topk_ids, _ = torch.ops.custom.npu_moe_gating_top_k( + x=router_logits.to(torch.float32), + k=topk_config.top_k, + bias=( + correction_bias.to(torch.float32) + if correction_bias is not None + else None + ), + input_ids=None, + tid2eid=None, + routed_scaling_factor=float(routed_scaling_factor), + norm_type=2, + ) + topk_weights = topk_weights.to(torch.float32) + # Fast path: simple top-k without grouped routing and bias - if not use_grouped_topk and correction_bias is None: + elif not use_grouped_topk and correction_bias is None: topk_weights, topk_ids, _ = torch.ops.npu.npu_moe_gating_top_k_softmax( router_logits, k=topk_config.top_k, @@ -58,26 +81,6 @@ def fused_topk_npu( ) topk_weights = topk_weights.to(torch.float32) - # sqrtsoftplus (DSV4 noaux_tc): the NPU op only scores sigmoid/softmax, so use - # a torch path. top-k over (scores + bias); weights from un-biased scores. - elif topk_config.scoring_func == "sqrtsoftplus": - scores = torch.nn.functional.softplus(router_logits.float()).sqrt() - scores_for_choice = ( - scores + correction_bias.unsqueeze(0).float() - if correction_bias is not None - else scores - ) - _, topk_ids = torch.topk( - scores_for_choice, k=topk_config.top_k, dim=-1, sorted=False - ) - topk_ids = topk_ids.to(torch.int32) - topk_weights = scores.gather(1, topk_ids) - if renormalize: - topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) - else: - topk_weights = topk_weights * topk_config.routed_scaling_factor - topk_weights = topk_weights.to(torch.float32) - # Support grouped top-k or correction bias or sigmoid or routed_scaling_factor elif ( correction_bias is not None diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index 9ab8b166c..f33c74f77 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -38,12 +38,13 @@ from sglang.srt.utils import ( ) _is_npu = is_npu() +_use_zbal = _is_npu and envs.SGLANG_ZBAL_LOCAL_MEM_SIZE.get() > 0 if TYPE_CHECKING: from sglang.srt.batch_overlap.single_batch_overlap import CombineOverlapArgs try: - if _is_npu and envs.SGLANG_ZBAL_LOCAL_MEM_SIZE.get() > 0: + if _use_zbal: from zbal.zbal.deepep_adaptor import Config from zbal.zbal_buffer import Buffer else: @@ -753,7 +754,11 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): self.num_max_dispatch_tokens_per_rank, self.num_experts, use_fp8=self.use_fp8, - **(dict(topk_weights=topk_weights) if _is_npu else dict()), + **( + dict(topk_weights=topk_weights) + if _is_npu and not _use_zbal + else dict() + ), **(dict(use_nvfp4=True) if self.use_nvfp4 else dict()), **( dict(x_global_scale=input_global_scale) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index c5ac9a6e7..892c20305 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -41,7 +41,12 @@ from torch.distributed import barrier from sglang.kernels.ops.mamba.triton_ops import ( initialize_mamba_selective_state_update_backend, ) -from sglang.srt.configs.model_config import ModelConfig, ModelImpl, is_minimax_sparse +from sglang.srt.configs.model_config import ( + ModelConfig, + ModelImpl, + is_deepseek_v4, + is_minimax_sparse, +) from sglang.srt.constrained.grammar_manager import GrammarManager from sglang.srt.debug_utils.pr_fix_toggle import maybe_revert_pr_fix from sglang.srt.disaggregation.decode import ( @@ -291,6 +296,13 @@ else: logger = logging.getLogger(__name__) + +def _prewarm_hccl_group(device, group, device_module): + warmup_tensor = torch.zeros(1, dtype=torch.int32, device=device) + torch.distributed.all_reduce(warmup_tensor, group=group) + device_module.synchronize() + + # Test retract decode for debugging purposes TEST_RETRACT = envs.SGLANG_TEST_RETRACT.get() TEST_RETRACT_INTERVAL = envs.SGLANG_TEST_RETRACT_INTERVAL.get() @@ -488,6 +500,22 @@ class Scheduler( self.disable_radix_cache = result.disable_radix_cache self.tree_cache = result.tree_cache + if _is_npu and is_deepseek_v4( + self.tp_worker.model_runner.model_config.hf_config + ): + rank = ( + self.ps.dp_rank + if self.ps.dp_rank is not None + else self.tp_group.rank_in_group + ) + logger.info("HCCL DP prewarm start: rank=%s", rank) + _prewarm_hccl_group( + device=self.tp_group.device, + group=self.tp_group.device_group, + device_module=self.tp_group.device_module, + ) + logger.info("HCCL DP prewarm done: rank=%s", rank) + if (c := self.tp_worker.model_runner.canary_manager) is not None: c.attach_radix_cache(self.tree_cache) diff --git a/python/sglang/srt/mem_cache/allocation.py b/python/sglang/srt/mem_cache/allocation.py index 3d6c8cf86..28c17b9b2 100644 --- a/python/sglang/srt/mem_cache/allocation.py +++ b/python/sglang/srt/mem_cache/allocation.py @@ -178,10 +178,20 @@ def _compute_dsv4_state_lens(batch, *, is_decode: bool): 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() + batch.reqs, batch.seq_lens_cpu.tolist(), prefix_lens ) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 3835f5fc7..264cbc337 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -682,7 +682,12 @@ class KVCacheConfigurator: extra_max_context_len: int, pre_alloc_size: int, ) -> ReqToTokenPool: - from sglang.srt.disaggregation.decode import DecodeReqToTokenPool + if _is_npu and is_deepseek_v4(self.model_config.hf_config): + from sglang.srt.hardware_backend.npu.dsv4.dsv4_req_to_token_pool import ( + DSV4NPUDecodeReqToTokenPool as DecodeReqToTokenPool, + ) + else: + from sglang.srt.disaggregation.decode import DecodeReqToTokenPool req_to_token_pool = DecodeReqToTokenPool( size=max_num_reqs, diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 4014e4e6c..a21936d73 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -23,9 +23,6 @@ import torch.nn as nn import torch.nn.functional as F import sglang.srt.models.deepseek_v2 as deepseek_v2 -from sglang.kernels.ops.attention.deepseek_v4_rope import ( - v4_rope_inplace_npu, -) from sglang.kernels.ops.attention.dsv4 import ( fused_norm_rope_inplace, fused_q_norm_rope, @@ -47,6 +44,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.environ import envs from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation +from sglang.srt.hardware_backend.npu.dsv4.dsv4_rope import Dsv4NpuRoPE from sglang.srt.layers.attention.dsa.utils import ( can_dsa_cp_split, dsa_use_prefill_cp, @@ -622,6 +620,11 @@ class MQALayer(MqaAttentionBase): device=get_server_args().device, ) + if _is_npu: + Dsv4NpuRoPE.for_freqs( + self.freqs_cis, getattr(self, "rotary_emb", None) + ).ensure_tables(torch.float32) + if _is_hip: cos_cache = ( self.freqs_cis.real.to(torch.bfloat16).unsqueeze(-2).unsqueeze(-2) @@ -686,6 +689,25 @@ class MQALayer(MqaAttentionBase): # (`_compute_kv_to_cache`), so the legacy "overlap store cache" flag # has no effect here -- the fused path is on by default. + def _get_npu_rope_position_cache( + self, positions: torch.Tensor, dtype: torch.dtype, inverse: bool = False + ) -> Tuple[torch.Tensor, torch.Tensor]: + # ``rotary_emb`` is shared by layers with the same RoPE configuration and + # can also be shared by the target and NextN models. Only cache the + # immutable full table on it. A position-gathered tensor is specific to + # this forward and reusing it based on shape alone gives MTP decode the + # previous step's RoPE values when positions change but batch size does not. + return Dsv4NpuRoPE.for_freqs( + self.freqs_cis, getattr(self, "rotary_emb", None) + ).get_cos_sin( + positions, + dtype, + view_4d=True, + inverse=inverse, + allow_build=False, + cache_dtype=torch.float32, + ) + def _compute_q_a( self, x: torch.Tensor, @@ -1059,11 +1081,15 @@ class MQALayer(MqaAttentionBase): kv, _ = self.wkv(x) kv = self.kv_norm(kv) - v4_rope_inplace_npu( - q[..., -self.qk_rope_head_dim :], - kv[..., -self.qk_rope_head_dim :].unsqueeze(1), - self.freqs_cis, - positions, + cos4, sin4 = self._get_npu_rope_position_cache( + positions, q.dtype, inverse=False + ) + Dsv4NpuRoPE.apply_rotary_mul_inplace( + q, + kv.unsqueeze(1), + cos4, + sin4, + qk_nope_dim=self.qk_nope_head_dim, ) attn_backend.store_cache( layer_id=self.layer_id, @@ -1259,12 +1285,15 @@ class MQALayer(MqaAttentionBase): ) o = o[:, tp_slice, :] if _is_npu: - v4_rope_inplace_npu( - o[..., -self.qk_rope_head_dim :], + cos4, sin4 = self._get_npu_rope_position_cache( + positions, o.dtype, inverse=True + ) + Dsv4NpuRoPE.apply_rotary_mul_inplace( + o, None, - self.freqs_cis, - positions, - inverse=True, + cos4, + sin4, + qk_nope_dim=self.qk_nope_head_dim, ) else: fused_rope_inplace( @@ -1354,14 +1383,18 @@ class DeepseekV4DecoderLayer(nn.Module): layer_id=layer_id, quant_config=quant_config, prefix=add_prefix("self_attn", prefix), - alt_streams=alt_streams, + alt_streams=None if _is_npu else alt_streams, compress_ratio_override=compress_ratio_override, ) moe_alt_stream = ( alt_streams[0] if ( alt_streams is not None - and (_is_cuda or envs.SGLANG_ROCM_USE_MULTI_STREAM.get()) + and ( + _is_cuda + or envs.SGLANG_ROCM_USE_MULTI_STREAM.get() + or envs.SGLANG_NPU_USE_MULTI_STREAM.get() + ) ) else None ) @@ -2105,16 +2138,21 @@ class DeepseekV4Model(nn.Module): else: self.embed_tokens = PPMissingLayer() self.rms_norm_eps = config.rms_norm_eps - use_stream_pool = _is_cuda or ( - _is_hip - and ( - envs.SGLANG_ROCM_USE_MULTI_STREAM.get() - or envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get() + use_stream_pool = ( + _is_cuda + or ( + _is_hip + and ( + envs.SGLANG_ROCM_USE_MULTI_STREAM.get() + or envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get() + ) ) + 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 self.alt_streams = ( - [torch.cuda.Stream() for _ in range(num_alt_streams)] + [device_module.Stream() for _ in range(num_alt_streams)] if use_stream_pool else None ) @@ -2339,7 +2377,6 @@ class DeepseekV4Model(nn.Module): for _attr in ("freqs_cis_c4", "freqs_cis_c128"): if hasattr(forward_batch, _attr): delattr(forward_batch, _attr) - capture_dspark = self.dspark_layers_to_capture is not None if capture_dspark and dsa_use_prefill_cp(forward_batch): raise NotImplementedError(