From d2f291c934a4e532160888533e2e7e670620c9c1 Mon Sep 17 00:00:00 2001 From: BourneSun0527 Date: Sun, 20 Sep 2026 09:17:27 +0800 Subject: [PATCH] [NPU][DSV4]dsv4 enable cpp (#39820) Co-authored-by: Even Zhou --- .../npu/attention/ascend_dsv4_backend.py | 5 +++-- .../hardware_backend/npu/dsv4/dsv4_memory_pool.py | 13 +++++++------ .../npu/attention/test_npu_ascend_dsv4_backend.py | 2 +- 3 files changed, 11 insertions(+), 9 deletions(-) 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 b370372ec..ef8eeef51 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 @@ -51,7 +51,7 @@ def _sparse_attn_ops(): ) return ( torch.ops.custom.npu_sparse_attn_sharedkv_metadata, - torch.ops.npu.sparse_attn_sharedkv, + torch.ops.custom.npu_sparse_attn_sharedkv, ) @@ -2104,13 +2104,14 @@ class DeepseekV4AscendAttnBackend( softmax_scale=layer.scaling, cmp_ratio=1, ) + _, attn_op = _sparse_attn_ops() if self._is_dspark_draft_worker: attn_kwargs["cu_seqlens_ori_kv"] = fm.actual_seq_lengths_q_pa + attn_op = torch.ops.npu.sparse_attn_sharedkv ori_sparse_indices = getattr(fm, "ori_sparse_indices", None) if ori_sparse_indices is not None: attn_kwargs["ori_sparse_indices"] = ori_sparse_indices q_arg = attn_kwargs.pop("q") - _, attn_op = _sparse_attn_ops() out, _ = attn_op(q_arg, **attn_kwargs) return out 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 fb0609a57..c165b4400 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 @@ -483,7 +483,7 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): item = self.layer_mapping[layer_id] ratio = item.compress_ratio if ratio == 0: - return self.swa_kv_pool.kv_buffer[item.compress_layer_id] + return self.swa_kv_pool.kv_buffer[self._swa_local_layer_id(layer_id)] if ratio == 4: return self.c4_kv_pool.kv_buffer[item.compress_layer_id] if ratio == 128: @@ -507,9 +507,10 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): shape becomes (num_tokens, 1, dim). """ self.wait_layer_transfer(layer_id) - # Index by RAW layer_id, not compress_layer_id (a per-bucket counter that - # would collide across ratios). swa_kv_pool is sized layer_num=total_layers. - kv = self.swa_kv_pool.kv_buffer[layer_id] + # Index by PP-stage-local layer_id, not compress_layer_id (a per-bucket + # counter that would collide across ratios). swa_kv_pool is sized + # layer_num=stage_layer_num (only this PP stage's layers). + kv = self.swa_kv_pool.kv_buffer[self._swa_local_layer_id(layer_id)] if loc is not None: kv = kv.flatten(0, 1)[loc] return kv @@ -554,8 +555,8 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): (num_pages, page_size, 1, dim) so we flatten the first two dims and index_put. """ - # Index by raw layer_id (see get_swa_buffer) to avoid bucket collision. - buf = self.swa_kv_pool.kv_buffer[layer_id] + # Index by PP-stage-local layer_id (see get_swa_buffer). + buf = self.swa_kv_pool.kv_buffer[self._swa_local_layer_id(layer_id)] if is_npu_arch35(): self._write_a5_packed_kv(buf=buf, loc=loc, cache=cache) return 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 99a7cf35a..e18a08ae9 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 @@ -712,7 +712,7 @@ class TestArch35SparseAttentionDispatch(unittest.TestCase): 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.assertIs(attention_op, custom_ops.npu_sparse_attn_sharedkv) self.assertEqual(kwargs, {})