A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[fix][cpu] fix prefill attention in CPU attention backend (#27035)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>
Fadi Arafeh committed
Oct 18, 2025 at 14:30 UTC
ab4be40fc5bde2dbe76e3fb081dc3721219b5a2b
2 files changed
+16
-4
vllm/engine/arg_utils.py
+9
-1
@@ -1293,7 +1293,8 @@ class EngineArgs:
1293
1294
# Set default arguments for V1 Engine.
1295
self._set_default_args(usage_context, model_config)
1296
- # Disable chunked prefill for POWER (ppc64le)/ARM/s390x/RISCV CPUs in V1
1296
+ # Disable chunked prefill and prefix caching for:
1297
+ # POWER (ppc64le)/ARM/s390x/RISCV CPUs in V1
1298
if current_platform.is_cpu() and current_platform.get_cpu_architecture() in (
1299
CpuArchEnum.POWERPC,
1300
CpuArchEnum.S390X,
@@ -1306,6 +1307,13 @@ class EngineArgs:
1307
"disabling it for V1 backend."
1308
)
1309
self.enable_chunked_prefill = False
1310
+ logger.info(
1311
+ "Prefix caching is not supported for ARM and POWER, "
1312
+ "S390X and RISC-V CPUs; "
1313
+ "disabling it for V1 backend."
1314
+ )
1315
+ self.enable_prefix_caching = False
1316
+
1317
assert self.enable_chunked_prefill is not None
1318
1319
sliding_window: int | None = None
vllm/v1/attention/backends/cpu_attn.py
+7
-3
@@ -412,7 +412,7 @@ class TorchSDPAMetadataBuilderV1(AttentionMetadataBuilder[TorchSDPAMetadata]):
412
num_decode_tokens=num_decode_tokens,
413
slot_mapping=slot_mapping,
414
# to ensure inference when chunked_prefill is disabled
415
- seq_lens=seq_lens_cpu.tolist(),
415
+ seq_lens=seq_lens_cpu.tolist()[num_decodes:], # prefill
416
decode_seq_lens_tensor=seq_lens_cpu[:num_decodes], # decode
417
decode_max_seq_len=max_decode_seq_len, # decode
418
decode_block_tables=block_table_tensor[:num_decodes], # decode
@@ -617,7 +617,6 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
617
prefill_meta.prefill_block_tables,
618
self.alibi_slopes,
619
)
620
-
620
if decode_meta := attn_metadata.decode_metadata:
621
assert attn_type != AttentionType.ENCODER_ONLY, (
622
"Encoder-only models should not have decode metadata."
@@ -686,7 +685,12 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
685
causal_attn = attn_type == AttentionType.DECODER
686
687
seq_lens_q, seq_lens_kv = attn_metadata.get_seq_lens(attn_type)
689
- start_q, start_kv = 0, 0
688
+ # Incoming Q and KV contain decoded tokens as well, hence start at an offset
689
+ # equal to num_decode_tokens since decode requests appear first
690
+ start_q, start_kv = (
691
+ attn_metadata.num_decode_tokens,
692
+ attn_metadata.num_decode_tokens,
693
+ )
694
for seq_len_q, seq_len_kv, mask in zip(seq_lens_q, seq_lens_kv, attn_masks):
695
end_q = start_q + seq_len_q
696
end_kv = start_kv + seq_len_kv