Upload inference/vllm-ascend_v0.11.0rc0.patch with huggingface_hub

inference/vllm-ascend_v0.11.0rc0.patch

@@ -0,0 +1,847 @@
+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