@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8760f972ca
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[CPU] Refine CPU attention frontend (#45391)

Signed-off-by: jiang1.li <jiang1.li@intel.com>

Li, Jiang committed Jun 15, 2026 at 10:26 UTC 8760f972caf57f592ae2c118cecd7f105891ba13
4 files changed +384 -217
csrc/cpu/cpu_attn_impl.hpp
+8 -7
@@ -822,8 +822,8 @@ struct AttentionInput {
822 logits_buffer_t *__restrict__ logits_buffer, \
823 float *__restrict__ partial_q_buffer, float *__restrict__ max_buffer, \
824 float *__restrict__ sum_buffer, int32_t *__restrict__ block_table, \
825 - const int32_t kv_tile_start_pos, const int32_t kv_tile_end_pos, \
826 - const int32_t kv_tile_token_num, \
825 + const int32_t kv_end_pos, const int32_t kv_tile_start_pos, \
826 + const int32_t kv_tile_end_pos, const int32_t kv_tile_token_num, \
827 const int64_t kv_cache_num_blocks_stride, const int32_t q_head_num, \
828 const int32_t q_token_num, const int32_t q_tile_start_pos, \
829 const int32_t q_heads_per_kv, const int32_t block_size, \
@@ -834,7 +834,7 @@ struct AttentionInput {
834
835 #define CPU_ATTENTION_PARAMS \
836 q_heads_buffer, k_head_cache_ptr, v_head_cache_ptr, logits_buffer, \
837 - partial_q_buffer, max_buffer, sum_buffer, block_table, \
837 + partial_q_buffer, max_buffer, sum_buffer, block_table, kv_end_pos, \
838 kv_tile_start_pos, kv_tile_end_pos, kv_tile_token_num, \
839 kv_cache_num_blocks_stride, q_head_num, q_token_num, q_tile_start_pos, \
840 q_heads_per_kv, block_size, left_window_size, right_window_size, scale, \
@@ -917,6 +917,7 @@ class AttentionMainLoop {
917 // - max_buffer: [MaxQHeadNumPerIteration, 1], store max logits
918 // - sum_buffer: [MaxQHeadNumPerIteration, 1], store sum of exp
919 // - block_table
920 + // - kv_end_pos: un-aligned end position of KV cache
921 // - kv_tile_start_pos: start position of KV cache, aligned to
922 // BlockSizeAlignment
923 // - kv_tile_end_pos: end position of KV cache, aligned to
@@ -1043,7 +1044,7 @@ class AttentionMainLoop {
1044 }
1045
1046 apply_mask(logits_buffer, kv_tile_token_num, q_tile_start_pos,
1046 - kv_tile_start_pos, kv_tile_end_pos, q_token_num,
1047 + kv_end_pos, kv_tile_start_pos, kv_tile_end_pos, q_token_num,
1048 q_heads_per_kv, left_window_size, right_window_size);
1049
1050 // if (debug_info){
@@ -1126,7 +1127,7 @@ class AttentionMainLoop {
1127
1128 void apply_mask(logits_buffer_t* __restrict__ logits_buffer,
1129 const int64_t logits_buffer_stride,
1129 - const int32_t q_tile_start_pos,
1130 + const int32_t q_tile_start_pos, const int32_t kv_end_pos,
1131 const int32_t kv_tile_start_pos,
1132 const int32_t kv_tile_end_pos, const int32_t q_token_num,
1133 const int32_t q_heads_per_kv,
@@ -1154,7 +1155,7 @@ class AttentionMainLoop {
1155 std::max(kv_tile_start_pos,
1156 curr_token_pos + sliding_window_right + 1));
1157 }
1157 - return pos;
1158 + return std::min(pos, kv_end_pos);
1159 }();
1160
1161 int32_t left_invalid_token_num = left_kv_pos - kv_tile_start_pos;
@@ -1789,7 +1790,7 @@ class AttentionMainLoop {
1790 attn_impl.template execute_attention<Attention>(
1791 curr_q_heads_buffer, curr_k_cache, curr_v_cache,
1792 logits_buffer, curr_partial_q_buffer, curr_max_buffer,
1792 - curr_sum_buffer, curr_block_table,
1793 + curr_sum_buffer, curr_block_table, kv_end_pos,
1794 aligned_actual_kv_tile_pos_left,
1795 aligned_actual_kv_tile_pos_right, actual_kv_token_num,
1796 kv_cache_block_num_stride, q_tile_head_num,
csrc/cpu/generate_cpu_attn_dispatch.py
+1 -1
@@ -11,7 +11,7 @@ import os
11 HEAD_DIMS_32 = [32, 64, 96, 128, 160, 192, 224, 256, 512]
12
13 # Head dimensions divisible by 16 but not 32 (VEC16 only)
14 -HEAD_DIMS_16 = [80, 112]
14 +HEAD_DIMS_16 = [48, 80, 112]
15
16 # ISA types
17 ISA_TYPES = {
tests/kernels/attention/test_cpu_attn.py
+287 -7
@@ -25,7 +25,6 @@ from vllm._custom_ops import (
25 if torch.cpu._is_amx_tile_supported():
26 torch.cpu._init_amx()
27
28 -
28 NUM_HEADS = [
29 (4, 4),
30 (8, 2),
@@ -43,6 +42,11 @@ SEQ_LENS = [ # (q_len, kv_len)
42 [(2345, 2345), (5, 5), (3, 16), (134, 5131)], # prefill batch
43 [(992, 2456), (1, 1234), (98, 1145), (1, 4162), (2345, 2345)], # mixed batch
44 ]
45 +_FP8_ATOL = {"fp8_e4m3": 0.2, "fp8_e5m2": 0.3}
46 +_FP8_RTOL = 0.1
47 +ENCODER_SEQ_LENS = [
48 + [1, 678, 2367, 145, 4162, 36, 7812],
49 +]
50
51
52 def get_attn_isa(
@@ -61,10 +65,7 @@ def get_attn_isa(
65
66 # rand number generation takes too much time, cache rand tensors
67 @functools.lru_cache(maxsize=128, typed=False)
64 -def tensor_cache(
65 - elem_num: int,
66 - dtype: torch.dtype,
67 -) -> torch.Tensor:
68 +def tensor_cache(elem_num: int, dtype: torch.dtype, tag: str = "none") -> torch.Tensor:
69 tensor = torch.randn(elem_num, dtype=dtype)
70
71 return tensor
@@ -183,8 +184,222 @@ def ref_paged_attn(
184 return torch.cat(outputs, dim=0)
185
186
186 -_FP8_ATOL = {"fp8_e4m3": 0.2, "fp8_e5m2": 0.3}
187 -_FP8_RTOL = 0.1
187 +def ref_varlen_encoder_attn(
188 + query: torch.Tensor, # [token, q_head_num, head_dim]
189 + key: torch.Tensor, # [token, kv_head_num, head_dim]
190 + value: torch.Tensor,
191 + seq_lens: list[int],
192 + scale: float,
193 + sliding_window: int | None = None,
194 +) -> torch.Tensor:
195 + num_seqs = len(seq_lens)
196 + dtype = query.dtype
197 +
198 + output = torch.empty_like(query)
199 +
200 + start_idx = 0
201 + for i in range(num_seqs):
202 + seq_len = seq_lens[i]
203 + q = query[start_idx : start_idx + seq_len].float()
204 + k = key[start_idx : start_idx + seq_len].float()
205 + v = value[start_idx : start_idx + seq_len].float()
206 + q *= scale
207 +
208 + if q.shape[1] != k.shape[1]:
209 + k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1)
210 + v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
211 + attn = torch.einsum("qhd,khd->hqk", q, k).float()
212 + empty_mask = torch.ones(seq_len, seq_len)
213 + if sliding_window is not None:
214 + mask = (
215 + torch.triu(empty_mask, diagonal=1 - sliding_window).bool()
216 + ^ torch.triu(empty_mask, diagonal=sliding_window).bool()
217 + ).logical_not()
218 + else:
219 + mask = empty_mask.logical_not()
220 +
221 + attn.masked_fill_(mask, float("-inf"))
222 + attn = torch.softmax(attn, dim=-1)
223 + out = torch.einsum("hqk,khd->qhd", attn, v).to(dtype=dtype)
224 + output[start_idx : start_idx + seq_len].copy_(out)
225 +
226 + start_idx += seq_len
227 +
228 + return output
229 +
230 +
231 +@torch.inference_mode()
232 +def varlen_encoder_attention(
233 + seq_lens: list[int],
234 + num_heads: tuple[int, int],
235 + head_size: int,
236 + sliding_window: int | None,
237 + dtype: torch.dtype,
238 + block_size: int,
239 + isa: str,
240 +) -> None:
241 + set_random_seed(0)
242 + num_seqs = len(seq_lens)
243 + num_query_heads = num_heads[0]
244 + num_kv_heads = num_heads[1]
245 + assert num_query_heads % num_kv_heads == 0
246 + window_size = (
247 + (sliding_window - 1, sliding_window - 1)
248 + if sliding_window is not None
249 + else (-1, -1)
250 + )
251 + scale = head_size**-0.5
252 + token_num = sum(seq_lens)
253 +
254 + seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32)
255 + query_start_loc = torch.zeros(num_seqs, dtype=torch.int32)
256 + torch.cumsum(seq_lens_tensor[:-1], 0, out=query_start_loc[1:])
257 + block_nums = (seq_lens_tensor + block_size - 1) // block_size
258 + start_block_ids = torch.zeros_like(seq_lens_tensor)
259 + torch.cumsum(block_nums[:-1], 0, out=start_block_ids[1:])
260 + total_block_num: int = block_nums.sum().item()
261 + max_block_num = block_nums.max().item()
262 + block_offsets = torch.arange(0, max_block_num, dtype=torch.int32)
263 + encoder_block_table = start_block_ids[:, None] + block_offsets[None, :]
264 + slot_mapping_list = []
265 + slot_start_idx = 0
266 + for i in range(num_seqs):
267 + block_num = block_nums[i].item()
268 + seq_len = seq_lens[i]
269 + slot_mapping_list.append(torch.arange(slot_start_idx, slot_start_idx + seq_len))
270 + slot_start_idx += block_num * block_size
271 + slot_mapping = torch.cat(slot_mapping_list)
272 +
273 + query = tensor_cache(
274 + elem_num=token_num * num_query_heads * head_size,
275 + dtype=dtype,
276 + tag="query",
277 + )
278 + query = query.view(
279 + token_num,
280 + num_query_heads,
281 + head_size,
282 + )
283 +
284 + key_value = tensor_cache(
285 + elem_num=2 * token_num * num_kv_heads * head_size,
286 + dtype=dtype,
287 + tag="kv",
288 + )
289 + key_value = key_value.view(
290 + 2,
291 + token_num,
292 + num_kv_heads,
293 + head_size,
294 + )
295 + key, value = key_value.unbind(0)
296 +
297 + # KV cache for CPU attention
298 + packed_key_value_cache = torch.zeros(
299 + total_block_num, num_kv_heads, block_size, head_size * 2, dtype=dtype
300 + )
301 + packed_key_value_cache = packed_key_value_cache.view(
302 + (total_block_num, num_kv_heads, block_size * 2, -1)
303 + )
304 + packed_key_cache, packed_value_cache = packed_key_value_cache.chunk(2, dim=2)
305 +
306 + cu_query_lens = torch.tensor([0] + seq_lens, dtype=torch.int32).cumsum(
307 + dim=0, dtype=torch.int32
308 + )
309 + kv_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32)
310 +
311 + # use reshape_and_cache to pack key_cache and value_cache
312 + cpu_attn_reshape_and_cache(
313 + key=key.view(-1, num_kv_heads, head_size),
314 + value=value.view(-1, num_kv_heads, head_size),
315 + key_cache=packed_key_cache,
316 + value_cache=packed_value_cache,
317 + slot_mapping=slot_mapping,
318 + isa=isa,
319 + )
320 +
321 + metadata = cpu_attn_get_scheduler_metadata(
322 + num_reqs=num_seqs,
323 + num_heads=num_query_heads,
324 + num_kv_heads=num_kv_heads,
325 + head_dim=head_size,
326 + seq_lens=kv_lens_tensor,
327 + dtype=dtype,
328 + query_start_loc=cu_query_lens,
329 + causal=False,
330 + sliding_window_size=sliding_window if sliding_window is not None else -1,
331 + isa=isa,
332 + enable_kv_split=False,
333 + )
334 +
335 + out_without_split = torch.empty_like(query)
336 + cpu_attention_with_kv_cache(
337 + query=query,
338 + key_cache=packed_key_cache,
339 + value_cache=packed_value_cache,
340 + output=out_without_split,
341 + query_start_loc=cu_query_lens,
342 + seq_lens=kv_lens_tensor,
343 + scale=scale,
344 + causal=False,
345 + alibi_slopes=None,
346 + sliding_window=window_size,
347 + block_table=encoder_block_table,
348 + softcap=0,
349 + scheduler_metadata=metadata,
350 + s_aux=None,
351 + )
352 +
353 + metadata = cpu_attn_get_scheduler_metadata(
354 + num_reqs=num_seqs,
355 + num_heads=num_query_heads,
356 + num_kv_heads=num_kv_heads,
357 + head_dim=head_size,
358 + seq_lens=kv_lens_tensor,
359 + dtype=dtype,
360 + query_start_loc=cu_query_lens,
361 + causal=False,
362 + sliding_window_size=sliding_window if sliding_window is not None else -1,
363 + isa=isa,
364 + enable_kv_split=True,
365 + )
366 +
367 + out_with_split = torch.empty_like(query)
368 + cpu_attention_with_kv_cache(
369 + query=query,
370 + key_cache=packed_key_cache,
371 + value_cache=packed_value_cache,
372 + output=out_with_split,
373 + query_start_loc=cu_query_lens,
374 + seq_lens=kv_lens_tensor,
375 + scale=scale,
376 + causal=False,
377 + alibi_slopes=None,
378 + sliding_window=window_size,
379 + block_table=encoder_block_table,
380 + softcap=0,
381 + scheduler_metadata=metadata,
382 + s_aux=None,
383 + )
384 +
385 + ref_output = ref_varlen_encoder_attn(
386 + query=query,
387 + key=key,
388 + value=value,
389 + seq_lens=seq_lens,
390 + scale=scale,
391 + sliding_window=sliding_window,
392 + )
393 + atol, rtol = 1.5e-2, 1e-2
394 +
395 + (
396 + torch.testing.assert_close(out_with_split, ref_output, atol=atol, rtol=rtol),
397 + f"{torch.max(torch.abs(out_with_split - ref_output))}",
398 + )
399 + (
400 + torch.testing.assert_close(out_without_split, ref_output, atol=atol, rtol=rtol),
401 + f"{torch.max(torch.abs(out_without_split - ref_output))}",
402 + )
403
404
405 @torch.inference_mode()
@@ -418,6 +633,71 @@ def varlen_with_paged_kv(
633 )
634
635
636 +@pytest.mark.parametrize("seq_lens", ENCODER_SEQ_LENS)
637 +@pytest.mark.parametrize("num_heads", NUM_HEADS)
638 +@pytest.mark.parametrize("head_size", HEAD_SIZES)
639 +@pytest.mark.parametrize(
640 + "block_size",
641 + [
642 + 128,
643 + ],
644 +)
645 +@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
646 +@pytest.mark.parametrize("dtype", QTYPES)
647 +@pytest.mark.parametrize("isa", ["vec"])
648 +def test_varlen_encoder_attention_vec(
649 + seq_lens: list[int],
650 + num_heads: tuple[int, int],
651 + head_size: int,
652 + sliding_window: int | None,
653 + dtype: torch.dtype,
654 + block_size: int,
655 + isa: str,
656 +) -> None:
657 + varlen_encoder_attention(
658 + seq_lens=seq_lens,
659 + num_heads=num_heads,
660 + head_size=head_size,
661 + sliding_window=sliding_window,
662 + dtype=dtype,
663 + block_size=block_size,
664 + isa=isa,
665 + )
666 +
667 +
668 +@pytest.mark.parametrize("seq_lens", ENCODER_SEQ_LENS)
669 +@pytest.mark.parametrize("num_heads", NUM_HEADS)
670 +@pytest.mark.parametrize("head_size", HEAD_SIZES)
671 +@pytest.mark.parametrize(
672 + "block_size",
673 + [
674 + 128,
675 + ],
676 +)
677 +@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
678 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
679 +@pytest.mark.parametrize("isa", ["amx"])
680 +@pytest.mark.skipif(not torch.cpu._is_amx_tile_supported(), reason="no AMX support.")
681 +def test_varlen_encoder_attention_amx(
682 + seq_lens: list[int],
683 + num_heads: tuple[int, int],
684 + head_size: int,
685 + sliding_window: int | None,
686 + dtype: torch.dtype,
687 + block_size: int,
688 + isa: str,
689 +) -> None:
690 + varlen_encoder_attention(
691 + seq_lens=seq_lens,
692 + num_heads=num_heads,
693 + head_size=head_size,
694 + sliding_window=sliding_window,
695 + dtype=dtype,
696 + block_size=block_size,
697 + isa=isa,
698 + )
699 +
700 +
701 @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3", "fp8_e5m2"])
702 @pytest.mark.parametrize("seq_lens", SEQ_LENS)
703 @pytest.mark.parametrize("num_heads", NUM_HEADS)
vllm/v1/attention/backends/cpu_attn.py
+88 -202
@@ -11,7 +11,7 @@ import torch
11
12 from vllm import _custom_ops as ops
13 from vllm import envs
14 -from vllm.config import VllmConfig
14 +from vllm.config import VllmConfig, get_current_vllm_config
15 from vllm.logger import init_logger
16 from vllm.platforms import CpuArchEnum, current_platform
17 from vllm.utils.torch_utils import is_quantized_kv_cache
@@ -26,20 +26,15 @@ from vllm.v1.attention.backend import (
26 )
27 from vllm.v1.attention.backends.utils import (
28 KVCacheLayoutType,
29 - split_decodes_and_prefills,
29 )
31 -from vllm.v1.kv_cache_interface import AttentionSpec, CrossAttentionSpec
30 +from vllm.v1.kv_cache_interface import (
31 + AttentionSpec,
32 + CrossAttentionSpec,
33 + EncoderOnlyAttentionSpec,
34 +)
35
36 logger = init_logger(__name__)
37
35 -_CPU_ARCH_PREFER_MIXED_BATCH = (
36 - CpuArchEnum.X86,
37 - CpuArchEnum.ARM,
38 - CpuArchEnum.S390X,
39 - CpuArchEnum.RISCV,
40 - CpuArchEnum.POWERPC,
41 -)
42 -
38
39 class CPUAttentionBackend(AttentionBackend):
40 forward_includes_kv_cache_update: bool = False
@@ -124,6 +119,8 @@ class CPUAttentionMetadata:
119 sdpa_attn_masks: list[torch.Tensor | None] | None = None
120 sdpa_start_loc: torch.Tensor | None = None
121
122 + encoder_cache: torch.Tensor | None = None
123 +
124
125 class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]):
126 def __init__(
@@ -135,17 +132,6 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
132 ) -> None:
133 super().__init__(kv_cache_spec, layer_names, vllm_config, device)
134
138 - self.use_sdpa_prefill = False
139 - reorder_batch_threshold = None
140 - if current_platform.get_cpu_architecture() not in _CPU_ARCH_PREFER_MIXED_BATCH:
141 - # in this case, decode seqs are reordered to the front of prefill seqs
142 - # to split decode and prefill. Then use SDPA for prefill and
143 - # cpu_attention_with_kv_cache for decode
144 - reorder_batch_threshold = 1
145 - self.use_sdpa_prefill = True
146 -
147 - self._init_reorder_batch_threshold(reorder_batch_threshold, False)
148 -
135 self.kv_cache_spec = kv_cache_spec
136 self.vllm_config = vllm_config
137
@@ -168,6 +154,9 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
154 kv_cache_dtype_str,
155 )
156 self.is_cross_attention = isinstance(kv_cache_spec, CrossAttentionSpec)
157 + self.is_encoder_only_attention = isinstance(
158 + kv_cache_spec, EncoderOnlyAttentionSpec
159 + )
160
161 def build(
162 self,
@@ -185,23 +174,34 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
174 slot_mapping = common_attn_metadata.slot_mapping
175 causal = False if self.is_cross_attention else common_attn_metadata.causal
176
188 - sdpa_start_loc = query_start_loc
189 - num_decode_tokens = 0
190 - if self.use_sdpa_prefill and causal:
191 - # Decoder, need reorder and truncate
192 - assert self.reorder_batch_threshold
193 - (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens) = (
194 - split_decodes_and_prefills(
195 - common_attn_metadata,
196 - decode_threshold=self.reorder_batch_threshold,
197 - require_uniform=True,
198 - )
177 + encoder_cache_tensor = None
178 + if self.is_encoder_only_attention:
179 + block_nums = (seq_lens + self.block_size - 1) // self.block_size
180 + start_block_ids = torch.zeros_like(seq_lens)
181 + torch.cumsum(block_nums[:-1], 0, out=start_block_ids[1:])
182 + total_block_num: int = block_nums.sum().item()
183 + max_block_num = block_nums.max().item()
184 + block_offsets = torch.arange(
185 + 0, max_block_num, dtype=block_table_tensor.dtype
186 + )
187 + encoder_block_table = start_block_ids[:, None] + block_offsets[None, :]
188 + torch.ops._C.compute_slot_mapping_kernel_impl(
189 + query_start_loc,
190 + common_attn_metadata.positions,
191 + encoder_block_table,
192 + slot_mapping,
193 + self.block_size,
194 + )
195 + encoder_cache_tensor = torch.zeros(
196 + (
197 + total_block_num,
198 + self.num_kv_heads,
199 + self.block_size,
200 + 2 * self.head_dim,
201 + ),
202 + dtype=self.dtype,
203 )
200 - num_reqs = num_decodes
201 - sdpa_start_loc = sdpa_start_loc[num_decodes:] - num_decode_tokens
202 - seq_lens = seq_lens[:num_decodes]
203 - query_start_loc = query_start_loc[: num_decodes + 1]
204 - block_table_tensor = block_table_tensor[:num_decodes]
204 + block_table_tensor = encoder_block_table
205
206 scheduler_metadata = ops.cpu_attn_get_scheduler_metadata(
207 num_reqs=num_reqs,
@@ -227,9 +227,7 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
227 slot_mapping=slot_mapping,
228 scheduler_metadata=scheduler_metadata,
229 causal=causal,
230 - use_sdpa_prefill=self.use_sdpa_prefill,
231 - num_decode_tokens=num_decode_tokens,
232 - sdpa_start_loc=sdpa_start_loc,
230 + encoder_cache=encoder_cache_tensor,
231 )
232
233 return attn_metadata
@@ -289,6 +287,14 @@ class CPUAttentionBackendImpl(AttentionImpl):
287 "heads in the layer"
288 )
289
290 + vllm_config = get_current_vllm_config()
291 + self.isa = _get_attn_isa(
292 + vllm_config.model_config.dtype,
293 + vllm_config.cache_config.block_size,
294 + self.head_size,
295 + self.kv_cache_dtype,
296 + )
297 +
298 def forward(
299 self,
300 layer: AttentionLayer,
@@ -325,60 +331,58 @@ class CPUAttentionBackendImpl(AttentionImpl):
331
332 num_actual_tokens = attn_metadata.num_actual_tokens
333
328 - # Handle encoder attention differently - no KV cache needed
334 + # For encoder attention
335 if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
336 # For encoder attention,
331 - return self._run_sdpa_forward(
332 - query[:num_actual_tokens],
333 - key[:num_actual_tokens],
334 - value[:num_actual_tokens],
335 - output[:num_actual_tokens],
336 - attn_metadata,
337 - self.attn_type,
338 - )
337 + kv_cache = attn_metadata.encoder_cache
338
340 - # For decoder and cross-attention, use KV cache, size are
341 - # [num_blocks, num_kv_heads, block_size, 2 * head_size]
342 - # Make a view [num_blocks, num_kv_heads, block_size * 2, head_size]
343 - # Then slice KV at dim 2
339 + # KV cache size are [num_blocks, num_kv_heads, block_size,
340 + # 2 * head_size]. Make a view [num_blocks, num_kv_heads,
341 + # block_size * 2, head_size]. Then slice KV at dim 2
342 num_blocks, num_kv_heads, block_size, _ = kv_cache.size()
343 kv_cache = kv_cache.view((num_blocks, num_kv_heads, block_size * 2, -1))
344 key_cache, value_cache = kv_cache.chunk(2, dim=2)
345
348 - if attn_metadata.use_sdpa_prefill:
349 - assert self.sinks is None, "Attention sink is unsupported in SDPA prefill"
350 - num_decode_tokens = attn_metadata.num_decode_tokens
351 - self._run_sdpa_forward(
352 - query[num_decode_tokens:num_actual_tokens],
353 - key[num_decode_tokens:num_actual_tokens],
354 - value[num_decode_tokens:num_actual_tokens],
355 - output[num_decode_tokens:num_actual_tokens],
356 - attn_metadata,
357 - self.attn_type,
358 - )
359 - num_actual_tokens = num_decode_tokens
360 -
361 - if num_actual_tokens > 0:
362 - ops.cpu_attention_with_kv_cache(
363 - query=query[:num_actual_tokens],
364 - key_cache=key_cache,
365 - value_cache=value_cache,
366 - output=output[:num_actual_tokens], # type: ignore
367 - query_start_loc=attn_metadata.query_start_loc,
368 - seq_lens=attn_metadata.seq_lens,
369 - scale=self.scale,
370 - causal=attn_metadata.causal,
371 - alibi_slopes=self.alibi_slopes, # type: ignore
372 - sliding_window=self.sliding_window,
373 - block_table=attn_metadata.block_table,
374 - softcap=self.logits_soft_cap,
375 - scheduler_metadata=attn_metadata.scheduler_metadata,
376 - s_aux=self.sinks,
346 + # key and value may be None in the case of cross attention. They are
347 + # calculated once based on the output from the encoder and then cached
348 + # in KV cache.
349 + if (
350 + self.kv_sharing_target_layer_name is None
351 + and key is not None
352 + and value is not None
353 + ):
354 + ops.cpu_attn_reshape_and_cache(
355 + key,
356 + value,
357 + key_cache,
358 + value_cache,
359 + attn_metadata.slot_mapping,
360 + self.isa,
361 k_scale=layer._k_scale_float,
362 v_scale=layer._v_scale_float,
363 kv_cache_dtype=self.kv_cache_dtype,
364 )
365
366 + ops.cpu_attention_with_kv_cache(
367 + query=query[:num_actual_tokens],
368 + key_cache=key_cache,
369 + value_cache=value_cache,
370 + output=output[:num_actual_tokens], # type: ignore
371 + query_start_loc=attn_metadata.query_start_loc,
372 + seq_lens=attn_metadata.seq_lens,
373 + scale=self.scale,
374 + causal=attn_metadata.causal,
375 + alibi_slopes=self.alibi_slopes, # type: ignore
376 + sliding_window=self.sliding_window,
377 + block_table=attn_metadata.block_table,
378 + softcap=self.logits_soft_cap,
379 + scheduler_metadata=attn_metadata.scheduler_metadata,
380 + s_aux=self.sinks,
381 + k_scale=layer._k_scale_float,
382 + v_scale=layer._v_scale_float,
383 + kv_cache_dtype=self.kv_cache_dtype,
384 + )
385 +
386 return output
387
388 def do_kv_cache_update(
@@ -395,136 +399,18 @@ class CPUAttentionBackendImpl(AttentionImpl):
399 num_blocks, num_kv_heads, block_size, _ = kv_cache.size()
400 kv_cache = kv_cache.view((num_blocks, num_kv_heads, block_size * 2, -1))
401 key_cache, value_cache = kv_cache.chunk(2, dim=2)
398 - isa = _get_attn_isa(
399 - key.dtype, key_cache.shape[2], self.head_size, self.kv_cache_dtype
400 - )
402 ops.cpu_attn_reshape_and_cache(
403 key,
404 value,
405 key_cache,
406 value_cache,
407 slot_mapping,
407 - isa,
408 + self.isa,
409 k_scale=layer._k_scale_float,
410 v_scale=layer._v_scale_float,
411 kv_cache_dtype=self.kv_cache_dtype,
412 )
413
413 - def _run_sdpa_forward(
414 - self,
415 - query: torch.Tensor,
416 - key: torch.Tensor,
417 - value: torch.Tensor,
418 - output: torch.Tensor,
419 - attn_metadata: CPUAttentionMetadata,
420 - attn_type: str,
421 - ) -> torch.Tensor:
422 - attn_masks = attn_metadata.sdpa_attn_masks
423 - if attn_masks is None:
424 - if self.alibi_slopes is not None:
425 - attn_masks = _make_alibi_bias(
426 - self.alibi_slopes,
427 - query.dtype,
428 - attn_metadata.sdpa_start_loc,
429 - )
430 - elif self.sliding_window[0] != -1 or self.sliding_window[1] != -1:
431 - assert attn_metadata.seq_lens is not None
432 - attn_masks = _make_sliding_window_bias(
433 - attn_metadata.sdpa_start_loc,
434 - self.sliding_window[0],
435 - self.sliding_window[1],
436 - query.dtype,
437 - )
438 - else:
439 - attn_masks = [None] * (attn_metadata.sdpa_start_loc.size(0) - 1) # type: ignore
440 - attn_metadata.sdpa_attn_masks = attn_masks
441 -
442 - query = query.movedim(0, query.dim() - 2)
443 - key = key.movedim(0, key.dim() - 2)
444 - value = value.movedim(0, value.dim() - 2)
445 -
446 - causal_attn = attn_type == AttentionType.DECODER
447 -
448 - sdpa_start_loc = attn_metadata.sdpa_start_loc.numpy() # type: ignore
449 - for i in range(len(attn_masks)):
450 - mask = attn_masks[i]
451 - start_q = sdpa_start_loc[i]
452 - end_q = sdpa_start_loc[i + 1]
453 - sub_out = (
454 - torch.nn.functional.scaled_dot_product_attention(
455 - query[None, :, start_q:end_q, :],
456 - key[None, :, start_q:end_q, :],
457 - value[None, :, start_q:end_q, :],
458 - attn_mask=mask,
459 - dropout_p=0.0,
460 - is_causal=causal_attn and mask is None,
461 - scale=self.scale,
462 - enable_gqa=self.num_heads > self.num_kv_heads,
463 - )
464 - .squeeze(0)
465 - .movedim(query.dim() - 2, 0)
466 - )
467 - output[start_q:end_q, :, :] = sub_out
468 - return output
469 -
470 -
471 -def _make_alibi_bias(
472 - alibi_slopes: torch.Tensor,
473 - dtype: torch.dtype,
474 - sdpa_start_loc: torch.Tensor,
475 -) -> list[torch.Tensor]:
476 - attn_biases: list[torch.Tensor] = []
477 - seq_num = sdpa_start_loc.size(0) - 1
478 - sdpa_start_loc = sdpa_start_loc.numpy() # type: ignore
479 - for i in range(seq_num):
480 - seq_len = sdpa_start_loc[i + 1] - sdpa_start_loc[i]
481 - bias = torch.arange(seq_len, dtype=dtype) # type: ignore
482 - # NOTE(zhuohan): HF uses
483 - # `bias = bias[None, :].repeat(seq_len, 1)`
484 - # here. We find that both biases give the same results, but
485 - # the bias below more accurately follows the original ALiBi
486 - # paper.
487 - bias = bias[None, :] - bias[:, None]
488 -
489 - num_heads = alibi_slopes.shape[0]
490 - bias = bias[None, :].repeat((num_heads, 1, 1))
491 - bias.mul_(alibi_slopes[:, None, None]).unsqueeze_(0)
492 - inf_mask = (
493 - torch.empty((1, seq_len, seq_len), dtype=bias.dtype) # type: ignore
494 - .fill_(-torch.inf)
495 - .triu_(diagonal=1)
496 - )
497 - attn_biases.append((bias + inf_mask).to(dtype))
498 -
499 - return attn_biases
500 -
501 -
502 -def _make_sliding_window_bias(
503 - sdpa_start_loc: torch.Tensor,
504 - left_window_size: int,
505 - right_window_size: int,
506 - dtype: torch.dtype,
507 -) -> list[torch.Tensor]:
508 - attn_biases: list[torch.Tensor] = []
509 - seq_num = sdpa_start_loc.size(0) - 1
510 - sdpa_start_loc = sdpa_start_loc.numpy() # type: ignore
511 - for i in range(seq_num):
512 - seq_len = sdpa_start_loc[i + 1] - sdpa_start_loc[i]
513 - mask = torch.full( # type: ignore
514 - (1, seq_len, seq_len), # type: ignore
515 - fill_value=1,
516 - dtype=dtype,
517 - )
518 -
519 - if right_window_size != -1:
520 - mask = torch.tril(mask, diagonal=right_window_size)
521 - if left_window_size != -1:
522 - mask = torch.triu(mask, diagonal=-left_window_size)
523 - mask = torch.log(mask)
524 - attn_biases.append(mask)
525 -
526 - return attn_biases
527 -
414
415 @functools.lru_cache(maxsize=1)
416 def _riscv_supports_rvv() -> bool: