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
[Kernel] [CPU] refactor `cpu_attn.py:_run_sdpa_forward` for better memory access (#24701)
Signed-off-by: ignaciosica <mignacio.sica@gmail.com>
Ignacio Sica committed
Sep 12, 2025 at 08:23 UTC
7a1c4025f1e2879fa398888d70c596e5818026cb
1 file changed
+4
-4
vllm/v1/attention/backends/cpu_attn.py
+4
-4
@@ -641,10 +641,6 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
641
attn_metadata: TorchSDPAMetadata,
642
attn_type: str = AttentionType.DECODER,
643
) -> None:
644
- if self.num_kv_heads != self.num_heads:
645
- key = key.repeat_interleave(self.num_queries_per_kv, dim=1)
646
- value = value.repeat_interleave(self.num_queries_per_kv, dim=1)
647
-
644
attn_masks = attn_metadata.get_attn_bias(attn_type)
645
if attn_masks is None:
646
if self.alibi_slopes is not None:
@@ -665,6 +661,10 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
661
key = key.movedim(0, key.dim() - 2)
662
value = value.movedim(0, value.dim() - 2)
663
664
+ if self.num_kv_heads != self.num_heads:
665
+ key = key.repeat_interleave(self.num_queries_per_kv, dim=-3)
666
+ value = value.repeat_interleave(self.num_queries_per_kv, dim=-3)
667
+
668
causal_attn = (attn_type == AttentionType.DECODER)
669
670
seq_lens_q, seq_lens_kv = attn_metadata.get_seq_lens(attn_type)