[Hardware][CPU] support cpu in v1 engine

vllm/attention/backends/torch_sdpa.py

@@ -297,26 +297,33 @@
                 dtype=torch.int32,
                 device="cpu",
             )
-            query_lens_tensor = torch.tensor(prefill_query_lens,
-                                             dtype=torch.int32,
-                                             device="cpu")
-            kv_lens_tensor = torch.tensor(prefill_seq_lens,
-                                          dtype=torch.int32,
-                                          device="cpu")
-            query_start_loc = torch.zeros(input_data.num_prefills + 1,
-                                          dtype=torch.int32,
-                                          device="cpu")
-            kv_start_loc = torch.zeros(input_data.num_prefills + 1,
-                                       dtype=torch.int32,
-                                       device="cpu")
-            torch.cumsum(query_lens_tensor,
-                         dim=0,
-                         dtype=torch.int32,
-                         out=query_start_loc[1:])
-            torch.cumsum(kv_lens_tensor,
-                         dim=0,
-                         dtype=torch.int32,
-                         out=kv_start_loc[1:])
+            query_start_loc: torch.Tensor
+            kv_start_loc: torch.Tensor
+            if input_data.seq_start_loc is not None and input_data.seq_start_loc is not None:
+                query_start_loc = input_data.query_start_loc[input_data.num_prefills + 1:]
+                kv_start_loc = input_data.seq_start_loc
+            else:
+                query_lens_tensor = torch.tensor(prefill_query_lens,
+                                                 dtype=torch.int32,
+                                                 device="cpu")
+                kv_lens_tensor = torch.tensor(prefill_seq_lens,
+                                              dtype=torch.int32,
+                                              device="cpu")
+
+                query_start_loc = torch.zeros(input_data.num_prefills + 1,
+                                              dtype=torch.int32,
+                                              device="cpu")
+                kv_start_loc = torch.zeros(input_data.num_prefills + 1,
+                                           dtype=torch.int32,
+                                           device="cpu")
+                torch.cumsum(query_lens_tensor,
+                             dim=0,
+                             dtype=torch.int32,
+                             out=query_start_loc[1:])
+                torch.cumsum(kv_lens_tensor,
+                             dim=0,
+                             dtype=torch.int32,
+                             out=kv_start_loc[1:])
             max_query_len = max(prefill_query_lens)
             max_kv_len = max(prefill_seq_lens)
         else:

vllm/v1/attention/backends/flash_attn.py

@@ -6,7 +6,7 @@
 
 from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
                                               AttentionMetadata, AttentionType)
-from vllm.vllm_flash_attn import flash_attn_varlen_func
+# from vllm.vllm_flash_attn import flash_attn_varlen_func
 
 
 class FlashAttentionBackend(AttentionBackend):
@@ -166,6 +166,7 @@ def forward(
         )
 
         # Compute attention and update output up to `num_actual_tokens`.
+        from vllm.vllm_flash_attn import flash_attn_varlen_func
         flash_attn_varlen_func(
             q=query[:num_actual_tokens],
             k=key_cache,

vllm/v1/core/scheduler.py

@@ -203,6 +203,9 @@
                     num_computed_tokens -= 1
                     num_new_tokens = 1
                     computed_blocks.pop()
+                # if current request can't be fully scheduled, skip and don't schedule it
+                if num_new_tokens > token_budget:
+                    break
                 num_new_tokens = min(num_new_tokens, token_budget)
                 assert num_new_tokens > 0
 

vllm/v1/engine/async_llm.py

@@ -10,6 +10,7 @@
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
 from vllm.outputs import PoolingRequestOutput, RequestOutput
+from vllm.platforms import current_platform
 from vllm.pooling_params import PoolingParams
 from vllm.prompt_adapter.request import PromptAdapterRequest
 from vllm.sampling_params import SamplingParams
@@ -21,6 +22,7 @@
 from vllm.v1.engine.detokenizer import Detokenizer
 from vllm.v1.engine.processor import Processor
 from vllm.v1.executor.abstract import Executor
+from vllm.v1.executor.cpu_executor import CPUExecutor
 
 logger = init_logger(__name__)
 
@@ -127,6 +129,8 @@ def shutdown(self):
 
     @classmethod
     def _get_executor_cls(cls, vllm_config: VllmConfig):
+        if current_platform.is_cpu():
+            return CPUExecutor
         distributed_executor_backend = (
             vllm_config.parallel_config.distributed_executor_backend)
         if distributed_executor_backend == "mp":

vllm/v1/engine/llm_engine.py

@@ -11,6 +11,7 @@
 from vllm.lora.request import LoRARequest
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
 from vllm.outputs import RequestOutput
+from vllm.platforms import current_platform
 from vllm.pooling_params import PoolingParams
 from vllm.prompt_adapter.request import PromptAdapterRequest
 from vllm.sampling_params import SamplingParams
@@ -21,6 +22,7 @@
 from vllm.v1.engine.detokenizer import Detokenizer
 from vllm.v1.engine.processor import Processor
 from vllm.v1.executor.abstract import Executor
+from vllm.v1.executor.cpu_executor import CPUExecutor
 
 logger = init_logger(__name__)
 
@@ -104,6 +106,8 @@ def from_engine_args(
 
     @classmethod
     def _get_executor_cls(cls, vllm_config: VllmConfig):
+        if current_platform.is_cpu():
+            return CPUExecutor
         distributed_executor_backend = (
             vllm_config.parallel_config.distributed_executor_backend)
         if distributed_executor_backend == "mp":