minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[NPU] Optimize DeepSeek-V4 performance (#31931)

Browse Source
This commit is contained in:
Talantan1102
2026-07-28 19:45:46 +08:00
committed by GitHub
parent 51397af885
commit 5558dbad00
21 changed files with 1311 additions and 523 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/deepseek_v4_rope.py
-120
View File
@@ -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],
)
python/sglang/srt/disaggregation/ascend/conn.py
+44 -8
View File
@@ -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
python/sglang/srt/disaggregation/decode.py
+75 -19
View File
@@ -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
python/sglang/srt/disaggregation/mooncake/conn.py
+18 -7
View File
@@ -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
python/sglang/srt/disaggregation/prefill.py
+34 -18
View File
@@ -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
python/sglang/srt/disaggregation/utils.py
+16 -1
View File
@@ -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 "
python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
+3 -3
View File
@@ -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
)
python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py
+362 -231
View File
@@ -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)
python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py
+85 -10
View File
@@ -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
python/sglang/srt/hardware_backend/npu/dsv4/dsv4_common_hooks.py
+201 -25
View File
@@ -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)
python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py
+77 -1
View File
@@ -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
python/sglang/srt/hardware_backend/npu/dsv4/dsv4_req_to_token_pool.py
+69 -28
View File
@@ -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)
python/sglang/srt/hardware_backend/npu/dsv4/dsv4_rope.py
+181
View File
@@ -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],
)
python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py
+3 -2
View File
@@ -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
)
python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py
+7 -2
View File
@@ -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 = (
python/sglang/srt/hardware_backend/npu/moe/topk.py
+24 -21
View File
@@ -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
python/sglang/srt/layers/moe/token_dispatcher/deepep.py
+7 -2
View File
@@ -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)
python/sglang/srt/managers/scheduler.py
+29 -1
View File
@@ -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)
python/sglang/srt/mem_cache/allocation.py
+11 -1
View File
@@ -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
)
python/sglang/srt/mem_cache/kv_cache_configurator.py
+6 -1
View File
@@ -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,
python/sglang/srt/models/deepseek_v4.py
+59 -22
View File
@@ -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(