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

[NPU] Add NPU arch35 support and enhance DSV4 processing in DeepSeek-V4 (#37373)

Browse Source 
Co-authored-by: AndyLi429 <AndyLi429@noreply.gitcode.com>
Co-authored-by: Kailong Lu <kelonlu@163.com>
Co-authored-by: cx <chengxin65@huawei.com>
Co-authored-by: ranjiewen <ranjiewen@huawei.com>
Co-authored-by: HEX1A0A <1a0ahex@gmail.com>
Co-authored-by: vstone-w <374330057@qq.com>
Co-authored-by: Even Zhou <even.y.zhou@outlook.com>
Co-authored-by: ClownBin <chaobin1993@126.com>
Co-authored-by: sglang-npu-bot <sglangnpu@163.com>
This commit is contained in:
AndyLi429
2026-09-07 21:08:06 +08:00
committed by GitHub
co-authored by AndyLi429 Kailong Lu cx ranjiewen HEX1A0A vstone-w Even Zhou ClownBin sglang-npu-bot
parent df623d3cbd
commit 62a4a6ea0e
25 changed files with 3296 additions and 192 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/ascend/conn.py
+24
View File
@@ -24,6 +24,10 @@ class AscendStateType(str, enum.Enum):
"""DSV4-on-NPU PD components without a cross-hardware equivalent."""
DSV4_C128 = "dsv4_c128"
# C4 compress-state rows (attention + indexer) addressed within each
# req_pool_idx bank on A5 (CYCLE cache_mode). Separate from StateType.SWA
# because each peer maps logical positions into its own local ring.
DSV4_C4_STATE = "dsv4_c4_state"
_DSV4_KVCACHE_STATE_TYPES = tuple(AscendStateType)
@@ -82,6 +86,15 @@ class AscendKVManager(MooncakeKVManager):
dst = dst_kv_ptrs[c128_start:c128_end]
return src_kv_ptrs, dst, len(src_kv_ptrs)
if state_type == AscendStateType.DSV4_C4_STATE:
# Layout: [attn_state_0..attn_{c4_full-1},
# idx_state_0..idx_{c4_full-1}]
# Two groups, each c4_full entries; slice both by PP stage.
dst = []
for offset in (0, c4_full):
dst.extend(dst_kv_ptrs[offset + c4_start : offset + c4_end])
return src_kv_ptrs, dst, len(src_kv_ptrs)
# NPU main KV layout: [C4 KV, index K, index scale].
if state_type is None and len(dst_kv_ptrs) == 3 * c4_full:
dst = []
@@ -89,6 +102,17 @@ class AscendKVManager(MooncakeKVManager):
dst.extend(dst_kv_ptrs[offset + c4_start : offset + c4_end])
return src_kv_ptrs, dst, len(src_kv_ptrs)
# On A5 (CYCLE cache_mode), StateType.SWA only contains SWA KV
# buffers (C4 compress state is registered separately as
# DSV4_C4_STATE). The common _mla_slice_ptrs_for_pp assumes
# SWA + C4 state are bundled (swa_L + 2*c4_full), so intercept
# here and slice SWA KV by layer index directly.
if state_type == StateType.SWA and AscendStateType.DSV4_C4_STATE in (
self.kv_args.state_types or []
):
dst = list(dst_kv_ptrs[start_layer:end_layer])
return src_kv_ptrs, dst, len(src_kv_ptrs)
return super().get_mla_kv_ptrs_with_pp(src_kv_ptrs, dst_kv_ptrs, state_type)
# src_kv_ptrs: k_data, v_data, index_k_data(optional)
python/sglang/srt/disaggregation/ascend/transfer_engine.py
+5 -1
View File
@@ -65,6 +65,10 @@ class AscendTransferEngine(MooncakeTransferEngine):
transfer_protocol = self._get_transfer_protocol()
if transfer_protocol is None or transfer_protocol == "sdma":
trans_op_type = TransferEngine.TransDataOpType.SDMA
elif transfer_protocol == "device_urma":
trans_op_type = TransferEngine.TransDataOpType.DEVICE_URMA
elif transfer_protocol == "device_uboe":
trans_op_type = TransferEngine.TransDataOpType.DEVICE_UBOE
else:
trans_op_type = TransferEngine.TransDataOpType.DEVICE_RDMA
"""with device RDMA for PD transfer"""
@@ -100,7 +104,7 @@ class AscendTransferEngine(MooncakeTransferEngine):
@staticmethod
def _get_transfer_protocol():
protocol = os.getenv("ASCEND_MF_TRANSFER_PROTOCOL")
allowed_protocols = {"device_rdma", "sdma"}
allowed_protocols = {"device_rdma", "sdma", "device_urma", "device_uboe"}
if protocol and protocol.lower() in allowed_protocols:
return protocol.lower()
else:
python/sglang/srt/disaggregation/utils.py
+41
View File
@@ -79,6 +79,31 @@ def is_dsv4_c128_online_enabled() -> bool:
return not _IS_HIP and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get()
def get_dsv4_c4_state_indices(
req_pool_idx: int,
seq_len: int,
*,
ring_size: int,
) -> np.ndarray:
"""Return physical rows for the live C4 compressor history.
Prefill and decode may use different C4 ring sizes (8 without speculative
decoding and 16 with EAGLE/MTP). State transfer must therefore pair rows
by logical token position instead of copying a whole request-local bank.
The C4 overlap compressor keeps ``seq_len % 4 + 4`` live rows.
"""
if ring_size < 8 or ring_size % 4 != 0:
raise ValueError(
f"C4 ring_size must be a multiple of 4 and at least 8, got {ring_size}"
)
seq_len = max(0, int(seq_len))
state_len = seq_len % 4 + 4
positions = np.arange(max(0, seq_len - state_len), seq_len, dtype=np.int64)
rows = int(req_pool_idx) * int(ring_size) + positions % int(ring_size)
return rows.astype(np.int32)
def get_dsv4_c128_state_indices(
req_pool_idx: int,
seq_len: int,
@@ -1468,6 +1493,22 @@ def setup_state_kv_args(
c128_item_lens,
)
# On A5 (CYCLE cache_mode), C4 state uses request-local ring rows rather
# than SWA pages. Register it separately so P and D can independently
# map logical positions when their local ring sizes differ.
from sglang.srt.hardware_backend.npu.utils import is_npu_arch35
if is_npu_arch35():
c4_ptrs, c4_lens, c4_item_lens = token_to_kv_pool.get_c4_state_buf_infos()
if c4_ptrs:
append_state_component(
kv_args,
AscendStateType.DSV4_C4_STATE,
c4_ptrs,
c4_lens,
c4_item_lens,
)
# DSV4 NextN shares the target allocator, so target and draft use the same
# local SWA indices. Keep draft buffers in a separate positional component
# to avoid mixing them into the target's heterogeneous state layout, while
python/sglang/srt/environ.py
+4
View File
@@ -1138,6 +1138,10 @@ class Envs:
# 0 lets ElasticBuffer select its theoretical communication SM/QP counts.
SGLANG_DEEPEP_V2_NUM_SMS = EnvInt(0)
SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS = EnvInt(32)
# A5 DSV4 FP4 + DeepEP low-latency dispatch wire format. This is read only
# by the model-specific dispatcher configuration; all other paths retain
# their existing behavior.
SGLANG_NPU_DSV4_DEEPEP_LL_DISPATCH_QUANT_MODE = EnvStr("mxfp8")
SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO = EnvBool(False)
SGLANG_ENABLE_QWEN_DEEPEP_SHARED_OVERLAP = EnvBool(True)
# Force dynamic Waterfill with runtime EP all-reduce instead of the default
python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py
+333 -109
View File
@@ -16,7 +16,8 @@ from sglang.kernels.ops.speculative.dspark.dspark_attn_metadata import (
)
from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.attention.ascend_backend import AscendAttnBackend
from sglang.srt.hardware_backend.npu.dsv4.dsv4_rope import Dsv4NpuRoPE
from sglang.srt.hardware_backend.npu.dsv4.dsv4_rope import Dsv4NpuRoPE, rope_cos_sin
from sglang.srt.hardware_backend.npu.utils import is_npu_arch35
from sglang.srt.model_executor.forward_batch_info import DSV4OutCacheLoc, ForwardMode
from sglang.srt.model_executor.forward_context import get_attn_backend
from sglang.srt.runtime_context import get_parallel
@@ -29,6 +30,40 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
# A5 kv-quant KV layout: nope is quantized in groups of 64 and the RoPE half is
# stored unquantized, so the kernels need both dimensions spelled out.
_NPU_ARCH35_KV_TILE_SIZE = 64
_NPU_ARCH35_KV_ROPE_HEAD_DIM = 64
def _sparse_attn_ops():
"""(metadata op, attention op) for the DSV4 shared-KV sparse attention.
A5 reads a quantized KV cache, which is a different kernel rather than a
flag on the pre-A5 one.
"""
if is_npu_arch35():
return (
torch.ops.custom.npu_kv_quant_sparse_attn_sharedkv_metadata,
torch.ops.custom.npu_kv_quant_sparse_attn_sharedkv,
)
return (
torch.ops.custom.npu_sparse_attn_sharedkv_metadata,
torch.ops.npu.sparse_attn_sharedkv,
)
def _sparse_attn_kv_quant_kwargs() -> dict:
"""Extra kwargs the A5 kv-quant kernels need to interpret the KV layout."""
if not is_npu_arch35():
return {}
return {
"kv_quant_mode": 1,
"tile_size": _NPU_ARCH35_KV_TILE_SIZE,
"rope_head_dim": _NPU_ARCH35_KV_ROPE_HEAD_DIM,
}
def _walsh_hadamard_matrix(n: int, dtype: torch.dtype, device) -> torch.Tensor:
# n**-0.5 norm is baked in via the sqrt(2) division per doubling; _apply_hadamard is a plain matmul
cache = _walsh_hadamard_matrix._cache
@@ -97,6 +132,21 @@ def _build_explicit_state_block_table(
).contiguous()
def _build_cycle_state_block_table(req_pool_indices: torch.Tensor) -> torch.Tensor:
"""Build the Atlas A5 cache_mode=2 request-bank table.
A5 interprets this input as one bank id per request and computes the
in-bank ring offset itself. It must never receive the A3 explicit
per-token location table.
"""
if req_pool_indices.ndim != 1:
raise ValueError(
"Atlas A5 compressor requires a 1-D request-bank table, got "
f"shape={tuple(req_pool_indices.shape)}"
)
return req_pool_indices.to(dtype=torch.int32).contiguous()
class CompressorAscendBackendMixin:
@staticmethod
def _to_cpu_int_list(values) -> Optional[list[int]]:
@@ -128,6 +178,11 @@ class CompressorAscendBackendMixin:
def _build_npu_compress_metadata(self, forward_batch: ForwardBatch) -> None:
fm = self.forward_metadata
fm.dsv4_cycle_state_block_table = (
_build_cycle_state_block_table(forward_batch.req_pool_indices)
if is_npu_arch35()
else None
)
is_decode = forward_batch.forward_mode.is_decode()
is_verify = forward_batch.forward_mode.is_target_verify()
fm.dsv4_explicit_state_block_tables = {}
@@ -313,10 +368,16 @@ class CompressorAscendBackendMixin:
else:
n_c_tokens = max(1, seq_lens_max // ratio)
if ratio == 4:
slots = req_to_token[req_pool_64, : n_c_tokens * ratio]
c_page_table = (slots[:, :: self.page_size] // self.page_size).to(
torch.int32
col_idx = torch.arange(
0,
n_c_tokens * ratio,
self.page_size,
device=req_to_token.device,
)
slots = torch.index_select(
torch.index_select(req_to_token, 1, col_idx), 0, req_pool_64
)
c_page_table = (slots // self.page_size).to(torch.int32)
else:
c128_page_size = req_to_token_pool.c128_page_size
n_groups = (n_c_tokens + c128_page_size - 1) // c128_page_size
@@ -372,21 +433,27 @@ class CompressorAscendBackendMixin:
pool = self.token_to_kv_pool
state_pool = pool._get_state_pool(compressor.layer_id, compressor.is_in_indexer)
state_cache = state_pool.state_cache_3d
table_cache = fm.dsv4_explicit_state_block_tables
if ratio not in table_cache:
table_cache[ratio] = _build_explicit_state_block_table(
compress_ratio=ratio,
coff=coff,
state_pool=state_pool,
token_to_kv_pool=pool,
req_to_token=self.req_to_token,
req_pool_indices=forward_batch.req_pool_indices,
start_pos=fm.start_pos,
cu_seqlens=fm.actual_seq_lengths_q_pa,
seqused=fm.seqused,
max_input_capacity=fm.dsv4_max_input_capacity,
)
state_block_table = table_cache[ratio]
if is_npu_arch35():
# A5 cache_mode=2 is CYCLE: one request bank per row. The
# compressor derives the in-bank offset from start_pos; passing
# the A3 explicit [B, width] table here would be an ABI violation.
state_block_table = fm.dsv4_cycle_state_block_table
else:
table_cache = fm.dsv4_explicit_state_block_tables
if ratio not in table_cache:
table_cache[ratio] = _build_explicit_state_block_table(
compress_ratio=ratio,
coff=coff,
state_pool=state_pool,
token_to_kv_pool=pool,
req_to_token=self.req_to_token,
req_pool_indices=forward_batch.req_pool_indices,
start_pos=fm.start_pos,
cu_seqlens=fm.actual_seq_lengths_q_pa,
seqused=fm.seqused,
max_input_capacity=fm.dsv4_max_input_capacity,
)
state_block_table = table_cache[ratio]
cos, sin = Dsv4NpuRoPE.for_freqs(
compressor.freqs_cis, getattr(compressor, "rotary_emb", None)
@@ -397,7 +464,11 @@ class CompressorAscendBackendMixin:
allow_build=False,
)
cmp_kv = torch.ops.npu.compressor(
# TODO: torch.ops.npu.compressor does not support Atlas A5 yet.
compressor_op = (
torch.ops.custom.compressor if is_npu_arch35() else torch.ops.npu.compressor
)
cmp_kv = compressor_op(
x,
compressor._fused_wkv_w,
compressor._fused_wgate_w,
@@ -466,6 +537,9 @@ class CompressorAscendBackendMixin:
) -> None:
kv_scale: Optional[torch.Tensor] = None
li_kv_dtype = getattr(compressor, "li_kv_dtype", "bf16")
# A5 quantizes and scatters in one fused kernel, so the dequant scale is
# produced inside indexer_compress_epilog rather than here.
fused_fp8_indexer_write = li_kv_dtype == "float8" and compressor.is_in_indexer
if li_kv_dtype == "int8" and compressor.is_in_indexer:
kv, kv_scale = torch_npu.npu_dynamic_quant(kv)
kv_scale = kv_scale.to(torch.float16)
@@ -499,6 +573,36 @@ class CompressorAscendBackendMixin:
kv = kv[valid]
if kv_scale is not None:
kv_scale = kv_scale[valid]
# Eager verify keeps no row when no request completed a compression block
# this step (loc is then all-zero, the skip sentinel), and prefill can hand
# us an empty chunk. Nothing to write: the pre-A5 scatter treated that as a
# no-op, while both A5 fused epilog kernels reject a zero-row input. Unlike
# the `loc is None` check below (missing metadata = a bug), an empty write is
# a legitimate step outcome. Static shape read, so graph capture is unaffected.
if kv.shape[0] == 0:
return
if fused_fp8_indexer_write:
if loc is None:
raise RuntimeError(
"DSV4 A5 fused indexer epilog needs a slot mapping, but "
f"loc is None (mode={forward_batch.forward_mode}, "
f"ratio={compressor.ratio}). Writing nothing here would "
"leave the indexer KV cache stale."
)
torch.ops.custom.indexer_compress_epilog(
indexer_compress_cache=self.token_to_kv_pool.get_compress_buffer(
compressor.layer_id, True
),
indexer_compress_scale=self.token_to_kv_pool.get_compress_dequant_scale_buffer(
compressor.layer_id, True
),
x=kv,
slot_mapping=loc.to(torch.int32),
)
return
self.token_to_kv_pool.set_compress_buffer(
compressor.layer_id,
loc,
@@ -520,7 +624,7 @@ class C4IndexerAscendBackendMixin:
q_lora: torch.Tensor,
forward_batch: ForwardBatch,
) -> tuple[torch.Tensor, torch.Tensor]:
q = self._compute_q_npu(c4_indexer, q_lora, forward_batch.positions)
q = self._compute_q_npu(c4_indexer, q_lora, forward_batch)
weights, _ = c4_indexer.weights_proj(x)
weights = weights * (c4_indexer.softmax_scale * c4_indexer.n_heads**-0.5)
c4_indexer.compressor(x, forward_batch)
@@ -569,7 +673,7 @@ class C4IndexerAscendBackendMixin:
with torch.npu.stream(stream_q):
if q_lora_ready is not None:
stream_q.wait_event(q_lora_ready)
q = self._compute_q_npu(c4_indexer, q_lora, forward_batch.positions)
q = self._compute_q_npu(c4_indexer, q_lora, forward_batch)
q.record_stream(stream_q)
cur.wait_stream(stream_w)
@@ -593,7 +697,7 @@ class C4IndexerAscendBackendMixin:
)
li_kv_dtype = getattr(c4_indexer.compressor, "li_kv_dtype", "bf16")
if li_kv_dtype == "int8":
if li_kv_dtype in ("int8", "float8"):
# Empty/idle rank (T=0) must skip the indexer kernel; test is_idle
# rather than .item() since a host sync is illegal during capture.
if bs == 0 or forward_batch.forward_mode.is_idle():
@@ -680,30 +784,29 @@ class C4IndexerAscendBackendMixin:
return torch.cat(topk_idxs, dim=0).to(dtype=torch.int32)
def _ensure_npu_c4_indexer(self, c4_indexer, device: torch.device) -> None:
c4_indexer.compressor.li_kv_dtype = "int8"
# A5's lightning indexer consumes FP8 K + fp32 scales; pre-A5 stays int8.
c4_indexer.compressor.li_kv_dtype = "float8" if is_npu_arch35() else "int8"
if getattr(c4_indexer, "hadamard_matrix", None) is None:
H = _walsh_hadamard_matrix(c4_indexer.head_dim, torch.float32, device)
c4_indexer.register_buffer("hadamard_matrix", H, persistent=False)
def _compute_q_npu(
self, c4_indexer, q_lora: torch.Tensor, positions: torch.Tensor
self, c4_indexer, q_lora: torch.Tensor, forward_batch: ForwardBatch
) -> torch.Tensor:
positions = forward_batch.positions
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)
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(
# Per-forward memo keyed on the c4 layers' freqs_cis (the indexer
# shares it), so every c4 layer reads one gather instead of its own.
cos4, sin4 = rope_cos_sin(
c4_indexer.freqs_cis,
getattr(c4_indexer, "rotary_emb", None),
forward_batch,
positions,
q.dtype,
view_4d=True,
allow_build=False,
cache_dtype=torch.float32,
)
Dsv4NpuRoPE.apply_rotary_mul_inplace(
q,
@@ -723,20 +826,33 @@ class C4IndexerAscendBackendMixin:
weights: torch.Tensor,
forward_batch: ForwardBatch,
) -> torch.Tensor:
q_int8, q_scale = torch_npu.npu_dynamic_quant(q)
import torch_npu
if k.dtype == torch.float8_e4m3fn:
# A5: block-quantize Q to FP8 so it matches the FP8 K buffer; scales
# stay fp32 and the kernel wants one scale per (token, head).
q_quant, q_scale = torch_npu.npu_dynamic_block_quant(
q.view(-1, q.shape[-1]), dst_type=k.dtype
)
q_quant = q_quant.view(-1, c4_indexer.n_heads, c4_indexer.head_dim)
q_scale = q_scale.view(-1, c4_indexer.n_heads)
else:
q_quant, q_scale = torch_npu.npu_dynamic_quant(q)
q_scale = q_scale.to(torch.float16)
fm = self.forward_metadata
li_quant_metadata = fm.kernel_metadata["li_quant_metadata"]
kwargs = dict(
query=q_int8,
query=q_quant,
key=k,
key_dequant_scale=k_scale.squeeze(-2),
key_dequant_scale=k_scale.squeeze(-2).to(q_scale.dtype),
actual_seq_lengths_query=fm.actual_seq_lengths_q,
actual_seq_lengths_key=fm.actual_seq_lengths_kv,
block_table=fm.c4_page_table,
layout_query="TND",
layout_key="PA_BSND",
weights=weights.to(torch.float16),
query_dequant_scale=q_scale.to(torch.float16),
weights=weights.to(q_scale.dtype),
query_dequant_scale=q_scale,
cmp_ratio=4,
query_quant_mode=0,
key_quant_mode=0,
@@ -810,6 +926,11 @@ class DeepseekV4AscendAttnBackend(
model_runner.spec_algorithm is not None
and model_runner.spec_algorithm.is_dspark()
)
self._is_eagle_algorithm = bool(
model_runner.spec_algorithm is not None
and model_runner.spec_algorithm.is_eagle()
and not model_runner.spec_algorithm.is_frozen_kv_mtp()
)
self._is_dspark_draft_worker = bool(
getattr(model_runner, "is_draft_worker", False)
and self._is_dspark_algorithm
@@ -820,6 +941,9 @@ class DeepseekV4AscendAttnBackend(
for pool in self.token_to_kv_pool.compress_state_pools
if pool is not None
}
# High-water mark of written page-table columns per shared graph
# buffer; see _copy_page_table_into_graph.
self._graph_table_high_water: dict[str, int] = {}
def _is_dspark_draft_block(self, forward_batch: ForwardBatch) -> bool:
spec_algorithm = forward_batch.spec_algorithm
@@ -1013,6 +1137,11 @@ class DeepseekV4AscendAttnBackend(
metadata.c4_loc = torch.zeros(c4_pad, dtype=torch.int64, device=device)
metadata.c128_loc = torch.zeros(c128_pad, dtype=torch.int64, device=device)
metadata.dsv4_max_input_capacity = tokens_per_req
metadata.dsv4_cycle_state_block_table = (
torch.zeros(bs, dtype=torch.int32, device=device)
if is_npu_arch35()
else None
)
metadata.dsv4_explicit_state_block_tables = {
ratio: torch.full(
(
@@ -1056,14 +1185,19 @@ class DeepseekV4AscendAttnBackend(
self.forward_metadata = metadata
@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.
def _copy_page_table_into_graph(self, key: str, src: torch.Tensor) -> None:
# Graph page tables live in buffers shared across bs buckets (each
# bucket's metadata holds a row slice), refreshed in place on replay.
full = self.graph_metadata[key]
r, c = src.shape
dst[:r, :c].copy_(src)
if c < dst.shape[1]:
dst[:, c:].fill_(val)
full[:r, :c].copy_(src)
high = self._graph_table_high_water.get(key, 0)
if c < high:
# Full height, not just this replay's slice: other buckets'
# replays may have written rows beyond this bucket's row count.
full[:, c:high].fill_(-1)
elif c > high:
self._graph_table_high_water[key] = c
@staticmethod
def _copy_1d_with_zero_tail(dst: torch.Tensor, src: Optional[torch.Tensor]) -> None:
@@ -1081,6 +1215,80 @@ class DeepseekV4AscendAttnBackend(
if n < dst.shape[0]:
dst[n:].fill_(0)
@staticmethod
def _stable_compact_1d(
dst: torch.Tensor, values: torch.Tensor, keep: torch.Tensor
) -> None:
"""Compact selected values into a fixed graph buffer without NonZero.
For every selected element, ``cumsum(keep) - 1`` is exactly its ordinal
in ``nonzero(keep)``. Selected elements therefore have unique scatter
destinations, while rejected elements contribute integer zero only.
This preserves the stable boolean-index order without a dynamic output
shape or a device-to-host size read.
"""
dst.zero_()
if dst.numel() == 0 or values.numel() == 0:
return
values = values.reshape(-1)
keep = keep.reshape(-1)
if values.numel() != keep.numel():
raise ValueError(
"stable compact requires value/mask size equality, got "
f"{values.numel()} and {keep.numel()}"
)
ranks = torch.cumsum(keep.to(torch.int64), dim=0) - 1
in_bounds = keep & (ranks < dst.numel())
safe_ranks = ranks.clamp(min=0, max=dst.numel() - 1)
compact_values = torch.where(
in_bounds,
values.to(dst.dtype),
torch.zeros_like(values, dtype=dst.dtype),
)
dst.scatter_add_(0, safe_ranks, compact_values)
def _fill_verify_positions_cmp_padding_one_device(
self,
positions: torch.Tensor,
dst: torch.Tensor,
ratio: int,
live_seq_lens: torch.Tensor,
n_draft: int,
) -> None:
"""Device-only fixed-shape equivalent of the eager CPU reference path."""
if ratio not in self._dsv4_compress_ratios or positions.numel() == 0:
dst.zero_()
return
n_draft = int(n_draft)
request_num = positions.shape[0] // n_draft
if request_num == 0:
dst.zero_()
return
if live_seq_lens.device != positions.device:
raise ValueError(
"device verify compression metadata requires live_seq_lens and "
"positions on the same device"
)
live_seq_lens = live_seq_lens[:request_num]
token_offsets = torch.arange(
1,
n_draft + 1,
dtype=live_seq_lens.dtype,
device=live_seq_lens.device,
)
absolute_lengths = live_seq_lens.view(-1, 1) + token_offsets.view(1, -1)
boundary_mask = ((absolute_lengths % ratio) == 0) & (
live_seq_lens.view(-1, 1) > 0
)
# Match the CPU reference exactly: select the boundary token from the
# request-major positions array, then move RoPE to the group's first token.
values = positions[: request_num * n_draft].reshape(-1) + (1 - ratio)
self._stable_compact_1d(dst, values, boundary_mask.reshape(-1))
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
@@ -1217,7 +1425,7 @@ class DeepseekV4AscendAttnBackend(
)
for key in ("c4_page_table", "c128_page_table"):
if key in result:
self._copy_2d_with_tail(getattr(ctx.fm, key), result[key], -1)
self._copy_page_table_into_graph(key, result[key])
def _refresh_graph_decode_compress_1d_direct(self, ctx) -> None:
fm = ctx.fm
@@ -1236,36 +1444,54 @@ class DeepseekV4AscendAttnBackend(
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
dst = getattr(fm, f"positions_cmp_padding_c{ratio}")
self._stable_compact_1d(
dst,
positions_last.to(torch.int64) + (1 - ratio),
should_compress,
)
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.final_seq_lens_cpu
verify_seq_lens_cpu = torch.where(
ctx.live_seq_lens_cpu > 0,
verify_seq_lens_cpu,
ctx.live_seq_lens_cpu,
)
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))
if self._is_eagle_algorithm:
self._fill_verify_positions_cmp_padding_one_device(
ctx.forward_batch.positions,
fm.positions_cmp_padding_c4,
4,
ctx.live_seq_lens,
n_draft=ctx.tokens_per_bs,
)
self._fill_verify_positions_cmp_padding_one_device(
ctx.forward_batch.positions,
fm.positions_cmp_padding_c128,
128,
ctx.live_seq_lens,
n_draft=ctx.tokens_per_bs,
)
else:
verify_seq_lens_cpu = ctx.final_seq_lens_cpu
verify_seq_lens_cpu = torch.where(
ctx.live_seq_lens_cpu > 0,
verify_seq_lens_cpu,
ctx.live_seq_lens_cpu,
)
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,
)
valid = ctx.live_seq_lens[: ctx.bs] > 0
fm.seqused.copy_(
(valid.to(torch.int32) * int(ctx.tokens_per_bs)).to(device=ctx.device)
@@ -1326,7 +1552,7 @@ class DeepseekV4AscendAttnBackend(
max_seq_pages = (max_len + self.page_size - 1) // self.page_size
if 0 < max_seq_pages < swa_src.shape[1]:
swa_src = swa_src[:, :max_seq_pages]
self._copy_2d_with_tail(fm.swa_page_table, swa_src, -1)
self._copy_page_table_into_graph("swa_page_table", swa_src)
def _refresh_graph_dspark_sparse_metadata(self, ctx) -> None:
if not (self._is_dspark_draft_worker and ctx.graph_mode.is_target_verify()):
@@ -1392,6 +1618,11 @@ class DeepseekV4AscendAttnBackend(
def _apply_dsv4_graph_metadata(self, forward_batch: ForwardBatch) -> None:
ctx = self._build_dsv4_graph_replay_ctx(forward_batch)
if is_npu_arch35():
ctx.fm.dsv4_cycle_state_block_table.copy_(
ctx.forward_batch.req_pool_indices[: ctx.bs]
)
self._refresh_graph_seq_metadata(ctx)
self._refresh_graph_compress_page_tables_direct(ctx)
@@ -1538,7 +1769,11 @@ class DeepseekV4AscendAttnBackend(
is_nextn: bool,
) -> dict:
fm = self.forward_metadata
metadata_op, _ = _sparse_attn_ops()
common = {
**_sparse_attn_kv_quant_kwargs(),
"cu_seqlens_q": actual_seq_lengths_q_pa,
"seqused_kv": actual_seq_lengths_kv,
"cmp_ratio": 1,
"ori_mask_mode": 4,
"cmp_mask_mode": 3,
@@ -1557,8 +1792,6 @@ class DeepseekV4AscendAttnBackend(
"has_ori_kv": True,
"has_cmp_kv": False,
}
# The host metadata op reads CPU int32 mirrors — never a D2H sync of the
# device tensors (that would drain the stream and stall overlapped prep).
c1a_kwargs = base_kwargs | common
if self._is_dspark_draft_worker:
cu_q_cpu = fm.actual_seq_lengths_q_pa_cpu
@@ -1570,14 +1803,13 @@ class DeepseekV4AscendAttnBackend(
c1a_kwargs = c1a_kwargs | host_inputs
metadata_op = torch.ops.npu.sparse_attn_sharedkv_metadata_host
else:
# The device-side op requires tensor args for backend dispatch; pass
# the device mirrors just like the pre-refactor call did.
c1a_kwargs = c1a_kwargs | {
"cu_seqlens_q": actual_seq_lengths_q_pa,
"seqused_kv": actual_seq_lengths_kv,
}
metadata_op = torch.ops.custom.npu_sparse_attn_sharedkv_metadata
kernel_metadata = {"c1a_metadata": metadata_op(**c1a_kwargs)}
metadata_op, _ = _sparse_attn_ops()
c1a_metadata = metadata_op(**c1a_kwargs)
kernel_metadata = {"c1a_metadata": c1a_metadata}
if self._dsv4_has_c4:
c4a_overrides = {
@@ -1586,9 +1818,8 @@ class DeepseekV4AscendAttnBackend(
"cmp_topk": self._dsv4_index_topk,
}
c4a_kwargs = c1a_kwargs | c4a_overrides
kernel_metadata["c4a_metadata"] = (
torch.ops.custom.npu_sparse_attn_sharedkv_metadata(**c4a_kwargs)
)
metadata_op, _ = _sparse_attn_ops()
kernel_metadata["c4a_metadata"] = metadata_op(**c4a_kwargs)
if actual_seq_lengths_q_pa is not None:
# the indexer metadata op wants a fresh contiguous tensor without the leading 0
@@ -1616,9 +1847,8 @@ class DeepseekV4AscendAttnBackend(
if self._dsv4_has_c128:
c128a_overrides = {"cmp_ratio": 128, "has_cmp_kv": True}
c128a_kwargs = c1a_kwargs | c128a_overrides
kernel_metadata["c128a_metadata"] = (
torch.ops.custom.npu_sparse_attn_sharedkv_metadata(**c128a_kwargs)
)
metadata_op, _ = _sparse_attn_ops()
kernel_metadata["c128a_metadata"] = metadata_op(**c128a_kwargs)
return kernel_metadata
@@ -1664,6 +1894,7 @@ class DeepseekV4AscendAttnBackend(
ori_kv = pool.get_swa_buffer(layer.layer_id)
attn_kwargs = dict(
**_sparse_attn_kv_quant_kwargs(),
cu_seqlens_q=fm.actual_seq_lengths_q_pa,
seqused_kv=fm.actual_seq_lengths_kv,
ori_mask_mode=4,
@@ -1687,7 +1918,8 @@ class DeepseekV4AscendAttnBackend(
if ori_sparse_indices is not None:
attn_kwargs["ori_sparse_indices"] = ori_sparse_indices
q_arg = attn_kwargs.pop("q")
out, _ = torch.ops.npu.sparse_attn_sharedkv(q_arg, **attn_kwargs)
_, attn_op = _sparse_attn_ops()
out, _ = attn_op(q_arg, **attn_kwargs)
return out
def _forward_compressed(
@@ -1733,6 +1965,7 @@ class DeepseekV4AscendAttnBackend(
)
attn_kwargs = dict(
**_sparse_attn_kv_quant_kwargs(),
cu_seqlens_q=fm.actual_seq_lengths_q_pa,
seqused_kv=fm.actual_seq_lengths_kv,
ori_mask_mode=4,
@@ -1758,7 +1991,8 @@ class DeepseekV4AscendAttnBackend(
else:
attn_kwargs["cmp_sparse_indices"] = None
q_arg = attn_kwargs.pop("q")
out, _ = torch.ops.npu.sparse_attn_sharedkv(q_arg, **attn_kwargs)
_, attn_op = _sparse_attn_ops()
out, _ = attn_op(q_arg, **attn_kwargs)
return out
def get_swa_out_cache_loc(self, forward_batch: ForwardBatch) -> torch.Tensor:
@@ -1860,7 +2094,7 @@ class DeepseekV4AscendAttnBackend(
abs_positions = start_positions.view(-1, 1) + torch.arange(
n_draft, dtype=start_positions.dtype
).view(1, -1)
boundary_mask = abs_positions % ratio == 0
boundary_mask = (abs_positions % ratio == 0) & (seq_lens_cpu.view(-1, 1) > 0)
indices = torch.nonzero(boundary_mask.flatten(), as_tuple=False).flatten()
if indices.numel() == 0:
@@ -1869,7 +2103,11 @@ class DeepseekV4AscendAttnBackend(
# on NPU, a non-blocking copy from a short-lived pinned CPU tensor can
# surface later as an unrelated CopyKernel stream failure.
indices = indices[: dst.numel()].to(device=positions.device)
dst[: indices.numel()].copy_(torch.gather(positions, 0, indices))
# ``indices`` selects the final token of each newly completed group.
# The compressed KV applies RoPE at the group's first token, matching
# the decode path and ``comp_pos = (position // ratio) * ratio``.
compressed_positions = torch.gather(positions, 0, indices) + (1 - ratio)
dst[: indices.numel()].copy_(compressed_positions)
def update_verify_buffers_to_fill_after_draft(
self, spec_info, cuda_graph_bs: Optional[int]
@@ -1992,31 +2230,17 @@ class DeepseekV4AscendMultiStepDraftBackend:
)
swa_steps = swa_steps.permute((2, 0, 1)).reshape(self.speculative_num_steps, -1)
def step_compress(loc, ratio: int):
if loc is None or loc.numel() == 0:
return loc
raw_bs = step_width // self.topk
seq_lens = forward_batch.seq_lens[:raw_bs].to(torch.int64)
positions = seq_lens[:, None, None] + torch.arange(
self.speculative_num_steps,
device=seq_lens.device,
dtype=seq_lens.dtype,
)
positions = positions.expand(-1, self.topk, -1)
should_compress = ((positions + 1) % ratio) == 0
counts = should_compress.reshape(-1).to(torch.int64)
offsets = torch.cumsum(counts, dim=0) - counts
step_mask = should_compress[:, :, step_id].reshape(-1)
step_offsets = offsets.reshape(
raw_bs, self.topk, self.speculative_num_steps
)[:, :, step_id].reshape(-1)
return loc[step_offsets[step_mask].to(torch.int64)]
return DSV4OutCacheLoc(
out_full_loc=full_steps[step_id],
out_swa_loc=swa_steps[step_id],
out_c4_loc=step_compress(bundle.out_c4_loc, 4),
out_c128_loc=step_compress(bundle.out_c128_loc, 128),
out_c4_loc=(
None if bundle.out_c4_loc is None else bundle.out_c4_loc.new_empty((0,))
),
out_c128_loc=(
None
if bundle.out_c128_loc is None
else bundle.out_c128_loc.new_empty((0,))
),
)
def _with_step_cache_locs(self, forward_batch: ForwardBatch, step_id: int, call_fn):
python/sglang/srt/hardware_backend/npu/dsv4/dsv4_common_hooks.py
+31 -4
View File
@@ -10,8 +10,9 @@ these hooks then:
2. Write newly allocated C128 page ids into the per-request sidecar.
Compressor state is fixed ring storage and does not participate in this
allocation/write path. PD reuses the public SWA/C128-state payloads and only
builds an NPU-specific payload for the independently addressed C128 KV pool.
allocation/write path. PD reuses the public SWA/C128-state payloads and builds
NPU-specific payloads for the independently addressed C128 KV pool and A5 C4
compress-state rows.
Non-DSV4 paths leave ``batch.out_cache_loc_dsv4`` None, so this module is a
no-op for them.
@@ -71,11 +72,24 @@ def dsv4_state_payloads(
*,
prefix_len: int = 0,
):
"""Build the only NPU-specific DSV4 PD payload: C128 KV pages."""
"""Build NPU-specific DSV4 PD payloads.
Returns payloads for components that are addressed differently from the
cross-hardware ``StateType.SWA`` / ``StateType.C128_STATE`` defaults:
* ``DSV4_C128`` — C128 KV pages from ``req_to_c128_sidecar``.
* ``DSV4_C4_STATE`` (A5 only) — live C4 compress-state rows. Prefill
and decode derive physical rows using their own local ring sizes, so
decode-only MTP can safely transfer from an 8-row ring to a 16-row ring.
Pre-A5 uses EXPLICIT cache_mode and the C4 state is handled by the
shared ``StateType.SWA`` payload.
"""
import numpy as np
from sglang.srt.disaggregation.ascend.conn import AscendStateType
from sglang.srt.disaggregation.utils import get_dsv4_c4_state_indices
from sglang.srt.hardware_backend.npu.utils import is_npu_arch35
seq_len = max(0, int(seq_len))
prefix_len = max(0, min(int(prefix_len), seq_len))
@@ -94,7 +108,20 @@ def dsv4_state_payloads(
)
return pages[pages > 0]
return {AscendStateType.DSV4_C128: c128_kv_pages}
payloads = {AscendStateType.DSV4_C128: c128_kv_pages}
if is_npu_arch35():
def c4_state_indices():
return get_dsv4_c4_state_indices(
req_pool_idx,
seq_len,
ring_size=req_to_token_pool.get_dsv4_c4_state_ring_size(),
)
payloads[AscendStateType.DSV4_C4_STATE] = c4_state_indices
return payloads
def dsv4_prealloc_kwargs(allocator, req, fill_len, req_to_token_pool, *, device):
python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py
+135 -29
View File
@@ -7,19 +7,22 @@ rules as the GPU implementation:
* C4A/C4Li state follows SWA physical pages.
* C128A state follows ``req_pool_idx`` and absolute position.
``NPUCompressStatePool`` only adds the contiguous 3-D view and positive dummy
location required by the Atlas A3 ``cache_mode=2`` operator. There is no paged
state allocator or ``cache_mode=1`` compatibility storage.
``NPUCompressStatePool`` adds the contiguous 3-D view and positive dummy
location required by the Atlas fused compressor operators. A3 uses explicit
locations; A5 uses the same ring storage through its request-bank (cycle) ABI.
There is no paged state allocator or ``cache_mode=1`` compatibility storage.
"""
from __future__ import annotations
import math
from typing import List, Optional, Tuple
import torch
import torch_npu
from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
from sglang.srt.hardware_backend.npu.utils import is_npu_arch35
from sglang.srt.mem_cache.deepseek_v4_compress_state import CompressStatePool
from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
ONLINE_C128,
@@ -29,13 +32,17 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
)
from sglang.srt.runtime_context import get_schedule
_NPU_ARCH35_KV_QUANT_GROUP_SIZE = 64
_NPU_ARCH35_KV_ROW_ALIGNMENT = 128
class NPUDeepSeekV4SingleKVPool(DeepSeekV4SingleKVPool):
"""NPU bf16 variant of the full / SWA / c4 / c128 single-KV pool.
"""NPU PA_ND variant of the full / SWA / c4 / c128 single-KV pool.
``npu_sparse_attn_sharedkv`` reads KV in PA_ND layout
``(num_pages, kernel_page_size, num_kv_heads=1, dim)`` with ``dim`` packing
K_nope + K_rope as bf16. C4 uses its native page so its physical page id can
K_nope + K_rope as bf16 before A5; A5 uses packed FP8 KV rows. C4 uses its
native page so its physical page id can
be shared with the corresponding full page. C128 uses its independently
configured physical page size; Full/SWA use the global page size.
The CUDA fp8-packed-bytes layout (the base ``create_buffer``) is untouched.
@@ -47,11 +54,27 @@ class NPUDeepSeekV4SingleKVPool(DeepSeekV4SingleKVPool):
self.kernel_page_size = kernel_page_size
super().__init__(*args, **kwargs)
@property
def a5_packed_kv_dim(self) -> int:
nope_dim = self.qk_nope_head_dim
rope_dim = self.qk_rope_head_dim
scale_dim = math.ceil(nope_dim / _NPU_ARCH35_KV_QUANT_GROUP_SIZE)
bytes_per_token = nope_dim + rope_dim * 2 + scale_dim
return (
math.ceil(bytes_per_token / _NPU_ARCH35_KV_ROW_ALIGNMENT)
* _NPU_ARCH35_KV_ROW_ALIGNMENT
)
def create_buffer(self, *, num_pages: int):
# Non-bf16 store dtype (shouldn't happen here) falls back to base layout.
if self.store_dtype != torch.bfloat16:
return super().create_buffer(num_pages=num_pages)
kv_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
if is_npu_arch35():
kv_dim = self.a5_packed_kv_dim
kv_dtype = torch.float8_e4m3fn
else:
kv_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
kv_dtype = torch.bfloat16
self.kv_cache_total_dim = kv_dim
# Writes are flat-indexed by loc; kernel_page_size controls the physical
# page layout exposed to the NPU operators.
@@ -61,18 +84,18 @@ class NPUDeepSeekV4SingleKVPool(DeepSeekV4SingleKVPool):
self.kernel_page_size,
1,
kv_dim,
dtype=torch.bfloat16,
dtype=kv_dtype,
device=self.device,
)
class NPUCompressStatePool(CompressStatePool):
"""Thin A3 adapter over the shared GPU-style ring state pool.
"""Thin Atlas adapter over the shared GPU-style ring state pool.
Allocation, sizing, ring ownership and address translation are inherited
from :class:`CompressStatePool`. NPU only requests a contiguous 3-D view,
enforces the A3 FP32 contract and replaces invalid locations with a cleared
positive dummy row.
enforces the FP32 state-cache contract and replaces invalid locations with
a cleared positive dummy row for explicit-location callers.
Location 0 is valid in explicit mode. Invalid/history-padding locations map
to the final cleared row instead of ``-1`` because the A3 kernel consumes
@@ -164,6 +187,10 @@ class NPUDeepSeekV4IndexerPool(DeepSeekV4IndexerPool):
super()._create_buffer()
kp = self._kernel_page_size
npu_num_pages = (self.size + kp + 1) // kp
if is_npu_arch35():
index_k_dtype, index_scale_dtype = torch.float8_e4m3fn, torch.float32
else:
index_k_dtype, index_scale_dtype = torch.int8, torch.float16
with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
self.index_k_buffer = [
torch.zeros(
@@ -171,7 +198,7 @@ class NPUDeepSeekV4IndexerPool(DeepSeekV4IndexerPool):
kp,
1,
self.index_head_dim,
dtype=torch.int8,
dtype=index_k_dtype,
device=self.device,
)
for _ in range(self.layer_num)
@@ -182,7 +209,7 @@ class NPUDeepSeekV4IndexerPool(DeepSeekV4IndexerPool):
kp,
1,
1,
dtype=torch.float16,
dtype=index_scale_dtype,
device=self.device,
)
for _ in range(self.layer_num)
@@ -205,20 +232,20 @@ class NPUDeepSeekV4IndexerPool(DeepSeekV4IndexerPool):
index_k: torch.Tensor,
index_k_scale: Optional[torch.Tensor],
) -> None:
# int8 K + fp16 scale come from _compressor_epilog_npu's npu_dynamic_quant
# output (index_k: int8 [T, D], index_k_scale: fp16 [T, 1]).
d = self.index_head_dim
loc_long = loc.view(-1, 1).long()
index_k_cache = self.index_k_buffer[layer_id]
torch_npu.npu_scatter_nd_update_(
self.index_k_buffer[layer_id].view(-1, 1, d),
index_k_cache.view(-1, 1, d),
loc_long,
index_k.to(torch.int8).view(-1, 1, d),
index_k.to(index_k_cache.dtype).view(-1, 1, d),
)
if index_k_scale is not None:
index_scale_cache = self.index_scale_buffer[layer_id]
torch_npu.npu_scatter_nd_update_(
self.index_scale_buffer[layer_id].view(-1, 1, 1),
index_scale_cache.view(-1, 1, 1),
loc_long,
index_k_scale.to(torch.float16).view(-1, 1, 1),
index_k_scale.to(index_scale_cache.dtype).view(-1, 1, 1),
)
@@ -303,9 +330,16 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool):
"SGLANG_OPT_USE_ONLINE_COMPRESS is incompatible with the "
"NPU fused compressor (no online mode in the kernel)."
)
ring_size = self.get_ring_size(ratio)
# A5 cache_mode=2 addresses one ring bank per request. The A3
# explicit-location path can share the smaller flat pool, but the A5
# cycle ABI needs enough physical banks for every req_pool_idx.
size = self._state_pool_size(ratio)
if is_npu_arch35():
size = max(size, self.num_req_slots * ring_size)
return NPUCompressStatePool(
size=self._state_pool_size(ratio),
ring_size=self.get_ring_size(ratio),
size=size,
ring_size=ring_size,
overlap=ratio == 4,
head_dim=self.qk_nope_head_dim + self.qk_rope_head_dim,
dtype=self.c4_state_dtype if ratio == 4 else self.c128_state_dtype,
@@ -320,9 +354,13 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool):
) -> NPUCompressStatePool:
# c4 indexer shares the c4 state pool size budget but has its own
# slot_dim (indexer_head_dim vs attention head_dim).
ring_size = self.get_ring_size(ratio)
size = self.c4_state_pool_size
if is_npu_arch35():
size = max(size, self.num_req_slots * ring_size)
return NPUCompressStatePool(
size=self.c4_state_pool_size,
ring_size=self.get_ring_size(ratio),
size=size,
ring_size=ring_size,
overlap=ratio == 4,
head_dim=self.indexer_head_dim,
device=self.device,
@@ -370,8 +408,11 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool):
def get_state_buf_infos(self) -> Tuple[List[int], List[int], List[int]]:
"""GPU-compatible ``StateType.SWA`` component.
SWA KV, C4 attention state and C4 indexer state retain separate buffers
but share the same SWA page/state index.
On pre-A5 (EXPLICIT cache_mode), SWA KV, C4 attention state and C4
indexer state retain separate buffers but share the same SWA page/state
index. On A5 (CYCLE cache_mode) the compressor addresses the C4 state
ring by ``req_pool_idx`` instead of SWA page, so C4 state is excluded
here and registered separately via :meth:`get_c4_state_buf_infos`.
"""
data_ptrs: List[int] = []
data_lens: List[int] = []
@@ -382,6 +423,33 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool):
data_lens.append(buf.nbytes)
item_lens.append(buf[0].nbytes)
if not is_npu_arch35():
for pools in (
self.compress_state_pools,
self.indexer_compress_state_pools,
):
for pool in pools:
if pool is None or pool.ratio != 4:
continue
state = pool.kv_score_buffer.kv_score
data_ptrs.append(state.data_ptr())
data_lens.append(state.nbytes)
item_lens.append(state[0].nbytes * pool.ring_size)
return data_ptrs, data_lens, item_lens
def get_c4_state_buf_infos(self) -> Tuple[List[int], List[int], List[int]]:
"""C4 compress state ring (attention + indexer).
Register one physical state row as an item. PD peers can have
different request-local ring sizes (for example prefill without MTP
and decode with MTP), so payload indices map the same logical token
positions into each peer's local ring independently.
"""
data_ptrs: List[int] = []
data_lens: List[int] = []
item_lens: List[int] = []
for pools in (self.compress_state_pools, self.indexer_compress_state_pools):
for pool in pools:
if pool is None or pool.ratio != 4:
@@ -389,7 +457,7 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool):
state = pool.kv_score_buffer.kv_score
data_ptrs.append(state.data_ptr())
data_lens.append(state.nbytes)
item_lens.append(state[0].nbytes * pool.ring_size)
item_lens.append(state[0].nbytes)
return data_ptrs, data_lens, item_lens
@@ -404,7 +472,8 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool):
def get_state_cache(self, layer_id: int, from_indexer: bool) -> torch.Tensor:
"""FP32 ``[block_num, ring_size, 2*coff*D]`` view of this layer's
kv+score buffer — the fused compressor op
(``torch.ops.npu.compressor``)'s ``state_cache`` argument."""
(``torch.ops.custom.compressor`` on A5 and ``torch.ops.npu.compressor``
elsewhere)'s ``state_cache`` argument."""
return self._get_state_pool(layer_id, from_indexer).state_cache_3d
# ------------------------------------------------------------------
@@ -457,7 +526,7 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool):
Routes to c4 / c128 kv_pool by layer compression ratio. Returns
``None`` for ratio == 0 (no compress KV exists). The
from_indexer=True branch returns the dedicated int8 K buffer that
from_indexer=True branch returns the dedicated quantized K buffer that
``torch.ops.custom.npu_quant_lightning_indexer`` consumes.
"""
item = self.layer_mapping[layer_id]
@@ -489,6 +558,9 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool):
"""
# Index by raw layer_id (see get_swa_buffer) to avoid bucket collision.
buf = self.swa_kv_pool.kv_buffer[layer_id]
if is_npu_arch35():
self._write_a5_packed_kv(buf=buf, loc=loc, cache=cache)
return
buf_flat = buf.flatten(0, 1) # (num_pages * page_size, 1, dim)
# Caller (V4 MQALayer) may hand us cache shaped (T, dim); the buffer has
# an explicit num_kv_heads=1 axis, so insert it.
@@ -496,6 +568,38 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool):
cache = cache.unsqueeze(1)
buf_flat[loc] = cache.to(buf_flat.dtype)
def _write_a5_packed_kv(
self,
*,
buf: torch.Tensor,
loc: torch.Tensor,
cache: torch.Tensor,
) -> None:
cache_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
if cache.shape[-1] != cache_dim:
raise RuntimeError(
f"DSV4 A5 KV cache expects input last dim {cache_dim}, "
f"got shape={tuple(cache.shape)}."
)
cache_2d = cache.reshape(-1, cache_dim).to(torch.bfloat16).contiguous()
slot_mapping = loc.reshape(-1).contiguous()
if cache_2d.shape[0] != slot_mapping.shape[0]:
raise RuntimeError(
"DSV4 A5 KV cache write expects one slot per token, got "
f"{cache_2d.shape[0]} rows and {slot_mapping.shape[0]} slots."
)
if cache_2d.shape[0] == 0:
return
torch.ops.npu.kv_compress_epilog(
buf.view(-1, 1, buf.shape[-1]),
cache_2d,
slot_mapping,
quant_group_size=_NPU_ARCH35_KV_QUANT_GROUP_SIZE,
quant_mode=2,
round_scale_flag=True,
layout=1,
)
def set_swa_key_buffer_radix_fused_norm_rope(
self,
layer_id: int,
@@ -568,6 +672,9 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool):
# PA_ND layout: kv_buffer[layer_id] shape = (num_pages, page_size,
# 1, kv_dim). Flatten (num_pages, page_size) and index by `loc`.
buf = compress_pool.kv_buffer[compress_layer_id]
if is_npu_arch35():
self._write_a5_packed_kv(buf=buf, loc=loc, cache=kv)
return
buf_flat = buf.flatten(0, 1)
kv_view = kv.to(buf_flat.dtype)
if kv_view.ndim == buf_flat.ndim - 1:
@@ -581,8 +688,7 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool):
layer_id: int,
from_indexer: bool,
) -> torch.Tensor:
# Returns the float16 dequant scale buffer (NPU indexer pool's dedicated
# scale buffer alongside the int8 K buffer).
# The indexer scale is fp16 on pre-A5 parts and fp32 on A5.
assert from_indexer, "only indexer compress pool has dequant scale"
compress_layer_id = self.layer_mapping[layer_id].compress_layer_id
return self.c4_indexer_kv_pool.get_index_scale(compress_layer_id)
python/sglang/srt/hardware_backend/npu/dsv4/dsv4_req_to_token_pool.py
+6
View File
@@ -77,6 +77,12 @@ class DSV4ReqToTokenTablesMixin:
release C128 KV pages."""
self._dsv4_allocator = allocator
def get_dsv4_c4_state_ring_size(self) -> int:
"""Return the local C4 state-ring size used by the NPU KV pool."""
if self._dsv4_allocator is None:
raise RuntimeError("DSV4 allocator is not registered")
return self._dsv4_allocator.get_kvcache().get_ring_size(4)
def set_c128_prefix_pages(self, req, page_ids: torch.Tensor) -> None:
"""Install pages returned by a Radix match.
python/sglang/srt/hardware_backend/npu/dsv4/dsv4_rope.py
+54
View File
@@ -179,3 +179,57 @@ class Dsv4NpuRoPE:
rotary_mode="interleave",
partial_slice=[qk_nope_dim, qk_nope_dim + rope_dim],
)
# Per-forward memo of position-gathered (cos, sin), stashed on the ForwardBatch
# under this attribute by prime_rope_cos_sin (the single writer).
_ROPE_MEMO_ATTR = "_dsv4_npu_rope_memo"
def prime_rope_cos_sin(attn_modules, forward_batch, positions) -> None:
memo: dict = {}
for attn in attn_modules:
freqs_cis = attn.freqs_cis
fwd_key = (id(freqs_cis), torch.bfloat16, False)
if fwd_key in memo:
continue
cos, sin = Dsv4NpuRoPE.for_freqs(
freqs_cis, getattr(attn, "rotary_emb", None)
).get_cos_sin(
positions,
torch.bfloat16,
view_4d=True,
inverse=False,
allow_build=False,
cache_dtype=torch.bfloat16,
)
memo[fwd_key] = (positions, cos, sin)
memo[(id(freqs_cis), torch.bfloat16, True)] = (positions, cos, -sin)
setattr(forward_batch, _ROPE_MEMO_ATTR, memo)
def rope_cos_sin(
freqs_cis: torch.Tensor,
rotary_emb,
forward_batch,
positions: torch.Tensor,
dtype: torch.dtype,
*,
inverse: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
memo = getattr(forward_batch, _ROPE_MEMO_ATTR, None)
entry = memo.get((id(freqs_cis), dtype, inverse)) if memo is not None else None
if entry is not None and entry[0] is positions:
return entry[1], entry[2]
# bf16 tables are ensured at layer init; gathering in the activation dtype
# skips the fp32-gather + cast pair. Bit-identical values: rounding the
# table once equals rounding each gathered element.
cache_dtype = dtype if dtype == torch.bfloat16 else torch.float32
return Dsv4NpuRoPE.for_freqs(freqs_cis, rotary_emb).get_cos_sin(
positions,
dtype,
view_4d=True,
inverse=inverse,
allow_build=False,
cache_dtype=cache_dtype,
)
python/sglang/srt/hardware_backend/npu/quantization/fp4_moe_methods.py
+630
View File
@@ -0,0 +1,630 @@
"""MXFP4 routed-expert MoE method for Ascend A5 (Ascend 950).
DeepSeek-V4's FP4 expert checkpoint stores block-32 MXFP4 weights with E8M0
scales. This module wires those weights to the A5 grouped-matmul kernels, both
for the plain (init-routing) path and for the DeepEP dispatch path.
"""
from typing import TYPE_CHECKING, Optional
import torch
from sgl_kernel_npu.activation.swiglu_mxfp8_quant import swiglu_quant
from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import (
_get_float4_e2m1fn_x2_dtype,
_get_float8_e8m0fnu_dtype,
)
from sglang.srt.hardware_backend.npu.utils import is_npu_arch35
from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase
from sglang.srt.utils import set_weight_attrs
if TYPE_CHECKING:
from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput
# MXFP4 group size, fixed at 32 by the msmodelslim export format.
MXFP4_BLOCK_SIZE = 32
def _configure_dsv4_deepep_dispatcher(layer: torch.nn.Module) -> None:
"""Select the DSV4 FP4 DeepEP wire format without changing other MoEs."""
dispatcher = getattr(layer, "dispatcher", None)
if dispatcher is None:
return
# This method is only instantiated for DSV4 FP4 experts on A5 today, but
# retain the former BF16 setting if that selection changes in the future.
if not is_npu_arch35():
dispatcher.set_quant_config({"dispatcher_output_dtype": "bf16"})
return
# Import lazily to avoid importing the MoE backend during quant method
# module initialization.
from sglang.srt.layers.moe import get_moe_a2a_backend
if not get_moe_a2a_backend().is_deepep():
dispatcher.set_quant_config({"dispatcher_output_dtype": "bf16"})
return
low_latency_dtype = envs.SGLANG_NPU_DSV4_DEEPEP_LL_DISPATCH_QUANT_MODE.get()
if low_latency_dtype not in {"mxfp8", "bf16"}:
raise ValueError(
"SGLANG_NPU_DSV4_DEEPEP_LL_DISPATCH_QUANT_MODE must be one of "
"'mxfp8' or 'bf16' for A5 DSV4 DeepEP low-latency dispatch; "
f"got {low_latency_dtype!r}."
)
# The concrete dispatcher selects one mode-specific value. Normal (prefill)
# remains BF16, while low-latency (decode) defaults to MXFP8.
dispatcher.set_quant_config(
{
"normal_dispatcher_output_dtype": "bf16",
"low_latency_dispatcher_output_dtype": low_latency_dtype,
}
)
def _wrap_mxfp4_scale_weight_loader(weight_loader):
def load_scale(param, loaded_weight, *args, **kwargs):
if param.dtype == torch.uint8 and loaded_weight.dtype == torch.float8_e8m0fnu:
loaded_weight = loaded_weight.view(torch.uint8)
return weight_loader(param, loaded_weight, *args, **kwargs)
return load_scale
class NPUW4A4Fp4MoEMethod(FusedMoEMethodBase):
"""DeepSeek-V4 routed experts on Ascend A5: W4A8 MXFP weights.
Delegates nothing to ``fp8_method`` except the shared runner config; it is
held so the FP8 method sees the same ``moe_runner_config`` the layer built.
"""
def __init__(self, fp8_method, prefix: str = ""):
self._fp8 = fp8_method
self.prefix = prefix
self.moe_runner_config = None
def create_weights(
self,
layer: torch.nn.Module,
num_experts: int,
hidden_size: int,
intermediate_size_per_partition: int,
params_dtype: torch.dtype,
**extra_weight_attrs,
):
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
# Two FP4 values per stored byte, hence the // 2 on the K dimension.
w13_weight = torch.nn.Parameter(
torch.empty(
(num_experts, 2 * intermediate_size_per_partition, hidden_size // 2),
dtype=torch.uint8,
),
requires_grad=False,
)
w2_weight = torch.nn.Parameter(
torch.empty(
(num_experts, hidden_size, intermediate_size_per_partition // 2),
dtype=torch.uint8,
),
requires_grad=False,
)
layer.register_parameter("w13_weight", w13_weight)
set_weight_attrs(w13_weight, extra_weight_attrs)
layer.register_parameter("w2_weight", w2_weight)
set_weight_attrs(w2_weight, extra_weight_attrs)
scale_attrs = dict(extra_weight_attrs)
scale_attrs["quant_method"] = FusedMoeWeightScaleSupported.BLOCK.value
if weight_loader := scale_attrs.get("weight_loader"):
scale_attrs["weight_loader"] = _wrap_mxfp4_scale_weight_loader(
weight_loader
)
w13_weight_scale = torch.nn.Parameter(
torch.zeros(
(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // MXFP4_BLOCK_SIZE,
),
dtype=torch.uint8,
),
requires_grad=False,
)
w2_weight_scale = torch.nn.Parameter(
torch.zeros(
(
num_experts,
hidden_size,
intermediate_size_per_partition // MXFP4_BLOCK_SIZE,
),
dtype=torch.uint8,
),
requires_grad=False,
)
# Scales ship as raw E8M0 exponent bytes; no ue8m0 requantization here.
w13_weight_scale.format_ue8m0 = False
w2_weight_scale.format_ue8m0 = False
layer.register_parameter("w13_weight_scale_inv", w13_weight_scale)
set_weight_attrs(w13_weight_scale, scale_attrs)
layer.register_parameter("w2_weight_scale_inv", w2_weight_scale)
set_weight_attrs(w2_weight_scale, scale_attrs)
def create_moe_runner(self, layer: torch.nn.Module, moe_runner_config):
self.moe_runner_config = moe_runner_config
self._fp8.moe_runner_config = moe_runner_config
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
from sglang.srt.hardware_backend.npu.utils import NPUACLFormat, npu_format_cast
if layer.w13_weight_scale_inv.data.max() == 0:
raise RuntimeError(
f"FP4 expert weight scales are all zero (never loaded) for "
f"prefix={self.prefix!r}; the checkpoint scale names likely did "
"not match w13_weight_scale_inv."
)
if layer.w2_weight_scale_inv.data.max() == 0:
raise RuntimeError(
f"FP4 expert weight scales are all zero (never loaded) for "
f"prefix={self.prefix!r}; the checkpoint scale names likely did "
"not match w2_weight_scale_inv."
)
nz_kwargs = {
"customize_dtype": torch.float8_e4m3fn,
"input_dtype": _get_float4_e2m1fn_x2_dtype(),
}
nz_format = NPUACLFormat.ACL_FORMAT_FRACTAL_NZ
layer.w13_weight.data = npu_format_cast(
layer.w13_weight.data.view(torch.uint8), nz_format, **nz_kwargs
).transpose(1, 2)
layer.w2_weight.data = npu_format_cast(
layer.w2_weight.data.view(torch.uint8), nz_format, **nz_kwargs
).transpose(1, 2)
layer.w13_weight_scale_inv = torch.nn.Parameter(
_reshape_mxfp4_scale_for_npu(layer.w13_weight_scale_inv.data),
requires_grad=False,
)
layer.w2_weight_scale_inv = torch.nn.Parameter(
_reshape_mxfp4_scale_for_npu(layer.w2_weight_scale_inv.data),
requires_grad=False,
)
_configure_dsv4_deepep_dispatcher(layer)
def apply(
self,
layer: torch.nn.Module,
dispatch_output: "DispatchOutput",
) -> "CombineInput":
combine_input = npu_apply_w4a8_mxfp_moe_deepep(layer, dispatch_output)
if combine_input is not None:
return combine_input
combine_input = npu_apply_w4a4_mxfp_moe_ascend_tp(layer, dispatch_output)
if combine_input is not None:
return combine_input
# Standard dispatch. Unreachable on NPU today — create_moe_dispatcher
# picks AscendTPDispatcher whenever is_npu() and no a2a backend is set —
# but kept so this method is not silently wrong if that changes.
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
hidden_states = dispatch_output.hidden_states
topk_weights, topk_ids, _ = dispatch_output.topk_output
topk_ids = topk_ids.to(torch.int32)
topk_weights = topk_weights.to(hidden_states.dtype)
moe_runner_config = layer.moe_runner_config
output = npu_fused_experts_w4a4_mxfp(
hidden_states,
layer.w13_weight,
layer.w13_weight_scale_inv,
layer.w2_weight,
layer.w2_weight_scale_inv,
topk_weights,
topk_ids,
moe_runner_config.top_k,
swiglu_limit=moe_runner_config.swiglu_limit,
)
return StandardCombineInput(hidden_states=output)
def _reshape_mxfp4_scale_for_npu(scale: torch.Tensor) -> torch.Tensor:
"""``[E, N, K/32] -> [E, K/64, N, 2]``, the packed-pair layout the GMM wants."""
if scale.dim() != 3:
return scale
num_experts, n, k32 = scale.shape
if k32 % 2 != 0:
raise ValueError(
"MXFP4 scale K dimension must be divisible by 2 for the "
f"[E, K/64, N, 2] layout, got {tuple(scale.shape)}."
)
return scale.view(num_experts, n, k32 // 2, 2).transpose(1, 2)
def _apply_swiglu_limit_npu(
gate_up: torch.Tensor, swiglu_limit: Optional[float]
) -> None:
"""Clamp the SwiGLU input in place before ``npu_swiglu`` (DeepSeek-V4).
gate (first half) <= limit; up (second half) in
[-limit, limit]. ``chunk`` returns views, so the in-place clamps mutate
``gate_up`` directly. No-op when ``swiglu_limit`` is unset or <= 0.
"""
if swiglu_limit is None or swiglu_limit <= 0:
return
gate, up = gate_up.chunk(2, dim=-1)
gate.clamp_(max=swiglu_limit)
up.clamp_(min=-swiglu_limit, max=swiglu_limit)
def npu_fused_experts_w4a4_mxfp(
hidden_states: torch.Tensor,
w13: torch.Tensor,
w13_weight_scale_inv: torch.Tensor,
w2: torch.Tensor,
w2_weight_scale_inv: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
top_k: int,
swiglu_limit: Optional[float] = None,
**kwargs,
):
if torch.npu.is_current_stream_capturing():
return npu_fused_experts_w4a4_mxfp_decode(
hidden_states=hidden_states,
w13=w13,
w13_weight_scale_inv=w13_weight_scale_inv,
w2=w2,
w2_weight_scale_inv=w2_weight_scale_inv,
topk_weights=topk_weights,
topk_ids=topk_ids,
top_k=top_k,
swiglu_limit=swiglu_limit,
**kwargs,
)
original_shape = hidden_states.shape
original_dtype = hidden_states.dtype
if len(original_shape) == 3:
hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
num_tokens = hidden_states.shape[0]
num_experts = w13.shape[0]
row_idx = (
torch.arange(
0, num_tokens * top_k, dtype=torch.int32, device=topk_weights.device
)
.view(top_k, -1)
.permute(1, 0)
.contiguous()
)
hidden_states, expanded_row_idx, expanded_expert_idx = (
torch.ops.npu.npu_moe_init_routing(
hidden_states,
row_idx=row_idx,
expert_idx=topk_ids,
active_num=num_tokens,
)
)
expert_tokens = torch.ops.npu.npu_moe_compute_expert_tokens(
expanded_expert_idx, num_experts
).to(torch.int64)
# npu_moe_init_routing pads its output to the worst case; rows past the last
# expert boundary hold garbage and must not reach finalize_routing.
row_ids = torch.arange(
hidden_states.shape[0], device=hidden_states.device, dtype=torch.int64
)
valid_mask_2d = (row_ids < expert_tokens[-1]).unsqueeze(1)
hidden_states = w4a8_mxfp_gmm(
input=hidden_states,
input_scale=None,
weight=w13,
weight_scale=w13_weight_scale_inv,
group_list_type=0,
group_list=expert_tokens,
output_dtype=original_dtype,
)
assert swiglu_limit is not None
hidden_states, hidden_states_scale = swiglu_quant(
hidden_states,
group_list=expert_tokens,
group_list_type=0,
need_quant=True,
do_limit=True,
limit=swiglu_limit,
)
hidden_states = w4a8_mxfp_gmm(
input=hidden_states,
input_scale=hidden_states_scale,
weight=w2,
weight_scale=w2_weight_scale_inv,
group_list_type=0,
group_list=expert_tokens,
output_dtype=original_dtype,
)
hidden_states = hidden_states * valid_mask_2d.to(hidden_states.dtype)
final_hidden_states = torch.ops.npu.npu_moe_finalize_routing(
hidden_states,
skip1=None,
skip2=None,
bias=None,
scales=topk_weights,
expanded_src_to_dst_row=expanded_row_idx,
export_for_source_row=topk_ids,
)
if len(original_shape) == 3:
final_hidden_states = final_hidden_states.view(original_shape)
return final_hidden_states
def npu_fused_experts_w4a4_mxfp_decode(
hidden_states: torch.Tensor,
w13: torch.Tensor,
w13_weight_scale_inv: torch.Tensor,
w2: torch.Tensor,
w2_weight_scale_inv: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
top_k: int,
swiglu_limit: Optional[float] = None,
**kwargs,
):
"""Graph-capturable variant: routing v2 + token_unpermute, no host syncs."""
num_tokens = hidden_states.shape[:-1].numel()
global_num_experts = w13.shape[0]
original_shape = hidden_states.shape
original_dtype = hidden_states.dtype
group_list_type = 1
hidden_states, expanded_row_idx, expert_tokens, _ = (
torch.ops.npu.npu_moe_init_routing_v2(
hidden_states,
topk_ids,
active_num=num_tokens * top_k,
expert_num=global_num_experts,
expert_tokens_num_type=group_list_type,
expert_tokens_num_flag=True,
active_expert_range=[0, global_num_experts],
quant_mode=-1,
)
)
expert_tokens = expert_tokens.to(torch.int64)
hidden_states = w4a8_mxfp_gmm(
input=hidden_states,
input_scale=None,
weight=w13,
weight_scale=w13_weight_scale_inv,
group_list_type=group_list_type,
group_list=expert_tokens,
output_dtype=original_dtype,
)
assert swiglu_limit is not None
hidden_states, hidden_states_scale = swiglu_quant(
hidden_states,
group_list=expert_tokens,
group_list_type=group_list_type,
need_quant=True,
do_limit=True,
limit=swiglu_limit,
)
hidden_states = w4a8_mxfp_gmm(
input=hidden_states,
input_scale=hidden_states_scale,
weight=w2,
weight_scale=w2_weight_scale_inv,
group_list_type=group_list_type,
group_list=expert_tokens,
output_dtype=original_dtype,
)
final_hidden_states = torch.ops.npu.npu_moe_token_unpermute(
permuted_tokens=hidden_states,
sorted_indices=torch.abs(expanded_row_idx),
probs=topk_weights,
)
if len(original_shape) == 3:
final_hidden_states = final_hidden_states.view(original_shape)
return final_hidden_states
def npu_apply_w4a4_mxfp_moe_ascend_tp(
layer: torch.nn.Module,
dispatch_output: "DispatchOutput",
) -> Optional["CombineInput"]:
"""Ascend TP path. Returns ``None`` when the dispatch is not an Ascend TP one.
AscendTPDispatcher already ran npu_moe_init_routing_v2 on dispatch and runs
npu_moe_finalize_routing (with topk_weights) on combine, so this only owns
the grouped-matmul chain in between — no permute, no routing-weight apply.
"""
from sglang.srt.layers.moe.token_dispatcher import AscendTPCombineInput
from sglang.srt.layers.moe.token_dispatcher.base import DispatchOutputChecker
if not DispatchOutputChecker.format_is_ascend_tp(dispatch_output):
return None
hidden_states = npu_apply_without_routing_weights_w4a4_mxfp(
layer,
dispatch_output.hidden_states,
dispatch_output.hidden_states_scale,
group_list_type=dispatch_output.group_list_type,
group_list=dispatch_output.expert_tokens,
output_dtype=torch.bfloat16,
)
return AscendTPCombineInput(hidden_states=hidden_states)
def npu_apply_w4a8_mxfp_moe_deepep(
layer: torch.nn.Module,
dispatch_output: "DispatchOutput",
) -> Optional["CombineInput"]:
"""DeepEP path. Returns ``None`` when the dispatch is not a DeepEP one."""
from sglang.srt.layers.moe.token_dispatcher import (
DeepEPLLCombineInput,
DeepEPNormalCombineInput,
)
from sglang.srt.layers.moe.token_dispatcher.base import DispatchOutputChecker
if not dispatch_output.format.is_deepep():
return None
if DispatchOutputChecker.format_is_deepep_normal(dispatch_output):
hidden_states, hidden_states_scale, _, _, num_recv_tokens_per_expert = (
dispatch_output
)
group_list = torch.tensor(
num_recv_tokens_per_expert, dtype=torch.int64, device=hidden_states.device
)
combine_cls = DeepEPNormalCombineInput
else:
hidden_states, hidden_states_scale, _, _, group_list, _ = dispatch_output
group_list = group_list.to(torch.int64)
combine_cls = DeepEPLLCombineInput
hidden_states = npu_apply_without_routing_weights_w4a4_mxfp(
layer,
hidden_states,
hidden_states_scale,
group_list_type=1,
group_list=group_list,
output_dtype=torch.bfloat16,
)
return combine_cls(
hidden_states=hidden_states,
topk_ids=dispatch_output.topk_ids,
topk_weights=dispatch_output.topk_weights,
)
def npu_apply_without_routing_weights_w4a4_mxfp(
layer,
hidden_states,
hidden_states_scale,
*,
group_list_type,
group_list,
output_dtype,
):
hidden_states = w4a8_mxfp_gmm(
input=hidden_states,
input_scale=hidden_states_scale,
weight=layer.w13_weight,
weight_scale=layer.w13_weight_scale_inv,
group_list_type=group_list_type,
group_list=group_list,
output_dtype=output_dtype,
)
assert layer.moe_runner_config.swiglu_limit is not None
hidden_states, hidden_states_scale = swiglu_quant(
hidden_states,
group_list=group_list,
group_list_type=group_list_type,
need_quant=True,
do_limit=True,
limit=layer.moe_runner_config.swiglu_limit,
)
return w4a8_mxfp_gmm(
input=hidden_states,
input_scale=hidden_states_scale,
weight=layer.w2_weight,
weight_scale=layer.w2_weight_scale_inv,
group_list_type=group_list_type,
group_list=group_list,
output_dtype=output_dtype,
)
def _pair_pack_mxfp_act_scale(
scale: torch.Tensor, input_shape: Optional[tuple[int, int]] = None
) -> torch.Tensor:
"""Adapt MXFP activation scales to the A5 GMM ``[M, K/64, 2]`` layout.
Low-latency DeepEP MXFP8 returns a flat E8M0 scale buffer, one byte for
every 32 activation elements. The grouped-matmul kernel expects those
bytes paired on the final dimension instead.
"""
if scale.ndim == 1:
if input_shape is None or len(input_shape) != 2:
raise ValueError(
"A flat MXFP activation scale requires its two-dimensional "
"activation input shape."
)
num_tokens, hidden_size = input_shape
if hidden_size % (2 * MXFP4_BLOCK_SIZE) != 0:
raise ValueError(
"MXFP activation hidden size must be divisible by "
f"{2 * MXFP4_BLOCK_SIZE}; got {hidden_size}."
)
expected_num_scales = num_tokens * (hidden_size // MXFP4_BLOCK_SIZE)
if scale.numel() != expected_num_scales:
raise ValueError(
"Invalid flat MXFP activation scale length: expected "
f"{expected_num_scales} for input shape {input_shape}, got "
f"{scale.numel()}."
)
scale = scale.reshape(num_tokens, hidden_size // MXFP4_BLOCK_SIZE)
# ``[M, K/32] -> [M, K/64, 2]`` MX per-token scale layout for the A5 GMM.
if scale.ndim != 2:
return scale
if scale.shape[-1] % 2 != 0:
raise ValueError(f"Invalid MXFP per-token scale shape: {tuple(scale.shape)}")
return scale.reshape(scale.shape[0], scale.shape[1] // 2, 2)
def w4a8_mxfp_gmm(
*,
input: torch.Tensor,
input_scale: Optional[torch.Tensor],
weight: torch.Tensor,
weight_scale: torch.Tensor,
group_list_type: int,
group_list: torch.Tensor,
output_dtype: torch.dtype,
scale_alg=None,
) -> torch.Tensor:
"""FP4 weight x FP8-e4m3 activation (the checkpoint's W4A8_MXFP scheme).
W4A8MXFP GMM call: FP8 ``x_dtype``, FP4
``weight_dtype``, and the weight block scales fed through ``antiquant_scale``
with ``scale=None`` — the ``scale=`` + ``scale_dtype=`` form belongs to
W4A4_MXFP4 and dequantizes differently.
"""
group_list = group_list.to(torch.int64)
if input_scale is None:
x, x_scale = torch.ops.npu.npu_dynamic_mx_quant(
input,
axis=1,
round_mode="rint",
dst_type=torch.float8_e4m3fn,
block_size=MXFP4_BLOCK_SIZE,
scale_alg=scale_alg,
)
else:
x, x_scale = input, input_scale
return torch.ops.npu.npu_grouped_matmul(
[x],
[weight],
scale=None,
antiquant_scale=[weight_scale],
scale_dtype=None,
per_token_scale=[
_pair_pack_mxfp_act_scale(x_scale, input_shape=tuple(x.shape))
],
split_item=2,
group_type=0,
group_list=group_list,
group_list_type=group_list_type,
output_dtype=output_dtype,
x_dtype=torch.float8_e4m3fn,
weight_dtype=_get_float4_e2m1fn_x2_dtype(),
per_token_scale_dtype=_get_float8_e8m0fnu_dtype(),
)[0]
python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py
+54 -1
View File
@@ -1,5 +1,5 @@
import logging
from typing import TYPE_CHECKING, Optional
from typing import TYPE_CHECKING, List, Optional
import torch
from torch.nn.parameter import Parameter
@@ -304,6 +304,59 @@ class NPUMXFP8LinearMethod(_NPULinearMethodBase):
return output.reshape(output_shape)
def npu_w8a8_mxfp8_linear(
input: torch.Tensor,
weight: torch.Tensor,
block_size: List[int],
weight_scale: torch.Tensor,
input_scale: Optional[torch.Tensor] = None,
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Block-FP8 linear on Atlas A5, used as the ``w8a8_block_fp8_linear``
backend on NPU (see ``fp8_utils._dispatch_auto_backend``).
The loading path requantizes block-FP8 weights into the A5 MXFP8 layout;
activations are quantized per call. ``block_size`` is retained for the shared
block-FP8 backend interface and ``input_scale`` is unused because activation
scales are always dynamic here.
"""
if weight.dtype != torch.float8_e4m3fn:
raise ValueError(
f"npu_w8a8_mxfp8_linear expects float8_e4m3fn weights, got {weight.dtype}"
)
original_dtype = input.dtype
if original_dtype not in (torch.float16, torch.bfloat16):
input = input.to(torch.bfloat16)
original_dtype = torch.bfloat16
orig_shape = input.shape
input_2d = input.view(-1, orig_shape[-1]).contiguous()
x_fp8, x_scale = torch.ops.npu.npu_dynamic_mx_quant(
input_2d, dst_type=torch.float8_e4m3fn
)
e8m0_dtype = _get_float8_e8m0fnu_dtype()
quant_bias = (
bias.to(torch.float32)
if bias is not None and bias.dtype != torch.float32
else bias
)
output_2d = torch.ops.npu.npu_quant_matmul(
x_fp8,
weight,
scale=weight_scale,
scale_dtype=e8m0_dtype,
pertoken_scale=x_scale,
pertoken_scale_dtype=e8m0_dtype,
bias=quant_bias,
output_dtype=original_dtype,
group_sizes=(1, 1, MXFP8_BLOCK_SIZE),
)
return output_2d.reshape(*orig_shape[:-1], output_2d.shape[-1])
class NPU_W4A4DynamicLinearMethod(_NPULinearMethodBase):
def process_weights_after_loading(self, layer):
layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
python/sglang/srt/hardware_backend/npu/quantization/w8a8_mxfp8.py
+120
View File
@@ -0,0 +1,120 @@
from typing import List
import torch
from torch.nn import Module
_NPU_ARCH35_MXFP8_BLOCK_SIZE = 32
def process_npu_arch35_mxfp8_linear_weights(
layer: Module, weight_block_size: List[int], scale_fmt: str
) -> None:
"""Convert UE8M0 block-FP8 weights to the NPU arch35 MXFP8 layout."""
if scale_fmt != "ue8m0":
raise ValueError(
"NPU arch35 MXFP8 weight loading requires scale_fmt='ue8m0', "
f"got {scale_fmt!r}."
)
_layout_npu_arch35_ue8m0_weights(layer, weight_block_size)
def _layout_npu_arch35_ue8m0_weights(
layer: Module, weight_block_size: List[int]
) -> None:
"""Reinterpret UE8M0 block scales and transpose weights without requantizing."""
block_n, block_k = weight_block_size
group_size = _NPU_ARCH35_MXFP8_BLOCK_SIZE
n_dim, k_dim = layer.weight.shape
if block_k % group_size != 0:
raise ValueError(
f"UE8M0 block K size must be divisible by {group_size}, got {block_k}."
)
if k_dim % (2 * group_size) != 0:
raise ValueError(
"NPU arch35 MXFP8 linear requires K to be divisible by "
f"{2 * group_size}, got {k_dim}."
)
expected_scale_shape = (
(n_dim + block_n - 1) // block_n,
(k_dim + block_k - 1) // block_k,
)
checkpoint_scale = layer.weight_scale_inv.data
if tuple(checkpoint_scale.shape) != expected_scale_shape:
raise ValueError(
"Unexpected UE8M0 scale shape: "
f"got {tuple(checkpoint_scale.shape)}, expected {expected_scale_shape}."
)
if checkpoint_scale.dtype == torch.float8_e8m0fnu:
scale_u8 = checkpoint_scale.view(torch.uint8)
elif checkpoint_scale.dtype == torch.uint8:
scale_u8 = checkpoint_scale
elif checkpoint_scale.dtype == torch.float32:
# SGLang's block scale parameter is currently allocated as FP32. The
# loader converts F8_E8M0 values to exact powers of two, so recover the
# original exponent byte without materializing the weight in FP32.
scale_u8 = ((checkpoint_scale.view(torch.int32) >> 23) & 0xFF).to(torch.uint8)
else:
raise TypeError(
"UE8M0 checkpoint scales must be float8_e8m0fnu, uint8, or float32, "
f"got {checkpoint_scale.dtype}."
)
scale_u8 = scale_u8.repeat_interleave(block_n, dim=0)[:n_dim]
scale_u8 = scale_u8.repeat_interleave(block_k // group_size, dim=1)
scale_u8 = scale_u8[:, : k_dim // group_size]
# Keep transpose views: the A5 kernel expects the original row-major
# storage scanned in K-major logical order.
layer.weight.data = layer.weight.data.transpose(0, 1)
layer.weight_scale_inv.data = scale_u8.reshape(
n_dim, k_dim // (2 * group_size), 2
).transpose(0, 1)
layer.weight_scale_inv.format_ue8m0 = True
if getattr(layer, "_dsv4_npu_arch35_mxfp8_wo_a", False):
batch_npu_arch35_wo_a_weights(layer)
def batch_npu_arch35_wo_a_weights(layer: Module) -> None:
"""Reshape DSV4's ``wo_a`` for arch35 batched MXFP8 matmul.
``npu_transpose_quant_batchmatmul`` expects weight
``[D, G*R] -> [G, D, R]`` and scale
``[D/64, G*R, 2] -> [G, D/64, R, 2]``.
"""
num_groups = layer._dsv4_num_groups
rank = layer._dsv4_o_lora_rank
hidden_dim = layer.weight.shape[0]
scale_k64 = layer.weight_scale_inv.shape[0]
output_dim = num_groups * rank
if layer.weight.shape != (hidden_dim, output_dim):
raise ValueError(
"Unexpected NPU arch35 wo_a weight layout after FP8 post-processing: "
f"got {tuple(layer.weight.shape)}, expected ({hidden_dim}, {output_dim})."
)
if layer.weight_scale_inv.shape != (scale_k64, output_dim, 2):
raise ValueError(
"Unexpected NPU arch35 wo_a scale layout after FP8 post-processing: "
f"got {tuple(layer.weight_scale_inv.shape)}, expected "
f"({scale_k64}, {output_dim}, 2)."
)
if scale_k64 * 64 != hidden_dim:
raise ValueError(
"Unexpected NPU arch35 wo_a scale K dimension: "
f"{scale_k64} packed pairs for hidden dim {hidden_dim}."
)
layer.weight.data = (
layer.weight.data.T.reshape(num_groups, rank, hidden_dim)
.transpose(1, 2)
.contiguous()
)
layer.weight_scale_inv.data = (
layer.weight_scale_inv.data.transpose(0, 1)
.reshape(num_groups, rank, scale_k64, 2)
.transpose(1, 2)
.contiguous()
)
python/sglang/srt/hardware_backend/npu/utils.py
+26
View File
@@ -24,6 +24,32 @@ indexer_weight_stream = None
gva_is_inited = False
@functools.lru_cache(maxsize=1)
def is_npu_arch35() -> bool:
"""Whether the runtime is on NPU architecture 35."""
if not is_npu():
return False
import acl
return acl.rt.get_device_info(0, 601) == (3510, 0)
def use_npu_arch35_mxfp8_wo_a(quant_config) -> bool:
"""Whether wo_a runs the native NPU arch35 MXFP8 GEMM.
Only for serialized DeepSeek block-FP8 checkpoints — those are the ones
``Fp8LinearMethod.process_weights_after_loading`` can reinterpret into the
NPU arch35 MXFP8 scale layout.
"""
if not _is_npu or not is_npu_arch35() or quant_config is None:
return False
if not getattr(quant_config, "is_checkpoint_fp8_serialized", False):
return False
weight_block_size = getattr(quant_config, "weight_block_size", None)
return tuple(weight_block_size or ()) == (128, 128)
class NPUACLFormat(IntEnum):
ACL_FORMAT_UNDEFINED = -1
ACL_FORMAT_ND = 2
python/sglang/srt/layers/moe/token_dispatcher/deepep.py
+69 -1
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import inspect
import logging
import os
from contextlib import nullcontext
@@ -457,6 +458,10 @@ class _DeepEPDispatcherImplBase:
"use_fp8": False,
"use_nvfp4": True,
},
DispatcherOutputDtype.MXFP8: {
"use_fp8": False,
"use_nvfp4": False,
},
}
# Validate and apply hardware-specific adjustments
@@ -473,6 +478,25 @@ class _DeepEPDispatcherImplBase:
def _validate_and_adjust_dtype(self) -> None:
"""Validate dtype against hardware and adjust if necessary."""
self.low_latency_quant_mode = None
self._low_latency_quant_mode_runtime_checked = False
if self.deepep_output_dtype == DispatcherOutputDtype.MXFP8:
if not _is_npu or self.dispatch_mode != DeepEPMode.LOW_LATENCY:
raise RuntimeError(
"MXFP8 DeepEP dispatch is supported only for A5 "
"low-latency dispatch."
)
from sglang.srt.hardware_backend.npu.utils import is_npu_arch35
if not is_npu_arch35():
raise RuntimeError(
"MXFP8 DeepEP dispatch is supported only on Ascend A5 "
"in low-latency mode."
)
self.low_latency_quant_mode = "mx_fp8_e4m3"
return
if _is_npu:
if self.deepep_output_dtype == DispatcherOutputDtype.FP8:
logger.warning_once(
@@ -765,6 +789,49 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
)
buffer = self._get_buffer()
if (
self.low_latency_quant_mode is not None
and not self._low_latency_quant_mode_runtime_checked
):
try:
dispatch_signature = inspect.signature(buffer.low_latency_dispatch)
except (TypeError, ValueError) as exc:
raise RuntimeError(
"A5 MXFP8 DeepEP dispatch requires a recent "
"sgl-kernel-npu/DeepEP runtime exposing "
"low_latency_dispatch(..., quant_mode=...)."
) from exc
if "quant_mode" not in dispatch_signature.parameters:
raise RuntimeError(
"A5 MXFP8 DeepEP dispatch requires a recent "
"sgl-kernel-npu/DeepEP runtime exposing "
"low_latency_dispatch(..., quant_mode=...)."
)
self._low_latency_quant_mode_runtime_checked = True
use_fp8 = self.use_fp8
low_latency_quant_kwargs = {}
if self.low_latency_quant_mode is not None:
deep_use_mode = os.environ.get("DEEP_USE_MODE", "default")
if deep_use_mode == "default":
low_latency_quant_kwargs = {
"quant_mode": self.low_latency_quant_mode,
}
elif deep_use_mode == "ops":
# The ops strategy ignores quant_mode and uses the legacy
# flags. Pass both forms so the request is explicit and the
# strategy still produces E4M3 + E8M0 MXFP8 tensors.
use_fp8 = True
low_latency_quant_kwargs = {
"quant_mode": self.low_latency_quant_mode,
"use_ue8m0": True,
}
else:
raise RuntimeError(
"A5 MXFP8 DeepEP dispatch supports only "
"DEEP_USE_MODE=default or DEEP_USE_MODE=ops; got "
f"{deep_use_mode!r}."
)
_deepep_precompile_tp_barrier()
packed_recv_hidden, self.packed_recv_count, self.handle, event, hook = (
buffer.low_latency_dispatch(
@@ -772,7 +839,8 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
topk_ids,
self.num_max_dispatch_tokens_per_rank,
self.num_experts,
use_fp8=self.use_fp8,
use_fp8=use_fp8,
**low_latency_quant_kwargs,
**(
dict(topk_weights=topk_weights)
if _is_npu and not _use_zbal
python/sglang/srt/layers/quantization/fp8.py
+23
View File
@@ -23,6 +23,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.utils import is_npu_arch35
from sglang.srt.layers.amx_utils import (
CPUQuantMethod,
_amx_process_weight_after_loading,
@@ -245,6 +246,7 @@ class Fp8Config(QuantizationConfig):
use_mxfp8: bool = False,
is_fp4_experts: bool = False,
kv_cache_quant_algo: Optional[str] = None,
scale_fmt: Optional[str] = None,
) -> None:
super().__init__()
# DSV4 mxfp4-packed (True) vs converted FP8 (False); injected by
@@ -269,6 +271,7 @@ class Fp8Config(QuantizationConfig):
self.packed_modules_mapping = packed_modules_mapping or {}
self.use_mxfp8 = use_mxfp8
self.kv_cache_quant_algo = kv_cache_quant_algo
self.scale_fmt = scale_fmt
if weight_block_size is not None:
if not is_checkpoint_fp8_serialized:
raise ValueError(
@@ -336,6 +339,7 @@ class Fp8Config(QuantizationConfig):
kv_cache_quant_algo = cls.get_from_keys_or(
config, ["kv_cache_quant_algo"], None
)
scale_fmt = cls.get_from_keys_or(config, ["scale_fmt"], None)
if use_mxfp8:
# MXFP8 (OCP) spec fixes block size to [1, 32]; ckpt field is metadata only.
if weight_block_size is not None and weight_block_size != [1, 32]:
@@ -352,6 +356,7 @@ class Fp8Config(QuantizationConfig):
packed_modules_mapping=packed_modules_mapping,
use_mxfp8=use_mxfp8,
kv_cache_quant_algo=kv_cache_quant_algo,
scale_fmt=scale_fmt,
)
def get_quant_method(
@@ -396,6 +401,13 @@ class Fp8Config(QuantizationConfig):
)
return fp8_method
if self.is_fp4_experts and is_npu_arch35():
from sglang.srt.hardware_backend.npu.quantization.fp4_moe_methods import (
NPUW4A4Fp4MoEMethod,
)
return NPUW4A4Fp4MoEMethod(fp8_method, prefix=prefix)
if self.is_fp4_experts and get_moe_runner_backend().is_marlin():
from sglang.srt.layers.quantization.mxfp4_marlin_moe import (
Mxfp4MarlinMoEMethod,
@@ -688,6 +700,17 @@ class Fp8LinearMethod(LinearMethodBase):
layer.weight_scale_inv.format_ue8m0 = True
self._process_mxfp8_linear_weight_scale(layer)
return
elif _is_npu and is_npu_arch35():
from sglang.srt.hardware_backend.npu.quantization.w8a8_mxfp8 import (
process_npu_arch35_mxfp8_linear_weights,
)
process_npu_arch35_mxfp8_linear_weights(
layer,
self.weight_block_size,
scale_fmt=getattr(self.quant_config, "scale_fmt", None),
)
return
# If ROCm, normalize the weights and scales to e4m3fnuz
if _is_fp8_fnuz:
# activation_scheme: dynamic
python/sglang/srt/layers/quantization/fp8_utils.py
+9 -1
View File
@@ -24,6 +24,7 @@ from sglang.kernels.ops.quantization.fp8_kernel import (
w8a8_block_fp8_matmul_triton,
)
from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.utils import is_npu_arch35
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.quantization.mxfp4_tensor import MXFP4QuantizeUtil
from sglang.srt.runtime_context import (
@@ -781,7 +782,8 @@ def _dispatch_auto_backend() -> Callable:
# 2. FlashInfer TRTLLM (if Blackwell GPU and FlashInfer available)
# 3. CUTLASS (if SM120 GPU and CUDA 12.8+)
# 4. AITER (if AMD GPU with AITER enabled)
# 5. Triton (fallback)
# 5. NPU (Ascend)
# 6. Triton (fallback)
if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM:
return deepgemm_w8a8_block_fp8_linear_with_fallback
@@ -791,6 +793,12 @@ def _dispatch_auto_backend() -> Callable:
return cutlass_w8a8_block_fp8_linear_with_fallback
elif _use_aiter:
return aiter_w8a8_block_fp8_linear
elif is_npu_arch35():
from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import (
npu_w8a8_mxfp8_linear,
)
return npu_w8a8_mxfp8_linear
else:
return triton_w8a8_block_fp8_linear
python/sglang/srt/models/deepseek_v2.py
+6 -6
View File
@@ -448,12 +448,13 @@ class DeepseekV2MLP(nn.Module):
# Fallback: fused silu+clamp kernel (still faster than unfused)
elif self.swiglu_limit is not None:
if _is_npu:
_g, _u = gate_up.chunk(2, dim=-1)
_lim = float(self.swiglu_limit)
gate_up = torch.cat(
[_g.clamp(max=_lim), _u.clamp(min=-_lim, max=_lim)], dim=-1
x = torch.ops.npu.npu_clipped_swiglu(
gate_up,
alpha=1,
limit=self.swiglu_limit,
bias=0,
interleaved=False,
)
x = self.act_fn(gate_up)
else:
M, N = gate_up.shape
x = gate_up.new_empty((M, N // 2))
@@ -485,7 +486,6 @@ class MoEGate(nn.Module):
),
)
)
if config.topk_method == "noaux_tc" and not is_hash_moe:
correction_bias_dtype = torch.float32
if quant_config is not None:
python/sglang/srt/models/deepseek_v4.py
+104 -30
View File
@@ -45,7 +45,15 @@ 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.hardware_backend.npu.dsv4.dsv4_rope import (
Dsv4NpuRoPE,
prime_rope_cos_sin,
rope_cos_sin,
)
from sglang.srt.hardware_backend.npu.utils import (
is_npu_arch35,
use_npu_arch35_mxfp8_wo_a,
)
from sglang.srt.layers.attention.dsa.utils import (
can_dsa_cp_split,
dsa_use_prefill_cp,
@@ -703,12 +711,16 @@ class MqaAttentionBase(nn.Module):
if wo_b_reduce_results is None
else wo_b_reduce_results
)
# NPU arch35 runs wo_a as a batched MXFP8 GEMM instead of deep_gemm's FP8 one,
# but it needs the same quantized weights.
self.use_npu_arch35_mxfp8_wo_a = use_npu_arch35_mxfp8_wo_a(quant_config)
quantize_wo_a = fp8 or self.use_npu_arch35_mxfp8_wo_a
if wo_a_keeps_quant_config is None:
keep_source_quant = (
quant_config is not None and quant_config.get_name() == "expert_pack"
)
wo_a_quant_config: Optional[QuantizationConfig] = (
quant_config if fp8 or keep_source_quant else None
quant_config if quantize_wo_a or keep_source_quant else None
)
elif wo_a_keeps_quant_config:
wo_a_quant_config = quant_config
@@ -761,14 +773,21 @@ class MqaAttentionBase(nn.Module):
prefix=add_prefix("wo_a", prefix),
tp_rank=self.attn_tp_rank,
tp_size=self.attn_tp_size,
**({} if fp8 else {"params_dtype": torch.bfloat16}),
**({} if quantize_wo_a else {"params_dtype": torch.bfloat16}),
)
if fp8:
from sglang.srt.layers import deep_gemm_wrapper
if quantize_wo_a:
assert hasattr(self.wo_a, "weight_scale_inv"), (
"FP8 quant_config must create weight_scale_inv"
)
if self.use_npu_arch35_mxfp8_wo_a:
# Read by the NPU arch35 MXFP8 weight processor to batch the
# weight/scale per attention group for npu_transpose_quant_batchmatmul.
self.wo_a._dsv4_npu_arch35_mxfp8_wo_a = True
self.wo_a._dsv4_num_groups = self.n_local_groups
self.wo_a._dsv4_o_lora_rank = self.o_lora_rank
elif fp8:
from sglang.srt.layers import deep_gemm_wrapper
self.wo_a.weight_scale_inv.format_ue8m0 = (
deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
)
@@ -905,9 +924,22 @@ class MQALayer(MqaAttentionBase):
)
if _is_npu:
Dsv4NpuRoPE.for_freqs(
rope = Dsv4NpuRoPE.for_freqs(
self.freqs_cis, getattr(self, "rotary_emb", None)
).ensure_tables(torch.float32)
)
# fp32 tables feed the compressor gather; bf16 tables make the
# activation-dtype gathers cast-free. Bit-identical values:
# rounding the table once equals rounding each gathered element.
rope.ensure_tables(torch.float32)
rope.ensure_tables(torch.bfloat16)
# npu_rms_norm has no weight-free overload; the per-head q norm
# reads this cached ones vector instead of paying a per-call
# alloc + fill.
self.register_buffer(
"q_rms_norm_ones",
torch.ones(self.head_dim, dtype=torch.bfloat16),
persistent=False,
)
if _is_hip:
cos_cache = (
@@ -990,22 +1022,25 @@ class MQALayer(MqaAttentionBase):
return result
def _get_npu_rope_position_cache(
self, positions: torch.Tensor, dtype: torch.dtype, inverse: bool = False
self,
forward_batch: ForwardBatch,
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(
# can also be shared by the target and NextN models. Only the immutable
# full table is cached on it; position-gathered tensors are memoized per
# forward (prime_rope_cos_sin / rope_cos_sin), never across forwards --
# reusing them based on shape alone gives MTP decode the previous step's
# RoPE values when positions change but batch size does not.
return rope_cos_sin(
self.freqs_cis,
getattr(self, "rotary_emb", None),
forward_batch,
positions,
dtype,
view_4d=True,
inverse=inverse,
allow_build=False,
cache_dtype=torch.float32,
)
def _compute_q_a(
@@ -1200,7 +1235,7 @@ class MQALayer(MqaAttentionBase):
kv, _ = self.wkv(x)
kv = self.kv_norm(kv)
cos4_k, sin4_k = self._get_npu_rope_position_cache(
positions, kv.dtype, inverse=False
forward_batch, positions, kv.dtype, inverse=False
)
Dsv4NpuRoPE.apply_rotary_mul_inplace(
kv.unsqueeze(1),
@@ -1220,10 +1255,9 @@ class MQALayer(MqaAttentionBase):
stream_q.wait_event(q_lora_ready)
q, _ = self.wq_b(q_lora)
q = q.view(-1, self.n_local_heads, self.head_dim)
_dummy = q.new_ones(q.shape[-1])
q = torch_npu.npu_rms_norm(q, _dummy, self.eps)[0]
q = torch_npu.npu_rms_norm(q, self.q_rms_norm_ones, self.eps)[0]
cos4_q, sin4_q = self._get_npu_rope_position_cache(
positions, q.dtype, inverse=False
forward_batch, positions, q.dtype, inverse=False
)
Dsv4NpuRoPE.apply_rotary_mul_inplace(
q,
@@ -1515,8 +1549,7 @@ class MQALayer(MqaAttentionBase):
q_lora = self.q_norm(q_lora)
q, _ = self.wq_b(q_lora)
q = q.view(-1, self.n_local_heads, self.head_dim)
_dummy = q.new_ones(q.shape[-1])
q = torch_npu.npu_rms_norm(q, _dummy, self.eps)[0]
q = torch_npu.npu_rms_norm(q, self.q_rms_norm_ones, self.eps)[0]
if qkv_a is not None:
kv = qkv_a[..., self.q_lora_rank :]
@@ -1525,7 +1558,7 @@ class MQALayer(MqaAttentionBase):
kv = self.kv_norm(kv)
cos4, sin4 = self._get_npu_rope_position_cache(
positions, q.dtype, inverse=False
forward_batch, positions, q.dtype, inverse=False
)
Dsv4NpuRoPE.apply_rotary_mul_inplace(
q,
@@ -1790,7 +1823,7 @@ class MQALayer(MqaAttentionBase):
else:
if _is_npu:
cos4, sin4 = self._get_npu_rope_position_cache(
positions, o.dtype, inverse=True
forward_batch, positions, o.dtype, inverse=True
)
Dsv4NpuRoPE.apply_rotary_mul_inplace(
o,
@@ -1810,7 +1843,23 @@ class MQALayer(MqaAttentionBase):
o = o.view(o.shape[0], self.n_local_groups, -1)
if _FP8_WO_A_GEMM and _wo_a_fp8_mxscale is not None:
if self.use_npu_arch35_mxfp8_wo_a:
o, o_scale = torch_npu.npu_dynamic_mx_quant(
o, dst_type=torch.float8_e4m3fn
)
o = torch_npu.npu_transpose_quant_batchmatmul(
o,
self.wo_a.weight,
dtype=torch.bfloat16,
bias=None,
group_sizes=(0, 0, 32),
x1_scale=o_scale.view(torch.float8_e8m0fnu),
x2_scale=self.wo_a.weight_scale_inv.view(torch.float8_e8m0fnu),
perm_x1=(1, 0, 2),
perm_x2=(0, 1, 2),
perm_y=(1, 0, 2),
)
elif _FP8_WO_A_GEMM and _wo_a_fp8_mxscale is not None:
# ROCm gfx950: same fp8 absorb GEMM as the DeepGEMM path below,
# but through aiter's e8m0 block-scale batched GEMM. The
# activation is quantized per token-group inside the helper.
@@ -2162,7 +2211,16 @@ class DeepseekV4DecoderLayer(nn.Module):
)
if _is_npu:
return torch.ops.custom.npu_hc_post(x, residual, post, comb)
if not is_npu_arch35():
return torch.ops.custom.npu_hc_post(x, residual, post, comb)
# The A5 build of npu_hc_post is batched — it requires a leading
# batch axis on every operand.
return torch.ops.custom.npu_hc_post(
x.unsqueeze(0),
residual.unsqueeze(0),
post.unsqueeze(0),
comb.unsqueeze(0),
).squeeze(0)
if _is_xpu:
return _get_mhc_ops().mhc_post(x, residual, post, comb)
@@ -3276,6 +3334,19 @@ class DeepseekV4Model(nn.Module):
for _attr in ("freqs_cis_c4", "freqs_cis_c128"):
if hasattr(forward_batch, _attr):
delattr(forward_batch, _attr)
if _is_npu and not run_tbo:
# Rope cos/sin for the whole forward: one bf16 gather per rope
# config on the current stream, before the layer loop forks the
# KV/Q side streams. TBO children carry their own positions and
# recompute per layer.
prime_rope_cos_sin(
(
self.layers[i].self_attn
for i in range(self.start_layer, self.end_layer)
),
forward_batch,
positions,
)
if run_tbo:
# Two-batch-overlap prefill (EP / mori). Cross-layer mHC fusion is
# disabled here (each layer self-contained), so no trailing hc_post.
@@ -3749,7 +3820,10 @@ class DeepseekV4ForCausalLM(nn.Module):
else:
raise ValueError("num_nextn_predict_layers is not in the config")
if not _FP8_WO_A_GEMM:
# Must mirror MQALayer.__init__'s `quantize_wo_a`: dequantizing wo_a here
# while the layer allocated an FP8 parameter (or vice versa) fails the
# weight loader's dtype check.
if not (_FP8_WO_A_GEMM or use_npu_arch35_mxfp8_wo_a(self.quant_config)):
weights = _prepare_deepseek_v4_weights(weights, self.quant_config)
stacked_params_mapping = DEEPSEEK_V4_STACKED_PARAMS_MAPPING
python/sglang/srt/models/deepseek_v4_nextn.py
+11 -1
View File
@@ -7,6 +7,7 @@ from torch import nn
from transformers import PretrainedConfig
from sglang.srt.distributed import get_pp_group
from sglang.srt.hardware_backend.npu.dsv4.dsv4_rope import prime_rope_cos_sin
from sglang.srt.layers.attention.dsa.utils import (
can_dsa_cp_split,
dsa_use_prefill_cp,
@@ -40,7 +41,11 @@ from sglang.srt.layers.vocab_parallel_embedding import (
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.forward_context import get_attn_backend
from sglang.srt.models.deepseek_v4 import DeepseekV4DecoderLayer, DeepseekV4ForCausalLM
from sglang.srt.models.deepseek_v4 import (
DeepseekV4DecoderLayer,
DeepseekV4ForCausalLM,
_is_npu,
)
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils import add_prefix
@@ -185,6 +190,11 @@ class DeepseekV4ModelNextN(nn.Module):
input_ids = cp_round_robin_input_ids(input_ids)
input_ids_global = input_ids
if _is_npu:
# Same per-forward rope prime as DeepseekV4Model.forward: the
# decoder layer reads the memoized gather instead of re-gathering.
prime_rope_cos_sin([self.decoder.self_attn], forward_batch, positions)
hidden_states, residual, post, comb = self.decoder(
positions=positions,
hidden_states=hidden_states,