[BugFix] isolate KV cache from sleep cleanup with dedicated memory tag for hadamard matrix (vllm-project#10775)

### What this PR does / why we need it?
Use a dedicated memory tag for KV cache used by DSA Hadamard, instead of
the shared tag that was previously being cleaned up during sleep(). This
ensures the Hadamard tensor is properly tagged and excluded from
sleep-mode memory release.
Fixes vllm-project#10684

### Does this PR introduce _any_ user-facing change?
No. 

### How was this patch tested?


- vLLM version: v0.23.0
- vLLM main:
vllm-project/vllm@967c5c3

---------

Signed-off-by: yangjinyang <yangjinyang5@huawei.com>
Signed-off-by: yjyang62 <yangjinyang5@huawei.com>
Co-authored-by: yjyang62 <yjyang62@users.noreply.github.com>

Author: yjyang62

vllm_ascend/attention/context_parallel/dsa_cp.py

@@ -189,9 +189,20 @@ def __init__(
                     ) from e
                 log_dim = math.ceil(math.log2(indexer_head_dim))
                 dim_padded = 2**log_dim
-                AscendDSACPMetadataBuilder.hadamard = torch.tensor(
-                    hadamard(dim_padded, dtype=float), dtype=torch.float, device=self.device
-                ).to(torch.bfloat16)
+                if self.vllm_config.model_config.enable_sleep_mode:
+                    # Sleep mode allocates KV inside CaMemAllocator; tag Hadamard so
+                    # sleep/wake does not treat it as KV cache.
+                    from vllm_ascend.device_allocator.camem import CaMemAllocator
+
+                    allocator = CaMemAllocator.get_instance()
+                    with allocator.use_allocation_tag(CaMemAllocator.sleep_persistent_tag):
+                        AscendDSACPMetadataBuilder.hadamard = torch.tensor(
+                            hadamard(dim_padded, dtype=float), dtype=torch.float, device=self.device
+                        ).to(torch.bfloat16)
+                else:
+                    AscendDSACPMetadataBuilder.hadamard = torch.tensor(
+                        hadamard(dim_padded, dtype=float), dtype=torch.float, device=self.device
+                    ).to(torch.bfloat16)
         self.start_pos_prefill = torch.zeros(scheduler_config.max_num_seqs, dtype=torch.int32, device=self.device)
         self.req_sas_metadata = torch.zeros(1024, dtype=torch.int32, device=self.device)
         self.req_qli_metadata = torch.zeros(1024, dtype=torch.int32, device=self.device)

vllm_ascend/attention/dsa_v1.py

@@ -424,9 +424,20 @@ def __init__(
                     raise ImportError("Please install scipy") from e
                 log_dim = math.ceil(math.log2(indexer_head_dim))
                 dim_padded = 2**log_dim
-                AscendDSAMetadataBuilder.hadamard = torch.tensor(
-                    hadamard(dim_padded, dtype=float), dtype=torch.float, device=self.device
-                ).to(torch.bfloat16)
+                if self.vllm_config.model_config.enable_sleep_mode:
+                    # Sleep mode allocates KV inside CaMemAllocator; tag Hadamard so
+                    # sleep/wake does not treat it as KV cache.
+                    from vllm_ascend.device_allocator.camem import CaMemAllocator
+
+                    allocator = CaMemAllocator.get_instance()
+                    with allocator.use_allocation_tag(CaMemAllocator.sleep_persistent_tag):
+                        AscendDSAMetadataBuilder.hadamard = torch.tensor(
+                            hadamard(dim_padded, dtype=float), dtype=torch.float, device=self.device
+                        ).to(torch.bfloat16)
+                else:
+                    AscendDSAMetadataBuilder.hadamard = torch.tensor(
+                        hadamard(dim_padded, dtype=float), dtype=torch.float, device=self.device
+                    ).to(torch.bfloat16)
         self.start_pos_prefill = torch.zeros(scheduler_config.max_num_seqs, dtype=torch.int32, device=self.device)
         self.start_pos_decode = torch.zeros(scheduler_config.max_num_seqs, dtype=torch.int32, device=self.device)
         self.decode_sas_metadata = torch.zeros(1024, dtype=torch.int32, device=self.device)

vllm_ascend/device_allocator/camem.py

@@ -134,6 +134,8 @@ class CaMemAllocator:
 
     instance = None
     default_tag: str = "default"
+    # Allocations with this tag stay mapped across sleep/wake cycles.
+    sleep_persistent_tag: str = "sleep_persistent"
 
     @staticmethod
     def get_instance() -> "CaMemAllocator":
@@ -200,6 +202,9 @@ def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None:
             offload_tags,
         )
         for ptr, data in self.pointer_to_data.items():
+            if data.tag == CaMemAllocator.sleep_persistent_tag:
+                # This memory is not offloaded or released during sleep.
+                continue
             handle = data.handle
             if data.tag in offload_tags:
                 size_in_bytes = handle[1]
@@ -227,6 +232,9 @@ def wake_up(self, tags: list[str] | None = None) -> None:
             tags or "all",
         )
         for ptr, data in self.pointer_to_data.items():
+            if data.tag == CaMemAllocator.sleep_persistent_tag:
+                # It was never released in sleep(), so there is nothing to remap.
+                continue
             if tags is None or data.tag in tags:
                 handle = data.handle
                 create_and_map(handle)
@@ -240,6 +248,16 @@ def wake_up(self, tags: list[str] | None = None) -> None:
                         memcpy(ptr, dest_max, cpu_ptr, size_in_bytes, ACL_MEMCPY_HOST_TO_DEVICE)
                         data.cpu_backup_tensor = None
 
+    @contextmanager
+    def use_allocation_tag(self, tag: str):
+        """Temporarily override the tag assigned to new allocations."""
+        old_tag = self.current_tag
+        self.current_tag = tag
+        try:
+            yield
+        finally:
+            self.current_tag = old_tag
+
     @contextmanager
     def use_memory_pool(self, tag: str | None = None):
         """