[NPU][DSV4]dsv4 enable cpp (#39820)

Co-authored-by: Even Zhou <even.y.zhou@outlook.com>

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
 

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
+        # Index by PP-stage-local layer_id, not compress_layer_id (a per-bucket
-        # would collide across ratios). swa_kv_pool is sized layer_num=total_layers.
+        # counter that would collide across ratios). swa_kv_pool is sized
-        kv = self.swa_kv_pool.kv_buffer[layer_id]
+        # 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.
+        # Index by PP-stage-local layer_id (see get_swa_buffer).
-        buf = self.swa_kv_pool.kv_buffer[layer_id]
+        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

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, {})