diff --git a/python/sglang/srt/disaggregation/ascend/conn.py b/python/sglang/srt/disaggregation/ascend/conn.py index 8ab6b06f7..bc44b5b95 100644 --- a/python/sglang/srt/disaggregation/ascend/conn.py +++ b/python/sglang/srt/disaggregation/ascend/conn.py @@ -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) diff --git a/python/sglang/srt/disaggregation/ascend/transfer_engine.py b/python/sglang/srt/disaggregation/ascend/transfer_engine.py index c4283552f..93d1a2956 100644 --- a/python/sglang/srt/disaggregation/ascend/transfer_engine.py +++ b/python/sglang/srt/disaggregation/ascend/transfer_engine.py @@ -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: diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 56aaec3d3..c9d8fb65e 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -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 diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index a5e613f08..cc262fee0 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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 diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py index 6db4fbdf1..9edff7606 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py @@ -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): diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_common_hooks.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_common_hooks.py index 3cda73847..0e0033efb 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_common_hooks.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_common_hooks.py @@ -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): diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py index 8aa279534..ec9fe9e70 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py @@ -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) diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_req_to_token_pool.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_req_to_token_pool.py index 86b3a8913..d51c551dc 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_req_to_token_pool.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_req_to_token_pool.py @@ -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. diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_rope.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_rope.py index aee422f3e..9b843b1f6 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_rope.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_rope.py @@ -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, + ) diff --git a/python/sglang/srt/hardware_backend/npu/quantization/fp4_moe_methods.py b/python/sglang/srt/hardware_backend/npu/quantization/fp4_moe_methods.py new file mode 100644 index 000000000..ab5e433b6 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/quantization/fp4_moe_methods.py @@ -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] diff --git a/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py b/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py index b761eb99f..ed2172b60 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py @@ -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() diff --git a/python/sglang/srt/hardware_backend/npu/quantization/w8a8_mxfp8.py b/python/sglang/srt/hardware_backend/npu/quantization/w8a8_mxfp8.py new file mode 100644 index 000000000..406638f0c --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/quantization/w8a8_mxfp8.py @@ -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() + ) diff --git a/python/sglang/srt/hardware_backend/npu/utils.py b/python/sglang/srt/hardware_backend/npu/utils.py index 3fd179443..32e4fe577 100644 --- a/python/sglang/srt/hardware_backend/npu/utils.py +++ b/python/sglang/srt/hardware_backend/npu/utils.py @@ -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 diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index 584511146..2e334c826 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index a4b0d9683..711397bba 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index e5a60f25f..bff62c44e 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -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 diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 27b3da146..13b664c41 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -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: diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 22626081c..937ab4152 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -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 diff --git a/python/sglang/srt/models/deepseek_v4_nextn.py b/python/sglang/srt/models/deepseek_v4_nextn.py index c9403fbfd..f54fd8d8a 100644 --- a/python/sglang/srt/models/deepseek_v4_nextn.py +++ b/python/sglang/srt/models/deepseek_v4_nextn.py @@ -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, diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index a2bdf5aab..a3a123276 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -32,6 +32,7 @@ from sglang.srt.disaggregation.mooncake.conn import ( ) from sglang.srt.disaggregation.utils import ( MetadataBuffers, + get_dsv4_c4_state_indices, get_dsv4_c128_state_indices, setup_state_kv_args, ) @@ -365,7 +366,12 @@ class TestEagleDsaSeedTransfer(unittest.TestCase): buffers.set_buf(self._make_req(seed)) buffers.set_buf(self._make_req(None, metadata_buffer_index=1)) - self.assertTrue(torch.equal(buffers.output_dsa_topk_indices[0], seed)) + self.assertTrue( + torch.equal( + buffers.output_dsa_topk_indices[0], + seed.to(buffers.output_dsa_topk_indices.device), + ) + ) self.assertEqual(buffers.output_dsa_topk_indices[1].tolist(), [-1, -1, -1]) ptrs, data_lens, item_lens = buffers.get_buf_infos() self.assertEqual(ptrs[-2], buffers.output_dsa_topk_indices.data_ptr()) @@ -519,6 +525,39 @@ class TestEagleDsaSeedTransfer(unittest.TestCase): self.assertEqual(future_map.topk_index_buf.shape, (4, 3)) +class TestDSV4C4StateIndices(unittest.TestCase): + def test_non_mtp_to_mtp_maps_the_same_logical_positions(self): + # seq_len=13 keeps logical positions [8, 13) for the overlap C4 state. + src = get_dsv4_c4_state_indices(2, 13, ring_size=8) + dst = get_dsv4_c4_state_indices(2, 13, ring_size=16) + + np.testing.assert_array_equal(src, np.array([16, 17, 18, 19, 20])) + np.testing.assert_array_equal(dst, np.array([40, 41, 42, 43, 44])) + self.assertEqual(src.size, dst.size) + + def test_ring_wrap_preserves_position_order(self): + np.testing.assert_array_equal( + get_dsv4_c4_state_indices(0, 10, ring_size=8), + np.array([4, 5, 6, 7, 0, 1], dtype=np.int32), + ) + + def test_short_and_empty_sequences(self): + np.testing.assert_array_equal( + get_dsv4_c4_state_indices(3, 3, ring_size=8), + np.array([24, 25, 26], dtype=np.int32), + ) + np.testing.assert_array_equal( + get_dsv4_c4_state_indices(3, 0, ring_size=8), + np.empty((0,), dtype=np.int32), + ) + + def test_invalid_ring_size_is_rejected(self): + with self.assertRaises(ValueError): + get_dsv4_c4_state_indices(0, 8, ring_size=4) + with self.assertRaises(ValueError): + get_dsv4_c4_state_indices(0, 8, ring_size=10) + + class TestDSV4C128StateIndices(unittest.TestCase): def test_online_aligned_boundary_has_no_partial_state(self): np.testing.assert_array_equal( diff --git a/test/registered/unit/managers/test_pp_cp_rank_offsets.py b/test/registered/unit/managers/test_pp_cp_rank_offsets.py index c0711b4db..4c6749c19 100644 --- a/test/registered/unit/managers/test_pp_cp_rank_offsets.py +++ b/test/registered/unit/managers/test_pp_cp_rank_offsets.py @@ -37,7 +37,10 @@ def _fake_group() -> SimpleNamespace: def _make_receiver(ps: ParallelState) -> SchedulerRequestReceiver: - group = _fake_group() + tp_group = _fake_group() + attn_tp_group = _fake_group() + attn_cp_group = _fake_group() + world_group = _fake_group() return SchedulerRequestReceiver( recv_from_tokenizer=None, recv_from_rpc=None, @@ -45,13 +48,13 @@ def _make_receiver(ps: ParallelState) -> SchedulerRequestReceiver: input_blocker=None, mm_receiver=None, ps=ps, - tp_group=group, - tp_cpu_group=group, - attn_tp_group=group, - attn_tp_cpu_group=group, - attn_cp_group=group, - attn_cp_cpu_group=group, - world_group=group, + tp_group=tp_group, + tp_cpu_group=tp_group, + attn_tp_group=attn_tp_group, + attn_tp_cpu_group=attn_tp_group, + attn_cp_group=attn_cp_group, + attn_cp_cpu_group=attn_cp_group, + world_group=world_group, server_args=SimpleNamespace( enable_dp_attention=True, enable_dp_attention_local_control_broadcast=False, @@ -63,6 +66,94 @@ def _make_receiver(ps: ParallelState) -> SchedulerRequestReceiver: ) +class TestRequestReceiverBroadcast(unittest.TestCase): + def test_local_control_skips_full_tp_broadcast_for_decode_dp(self): + # Decode uses pure DP attention (attn_tp=attn_cp=1). The DP controller + # sends control requests to every local leader, so no per-tick Gloo + # broadcast should remain in SchedulerRequestReceiver. + ps = SimpleNamespace( + attn_tp_rank=0, + attn_cp_rank=0, + attn_tp_size=1, + attn_cp_size=1, + tp_size=32, + ) + receiver = _make_receiver(ps) + control_req = SimpleNamespace(kind="control") + parallel = SimpleNamespace( + enable_dp_attention=True, + enable_dp_attention_local_control_broadcast=True, + ) + + with ( + patch( + "sglang.srt.managers.scheduler_components.request_receiver." + "get_parallel", + return_value=parallel, + ), + patch( + "sglang.srt.managers.scheduler_components.request_receiver." + "attn_cp_tp_broadcast_pyobj", + side_effect=lambda requests: requests, + ), + patch( + "sglang.srt.managers.scheduler_components.request_receiver." + "broadcast_pyobj" + ) as broadcast, + ): + result = receiver._broadcast_reqs_across_ranks([control_req]) + + self.assertEqual(result, [control_req]) + broadcast.assert_not_called() + + def test_default_control_uses_full_tp_broadcast(self): + ps = SimpleNamespace( + attn_tp_rank=0, + attn_cp_rank=0, + attn_tp_size=1, + attn_cp_size=1, + tp_size=32, + ) + receiver = _make_receiver(ps) + control_req = SimpleNamespace(kind="control") + parallel = SimpleNamespace( + enable_dp_attention=True, + enable_dp_attention_local_control_broadcast=False, + ) + + with ( + patch( + "sglang.srt.managers.scheduler_components.request_receiver." + "get_parallel", + return_value=parallel, + ), + patch( + "sglang.srt.managers.scheduler_components.request_receiver." + "is_ep_scale_joiner", + return_value=False, + ), + patch( + "sglang.srt.managers.scheduler_components.request_receiver." + "attn_cp_tp_broadcast_pyobj", + side_effect=lambda requests: requests, + ), + patch( + "sglang.srt.managers.scheduler_components.request_receiver." + "broadcast_pyobj", + side_effect=lambda requests, *_args, **_kwargs: requests, + ) as broadcast, + ): + result = receiver._broadcast_reqs_across_ranks([control_req]) + + self.assertEqual(result, [control_req]) + broadcast.assert_called_once_with( + [control_req], + receiver.tp_group.rank, + receiver.tp_cpu_group, + src=receiver.tp_group.ranks[0], + ) + + class TestPPCPRankOffsets(unittest.TestCase): def test_request_receiver_uses_cp_size_for_pp_recv_rank(self): ps = _make_ps() diff --git a/test/registered/unit/npu/attention/test_npu_ascend_dsv4_backend.py b/test/registered/unit/npu/attention/test_npu_ascend_dsv4_backend.py index fc834b971..99a7cf35a 100644 --- a/test/registered/unit/npu/attention/test_npu_ascend_dsv4_backend.py +++ b/test/registered/unit/npu/attention/test_npu_ascend_dsv4_backend.py @@ -56,11 +56,328 @@ sys.modules.setdefault("sglang.srt.speculative", ModuleType("sglang.srt.speculat sys.modules.setdefault("sglang.srt.speculative.eagle_utils", _eagle_stub) from sglang.srt.hardware_backend.npu.attention.ascend_dsv4_backend import ( + C4IndexerAscendBackendMixin, + CompressorAscendBackendMixin, + DeepseekV4AscendAttnBackend, DeepseekV4AscendMultiStepDraftBackend, _apply_hadamard, + _build_cycle_state_block_table, _get_kv_indices, + _sparse_attn_kv_quant_kwargs, + _sparse_attn_ops, _walsh_hadamard_matrix, ) +from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import ( + dsv4_state_payloads, +) +from sglang.srt.hardware_backend.npu.dsv4.dsv4_memory_pool import DSV4NPUTokenToKVPool +from sglang.srt.hardware_backend.npu.dsv4.dsv4_req_to_token_pool import ( + DSV4ReqToTokenTablesMixin, +) + + +class TestVerifyCompressPositions(unittest.TestCase): + @staticmethod + def _backend(): + backend = DeepseekV4AscendAttnBackend.__new__(DeepseekV4AscendAttnBackend) + backend._dsv4_compress_ratios = (4, 128) + return backend + + def _assert_device_path_matches_cpu_reference( + self, + *, + positions, + live_seq_lens, + n_draft, + ratio, + dst_size, + ): + backend = self._backend() + positions = torch.tensor(positions, dtype=torch.int64) + live_seq_lens = torch.tensor(live_seq_lens, dtype=torch.int32) + final_seq_lens = torch.where( + live_seq_lens > 0, + live_seq_lens + int(n_draft), + live_seq_lens, + ) + expected = torch.full((dst_size,), -1, dtype=torch.int64) + actual = torch.full((dst_size,), -2, dtype=torch.int64) + + backend._fill_verify_positions_cmp_padding_one( + positions, + expected, + ratio=ratio, + seq_lens_cpu=final_seq_lens, + n_draft=n_draft, + ) + backend._fill_verify_positions_cmp_padding_one_device( + positions, + actual, + ratio=ratio, + live_seq_lens=live_seq_lens, + n_draft=n_draft, + ) + + self.assertEqual(actual.tolist(), expected.tolist()) + + def test_uses_group_start_rope_position(self): + backend = DeepseekV4AscendAttnBackend.__new__(DeepseekV4AscendAttnBackend) + backend._dsv4_compress_ratios = (4, 128) + + # Two linear three-token verify trees. Their completed C4 groups end + # at token positions 7 and 11, whose compressed RoPE positions are the + # corresponding group starts 4 and 8. + positions = torch.tensor([7, 8, 9, 10, 11, 12], dtype=torch.int64) + final_seq_lens = torch.tensor([10, 13], dtype=torch.int32) + dst = torch.full((4,), -1, dtype=torch.int64) + + backend._fill_verify_positions_cmp_padding_one( + positions, + dst, + ratio=4, + seq_lens_cpu=final_seq_lens, + n_draft=3, + ) + + self.assertEqual(dst.tolist(), [4, 8, 0, 0]) + + def test_c128_uses_group_start_rope_position(self): + backend = DeepseekV4AscendAttnBackend.__new__(DeepseekV4AscendAttnBackend) + backend._dsv4_compress_ratios = (4, 128) + dst = torch.full((2,), -1, dtype=torch.int64) + + backend._fill_verify_positions_cmp_padding_one( + torch.tensor([126, 127, 128], dtype=torch.int64), + dst, + ratio=128, + seq_lens_cpu=torch.tensor([129], dtype=torch.int32), + n_draft=3, + ) + + self.assertEqual(dst.tolist(), [0, 0]) + + def test_no_completed_group_clears_destination(self): + backend = DeepseekV4AscendAttnBackend.__new__(DeepseekV4AscendAttnBackend) + backend._dsv4_compress_ratios = (4, 128) + dst = torch.full((2,), -1, dtype=torch.int64) + + backend._fill_verify_positions_cmp_padding_one( + torch.tensor([5, 6], dtype=torch.int64), + dst, + ratio=4, + seq_lens_cpu=torch.tensor([7], dtype=torch.int32), + n_draft=2, + ) + + self.assertEqual(dst.tolist(), [0, 0]) + + def test_zero_length_graph_padding_does_not_emit_a_position(self): + backend = DeepseekV4AscendAttnBackend.__new__(DeepseekV4AscendAttnBackend) + backend._dsv4_compress_ratios = (4, 128) + dst = torch.full((2,), -1, dtype=torch.int64) + + backend._fill_verify_positions_cmp_padding_one( + torch.tensor([0, 0, 0, 10, 11, 12], dtype=torch.int64), + dst, + ratio=4, + seq_lens_cpu=torch.tensor([0, 13], dtype=torch.int32), + n_draft=3, + ) + + self.assertEqual(dst.tolist(), [8, 0]) + + def test_device_path_matches_reference_across_boundaries_and_padding(self): + cases = ( + # Existing C4 and C128 boundary examples. + dict( + positions=[7, 8, 9, 10, 11, 12], + live_seq_lens=[7, 10], + n_draft=3, + ratio=4, + dst_size=4, + ), + dict( + positions=[126, 127, 128], + live_seq_lens=[126], + n_draft=3, + ratio=128, + dst_size=2, + ), + # A zero-length graph-padding row may appear before a live row. + dict( + positions=[0, 0, 0, 10, 11, 12], + live_seq_lens=[0, 10], + n_draft=3, + ratio=4, + dst_size=4, + ), + # More than one boundary per request and destination truncation. + dict( + positions=list(range(2, 11)) + list(range(7, 16)), + live_seq_lens=[2, 7], + n_draft=9, + ratio=4, + dst_size=3, + ), + # Preserve values from a non-linear tree position array rather than + # reconstructing them arithmetically from sequence lengths. + dict( + positions=[50, 90, 51, 52, 70, 71, 120, 72], + live_seq_lens=[3, 126], + n_draft=4, + ratio=4, + dst_size=4, + ), + # Long-context C128 boundaries around 128K. + dict( + positions=[131070, 131071, 131072, 131073], + live_seq_lens=[131070], + n_draft=4, + ratio=128, + dst_size=2, + ), + ) + for case in cases: + with self.subTest(case=case): + self._assert_device_path_matches_cpu_reference(**case) + + def test_stable_compact_preserves_boolean_index_order_and_zero_tail(self): + dst = torch.full((4,), -1, dtype=torch.int64) + values = torch.tensor([11, 22, 33, 44, 55, 66], dtype=torch.int64) + keep = torch.tensor([False, True, False, True, True, False]) + + DeepseekV4AscendAttnBackend._stable_compact_1d(dst, values, keep) + + self.assertEqual(dst.tolist(), [22, 44, 55, 0]) + + def test_stable_compact_matches_boolean_index_truncation(self): + dst = torch.full((2,), -1, dtype=torch.int64) + values = torch.tensor([5, 6, 7, 8, 9], dtype=torch.int64) + keep = torch.tensor([True, False, True, True, True]) + + DeepseekV4AscendAttnBackend._stable_compact_1d(dst, values, keep) + + self.assertEqual(dst.tolist(), values[keep][:2].tolist()) + + def test_stable_compact_matches_all_small_boolean_masks(self): + values = torch.arange(1, 7, dtype=torch.int64) + for mask_bits in range(1 << values.numel()): + keep = torch.tensor( + [(mask_bits >> index) & 1 for index in range(values.numel())], + dtype=torch.bool, + ) + for dst_size in range(1, values.numel() + 1): + dst = torch.full((dst_size,), -1, dtype=torch.int64) + DeepseekV4AscendAttnBackend._stable_compact_1d(dst, values, keep) + selected = values[keep][:dst_size] + expected = torch.zeros_like(dst) + expected[: selected.numel()].copy_(selected) + self.assertEqual(dst.tolist(), expected.tolist()) + + +class TestMultiStepDraftCompressedLocs(unittest.TestCase): + def test_draft_steps_skip_compressed_locs_but_preserve_full_and_swa(self): + backend = DeepseekV4AscendMultiStepDraftBackend.__new__( + DeepseekV4AscendMultiStepDraftBackend + ) + backend.topk = 1 + backend.speculative_num_steps = 3 + bundle = SimpleNamespace( + out_full_loc=torch.arange(6, dtype=torch.int64), + out_swa_loc=torch.arange(10, 16, dtype=torch.int64), + out_c4_loc=torch.tensor([101, 102], dtype=torch.int64), + out_c128_loc=torch.tensor([201], dtype=torch.int64), + ) + forward_batch = SimpleNamespace( + batch_size=2, + out_cache_loc=bundle.out_full_loc, + out_cache_loc_dsv4=bundle, + seq_lens=torch.tensor([7, 11], dtype=torch.int32), + ) + + with patch("torch.cumsum", side_effect=AssertionError("unexpected compaction")): + step = backend._step_out_cache_loc_dsv4(forward_batch, step_id=1) + + self.assertEqual(step.out_full_loc.tolist(), [1, 4]) + self.assertEqual(step.out_swa_loc.tolist(), [11, 14]) + self.assertEqual(step.out_c4_loc.numel(), 0) + self.assertEqual(step.out_c128_loc.numel(), 0) + self.assertEqual(step.out_c4_loc.dtype, bundle.out_c4_loc.dtype) + self.assertEqual(step.out_c128_loc.dtype, bundle.out_c128_loc.dtype) + + +class TestC4StateTransferLayout(unittest.TestCase): + @patch( + "sglang.srt.hardware_backend.npu.utils.is_npu_arch35", + return_value=True, + ) + def test_payload_uses_each_peers_private_ring_size(self, _): + def state_rows(ring_size): + req_pool = SimpleNamespace( + c128_page_size=1, + req_to_c128_sidecar=torch.zeros((4, 1), dtype=torch.int32), + get_dsv4_c4_state_ring_size=lambda: ring_size, + ) + payloads = dsv4_state_payloads(req_pool, 2, 13, page_size=1) + return next( + payload() + for state_type, payload in payloads.items() + if state_type.value == "dsv4_c4_state" + ) + + self.assertEqual(state_rows(8).tolist(), [16, 17, 18, 19, 20]) + self.assertEqual(state_rows(16).tolist(), [40, 41, 42, 43, 44]) + + def test_req_pool_reads_ring_size_from_registered_kv_pool(self): + req_pool = DSV4ReqToTokenTablesMixin.__new__(DSV4ReqToTokenTablesMixin) + req_pool._dsv4_allocator = MagicMock() + req_pool._dsv4_allocator.get_kvcache().get_ring_size.return_value = 16 + + self.assertEqual(req_pool.get_dsv4_c4_state_ring_size(), 16) + req_pool._dsv4_allocator.get_kvcache().get_ring_size.assert_called_once_with(4) + + def test_registers_single_rows_instead_of_request_banks(self): + attn_state = torch.empty((32, 5), dtype=torch.float32) + indexer_state = torch.empty((32, 7), dtype=torch.float32) + pool = DSV4NPUTokenToKVPool.__new__(DSV4NPUTokenToKVPool) + pool.compress_state_pools = [ + SimpleNamespace( + ratio=4, + ring_size=8, + kv_score_buffer=SimpleNamespace(kv_score=attn_state), + ) + ] + pool.indexer_compress_state_pools = [ + SimpleNamespace( + ratio=4, + ring_size=8, + kv_score_buffer=SimpleNamespace(kv_score=indexer_state), + ) + ] + + _, data_lens, item_lens = pool.get_c4_state_buf_infos() + + self.assertEqual(data_lens, [attn_state.nbytes, indexer_state.nbytes]) + self.assertEqual( + item_lens, + [attn_state[0].nbytes, indexer_state[0].nbytes], + ) + + +class TestC4IndexerInitialization(unittest.TestCase): + @patch( + "sglang.srt.hardware_backend.npu.attention.ascend_dsv4_backend.is_npu_arch35", + return_value=True, + ) + def test_arch35_indexer_uses_float8_kv(self, _): + indexer = torch.nn.Module() + indexer.head_dim = 8 + indexer.compressor = SimpleNamespace() + backend = C4IndexerAscendBackendMixin.__new__(C4IndexerAscendBackendMixin) + + backend._ensure_npu_c4_indexer(indexer, torch.device("cpu")) + + self.assertEqual(indexer.compressor.li_kv_dtype, "float8") class TestWalshHadamardMatrix(unittest.TestCase): @@ -181,6 +498,318 @@ class TestApplyHadamard(unittest.TestCase): self.assertTrue(torch.equal(out, expected)) +class TestCompressorStateTableABI(unittest.TestCase): + def test_arch35_cycle_table_is_one_bank_per_request(self): + req_pool_indices = torch.tensor([7, 3], dtype=torch.int64) + table = _build_cycle_state_block_table(req_pool_indices) + self.assertEqual(tuple(table.shape), (2,)) + self.assertEqual(table.dtype, torch.int32) + self.assertEqual(table.tolist(), [7, 3]) + + def test_arch35_cycle_table_rejects_explicit_shape(self): + with self.assertRaises(ValueError): + _build_cycle_state_block_table(torch.zeros((2, 8), dtype=torch.int32)) + + @patch( + "sglang.srt.hardware_backend.npu.attention.ascend_dsv4_backend.is_npu_arch35", + return_value=True, + ) + def test_arch35_eager_metadata_builds_cycle_table(self, _): + backend = CompressorAscendBackendMixin.__new__(CompressorAscendBackendMixin) + backend.forward_metadata = SimpleNamespace() + backend.token_to_kv_pool = MagicMock() + backend.req_to_token = torch.empty((0, 0), dtype=torch.int32) + backend.req_to_token_pool = MagicMock() + backend._dsv4_compress_ratios = () + backend._compute_compress_locs = MagicMock(return_value={}) + + forward_mode = MagicMock() + forward_mode.is_decode.return_value = True + forward_mode.is_target_verify.return_value = False + forward_batch = SimpleNamespace( + forward_mode=forward_mode, + req_pool_indices=torch.tensor([7, 3], dtype=torch.int64), + seq_lens=torch.tensor([5, 9], dtype=torch.int32), + out_cache_loc=torch.empty(0, dtype=torch.int64), + out_cache_loc_dsv4=None, + batch_size=2, + ) + + backend._build_npu_compress_metadata(forward_batch) + + table = getattr(backend.forward_metadata, "dsv4_cycle_state_block_table", None) + self.assertIsNotNone(table) + self.assertEqual(table.tolist(), [7, 3]) + self.assertEqual(table.dtype, torch.int32) + + @patch( + "sglang.srt.hardware_backend.npu.attention.ascend_dsv4_backend.is_npu_arch35", + return_value=True, + ) + def test_arch35_graph_replay_slices_static_req_pool_buffer_to_graph_bs(self, _): + backend = DeepseekV4AscendAttnBackend.__new__(DeepseekV4AscendAttnBackend) + table = torch.zeros(1, dtype=torch.int32) + graph_mode = MagicMock() + graph_mode.is_decode.return_value = False + graph_mode.is_target_verify.return_value = False + ctx = SimpleNamespace( + fm=SimpleNamespace(dsv4_cycle_state_block_table=table), + forward_batch=SimpleNamespace( + req_pool_indices=torch.arange(7, 19, dtype=torch.int64) + ), + graph_mode=graph_mode, + bs=1, + ) + backend._build_dsv4_graph_replay_ctx = MagicMock(return_value=ctx) + for name in ( + "_refresh_graph_seq_metadata", + "_refresh_graph_compress_page_tables_direct", + "_refresh_graph_explicit_state_block_tables", + "_refresh_graph_swa_metadata_direct", + "_refresh_graph_dspark_sparse_metadata", + "_refresh_graph_kernel_metadata", + ): + setattr(backend, name, MagicMock()) + + backend._apply_dsv4_graph_metadata(SimpleNamespace()) + + self.assertIs(ctx.fm.dsv4_cycle_state_block_table, table) + self.assertEqual(table.tolist(), [7]) + + @patch( + "sglang.srt.hardware_backend.npu.attention.ascend_dsv4_backend.is_npu_arch35", + return_value=True, + ) + def test_arch35_graph_capture_allocates_cycle_table_buffer(self, _): + backend = DeepseekV4AscendAttnBackend.__new__(DeepseekV4AscendAttnBackend) + metadata = SimpleNamespace() + backend.device = "cpu" + backend.graph_metadata = { + 2: metadata, + "swa_page_table": torch.full((2, 4), -1, dtype=torch.int32), + "c4_page_table": torch.full((2, 4), -1, dtype=torch.int32), + "c128_page_table": torch.full((2, 4), -1, dtype=torch.int32), + "kernel_metadata_c1a": torch.zeros(1024, dtype=torch.int32), + "kernel_metadata_c4a": torch.zeros(1024, dtype=torch.int32), + "kernel_metadata_c128a": torch.zeros(1024, dtype=torch.int32), + "kernel_metadata_li_quant": torch.zeros(1024, dtype=torch.int32), + "c4_topk_indices": torch.full((2, 1), -1, dtype=torch.int32), + } + backend._dsv4_graph_tokens_per_req = 1 + backend._dsv4_index_topk = 1 + backend._dsv4_state_pools_by_ratio = {} + backend._dsv4_sliding_window_size = 128 + backend._is_dspark_draft_worker = False + forward_mode = MagicMock() + forward_mode.is_target_verify.return_value = False + forward_mode.is_draft_extend_v2.return_value = False + + backend._init_dsv4_graph_metadata(2, forward_mode) + + table = getattr(metadata, "dsv4_cycle_state_block_table", None) + self.assertIsNotNone(table) + self.assertEqual(tuple(table.shape), (2,)) + self.assertEqual(table.dtype, torch.int32) + + @patch( + "sglang.srt.hardware_backend.npu.attention.ascend_dsv4_backend.is_npu_arch35", + return_value=True, + ) + def test_arch35_forward_reuses_batch_cycle_table(self, _): + table = torch.tensor([7, 3], dtype=torch.int32) + backend = CompressorAscendBackendMixin.__new__(CompressorAscendBackendMixin) + backend.graph_mode = False + backend.forward_metadata = SimpleNamespace( + dsv4_cycle_state_block_table=table, + positions_cmp_padding_c128=torch.empty(0, dtype=torch.int64), + actual_seq_lengths_q_pa=torch.tensor([0, 1, 2], dtype=torch.int32), + seqused=torch.ones(2, dtype=torch.int32), + start_pos=torch.zeros(2, dtype=torch.int32), + c128_loc=None, + ) + backend.token_to_kv_pool = MagicMock() + backend.token_to_kv_pool._get_state_pool.return_value = SimpleNamespace( + state_cache_3d=torch.empty(0) + ) + backend._ensure_compressor_hadamard = MagicMock() + backend._ensure_fused_caches = MagicMock() + backend._compressor_epilog_npu = MagicMock() + + compressor = SimpleNamespace( + ratio=128, + overlap=False, + layer_id=0, + is_in_indexer=False, + freqs_cis=None, + rotary_emb=None, + _fused_wkv_w=torch.empty(0), + _fused_wgate_w=torch.empty(0), + ape=torch.empty(0), + _fused_norm_weight_fp32=torch.empty(0), + rope_head_dim=64, + norm=SimpleNamespace(variance_epsilon=1e-6), + rotate=False, + ) + forward_mode = MagicMock() + forward_mode.is_prefill.return_value = False + forward_mode.is_target_verify.return_value = False + forward_batch = SimpleNamespace( + req_pool_indices=torch.tensor([7, 3], dtype=torch.int64), + forward_mode=forward_mode, + ) + rope = MagicMock() + rope.get_cos_sin.return_value = (torch.empty(0), torch.empty(0)) + + with ( + patch( + "sglang.srt.hardware_backend.npu.attention.ascend_dsv4_backend." + "Dsv4NpuRoPE.for_freqs", + return_value=rope, + ), + patch.object(torch.ops, "custom", MagicMock(), create=True) as custom_ops, + patch.object(torch.ops, "npu", MagicMock(), create=True) as npu_ops, + ): + custom_ops.compressor.return_value = torch.empty((0, 1)) + backend.forward_compress(compressor, torch.empty((2, 1)), forward_batch) + backend.forward_compress(compressor, torch.empty((2, 1)), forward_batch) + + self.assertEqual(npu_ops.compressor.call_count, 0) + self.assertIs( + custom_ops.compressor.call_args_list[0].kwargs["state_block_table"], table + ) + self.assertIs( + custom_ops.compressor.call_args_list[1].kwargs["state_block_table"], table + ) + + +class TestArch35SparseAttentionDispatch(unittest.TestCase): + _ARCH35_PATCH_TARGET = ( + "sglang.srt.hardware_backend.npu.attention.ascend_dsv4_backend.is_npu_arch35" + ) + + @patch(_ARCH35_PATCH_TARGET, return_value=True) + def test_arch35_uses_kv_quant_ops_and_layout_kwargs(self, _): + with patch("torch.ops.custom", MagicMock(), create=True) as custom_ops: + metadata_op, attention_op = _sparse_attn_ops() + kwargs = _sparse_attn_kv_quant_kwargs() + + self.assertIs( + metadata_op, custom_ops.npu_kv_quant_sparse_attn_sharedkv_metadata + ) + self.assertIs(attention_op, custom_ops.npu_kv_quant_sparse_attn_sharedkv) + self.assertEqual( + kwargs, + {"kv_quant_mode": 1, "tile_size": 64, "rope_head_dim": 64}, + ) + + @patch(_ARCH35_PATCH_TARGET, return_value=False) + def test_pre_arch35_keeps_legacy_ops_without_quant_kwargs(self, _): + with ( + patch("torch.ops.custom", MagicMock(), create=True) as custom_ops, + patch("torch.ops.npu", MagicMock(), create=True) as npu_ops, + ): + metadata_op, attention_op = _sparse_attn_ops() + kwargs = _sparse_attn_kv_quant_kwargs() + + self.assertIs(metadata_op, custom_ops.npu_sparse_attn_sharedkv_metadata) + self.assertIs(attention_op, npu_ops.sparse_attn_sharedkv) + self.assertEqual(kwargs, {}) + + +class TestSparseAttentionMetadata(unittest.TestCase): + _ARCH35_PATCH_TARGET = ( + "sglang.srt.hardware_backend.npu.attention.ascend_dsv4_backend.is_npu_arch35" + ) + + def test_device_metadata_receives_sequence_lengths(self): + cu_seqlens_q = torch.tensor([0, 2, 3], dtype=torch.int32) + seqused_kv = torch.tensor([8, 12], dtype=torch.int32) + + for is_arch35, metadata_op_name in ( + (False, "npu_sparse_attn_sharedkv_metadata"), + (True, "npu_kv_quant_sparse_attn_sharedkv_metadata"), + ): + with ( + self.subTest(is_arch35=is_arch35), + patch(self._ARCH35_PATCH_TARGET, return_value=is_arch35), + patch("torch.ops.custom", MagicMock(), create=True) as custom_ops, + patch("torch.ops.npu", MagicMock(), create=True), + ): + backend = DeepseekV4AscendAttnBackend.__new__( + DeepseekV4AscendAttnBackend + ) + backend.forward_metadata = SimpleNamespace() + backend._is_dspark_draft_worker = False + backend._dsv4_sliding_window_size = 128 + backend._dsv4_q_head_num = 64 + backend._dsv4_kv_head_num = 1 + backend._dsv4_head_dim = 512 + backend._dsv4_has_c4 = True + backend._dsv4_has_c128 = True + backend._dsv4_index_topk = 512 + backend._dsv4_index_n_heads = 16 + backend._dsv4_index_head_dim = 128 + + backend._kernel_metadata_from_parts( + bs=2, + actual_seq_lengths_q_pa=cu_seqlens_q, + actual_seq_lengths_kv=seqused_kv, + block_tables=torch.zeros((2, 1), dtype=torch.int32), + max_seqlen_q=2, + is_nextn=False, + ) + + metadata_op = getattr(custom_ops, metadata_op_name) + self.assertEqual(metadata_op.call_count, 3) + for call in metadata_op.call_args_list: + self.assertIs(call.kwargs["cu_seqlens_q"], cu_seqlens_q) + self.assertIs(call.kwargs["seqused_kv"], seqused_kv) + + def test_dspark_host_metadata_receives_host_sequence_lengths(self): + cu_seqlens_q = torch.tensor([0, 2, 3], dtype=torch.int32) + seqused_kv = torch.tensor([8, 12], dtype=torch.int32) + cu_seqlens_q_cpu = cu_seqlens_q.clone() + seqused_kv_cpu = seqused_kv.clone() + + with ( + patch("torch.ops.npu", MagicMock(), create=True) as npu_ops, + patch( + "sglang.srt.hardware_backend.npu.attention.ascend_dsv4_backend._sparse_attn_ops", + return_value=(MagicMock(), MagicMock()), + ), + ): + backend = DeepseekV4AscendAttnBackend.__new__(DeepseekV4AscendAttnBackend) + backend.forward_metadata = SimpleNamespace( + actual_seq_lengths_q_pa_cpu=cu_seqlens_q_cpu, + seq_lens_cpu_int=seqused_kv_cpu, + ) + backend._is_dspark_draft_worker = True + backend._dsv4_sliding_window_size = 128 + backend._dsv4_q_head_num = 64 + backend._dsv4_kv_head_num = 1 + backend._dsv4_head_dim = 512 + backend._dsv4_has_c4 = False + backend._dsv4_has_c128 = False + + kernel_metadata = backend._kernel_metadata_from_parts( + bs=2, + actual_seq_lengths_q_pa=cu_seqlens_q, + actual_seq_lengths_kv=seqused_kv, + block_tables=torch.zeros((2, 1), dtype=torch.int32), + max_seqlen_q=2, + is_nextn=False, + ) + + metadata_op = npu_ops.sparse_attn_sharedkv_metadata_host + metadata_op.assert_called_once() + self.assertIs(kernel_metadata["c1a_metadata"], metadata_op.return_value) + self.assertIs(metadata_op.call_args.kwargs["cu_seqlens_q"], cu_seqlens_q_cpu) + actual_seqused_kv = metadata_op.call_args.kwargs["seqused_kv"] + torch.testing.assert_close(actual_seqused_kv, seqused_kv_cpu[:2]) + self.assertEqual(actual_seqused_kv.dtype, torch.int32) + self.assertEqual(actual_seqused_kv.device.type, "cpu") + + class TestGetKvIndices(unittest.TestCase): _PATCH_TARGET = ( "sglang.srt.hardware_backend.npu.attention.ascend_dsv4_backend.get_attn_backend" @@ -383,5 +1012,69 @@ class TestCommonTemplate(unittest.TestCase): self.assertEqual(call_fn.call_count, 1) +class TestCompressorEpilogEmptyWrite(unittest.TestCase): + @staticmethod + def _backend(*, loc, graph_mode=False): + backend = CompressorAscendBackendMixin.__new__(CompressorAscendBackendMixin) + backend.graph_mode = graph_mode + backend.token_to_kv_pool = MagicMock() + backend.forward_metadata = SimpleNamespace(c4_loc=loc, c128_loc=loc) + return backend + + @staticmethod + def _compressor(*, li_kv_dtype="bf16", is_in_indexer=False): + return SimpleNamespace( + ratio=128, + layer_id=0, + is_in_indexer=is_in_indexer, + li_kv_dtype=li_kv_dtype, + ) + + @staticmethod + def _verify_batch(): + forward_mode = MagicMock() + forward_mode.is_target_verify.return_value = True + return SimpleNamespace(forward_mode=forward_mode) + + def test_all_slots_masked_skips_compress_write(self): + backend = self._backend(loc=torch.zeros(3, dtype=torch.int32)) + backend._compressor_epilog_npu( + self._compressor(), torch.zeros(3, 512), self._verify_batch() + ) + backend.token_to_kv_pool.set_compress_buffer.assert_not_called() + + def test_partially_masked_slots_writes_surviving_rows(self): + backend = self._backend(loc=torch.tensor([0, 7, 0], dtype=torch.int32)) + kv = torch.arange(12, dtype=torch.float32).view(3, 4) + backend._compressor_epilog_npu(self._compressor(), kv, self._verify_batch()) + + backend.token_to_kv_pool.set_compress_buffer.assert_called_once() + _, written_loc, written_kv, _, _ = ( + backend.token_to_kv_pool.set_compress_buffer.call_args.args + ) + self.assertEqual(written_loc.tolist(), [7]) + self.assertEqual(written_kv.tolist(), [kv[1].tolist()]) + + def test_graph_mode_keeps_static_shape_write(self): + backend = self._backend(loc=torch.zeros(3, dtype=torch.int32), graph_mode=True) + backend._compressor_epilog_npu( + self._compressor(), torch.ones(3, 4), self._verify_batch() + ) + + backend.token_to_kv_pool.set_compress_buffer.assert_called_once() + written_kv = backend.token_to_kv_pool.set_compress_buffer.call_args.args[2] + self.assertEqual(written_kv.shape[0], 3) + self.assertEqual(written_kv.abs().sum().item(), 0.0) + + def test_all_slots_masked_skips_fused_indexer_write(self): + backend = self._backend(loc=torch.zeros(3, dtype=torch.int32)) + compressor = self._compressor(li_kv_dtype="float8", is_in_indexer=True) + with patch("torch.ops.custom", MagicMock(), create=True) as custom_ops: + backend._compressor_epilog_npu( + compressor, torch.zeros(3, 512), self._verify_batch() + ) + custom_ops.indexer_compress_epilog.assert_not_called() + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/npu/quantization/test_fp4_moe_methods.py b/test/registered/unit/npu/quantization/test_fp4_moe_methods.py new file mode 100644 index 000000000..a21ac940f --- /dev/null +++ b/test/registered/unit/npu/quantization/test_fp4_moe_methods.py @@ -0,0 +1,558 @@ +import inspect +import os +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import torch + +from sglang.test.ci.ci_register import register_npu_ci + +register_npu_ci(est_time=1, suite="stage-a-unit-test-npu") + +# Load the quantization package first so `base_config`, `moe_methods`, and +# `linear_method_npu` initialize in dependency order. Importing `fp4_moe_methods` +# (or `linear_method_npu`) directly from a cold process triggers a circular +# import: linear_method_npu -> base_config -> quantization/__init__ -> +# gguf/unquant/gptq_moe -> moe_methods -> linear_method_npu (partially +# initialized, `_get_float8_e8m0fnu_dtype` not yet defined). Initializing the +# package first mirrors how the engine loads quantization at model-config time. +import sglang.srt.layers.quantization # noqa: F401 +from sglang.srt.environ import envs +from sglang.srt.hardware_backend.npu.quantization import fp4_moe_methods +from sglang.srt.hardware_backend.npu.quantization.fp4_moe_methods import ( + NPUW4A4Fp4MoEMethod, + _apply_swiglu_limit_npu, + _configure_dsv4_deepep_dispatcher, + _pair_pack_mxfp_act_scale, + _reshape_mxfp4_scale_for_npu, + npu_apply_without_routing_weights_w4a4_mxfp, + w4a8_mxfp_gmm, +) +from sglang.srt.layers.moe.fused_moe_triton import FusedMoE +from sglang.srt.layers.moe.token_dispatcher import deepep +from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8MoEMethod + +_NOT_PASSED = object() + + +class TestFP4MethodGate(unittest.TestCase): + def test_pre_arch35_keeps_fp8_moe_method(self): + config = Fp8Config(is_fp4_experts=True) + layer = FusedMoE.__new__(FusedMoE) + + with ( + patch("sglang.srt.layers.quantization.fp8.is_npu", return_value=True), + patch( + "sglang.srt.layers.quantization.fp8.is_npu_arch35", + return_value=False, + ), + ): + method = config.get_quant_method(layer, "model.layers.0.experts") + + self.assertIsInstance(method, Fp8MoEMethod) + + +class TestApplySwiGLULimitNpu(unittest.TestCase): + def test_clamps_gate_and_up_asymmetrically(self): + # DeepSeek-V4 clamps gate (first half) to <= limit but only the upper + # bound, while up (second half) is clamped symmetrically to [-limit, limit]. + # A regression that swapped these would silently change expert activations. + gate_up = torch.tensor([[8.0, -9.0, 9.0, -9.0]]) + _apply_swiglu_limit_npu(gate_up, 7.0) + self.assertTrue(torch.equal(gate_up, torch.tensor([[7.0, -9.0, 7.0, -7.0]]))) + + def test_noop_when_limit_none(self): + gate_up = torch.tensor([[8.0, -9.0]]) + _apply_swiglu_limit_npu(gate_up, None) + self.assertTrue(torch.equal(gate_up, torch.tensor([[8.0, -9.0]]))) + + def test_noop_when_limit_nonpositive(self): + gate_up = torch.tensor([[8.0, -9.0]]) + _apply_swiglu_limit_npu(gate_up, 0.0) + self.assertTrue(torch.equal(gate_up, torch.tensor([[8.0, -9.0]]))) + + +class TestReshapeMxfp4ScaleForNpu(unittest.TestCase): + def test_packs_scale_to_gmm_layout(self): + # [E, N, K/32] -> [E, K/64, N, 2] is the packed-pair layout the GMM reads; + # getting the transpose axis wrong silently dequantizes with the wrong scale. + scale = torch.arange(8, dtype=torch.uint8).view(1, 2, 4) + out = _reshape_mxfp4_scale_for_npu(scale) + self.assertEqual(tuple(out.shape), (1, 2, 2, 2)) + self.assertTrue(torch.equal(out, scale.view(1, 2, 2, 2).transpose(1, 2))) + + def test_rejects_odd_k_dim(self): + with self.assertRaises(ValueError): + _reshape_mxfp4_scale_for_npu(torch.zeros(1, 2, 3, dtype=torch.uint8)) + + +class TestMxfp4ScaleWeightLoader(unittest.TestCase): + def test_reinterprets_e8m0_scale_as_raw_uint8(self): + loaded = [] + + def weight_loader(param, loaded_weight, *args, **kwargs): + loaded.append(loaded_weight.clone()) + + layer = torch.nn.Module() + method = NPUW4A4Fp4MoEMethod(fp8_method=MagicMock(), prefix="test") + method.create_weights( + layer, + num_experts=1, + hidden_size=64, + intermediate_size_per_partition=64, + params_dtype=torch.bfloat16, + weight_loader=weight_loader, + ) + checkpoint_scale = torch.tensor([0.5, 0.25, 0.125], dtype=torch.float8_e8m0fnu) + + layer.w13_weight_scale_inv.weight_loader( + layer.w13_weight_scale_inv, + checkpoint_scale, + "model.layers.0.mlp.experts.0.gate_proj.weight_scale_inv", + "w1", + 0, + ) + + self.assertEqual(loaded[0].dtype, torch.uint8) + self.assertTrue(torch.equal(loaded[0], checkpoint_scale.view(torch.uint8))) + + +class TestPairPackMxfpActScale(unittest.TestCase): + def test_packs_as_view(self): + # The GMM expects a pair-packed *view* of the per-token scale, not a copy; + # materializing a copy here would break the kernel's aliasing contract. + flat = torch.arange(8).view(2, 4) + packed = _pair_pack_mxfp_act_scale(flat) + self.assertEqual(tuple(packed.shape), (2, 2, 2)) + self.assertEqual(packed.data_ptr(), flat.data_ptr()) + + def test_rejects_odd_scale_dim(self): + with self.assertRaises(ValueError): + _pair_pack_mxfp_act_scale(torch.zeros(2, 3)) + + def test_unflattens_low_latency_deepep_scale_as_view(self): + # DeepEP returns one flat E8M0 scale per 32-element block. Passing + # that flat buffer to GMM would use the wrong scale layout and either + # fail or dequantize activations incorrectly. + flat = torch.arange(4, dtype=torch.uint8) + packed = _pair_pack_mxfp_act_scale(flat, input_shape=(2, 64)) + + self.assertEqual(tuple(packed.shape), (2, 1, 2)) + self.assertEqual(packed.data_ptr(), flat.data_ptr()) + self.assertTrue(torch.equal(packed, torch.tensor([[[0, 1]], [[2, 3]]]))) + + def test_rejects_low_latency_deepep_scale_with_wrong_length(self): + with self.assertRaises(ValueError): + _pair_pack_mxfp_act_scale( + torch.zeros(3, dtype=torch.uint8), input_shape=(2, 64) + ) + + +class TestDsv4DeepEPMxfp8DispatcherConfig(unittest.TestCase): + @staticmethod + def _deepep_backend(): + return SimpleNamespace(is_deepep=lambda: True) + + def test_a5_deepep_defaults_low_latency_dispatch_to_mxfp8(self): + dispatcher = MagicMock() + layer = SimpleNamespace(dispatcher=dispatcher) + + with ( + patch.object(fp4_moe_methods, "is_npu_arch35", return_value=True), + patch( + "sglang.srt.layers.moe.get_moe_a2a_backend", + return_value=self._deepep_backend(), + ), + patch.dict(os.environ, {}, clear=True), + ): + _configure_dsv4_deepep_dispatcher(layer) + + dispatcher.set_quant_config.assert_called_once_with( + { + "normal_dispatcher_output_dtype": "bf16", + "low_latency_dispatcher_output_dtype": "mxfp8", + } + ) + + def test_non_deepep_ignores_the_low_latency_quant_environment(self): + dispatcher = MagicMock() + layer = SimpleNamespace(dispatcher=dispatcher) + + with ( + patch.object(fp4_moe_methods, "is_npu_arch35", return_value=True), + patch( + "sglang.srt.layers.moe.get_moe_a2a_backend", + return_value=SimpleNamespace(is_deepep=lambda: False), + ), + envs.SGLANG_NPU_DSV4_DEEPEP_LL_DISPATCH_QUANT_MODE.override("invalid"), + ): + _configure_dsv4_deepep_dispatcher(layer) + + dispatcher.set_quant_config.assert_called_once_with( + {"dispatcher_output_dtype": "bf16"} + ) + + def test_a5_deepep_allows_bf16_low_latency_fallback(self): + dispatcher = MagicMock() + layer = SimpleNamespace(dispatcher=dispatcher) + + with ( + patch.object(fp4_moe_methods, "is_npu_arch35", return_value=True), + patch( + "sglang.srt.layers.moe.get_moe_a2a_backend", + return_value=self._deepep_backend(), + ), + envs.SGLANG_NPU_DSV4_DEEPEP_LL_DISPATCH_QUANT_MODE.override("bf16"), + ): + _configure_dsv4_deepep_dispatcher(layer) + + dispatcher.set_quant_config.assert_called_once_with( + { + "normal_dispatcher_output_dtype": "bf16", + "low_latency_dispatcher_output_dtype": "bf16", + } + ) + + def test_a5_deepep_rejects_an_invalid_low_latency_quant_mode(self): + layer = SimpleNamespace(dispatcher=MagicMock()) + + with ( + patch.object(fp4_moe_methods, "is_npu_arch35", return_value=True), + patch( + "sglang.srt.layers.moe.get_moe_a2a_backend", + return_value=self._deepep_backend(), + ), + envs.SGLANG_NPU_DSV4_DEEPEP_LL_DISPATCH_QUANT_MODE.override("invalid"), + self.assertRaisesRegex(ValueError, "SGLANG_NPU_DSV4"), + ): + _configure_dsv4_deepep_dispatcher(layer) + + def test_non_a5_ignores_the_low_latency_quant_environment(self): + dispatcher = MagicMock() + layer = SimpleNamespace(dispatcher=dispatcher) + + with ( + patch.object(fp4_moe_methods, "is_npu_arch35", return_value=False), + patch( + "sglang.srt.layers.moe.get_moe_a2a_backend", + return_value=self._deepep_backend(), + ), + envs.SGLANG_NPU_DSV4_DEEPEP_LL_DISPATCH_QUANT_MODE.override("invalid"), + ): + _configure_dsv4_deepep_dispatcher(layer) + + dispatcher.set_quant_config.assert_called_once_with( + {"dispatcher_output_dtype": "bf16"} + ) + + +class _LowLatencyBuffer: + def __init__(self): + self.kwargs = None + + def low_latency_dispatch( + self, + hidden_states, + topk_ids, + num_max_dispatch_tokens_per_rank, + num_experts, + *, + use_fp8, + quant_mode=_NOT_PASSED, + **kwargs, + ): + self.kwargs = {"use_fp8": use_fp8, "quant_mode": quant_mode, **kwargs} + return torch.empty(0), torch.empty(0), object(), object(), object() + + +class _LegacyLowLatencyBuffer: + def low_latency_dispatch( + self, + hidden_states, + topk_ids, + num_max_dispatch_tokens_per_rank, + num_experts, + *, + use_fp8, + **kwargs, + ): + return torch.empty(0), torch.empty(0), object(), object(), object() + + +class TestDeepEPLowLatencyMxfp8Dispatch(unittest.TestCase): + @staticmethod + def _dispatcher(quant_mode, buffer): + dispatcher = object.__new__(deepep._DeepEPDispatcherImplLowLatency) + dispatcher.quant_config = {} + dispatcher.use_fp8 = False + dispatcher.use_nvfp4 = False + dispatcher.low_latency_quant_mode = quant_mode + dispatcher._low_latency_quant_mode_runtime_checked = False + dispatcher.num_max_dispatch_tokens_per_rank = 2 + dispatcher.num_experts = 2 + dispatcher.return_recv_hook = False + dispatcher._get_buffer = lambda: buffer + return dispatcher + + def test_mxfp8_passes_the_kernel_quant_mode(self): + buffer = _LowLatencyBuffer() + dispatcher = self._dispatcher("mx_fp8_e4m3", buffer) + + with ( + patch.dict(os.environ, {}, clear=True), + patch.object(deepep, "_deepep_precompile_tp_barrier"), + ): + dispatcher._dispatch_core( + torch.zeros(1, 64), + torch.zeros(1, 1, dtype=torch.int64), + torch.ones(1, 1), + ) + + self.assertEqual(buffer.kwargs["quant_mode"], "mx_fp8_e4m3") + + def test_mxfp8_ops_strategy_uses_legacy_mxfp8_flags(self): + buffer = _LowLatencyBuffer() + dispatcher = self._dispatcher("mx_fp8_e4m3", buffer) + + with ( + patch.dict(os.environ, {"DEEP_USE_MODE": "ops"}, clear=True), + patch.object(deepep, "_deepep_precompile_tp_barrier"), + ): + dispatcher._dispatch_core( + torch.zeros(1, 64), + torch.zeros(1, 1, dtype=torch.int64), + torch.ones(1, 1), + ) + + self.assertTrue(buffer.kwargs["use_fp8"]) + self.assertTrue(buffer.kwargs["use_ue8m0"]) + self.assertEqual(buffer.kwargs["quant_mode"], "mx_fp8_e4m3") + + def test_mxfp8_rejects_an_unsupported_low_latency_strategy(self): + dispatcher = self._dispatcher("mx_fp8_e4m3", _LowLatencyBuffer()) + + with ( + patch.dict(os.environ, {"DEEP_USE_MODE": "alltoall"}, clear=True), + self.assertRaisesRegex(RuntimeError, "DEEP_USE_MODE"), + ): + dispatcher._dispatch_core( + torch.zeros(1, 64), + torch.zeros(1, 1, dtype=torch.int64), + torch.ones(1, 1), + ) + + def test_mxfp8_checks_runtime_interface_once_per_dispatcher(self): + buffer = _LowLatencyBuffer() + dispatcher = self._dispatcher("mx_fp8_e4m3", buffer) + + with ( + patch.dict(os.environ, {}, clear=True), + patch.object(deepep, "_deepep_precompile_tp_barrier"), + patch.object( + deepep.inspect, "signature", wraps=inspect.signature + ) as signature, + ): + dispatcher._dispatch_core( + torch.zeros(1, 64), + torch.zeros(1, 1, dtype=torch.int64), + torch.ones(1, 1), + ) + dispatcher._dispatch_core( + torch.zeros(1, 64), + torch.zeros(1, 1, dtype=torch.int64), + torch.ones(1, 1), + ) + + self.assertEqual(signature.call_count, 1) + + def test_bf16_does_not_pass_a_quant_mode(self): + buffer = _LowLatencyBuffer() + dispatcher = self._dispatcher(None, buffer) + + with patch.object(deepep, "_deepep_precompile_tp_barrier"): + dispatcher._dispatch_core( + torch.zeros(1, 64), + torch.zeros(1, 1, dtype=torch.int64), + torch.ones(1, 1), + ) + + self.assertIs(buffer.kwargs["quant_mode"], _NOT_PASSED) + + def test_mxfp8_rejects_legacy_runtime_without_quant_mode(self): + dispatcher = self._dispatcher("mx_fp8_e4m3", _LegacyLowLatencyBuffer()) + + with self.assertRaisesRegex(RuntimeError, "quant_mode"): + dispatcher._dispatch_core( + torch.zeros(1, 64), + torch.zeros(1, 1, dtype=torch.int64), + torch.ones(1, 1), + ) + + +class TestW4A8MxfpGmmInputScale(unittest.TestCase): + def setUp(self): + self.input = torch.randn(2, 64) + self.input_scale = torch.ones(2, 1, 2) + self.weight = torch.empty(2, 64, 32, dtype=torch.uint8) + self.weight_scale = torch.ones(2, 1, 32, 2, dtype=torch.uint8) + self.group_list = torch.tensor([1, 1], dtype=torch.int32) + + def _call_gmm(self, input_scale): + return w4a8_mxfp_gmm( + input=self.input, + input_scale=input_scale, + weight=self.weight, + weight_scale=self.weight_scale, + group_list_type=1, + group_list=self.group_list, + output_dtype=torch.bfloat16, + ) + + def test_supplied_scale_skips_dynamic_quant(self): + expected = torch.randn(2, 32) + with ( + patch.object( + torch.ops.npu, "npu_dynamic_mx_quant", create=True + ) as dynamic_quant, + patch.object( + torch.ops.npu, + "npu_grouped_matmul", + return_value=[expected], + create=True, + ) as grouped_matmul, + ): + output = self._call_gmm(self.input_scale) + + dynamic_quant.assert_not_called() + self.assertIs(output, expected) + call_kwargs = grouped_matmul.call_args.kwargs + self.assertIs(call_kwargs["per_token_scale"][0], self.input_scale) + self.assertEqual(call_kwargs["group_list"].dtype, torch.int64) + self.assertTrue(torch.equal(call_kwargs["group_list"], self.group_list)) + + def test_flat_deepep_scale_skips_dynamic_quant_after_layout_adaptation(self): + flat_scale = torch.arange(4, dtype=torch.uint8) + expected = torch.randn(2, 32) + with ( + patch.object( + torch.ops.npu, "npu_dynamic_mx_quant", create=True + ) as dynamic_quant, + patch.object( + torch.ops.npu, + "npu_grouped_matmul", + return_value=[expected], + create=True, + ) as grouped_matmul, + ): + output = self._call_gmm(flat_scale) + + dynamic_quant.assert_not_called() + self.assertIs(output, expected) + packed_scale = grouped_matmul.call_args.kwargs["per_token_scale"][0] + self.assertEqual(tuple(packed_scale.shape), (2, 1, 2)) + self.assertEqual(packed_scale.data_ptr(), flat_scale.data_ptr()) + + def test_missing_scale_uses_dynamic_quant(self): + quantized = torch.empty(2, 64, dtype=torch.float8_e4m3fn) + quantized_scale = torch.ones(2, 1, 2) + expected = torch.randn(2, 32) + with ( + patch.object( + torch.ops.npu, + "npu_dynamic_mx_quant", + return_value=(quantized, quantized_scale), + create=True, + ) as dynamic_quant, + patch.object( + torch.ops.npu, + "npu_grouped_matmul", + return_value=[expected], + create=True, + ) as grouped_matmul, + ): + output = self._call_gmm(None) + + dynamic_quant.assert_called_once() + self.assertIs(output, expected) + self.assertIs( + grouped_matmul.call_args.kwargs["per_token_scale"][0], quantized_scale + ) + + +class TestW4A8MxfpGmmChain(unittest.TestCase): + def test_passes_swiglu_limit_to_quant(self): + gate_up = torch.randn(1, 64) + activated = torch.randn(1, 32) + activated_scale = torch.randn(1, 1) + expected = torch.randn(1, 32) + layer = SimpleNamespace( + w13_weight=MagicMock(), + w13_weight_scale_inv=MagicMock(), + w2_weight=MagicMock(), + w2_weight_scale_inv=MagicMock(), + moe_runner_config=SimpleNamespace(swiglu_limit=7.0), + ) + + with ( + patch.object( + fp4_moe_methods, "w4a8_mxfp_gmm", side_effect=[gate_up, expected] + ) as gmm, + patch.object( + fp4_moe_methods, + "swiglu_quant", + return_value=(activated, activated_scale), + ) as swiglu, + ): + output = npu_apply_without_routing_weights_w4a4_mxfp( + layer, + torch.randn(1, 4), + torch.ones(1, 1, 2), + group_list_type=1, + group_list=torch.tensor([1], dtype=torch.int64), + output_dtype=torch.bfloat16, + ) + + self.assertIs(output, expected) + self.assertTrue(torch.equal(swiglu.call_args.args[0], gate_up)) + self.assertTrue(swiglu.call_args.kwargs["do_limit"]) + self.assertEqual(swiglu.call_args.kwargs["limit"], 7.0) + self.assertIs(gmm.call_args_list[1].kwargs["input"], activated) + self.assertIs(gmm.call_args_list[1].kwargs["input_scale"], activated_scale) + + +class TestProcessWeightsAfterLoadingZeroScale(unittest.TestCase): + @staticmethod + def _method(): + return NPUW4A4Fp4MoEMethod(fp8_method=MagicMock(), prefix="test") + + def test_raises_when_w13_scales_never_loaded(self): + # An all-zero scale is the signature of a checkpoint whose scale names + # never matched; without this guard every routed expert computes silently + # as zero instead of failing loudly. + layer = SimpleNamespace( + w13_weight_scale_inv=torch.nn.Parameter( + torch.zeros(2, 2, 4, dtype=torch.uint8), requires_grad=False + ), + w2_weight_scale_inv=torch.nn.Parameter( + torch.zeros(2, 2, 4, dtype=torch.uint8), requires_grad=False + ), + ) + with self.assertRaises(RuntimeError): + self._method().process_weights_after_loading(layer) + + def test_raises_when_w2_scales_never_loaded(self): + layer = SimpleNamespace( + w13_weight_scale_inv=torch.nn.Parameter( + torch.ones(2, 2, 4, dtype=torch.uint8), requires_grad=False + ), + w2_weight_scale_inv=torch.nn.Parameter( + torch.zeros(2, 2, 4, dtype=torch.uint8), requires_grad=False + ), + ) + with self.assertRaises(RuntimeError): + self._method().process_weights_after_loading(layer) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/npu/quantization/test_npu_mxfp8_linear.py b/test/registered/unit/npu/quantization/test_npu_mxfp8_linear.py new file mode 100644 index 000000000..456ecc0d0 --- /dev/null +++ b/test/registered/unit/npu/quantization/test_npu_mxfp8_linear.py @@ -0,0 +1,171 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import torch + +from sglang.test.ci.ci_register import register_npu_ci + +register_npu_ci(est_time=1, suite="stage-a-unit-test-npu") + +# Load the quantization package first so `base_config`, `moe_methods`, and +# `linear_method_npu` initialize in dependency order. Importing +# `linear_method_npu` directly from a cold process triggers a circular import: +# linear_method_npu -> base_config -> quantization/__init__ -> +# gguf/unquant/gptq_moe -> moe_methods -> linear_method_npu (partially +# initialized, `_get_float8_e8m0fnu_dtype` not yet defined). Initializing the +# package first mirrors how the engine loads quantization at model-config time. +import sglang.srt.layers.quantization # noqa: F401 +from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import ( + npu_w8a8_mxfp8_linear, +) +from sglang.srt.hardware_backend.npu.quantization.w8a8_mxfp8 import ( + process_npu_arch35_mxfp8_linear_weights, +) +from sglang.srt.layers.quantization.fp8 import Fp8Config + + +class TestNPUW8A8BlockFP8Linear(unittest.TestCase): + def test_fp8_config_preserves_ue8m0_scale_format(self): + quant_config = Fp8Config.from_config( + { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [128, 128], + "scale_fmt": "ue8m0", + } + ) + + self.assertEqual(quant_config.scale_fmt, "ue8m0") + + def test_layout_only_ue8m0_conversion_preserves_fp8_weight(self): + original_weight = torch.randint(1, 255, (128, 64), dtype=torch.uint8).view( + torch.float8_e4m3fn + ) + layer = SimpleNamespace( + weight=torch.nn.Parameter(original_weight.clone(), requires_grad=False), + weight_scale_inv=torch.nn.Parameter( + torch.tensor([[2**-12]], dtype=torch.float32), requires_grad=False + ), + ) + + process_npu_arch35_mxfp8_linear_weights(layer, [128, 128], scale_fmt="ue8m0") + + self.assertEqual(layer.weight.shape, (64, 128)) + torch.testing.assert_close( + layer.weight.data.T.contiguous().view(torch.uint8), + original_weight.view(torch.uint8), + ) + self.assertEqual(layer.weight_scale_inv.shape, (1, 128, 2)) + self.assertTrue( + torch.equal( + layer.weight_scale_inv.data, + torch.full((1, 128, 2), 0x73, dtype=torch.uint8), + ) + ) + self.assertTrue(layer.weight_scale_inv.format_ue8m0) + + def test_rejects_non_ue8m0_scale_format(self): + with self.assertRaisesRegex(ValueError, "scale_fmt='ue8m0'"): + process_npu_arch35_mxfp8_linear_weights( + SimpleNamespace(), [128, 128], scale_fmt="float32" + ) + + def test_rejects_non_fp8_weight(self): + with self.assertRaisesRegex(ValueError, "expects float8_e4m3fn weights"): + npu_w8a8_mxfp8_linear( + torch.empty(1, 128, dtype=torch.bfloat16), + torch.empty(128, 64, dtype=torch.bfloat16), + [128, 128], + torch.empty(1), + ) + + def test_quantizes_flattened_input_and_restores_batch_shape(self): + input_tensor = torch.randn(2, 3, 128, dtype=torch.bfloat16) + weight = torch.empty(128, 64, dtype=torch.float8_e4m3fn) + weight_scale = torch.empty(2, 64, 2, dtype=torch.uint8) + bias = torch.randn(64, dtype=torch.float32) + quantized = torch.empty(6, 128, dtype=torch.float8_e4m3fn) + input_scale = torch.empty(6, 2, 2, dtype=torch.uint8) + matmul_output = torch.randn(6, 64, dtype=torch.bfloat16) + + npu_ops = MagicMock() + npu_ops.npu_dynamic_mx_quant.return_value = (quantized, input_scale) + npu_ops.npu_quant_matmul.return_value = matmul_output + with patch.object(torch.ops, "npu", npu_ops, create=True): + output = npu_w8a8_mxfp8_linear( + input_tensor, + weight, + [64, 128], + weight_scale, + bias=bias, + ) + + self.assertEqual(output.shape, (2, 3, 64)) + quant_call = npu_ops.npu_dynamic_mx_quant.call_args + self.assertEqual(quant_call.args[0].shape, (6, 128)) + self.assertEqual(quant_call.kwargs["dst_type"], torch.float8_e4m3fn) + + matmul_call = npu_ops.npu_quant_matmul.call_args + self.assertIs(matmul_call.args[0], quantized) + self.assertIs(matmul_call.args[1], weight) + self.assertIs(matmul_call.kwargs["scale"], weight_scale) + self.assertIs(matmul_call.kwargs["pertoken_scale"], input_scale) + self.assertIs(matmul_call.kwargs["bias"], bias) + self.assertEqual(matmul_call.kwargs["group_sizes"], (1, 1, 32)) + + def test_rejects_noncontiguous_input(self): + input_tensor = torch.randn(2, 3, 128, dtype=torch.bfloat16).transpose(0, 1) + weight = torch.empty(128, 64, dtype=torch.float8_e4m3fn) + weight_scale = torch.empty(2, 64, 2, dtype=torch.uint8) + + with self.assertRaisesRegex(RuntimeError, "view size is not compatible"): + npu_w8a8_mxfp8_linear(input_tensor, weight, [64, 128], weight_scale) + + def test_preserves_supported_input_dtype(self): + input_tensor = torch.randn(2, 128, dtype=torch.float16) + weight = torch.empty(128, 64, dtype=torch.float8_e4m3fn) + weight_scale = torch.empty(2, 64, 2, dtype=torch.uint8) + npu_ops = MagicMock() + npu_ops.npu_dynamic_mx_quant.return_value = ( + torch.empty(2, 128, dtype=torch.float8_e4m3fn), + torch.empty(2, 2, 2, dtype=torch.uint8), + ) + npu_ops.npu_quant_matmul.return_value = torch.empty(2, 64) + + with patch.object(torch.ops, "npu", npu_ops, create=True): + npu_w8a8_mxfp8_linear(input_tensor, weight, [128, 128], weight_scale) + + self.assertEqual( + npu_ops.npu_quant_matmul.call_args.kwargs["output_dtype"], + torch.float16, + ) + + def test_converts_bias_to_float32(self): + input_tensor = torch.randn(2, 128, dtype=torch.bfloat16) + weight = torch.empty(128, 64, dtype=torch.float8_e4m3fn) + weight_scale = torch.empty(2, 64, 2, dtype=torch.uint8) + bias = torch.randn(64, dtype=torch.bfloat16) + npu_ops = MagicMock() + npu_ops.npu_dynamic_mx_quant.return_value = ( + torch.empty(2, 128, dtype=torch.float8_e4m3fn), + torch.empty(2, 2, 2, dtype=torch.uint8), + ) + npu_ops.npu_quant_matmul.return_value = torch.empty(2, 64) + + with patch.object(torch.ops, "npu", npu_ops, create=True): + npu_w8a8_mxfp8_linear( + input_tensor, + weight, + [128, 128], + weight_scale, + bias=bias, + ) + + quant_bias = npu_ops.npu_quant_matmul.call_args.kwargs["bias"] + self.assertEqual(quant_bias.dtype, torch.float32) + torch.testing.assert_close(quant_bias, bias.float()) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/npu/test_npu_arch35_capability.py b/test/registered/unit/npu/test_npu_arch35_capability.py new file mode 100644 index 000000000..2fc1be21a --- /dev/null +++ b/test/registered/unit/npu/test_npu_arch35_capability.py @@ -0,0 +1,50 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from sglang.test.ci.ci_register import register_npu_ci + +register_npu_ci(est_time=1, suite="stage-a-unit-test-npu") + +from sglang.srt.hardware_backend.npu.utils import is_npu_arch35 + + +class TestArch35Capability(unittest.TestCase): + def tearDown(self): + is_npu_arch35.cache_clear() + + def test_arch35_is_supported_from_acl_device_info(self): + with ( + patch("sglang.srt.hardware_backend.npu.utils.is_npu", return_value=True), + patch.dict( + "sys.modules", + { + "acl": SimpleNamespace( + rt=SimpleNamespace(get_device_info=lambda *_: (3510, 0)) + ) + }, + ), + ): + self.assertTrue(is_npu_arch35()) + + def test_non_arch35_npu_is_not_supported(self): + with ( + patch("sglang.srt.hardware_backend.npu.utils.is_npu", return_value=True), + patch.dict( + "sys.modules", + { + "acl": SimpleNamespace( + rt=SimpleNamespace(get_device_info=lambda *_: (2901, 0)) + ) + }, + ), + ): + self.assertFalse(is_npu_arch35()) + + @patch("sglang.srt.hardware_backend.npu.utils.is_npu", return_value=False) + def test_non_npu_is_not_supported(self, _): + self.assertFalse(is_npu_arch35()) + + +if __name__ == "__main__": + unittest.main()