refine APIs

Co-authored-by: yan <yan.ma@intel.com>
Co-authored-by: mayuyuace <qiming1.zhang@intel.com>

Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>

Author: jikunshang

vllm/v1/attention/backends/ipex_attn.py

@@ -132,6 +132,7 @@ def __init__(
         else:
             self.sliding_window = (sliding_window - 1, 0)
         self.kv_cache_dtype = kv_cache_dtype
+        self.use_irope = use_irope
         if logits_soft_cap is None:
             # In flash-attn, setting logits_soft_cap as 0 means no soft cap.
             logits_soft_cap = 0
@@ -204,19 +205,45 @@ def forward(
             layer._k_scale_float,
             layer._v_scale_float,
         )
+        use_local_attn = \
+            (self.use_irope and attn_metadata.local_attn_metadata is not None)
+
+        if use_local_attn:
+            assert attn_metadata.local_attn_metadata is not None
+            local_metadata = attn_metadata.local_attn_metadata
+            cu_seqlens_q = local_metadata.local_query_start_loc
+            sequesd_k = local_metadata.local_seqused_k
+            max_seqlen_q = local_metadata.local_max_query_len
+            max_seqlen_k = local_metadata.local_max_seq_len
+            block_table = local_metadata.local_block_table
+        else:
+            cu_seqlens_q = attn_metadata.query_start_loc
+            sequesd_k = attn_metadata.seq_lens
+            max_seqlen_q = attn_metadata.max_query_len
+            max_seqlen_k = attn_metadata.max_seq_len
+            block_table = attn_metadata.block_table
+
+        if not hasattr(attn_metadata, "seq_start_loc"):
+            cumsum = torch.cumsum(sequesd_k, dim=0)
+            cu_seqlens_k = torch.cat([
+                torch.tensor([0], device=sequesd_k.device, dtype=torch.int32),
+                cumsum
+            ]).to(torch.int32)
+        else:
+            cu_seqlens_k = attn_metadata.seq_start_loc
 
         ipex_ops.flash_attn_varlen_func(
             output[:num_actual_tokens],
             query[:num_actual_tokens],
             key_cache,
             value_cache,
-            attn_metadata.query_start_loc,
-            attn_metadata.seq_start_loc,
-            attn_metadata.max_query_len,
-            attn_metadata.max_seq_len,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
             self.scale,
             is_casual=True,
-            block_table=attn_metadata.block_table,
+            block_table=block_table,
             alibi_slopes=self.alibi_slopes,
         )
         return output

vllm/v1/worker/xpu_model_runner.py

@@ -1,11 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Optional
 
 import numpy as np
 import torch
 
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
+from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
 if TYPE_CHECKING:
@@ -38,7 +39,9 @@ def _init_device_properties(self) -> None:
     def _sync_device(self) -> None:
         torch.xpu.synchronize()
 
-    def _prepare_inputs(self, scheduler_output: "SchedulerOutput"):
+    def _prepare_inputs(
+        self, scheduler_output: "SchedulerOutput"
+    ) -> tuple[dict[str, Any], torch.Tensor, Optional[SpecDecodeMetadata]]:
         total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
         assert total_num_scheduled_tokens > 0
         num_reqs = self.input_batch.num_reqs