inference/vllm-ascend_v0.11.0rc0.patch

diff --git a/vllm_ascend/attention/attention_v1.py b/vllm_ascend/attention/attention_v1.py
index d289bb4..0357b50 100644
--- a/vllm_ascend/attention/attention_v1.py
+++ b/vllm_ascend/attention/attention_v1.py
@@ -21,6 +21,7 @@ from typing import ClassVar, List, Optional, Tuple, Type
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 import torch_npu
 from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
                                               AttentionLayer, AttentionType)
@@ -30,6 +31,7 @@ from vllm.utils import cdiv, direct_register_custom_op
 from vllm.v1.attention.backends.utils import AttentionCGSupport
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.kv_cache_interface import AttentionSpec
+from vllm.model_executor.models.utils import extract_layer_index
 
 from vllm_ascend.attention.utils import (AscendCommonAttentionMetadata,
                                          maybe_save_kv_layer_to_connector,
@@ -39,6 +41,9 @@ from vllm_ascend.ops.attention import vanilla_chunked_prefill
 from vllm_ascend.utils import (ACL_FORMAT_FRACTAL_NZ, aligned_16, is_310p,
                                nd_to_nz_2d, nd_to_nz_spec)
 
+if torch.version.cann.startswith("8.3"):
+    import omni_custom_ops
+
 
 class AscendAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
@@ -115,6 +120,7 @@ class AscendAttentionBackend(AttentionBackend):
         return [64]
 
 
+
 class AscendAttentionState(Enum):
     PrefillNoCache = 0
     PrefillCacheHit = 1
@@ -135,8 +141,8 @@ class AscendMetadata:
     num_actual_tokens: int = 0
 
     # The sequence length per sequence. Sequence length means the computed
-    # tokens + new tokens (is None if it is a decoding).
-    # (batch_size,)
+    # tokens + new tokens (is None if it is a decoding).(batch_size,)
+
     seq_lens: torch.Tensor = None
 
     query_start_loc: torch.Tensor = None
@@ -145,20 +151,25 @@ class AscendMetadata:
     max_query_len: Optional[int] = None
 
     # ********************** KV Cache Related Properties ********************* #
-    # Block addresses per sequence (Seq id -> list of physical block).
-    # (batch_size, max_blocks_per_seq)
+    # Block addresses per sequence (Seq id -> list of physical block).(batch_size, max_blocks_per_seq)
+
     block_tables: torch.Tensor = None
 
     # The indices of the token slots that input tokens will be stored into.
     # E.g., if `slot_mapping` is [35, 2, 17] and the block size is 16, the
     # three tokens are stored in the 3rd slot in block 2, 2nd slot in block 0,
-    # and 1st slot in block 1, respectively.
-    # (num_tokens,)
+    # and 1st slot in block 1, respectively. (num_tokens,)
+
     slot_mapping: torch.Tensor = None
 
     # *************************** Other Properties *************************** #
     enable_dbo_across_dp: bool = False
 
+    # Patch for param sink
+    sink_block_tables: Optional[List[torch.Tensor]] = None
+    sink_attn_mask: Optional[torch.Tensor] = None
+    sink_seq_kvlens: torch.Tensor = None
+    swa_seq_qlens: torch.Tensor = None
 
 class AscendAttentionMetadataBuilder:
     # Does this backend/builder support ACL Graphs for attention (default: no).
@@ -182,6 +193,7 @@ class AscendAttentionMetadataBuilder:
         self.max_num_blocks_per_req = cdiv(
             self.model_config.max_model_len,
             AscendAttentionBackend.get_supported_block_size()[0])
+        self.param_sink_number = self.model_config.hf_config.param_sink_number
 
     def reorder_batch(self, input_batch,
                       scheduler_output: "SchedulerOutput") -> bool:
@@ -210,6 +222,33 @@ class AscendAttentionMetadataBuilder:
         query_start_loc = query_start_loc_cpu.to(self.device,
                                                  non_blocking=True)
 
+        num_input_tokens = common_attn_metadata.num_input_tokens
+
+
+        if num_input_tokens > num_reqs and attn_state == AscendAttentionState.DecodeOnly:
+            tokens_gap_num = num_input_tokens-num_reqs
+
+            sink_block_tables = F.pad(block_table, (1, 0, 0, tokens_gap_num), value=0)
+
+            sink_seq_kvlens = seq_lens + self.param_sink_number
+            sink_seq_kvlens = torch.cat([sink_seq_kvlens, torch.full((tokens_gap_num,), \
+                self.param_sink_number, dtype=torch.int32)], dim=0)
+
+            gap_query_lens = torch.cat([query_lens, torch.ones(tokens_gap_num, dtype=torch.int32)], dim=0)
+            swa_seq_qlens = torch.cumsum(gap_query_lens, dim=0).to(dtype=torch.int32)
+        else:
+            sink_block_tables = F.pad(block_table, (1, 0, 0, 0), value=0)
+            sink_seq_kvlens = seq_lens + self.param_sink_number
+            swa_seq_qlens = torch.cumsum(query_lens, dim=0).to(dtype=torch.int32)
+
+
+        if attn_mask is not None:
+            sink_attn_mask = F.pad(attn_mask, (self.param_sink_number, 0, 0, 0), value=0)
+        else:
+            sink_attn_mask = None
+
+
+
         if is_310p():
             if attn_state == AscendAttentionState.PrefillNoCache:
                 mask_nz = nd_to_nz_2d(attn_mask)
@@ -230,7 +269,12 @@ class AscendAttentionMetadataBuilder:
             slot_mapping=slot_mapping,
             attn_mask=attn_mask,
             attn_state=attn_state,
-            enable_dbo_across_dp=common_attn_metadata.enable_dbo_across_dp)
+            enable_dbo_across_dp=common_attn_metadata.enable_dbo_across_dp,
+            sink_block_tables=sink_block_tables,
+            sink_attn_mask=sink_attn_mask,
+            sink_seq_kvlens=sink_seq_kvlens,
+            swa_seq_qlens=swa_seq_qlens
+        )
         return attn_metadata
 
     def build_for_graph_capture(
@@ -265,6 +309,7 @@ class AscendAttentionBackendImpl(AttentionImpl):
         kv_cache_dtype: str,
         logits_soft_cap: Optional[float],
         attn_type: str,
+        layer_name: str,
         kv_sharing_target_layer_name: Optional[str],
         **kwargs,
     ) -> None:
@@ -287,6 +332,13 @@ class AscendAttentionBackendImpl(AttentionImpl):
         self.key_cache = None
         self.value_cache = None
 
+        self.layer_idx = extract_layer_index(layer_name)
+
+        # Patch for Sink
+        self.sink_cached = False
+        self.attn_mask = torch.ones((2048, 2048), dtype=torch.int8, device="npu").triu_(diagonal=1)
+        self.attn_mask = self.attn_mask.to(torch.bool)
+
     def _forward_prefill_no_cache(
         self,
         query: torch.Tensor,
@@ -295,6 +347,7 @@ class AscendAttentionBackendImpl(AttentionImpl):
         attn_metadata: AscendMetadata,
         output: Optional[torch.Tensor] = None,
         num_tokens=0,
+        param_sink_number: Optional[int] = 0
     ) -> torch.Tensor:
         assert attn_metadata is not None
         assert attn_metadata.attn_mask is not None
@@ -311,18 +364,72 @@ class AscendAttentionBackendImpl(AttentionImpl):
             mask = mask.repeat(attn_metadata.seq_lens.size(0), 1, 1, 1)
             mask = torch_npu.npu_format_cast(mask.contiguous(),
                                              ACL_FORMAT_FRACTAL_NZ)
+        if torch.version.cann.startswith("8.3"):
+            mask = torch.ones((2048, 2048), dtype=torch.int8, device=mask.device).triu_(diagonal=1)
+            # TODO: nocache swa
+            if param_sink_number > 0:
+                query_lens = attn_metadata.query_lens
+                seq_lens = attn_metadata.seq_lens + param_sink_number
+                output, _ = torch.ops.custom.npu_fused_infer_attention_sink(
+                    query,
+                    key,
+                    value,
+                    atten_mask=mask,
+                    actual_seq_qlen=query_lens,
+                    actual_seq_kvlen=seq_lens,
+                    num_query_heads=self.num_heads,
+                    num_key_value_heads=self.num_kv_heads,
+                    input_layout='TND',
+                    sparse_mode=3,
+                    sink_number=param_sink_number,
+                    softmax_scale=self.scale,
+                )
+            else:
+                output, _ = torch_npu.npu_fused_infer_attention_score(
+                    query=query,
+                    key=key,
+                    value=value,
+                    atten_mask=mask,
+                    input_layout="TND",
+                    actual_seq_lengths=attn_metadata.query_start_loc[1:],
+                    actual_seq_lengths_kv=attn_metadata.seq_lens,
+                    num_key_value_heads=self.num_kv_heads,
+                    num_heads=self.num_heads,
+                    scale=self.scale,
+                    sparse_mode=3,
+                )
+            return output
+        # Patch for sink on CANN8.2
+        if param_sink_number > 0: 
+            seq_lens = attn_metadata.seq_lens + param_sink_number
+            # TODO: _npu_flash_attention only allows qlen==kvlen， 
+            mask_elem = mask[0, -1]
+            sink_mask = torch.full((mask.size(0) + param_sink_number, 
+                                mask.size(1) + param_sink_number), 
+                                mask_elem, dtype=mask.dtype, device=mask.device)
+            sink_mask[param_sink_number:, :param_sink_number] = 0.0
+            sink_mask[param_sink_number:, param_sink_number:] = mask
+            sink_mask[:param_sink_number, :param_sink_number].triu_(diagonal=1)
+            mask = sink_mask
+
+            output = torch.zeros((output.size(0) + param_sink_number, 
+                                    output.size(1), output.size(2)),
+                                dtype=output.dtype,
+                                device=output.device)
+        else:
+            seq_lens = attn_metadata.seq_lens
 
         torch_npu._npu_flash_attention(query=query,
                                        key=key,
                                        value=value,
                                        mask=mask,
-                                       seq_len=attn_metadata.seq_lens,
+                                       seq_len=seq_lens,
                                        scale_value=self.scale,
                                        num_heads=self.num_heads,
                                        num_kv_heads=self.num_kv_heads,
                                        out=output)
         assert output is not None
-        return output[:num_tokens, :, :]
+        return output[param_sink_number:param_sink_number + num_tokens, :, :]
 
     def _forward_prefill_cache_hit(
         self,
@@ -356,6 +463,8 @@ class AscendAttentionBackendImpl(AttentionImpl):
         query: torch.Tensor,
         attn_metadata: AscendMetadata,
         output: Optional[torch.Tensor] = None,
+        layer: AttentionLayer = None,
+        param_sink_number: Optional[int] = 0
     ) -> torch.Tensor:
         if is_310p():
             # seq_lens_tensor needs to be transferred to the device for 310P.
@@ -426,16 +535,46 @@ class AscendAttentionBackendImpl(AttentionImpl):
                 handle = torch.npu.graph_task_group_end(stream)
                 graph_params.handles[num_tokens].append(handle)
             else:
-                torch_npu._npu_paged_attention(
-                    query=query,
-                    key_cache=self.key_cache,
-                    value_cache=self.value_cache,
-                    num_kv_heads=self.num_kv_heads,
-                    num_heads=self.num_heads,
-                    scale_value=self.scale,
-                    block_table=attn_metadata.block_tables,
-                    context_lens=attn_metadata.seq_lens,
-                    out=output)
+                # Patch for Sparse KV cache of SWA.
+                num_block, block_size, _, _ = self.key_cache.shape  # type: ignore
+                key = self.key_cache.view(  # type: ignore
+                    num_block, block_size, -1)
+                value = self.value_cache.view(  # type: ignore
+                    num_block, block_size, -1)
+                block_tables = attn_metadata.sink_block_tables
+                use_swa = (self.layer_idx % 2 == 0)
+                seq_kvlens = attn_metadata.sink_seq_kvlens
+                if use_swa:
+                    attn_mask = self.attn_mask.to(query.device, non_blocking=True)
+
+                    output, _ = torch.ops.custom.npu_fused_infer_attention_sink(
+                                query,
+                                key,
+                                value,
+                                atten_mask=attn_mask,
+                                actual_seq_qlen=attn_metadata.swa_seq_qlens,
+                                actual_seq_kvlen=seq_kvlens,
+                                block_table=block_tables,
+                                pre_tokens=128,
+                                next_tokens=0,
+                                num_query_heads=self.num_heads,
+                                num_key_value_heads=self.num_kv_heads,
+                                input_layout='TND',
+                                sparse_mode=4,
+                                block_size=block_size,
+                                sink_number=param_sink_number,
+                                softmax_scale=self.scale)
+                else:
+                    torch_npu._npu_paged_attention(
+                        query=query,
+                        key_cache=self.key_cache,
+                        value_cache=self.value_cache,
+                        num_kv_heads=self.num_kv_heads,
+                        num_heads=self.num_heads,
+                        scale_value=self.scale,
+                        block_table=block_tables,
+                        context_lens=seq_kvlens,
+                        out=output)
         return output
 
     def _forward_v1_style(
@@ -443,6 +582,7 @@ class AscendAttentionBackendImpl(AttentionImpl):
         query: torch.Tensor,
         attn_metadata: AscendMetadata,
         output: Optional[torch.Tensor] = None,
+        param_sink_number: Optional[int] = 0
     ) -> torch.Tensor:
         # Use chunked prefill for head size 192 scenario, like deepseek
         # paged_attention_splitfuse maybe crash at such scenario.
@@ -485,34 +625,87 @@ class AscendAttentionBackendImpl(AttentionImpl):
             value = self.value_cache.view(  # type: ignore
                 num_block, block_size, -1)
 
-            output, _ = torch_npu.npu_fused_infer_attention_score(
-                query=query,
-                key=key,
-                value=value,
-                atten_mask=attn_metadata.attn_mask,
-                block_table=attn_metadata.block_tables,
-                input_layout="TND",
-                block_size=block_size,
-                actual_seq_lengths=attn_metadata.query_start_loc[1:],
-                actual_seq_lengths_kv=attn_metadata.seq_lens,
-                num_key_value_heads=self.num_kv_heads,
-                num_heads=self.num_heads,
-                scale=self.scale,
-                sparse_mode=3,
-            )
+            #TODO: swa层，window长度 传参
+            use_swa = (self.layer_idx % 2 == 0)
+            sparse_mode = 4 if use_swa else 3
+            if param_sink_number > 0:
+                if sparse_mode == 4:
+                    output, _ = torch.ops.custom.npu_fused_infer_attention_sink(
+                                        query,
+                                        key,
+                                        value,
+                                        atten_mask=self.attn_mask,
+                                        actual_seq_qlen=attn_metadata.swa_seq_qlens,
+                                        actual_seq_kvlen=attn_metadata.sink_seq_kvlens,
+                                        block_table=attn_metadata.sink_block_tables,
+                                        pre_tokens=128,
+                                        next_tokens=0,
+                                        num_query_heads=self.num_heads,
+                                        num_key_value_heads=self.num_kv_heads,
+                                        input_layout='TND',
+                                        sparse_mode=4,
+                                        block_size=block_size,
+                                        sink_number=param_sink_number,
+                                        softmax_scale=self.scale
+                                    )
+                elif sparse_mode == 3:
+                    output, _ = torch.ops.custom.npu_fused_infer_attention_sink(
+                                        query,
+                                        key,
+                                        value,
+                                        atten_mask=self.attn_mask,
+                                        actual_seq_qlen=attn_metadata.swa_seq_qlens,
+                                        actual_seq_kvlen=attn_metadata.sink_seq_kvlens,
+                                        block_table=attn_metadata.sink_block_tables,
+                                        num_query_heads=self.num_heads,
+                                        num_key_value_heads=self.num_kv_heads,
+                                        input_layout='TND',
+                                        sparse_mode=3,
+                                        block_size=block_size,
+                                        sink_number=param_sink_number,
+                                        softmax_scale=self.scale
+                                    )
+
+            else:
+                output, _ = torch_npu.npu_fused_infer_attention_score(
+                    query=query,
+                    key=key,
+                    value=value,
+                    atten_mask=attn_metadata.attn_mask,
+                    block_table=attn_metadata.block_tables,
+                    input_layout="TND",
+                    block_size=block_size,
+                    actual_seq_lengths=attn_metadata.query_start_loc[1:],
+                    actual_seq_lengths_kv=attn_metadata.seq_lens,
+                    num_key_value_heads=self.num_kv_heads,
+                    num_heads=self.num_heads,
+                    scale=self.scale,
+                    sparse_mode=3,
+                )
         else:
+            # Patch for sink on CANN 8.2
+            if param_sink_number > 0:
+                seq_kvlens = attn_metadata.seq_lens + param_sink_number
+                block_tables = F.pad(attn_metadata.block_tables, (1, 0, 0, 0), value=0)
+                mask = F.pad(attn_metadata.attn_mask, (param_sink_number, 0, 0, 0), value=0)
+            else:
+                seq_kvlens = attn_metadata.seq_lens
+                block_tables = attn_metadata.block_tables
+                mask = attn_metadata.attn_mask
+
             torch_npu._npu_paged_attention_splitfuse(
                 query=query,
                 key_cache=self.key_cache,
                 value_cache=self.value_cache,
-                mask=attn_metadata.attn_mask,
-                block_table=attn_metadata.block_tables,
+                mask=mask,
+                block_table=block_tables,
                 seq_len=attn_metadata.query_lens,
-                context_lens=attn_metadata.seq_lens,
+                context_lens=seq_kvlens,
                 num_kv_heads=self.num_kv_heads,
                 num_heads=self.num_heads,
                 scale_value=self.scale,
                 out=output)
+
         return output
 
     def forward(
@@ -525,6 +718,10 @@ class AscendAttentionBackendImpl(AttentionImpl):
         attn_metadata: AscendMetadata,
         output: Optional[torch.Tensor] = None,
         trace_flag: bool = True,
+        sink_query: Optional[torch.Tensor] = None,
+        sink_key: Optional[torch.Tensor] = None,
+        sink_value: Optional[torch.Tensor] = None,
+        v_head_size: Optional[int] = None,
     ) -> torch.Tensor:
         """Forward pass with Ascend attention.
         Args:
@@ -556,7 +753,12 @@ class AscendAttentionBackendImpl(AttentionImpl):
                 key=key,
                 value=value,
                 output=output,
-                layer_name=layer.layer_name)
+                layer_name=layer.layer_name,
+                sink_query=sink_query,
+                sink_key=sink_key,
+                sink_value=sink_value,
+                v_head_size=v_head_size
+            )
 
         elif hasattr(layer, 'quant_method') and use_kv_cache_int8:
             output = layer.quant_method.apply(layer, query, key, value,
@@ -575,10 +777,13 @@ class AscendAttentionBackendImpl(AttentionImpl):
                                           "encoder/decoder cross-attention "
                                           "are not implemented for "
                                           "PallasAttentionBackendImpl")
+            sink_key_flag = (sink_key is not None)
+            param_sink_number = sink_key.shape[0] if sink_key_flag else 0
             # View q k v to BSH.
             query = query.view(-1, self.num_heads, self.head_size)
             key = key.view(-1, self.num_kv_heads, self.head_size)
-            value = value.view(-1, self.num_kv_heads, self.head_size)
+            value = value.view(-1, self.num_kv_heads, 
+                               v_head_size if v_head_size is not None else self.head_size)
             # TODO: Remove this contiguous in the future.
             value = value.contiguous()
 
@@ -586,33 +791,63 @@ class AscendAttentionBackendImpl(AttentionImpl):
                 if self.key_cache is None:
                     self.key_cache, self.value_cache = kv_cache[0], kv_cache[1]
                 slots = attn_metadata.slot_mapping
+
                 torch_npu._npu_reshape_and_cache(
                     key=key[:num_actual_tokens],
                     value=value[:num_actual_tokens],
                     key_cache=self.key_cache,
                     value_cache=self.value_cache,
                     slot_indices=slots)
-
+                if sink_key_flag and not self.sink_cached:
+                    # kv cache start from block 1 and slots 128, so we store sink in block 0.
+                    slots = torch.arange(0, param_sink_number, 
+                                        dtype=attn_metadata.slot_mapping.dtype,
+                                        device=attn_metadata.slot_mapping.device)
+                    torch_npu._npu_reshape_and_cache(
+                        key=sink_key,
+                        value=sink_value,
+                        key_cache=self.key_cache,
+                        value_cache=self.value_cache,
+                        slot_indices=slots)
+                    self.sink_cached = True
+
+            # TODO: 暂不进PrefillCacheHit分支，不更新sink实现
+            if attn_metadata.attn_state == AscendAttentionState.PrefillNoCache and sink_key_flag:
+                attn_metadata.attn_state = AscendAttentionState.ChunkedPrefill
             # V0-Style scheduler situation.
             if attn_metadata.attn_state == AscendAttentionState.PrefillNoCache:
+                if torch.version.cann.startswith("8.3"):
+                    # npu_fused_infer_attention_score and npu_fused_infer_attention_sink
+                    # does not support cases where query.shape[0] != actual_seq_lengths
+                    # Thus we need unpad it here.
+                    num_tokens = attn_metadata.query_start_loc[-1]
+                    query = query[:num_tokens]
+                    key = key[:num_tokens]
+                    value = value[:num_tokens]
+                elif sink_key_flag:
+                    query = torch.cat([sink_query, query], dim=0)
+                if sink_key_flag:
+                    key = torch.cat([sink_key, key], dim=0)
+                    value = torch.cat([sink_value, value], dim=0)
                 output = self._forward_prefill_no_cache(
-                    query, key, value, attn_metadata, output, num_tokens)
+                    query, key, value, attn_metadata, output, num_tokens,
+                    param_sink_number
+                )
             elif attn_metadata.attn_state == \
                 AscendAttentionState.PrefillCacheHit:
                 output = self._forward_prefill_cache_hit(
                     query, attn_metadata, output)
             elif attn_metadata.attn_state == AscendAttentionState.DecodeOnly:
                 output = self._forward_decode_only(query, attn_metadata,
-                                                   output)
+                                                   output, layer,
+                                                   param_sink_number)
             # Normal V1 situation.
             else:
                 if torch.version.cann.startswith("8.3"):
-                    # npu_fused_infer_attention_score does not support cases
-                    # where query.shape[0] != attn_metadata.query_start_loc[-1].
-                    # Thus we need unpad it here.
                     num_tokens = attn_metadata.query_start_loc[-1]
                     query = query[:num_tokens]
-                output = self._forward_v1_style(query, attn_metadata, output)
+                output = self._forward_v1_style(query, attn_metadata, output,
+                                                param_sink_number)
 
         # to make in-place change to the output tensor
         if hasattr(layer, 'quant_method') and use_kv_cache_int8:
@@ -627,6 +862,10 @@ def unified_ascend_attention_with_output(
     value: torch.Tensor,
     output: torch.Tensor,
     layer_name: str,
+    sink_query: Optional[torch.Tensor] = None,
+    sink_key: Optional[torch.Tensor] = None,
+    sink_value: Optional[torch.Tensor] = None,
+    v_head_size: Optional[int] = None,
 ) -> None:
     wait_for_kv_layer_from_connector(layer_name)
     forward_context: ForwardContext = get_forward_context()
@@ -642,7 +881,11 @@ def unified_ascend_attention_with_output(
                       kv_cache,
                       attn_metadata,
                       output,
-                      trace_flag=False)
+                      trace_flag=False,
+                      sink_query=sink_query,
+                      sink_key=sink_key,
+                      sink_value=sink_value,
+                      v_head_size=v_head_size)
     maybe_save_kv_layer_to_connector(layer_name, kv_cache)
     return
 
@@ -653,6 +896,11 @@ def unified_attention_with_output_fake(
     value: torch.Tensor,
     output: torch.Tensor,
     layer_name: str,
+    # patch for pangu with attention sink
+    sink_query: Optional[torch.Tensor] = None,
+    sink_key: Optional[torch.Tensor] = None,
+    sink_value: Optional[torch.Tensor] = None,
+    v_head_size: Optional[int] = None,
 ) -> None:
     return
 
diff --git a/vllm_ascend/attention/utils.py b/vllm_ascend/attention/utils.py
index 519cde0..93e1c95 100644
--- a/vllm_ascend/attention/utils.py
+++ b/vllm_ascend/attention/utils.py
@@ -63,6 +63,8 @@ class AscendCommonAttentionMetadata:
 
     graph_pad_size: int = -1
 
+    num_input_tokens: int = -1
+
 
 def split_decodes_and_prefills(
     common_attn_metadata: AscendCommonAttentionMetadata,
diff --git a/vllm_ascend/platform.py b/vllm_ascend/platform.py
index f1581df..b690bcb 100644
--- a/vllm_ascend/platform.py
+++ b/vllm_ascend/platform.py
@@ -216,6 +216,9 @@ class NPUPlatform(Platform):
             if compilation_config.cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE:
                 compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
 
+        if compilation_config.cudagraph_mode == CUDAGraphMode.PIECEWISE:
+            compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE_DECODE_ONLY
+
         if compilation_config.cudagraph_mode == CUDAGraphMode.NONE:
             compilation_config.level = CompilationLevel.NO_COMPILATION
         # TODO: Currently MLA does not support FULL_DECODE_ONLY, remove the second condition
@@ -223,7 +226,8 @@ class NPUPlatform(Platform):
         elif compilation_config.cudagraph_mode == CUDAGraphMode.PIECEWISE or (
                 compilation_config.cudagraph_mode
                 == CUDAGraphMode.FULL_DECODE_ONLY and model_config is not None
-                and model_config.use_mla):
+                and model_config.use_mla) or (
+                compilation_config.cudagraph_mode == CUDAGraphMode.PIECEWISE_DECODE_ONLY):
             logger.info(
                 "PIECEWISE compilation enabled on NPU. use_inductor not supported - "
                 "using only ACL Graph mode")
@@ -232,7 +236,8 @@ class NPUPlatform(Platform):
             compilation_config.set_splitting_ops_for_v1()
             compilation_config.use_inductor = False
             compilation_config.splitting_ops.extend([
-                "vllm.unified_ascend_attention_with_output", "vllm.mla_forward"
+                "vllm.unified_ascend_attention_with_output", "vllm.mla_forward",
+                "vllm.aggregate_hiddden",
             ])
             update_aclgraph_sizes(vllm_config)
         elif compilation_config.cudagraph_mode == CUDAGraphMode.FULL_DECODE_ONLY:
diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py
index 9281dd7..34808ec 100644
--- a/vllm_ascend/worker/model_runner_v1.py
+++ b/vllm_ascend/worker/model_runner_v1.py
@@ -281,6 +281,7 @@ class NPUModelRunner(LoRAModelRunnerMixin):
         self.encoder_cache: Dict[str, torch.Tensor] = {}
         self.attn_mask = None
         self.attn_state = None
+        self.with_prefill = False
         self.requests: Dict[str, CachedRequestState] = {}
         self.intermediate_tensors: Optional[IntermediateTensors] = None
         self.runner_only_attn_layers: set[str] = set()
@@ -509,6 +510,48 @@ class NPUModelRunner(LoRAModelRunnerMixin):
         self.num_draft_tokens = self._make_buffer(self.max_num_reqs,
                                                   dtype=torch.int32)
 
+        # Patch for conv cache
+        self.router_sliding_window = getattr(self.model_config.hf_text_config, "router_sliding_window", 0)
+        if self.router_sliding_window > 1:
+            self.cache_length = self.router_sliding_window - 1
+            self.req_cache_map = {}
+            self.occupied_cache = [0]*(self.max_num_reqs)
+            self.q_offsets = torch.arange(-self.cache_length, 0, device=self.device)
+            self.cache_slot_id = torch.empty(self.max_num_reqs, 
+                                    dtype=torch.long, device=self.device)
+            self.is_first_chunk = torch.empty(self.max_num_reqs, dtype=torch.bool, device=self.device) # For chunked prefill
+
+    def _build_conv_context(self, with_prefill:bool = False, dummy:bool = False, num_tokens:int = 0):
+        # conv cache slot & prefill hiddenstates loc
+        cache_slot_id = self.cache_slot_id[:self.input_batch.num_reqs]
+        query_start_loc = self.query_start_loc[:self.input_batch.num_reqs + 1]
+        is_first_chunk = self.is_first_chunk[:self.input_batch.num_reqs]
+
+        if with_prefill:
+            for idx, req_id in enumerate(self.input_batch.req_ids):
+                if req_id in self.req_cache_map:
+                    cache_id = self.req_cache_map[req_id]
+                    cache_slot_id[idx] = cache_id
+                    is_first_chunk[idx] = False
+                else:
+                    # new request with the first chunk
+                    new_cahce_id = self.occupied_cache.index(0)
+                    self.occupied_cache[new_cahce_id] = 1
+                    self.req_cache_map[req_id] = new_cahce_id
+                    cache_slot_id[idx] = new_cahce_id
+                    is_first_chunk[idx] = True
+        else:
+            for idx, req_id in enumerate(self.input_batch.req_ids):
+                cache_id = self.req_cache_map[req_id]
+                cache_slot_id[idx] = cache_id
+                is_first_chunk[idx] = False
+
+        forward_context = get_forward_context()
+        forward_context.cache_slot_id = cache_slot_id
+        forward_context.is_first_chunk = is_first_chunk
+        forward_context.query_start_loc = query_start_loc
+
+
     def _make_buffer(self,
                      *size: Union[int, torch.SymInt],
                      dtype: torch.dtype,
@@ -548,12 +591,16 @@ class NPUModelRunner(LoRAModelRunnerMixin):
             self.input_batch.num_accepted_tokens_cpu[i] = num_tokens
 
     def _use_aclgraph(self) -> bool:
-        return self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE and self.compilation_config.level == CompilationLevel.PIECEWISE and not self.model_config.enforce_eager
+        return self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE and \
+            self.compilation_config.level == CompilationLevel.PIECEWISE and not self.model_config.enforce_eager
 
     def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
         # Remove finished requests from the cached states.
         for req_id in scheduler_output.finished_req_ids:
-            self.requests.pop(req_id, None)
+            self.requests.pop(req_id, None) 
+            if self.router_sliding_window > 1 and req_id in self.req_cache_map:
+                cache_id = self.req_cache_map.pop(req_id)
+                self.occupied_cache[cache_id] = 0
 
         # Remove the finished requests from the persistent batch.
         # NOTE(woosuk): There could be an edge case where finished_req_ids and
@@ -891,7 +938,8 @@ class NPUModelRunner(LoRAModelRunnerMixin):
     def _make_attention_mask(self, seq_lens, position,
                              attn_state) -> torch.Tensor:
         # Chunk Prefill situation.
-        if attn_state == AscendAttentionState.ChunkedPrefill and not self.vllm_config.model_config.use_mla and not self.ascend_config.use_sfa:
+        if attn_state == AscendAttentionState.ChunkedPrefill and not self.vllm_config.model_config.use_mla and not \
+                self.ascend_config.use_sfa:
             if torch.version.cann.startswith("8.3"):
                 return self.attn_mask_builder.get_splitfuse_attn_mask()
             else:
@@ -942,7 +990,7 @@ class NPUModelRunner(LoRAModelRunnerMixin):
                 src_end = num_computed_tokens + prompt_part_len
 
                 self.mrope_positions_cpu[:, dst_start:dst_end] = \
-                    req.mrope_positions[:,src_start:src_end]
+                    req.mrope_positions[:, src_start:src_end]
 
                 mrope_pos_ptr += prompt_part_len
 
@@ -1126,9 +1174,9 @@ class NPUModelRunner(LoRAModelRunnerMixin):
         cumsum_dtype: Optional[np.dtype] = None,
     ) -> tuple[np.ndarray, np.ndarray]:
         """Get the cumulative sum and batched arange of the given array.
-        # E.g., [2, 5, 3] -> ([2, 7, 10], [0, 1, 0, 1, 2, 3, 4, 0, 1, 2])
-        # Equivalent to but faster than:
-        # np.concatenate([np.arange(n) for n in num_tokens])
+        E.g., [2, 5, 3] -> ([2, 7, 10], [0, 1, 0, 1, 2, 3, 4, 0, 1, 2])
+        Equivalent to but faster than:
+        np.concatenate([np.arange(n) for n in num_tokens])
         """
         # Step 1. [2, 5, 3] -> [2, 7, 10]
         cu_num_tokens = np.cumsum(num_tokens, dtype=cumsum_dtype)
@@ -1518,6 +1566,7 @@ class NPUModelRunner(LoRAModelRunnerMixin):
                 max_query_len=max_num_scheduled_tokens,
                 graph_pad_size=self.graph_pad_size,
                 decode_token_per_req=self.decode_token_per_req,
+                num_input_tokens=num_input_tokens
             )
 
             if self.speculative_config and \
@@ -1964,6 +2013,9 @@ class NPUModelRunner(LoRAModelRunnerMixin):
                     model_instance=self.model):
                 self.maybe_setup_kv_connector(scheduler_output)
 
+                if self.router_sliding_window > 1:
+                    self._build_conv_context(self.with_prefill)
+                
                 hidden_states = self._generate_process_reqs_hidden_states(
                     attn_metadata, self.with_prefill, maybe_padded_num_tokens,
                     input_ids, positions, intermediate_tensors, inputs_embeds)
@@ -2339,7 +2391,8 @@ class NPUModelRunner(LoRAModelRunnerMixin):
     ) -> torch.Tensor:
         # only support eager mode and piecewise graph now
         assert aclgraph_runtime_mode in {
-            CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL
+            CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL,
+            CUDAGraphMode.PIECEWISE_DECODE_ONLY
         }
 
         # Padding for DP
@@ -2472,6 +2525,8 @@ class NPUModelRunner(LoRAModelRunnerMixin):
                     batch_descriptor=batch_descriptor,
                     prefetch_stream=self.prefetch_stream,
                     model_instance=self.model):
+                if self.router_sliding_window > 1:
+                    self._build_conv_context(with_prefill, dummy=True, num_tokens=num_tokens)
                 hidden_states = self._generate_dummy_run_hidden_states(
                     with_prefill, is_torchair_compile, input_ids, positions,
                     attn_metadata, num_tokens, intermediate_tensors,
@@ -2789,8 +2844,8 @@ class NPUModelRunner(LoRAModelRunnerMixin):
 
                     # In order to transfer kv cache through the reigster_memory api from llmdatadist, the memory
                     # address should be aligned by 2M. In most case, torch_npu can allocate 2M aligned memory, but
-                    # we found there are also some exceptions during test, so we manual align those memory here, this part
-                    # of code may consume 2M * 2 * elem_size memory every layer.
+                    # we found there are also some exceptions during test, so we manual align those memory here, 
+                    # this part of code may consume 2M * 2 * elem_size memory every layer.
                     nope_allocate_shape = num_blocks * block_size * num_kv_heads * nope_dim
                     nope_allocate_shape_alignment = nope_allocate_shape + alignment
                     rope_allocate_shape = num_blocks * block_size * num_kv_heads * rope_dim
@@ -2888,8 +2943,8 @@ class NPUModelRunner(LoRAModelRunnerMixin):
 
                     # In order to transfer kv cache through the reigster_memory api from llmdatadist, the memory
                     # address should be aligned by 2M. In most case, torch_npu can allocate 2M aligned memory, but
-                    # we found there are also some exceptions during test, so we manual align those memory here, this part
-                    # of code may consume 2M * 2 * elem_size memory every layer.
+                    # we found there are also some exceptions during test, so we manual align those memory here,
+                    # this part of code may consume 2M * 2 * elem_size memory every layer.
                     nope_allocate_shape = num_blocks * block_size * num_kv_heads * nope_dim
                     nope_allocate_shape_alignment = nope_allocate_shape + alignment
                     rope_allocate_shape = num_blocks * block_size * num_kv_heads * rope_dim
@@ -3432,6 +3487,7 @@ class NPUModelRunner(LoRAModelRunnerMixin):
             and all(op in self.compilation_config.splitting_ops for op in [
                 "vllm.unified_ascend_attention_with_output",
                 "vllm.mla_forward",
+                "vllm.aggregate_hiddden", 
             ]))
 
         # Flexible resolve the aclgraph mode
@@ -3495,7 +3551,8 @@ class NPUModelRunner(LoRAModelRunnerMixin):
                            uniform_decode: bool):
         assert aclgraph_runtime_mode != CUDAGraphMode.NONE and \
             aclgraph_runtime_mode in [CUDAGraphMode.FULL,
-                                      CUDAGraphMode.PIECEWISE]
+                                      CUDAGraphMode.PIECEWISE,
+                                      CUDAGraphMode.PIECEWISE_DECODE_ONLY]
 
         # Only rank 0 should print progress bar during capture
         if is_global_first_rank():
@@ -3519,10 +3576,12 @@ class NPUModelRunner(LoRAModelRunnerMixin):
                 # attention while `PIECEWISE` implies no attention.
                 force_attention = (aclgraph_runtime_mode == CUDAGraphMode.FULL)
                 self._dummy_run(num_tokens,
+                                with_prefill = (uniform_decode == False),
                                 aclgraph_runtime_mode=CUDAGraphMode.NONE,
                                 force_attention=force_attention,
                                 uniform_decode=uniform_decode)
             self._dummy_run(num_tokens,
+                            with_prefill = (uniform_decode == False),
                             aclgraph_runtime_mode=aclgraph_runtime_mode,
                             force_attention=force_attention,
                             uniform_decode=uniform_decode)
@@ -3556,7 +3615,8 @@ class NPUModelRunner(LoRAModelRunnerMixin):
                     logger.error(
                         f"ACLgraph sizes capture fail: {type(e).__name__}:\n"
                         "ACLgraph has insufficient available streams to capture the configured number of sizes. "
-                        "Please verify both the availability of adequate streams and the appropriateness of the configured size count.\n\n"
+                        "Please verify both the availability of adequate streams "
+                        "and the appropriateness of the configured size count.\n\n"
                         "Recommended solutions:\n"
                         "1. Manually configure the compilation_config parameter "
                         "with a reduced set of sizes: '{\"cudagraph_capture_sizes\":[size1, size2, size3, ...]}'.\n"
@@ -3564,8 +3624,9 @@ class NPUModelRunner(LoRAModelRunnerMixin):
                         f"{str(e)}")
                     raise
 
-            if aclgraph_mode.decode_mode() == CUDAGraphMode.FULL and \
-                aclgraph_mode.separate_routine():
+            if aclgraph_mode.separate_routine() and \
+                (aclgraph_mode.decode_mode() == CUDAGraphMode.FULL or \
+                aclgraph_mode.decode_mode() == CUDAGraphMode.PIECEWISE):
                 max_num_tokens = self.scheduler_config.max_num_seqs * \
                         self.uniform_decode_query_len
                 decode_cudagraph_batch_sizes = [
@@ -3576,7 +3637,7 @@ class NPUModelRunner(LoRAModelRunnerMixin):
                     reversed(decode_cudagraph_batch_sizes))
                 self._capture_aclgraphs(
                     compilation_cases=compilation_cases_decode,
-                    aclgraph_runtime_mode=CUDAGraphMode.FULL,
+                    aclgraph_runtime_mode=aclgraph_mode.decode_mode(),
                     uniform_decode=True)
 
         # Disable aclgraph capturing globally, so any unexpected aclgraph