[CPU] Adds support for fp32 attention sinks (#56252)
Signed-off-by: Ankit Jaiswal <ankit.jaiswal@amd.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>
Ankit-Jaiswal-AMD committed
Sep 23, 2026 at 19:25 UTC
88afb9dcac7672f8e4a16590e67c2f09fe173dff
4 files changed
+95
-13
csrc/cpu/cpu_attn.cpp
+13
-1
@@ -243,7 +243,19 @@ void cpu_attention_with_kv_cache(
243
input.block_table = block_table.data_ptr<int32_t>();
244
input.alibi_slopes =
245
alibi_slopes.has_value() ? alibi_slopes->data_ptr<float>() : nullptr;
246
- input.s_aux = s_aux.has_value() ? s_aux->data_ptr<c10::BFloat16>() : nullptr;
246
+ // Attention sinks may be bf16 (native path) or fp32. Anything that is not
247
+ // bf16 must be provided as fp32 and is executed in full float precision.
248
+ if (s_aux.has_value()) {
249
+ TORCH_CHECK(s_aux->scalar_type() == at::ScalarType::BFloat16 ||
250
+ s_aux->scalar_type() == at::ScalarType::Float,
251
+ "cpu_attention_with_kv_cache: s_aux (attention sinks) dtype ",
252
+ s_aux->scalar_type(), " must be bfloat16 or float32");
253
+ input.s_aux = s_aux->data_ptr();
254
+ input.s_aux_is_bf16 = s_aux->scalar_type() == at::ScalarType::BFloat16;
255
+ } else {
256
+ input.s_aux = nullptr;
257
+ input.s_aux_is_bf16 = false;
258
+ }
259
input.dynamic_causal =
260
dynamic_causal.has_value() ? dynamic_causal->data_ptr<bool>() : nullptr;
261
input.scale = scale;
csrc/cpu/cpu_attn_impl.hpp
+24
-9
@@ -879,7 +879,11 @@ struct AttentionInput {
879
int32_t* seq_lens;
880
int32_t* block_table;
881
float* alibi_slopes;
882
- c10::BFloat16* s_aux;
882
+ // Attention sinks pointer. May reference bf16 or fp32 data depending on
883
+ // `s_aux_is_bf16`. bf16 sinks keep the native bf16 path; anything else is
884
+ // provided as fp32 and executed in full float precision.
885
+ const void* s_aux;
886
+ bool s_aux_is_bf16;
887
bool* dynamic_causal;
888
float scale;
889
bool causal;
@@ -1565,7 +1569,8 @@ class AttentionMainLoop {
1569
const float scale = input->scale;
1570
const float softcap_scale = input->softcap;
1571
const float* alibi_slopes = input->alibi_slopes;
1568
- const c10::BFloat16* s_aux = input->s_aux;
1572
+ const void* s_aux = input->s_aux;
1573
+ const bool s_aux_is_bf16 = input->s_aux_is_bf16;
1574
const bool* dynamic_causal = input->dynamic_causal;
1575
const bool is_dynamic_causal = dynamic_causal != nullptr;
1576
@@ -1773,9 +1778,6 @@ class AttentionMainLoop {
1778
const float* curr_alibi_slopes =
1779
(alibi_slopes != nullptr ? alibi_slopes + q_head_start_idx
1780
: nullptr);
1776
- const c10::BFloat16* curr_s_aux =
1777
- (s_aux != nullptr ? s_aux + q_head_start_idx : nullptr);
1778
-
1781
// copy the Q tile to q_buffer, the logical layout of q_buffer is
1782
// [actual_q_token_num, curr_q_heads_per_kv, head_dim]
1783
{
@@ -1787,10 +1789,23 @@ class AttentionMainLoop {
1789
1790
if (use_sink) {
1791
alignas(64) float s_aux_fp32[16];
1790
- // All other platforms have BF16Vec16 available
1791
- vec_op::BF16Vec16 vec_bf16(curr_s_aux);
1792
- vec_op::FP32Vec16 vec_fp32(vec_bf16);
1793
- vec_fp32.save(s_aux_fp32);
1792
+ // Sinks may be provided as bf16 or fp32. Either way they are
1793
+ // loaded into fp32: bf16 sinks are up-converted (native
1794
+ // BF16Vec16 path, available on all platforms), while fp32
1795
+ // sinks are loaded directly and executed in full precision.
1796
+ if (s_aux_is_bf16) {
1797
+ const c10::BFloat16* curr_s_aux =
1798
+ static_cast<const c10::BFloat16*>(s_aux) +
1799
+ q_head_start_idx;
1800
+ vec_op::BF16Vec16 vec_bf16(curr_s_aux);
1801
+ vec_op::FP32Vec16 vec_fp32(vec_bf16);
1802
+ vec_fp32.save(s_aux_fp32);
1803
+ } else {
1804
+ const float* curr_s_aux =
1805
+ static_cast<const float*>(s_aux) + q_head_start_idx;
1806
+ vec_op::FP32Vec16 vec_fp32(curr_s_aux);
1807
+ vec_fp32.save(s_aux_fp32);
1808
+ }
1809
1810
float* __restrict__ curr_sum_buffer = sum_buffer;
1811
float* __restrict__ curr_max_buffer = max_buffer;
tests/kernels/attention/test_cpu_attn.py
+49
-3
@@ -428,6 +428,7 @@ def varlen_with_paged_kv(
428
k_scale: float = 1.0,
429
v_scale: float = 1.0,
430
dynamic_causal: list[bool] | None = None,
431
+ s_aux_dtype: torch.dtype = torch.bfloat16,
432
) -> None:
433
set_random_seed(0)
434
num_seqs = len(seq_lens)
@@ -449,9 +450,7 @@ def varlen_with_paged_kv(
450
# 2^(-8/n)
451
alibi_slopes = _get_alibi_slopes(num_query_heads) if use_alibi else None
452
452
- s_aux = (
453
- 15 * torch.rand((num_query_heads,), dtype=torch.bfloat16) if use_sink else None
454
- )
453
+ s_aux = 15 * torch.rand((num_query_heads,), dtype=s_aux_dtype) if use_sink else None
454
455
is_fp8 = kv_cache_dtype != "auto"
456
if is_fp8 and current_platform.get_cpu_architecture() != CpuArchEnum.X86:
@@ -538,6 +537,8 @@ def varlen_with_paged_kv(
537
)
538
539
out_without_split = torch.empty_like(query)
540
+ if s_aux is not None and s_aux.dtype != torch.bfloat16:
541
+ s_aux = s_aux.to(torch.float32)
542
cpu_attention_with_kv_cache(
543
query=query,
544
key_cache=packed_key_cache,
@@ -1128,6 +1129,51 @@ def test_varlen_with_paged_kv_sink(
1129
)
1130
1131
1132
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3"])
1133
+@pytest.mark.parametrize("seq_lens", SEQ_LENS)
1134
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
1135
+@pytest.mark.parametrize("head_size", [96])
1136
+@pytest.mark.parametrize("block_size", [128])
1137
+@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
1138
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
1139
+@pytest.mark.parametrize("soft_cap", [None])
1140
+@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
1141
+@pytest.mark.parametrize("use_alibi", [False])
1142
+@pytest.mark.parametrize("use_sink", [True])
1143
+@pytest.mark.parametrize("isa", [get_attn_isa()])
1144
+@pytest.mark.parametrize("s_aux_dtype", [torch.float16])
1145
+def test_varlen_with_paged_kv_sink_fp16(
1146
+ seq_lens: list[tuple[int, int]],
1147
+ num_heads: tuple[int, int],
1148
+ head_size: int,
1149
+ sliding_window: int | None,
1150
+ dtype: torch.dtype,
1151
+ block_size: int,
1152
+ soft_cap: float | None,
1153
+ num_blocks: int,
1154
+ use_alibi: bool,
1155
+ use_sink: bool,
1156
+ isa: str,
1157
+ kv_cache_dtype: str,
1158
+ s_aux_dtype: torch.dtype,
1159
+) -> None:
1160
+ varlen_with_paged_kv(
1161
+ seq_lens=seq_lens,
1162
+ num_heads=num_heads,
1163
+ head_size=head_size,
1164
+ sliding_window=sliding_window,
1165
+ dtype=dtype,
1166
+ block_size=block_size,
1167
+ soft_cap=soft_cap,
1168
+ num_blocks=num_blocks,
1169
+ use_alibi=use_alibi,
1170
+ use_sink=use_sink,
1171
+ isa=isa,
1172
+ kv_cache_dtype=kv_cache_dtype,
1173
+ s_aux_dtype=s_aux_dtype,
1174
+ )
1175
+
1176
+
1177
@pytest.mark.parametrize(
1178
"kv_cache_dtype",
1179
[
vllm/v1/attention/backends/cpu_attn.py
+9
@@ -406,6 +406,15 @@ class CPUAttentionBackendImpl(AttentionImpl):
406
kv_cache_dtype=self.kv_cache_dtype,
407
)
408
409
+ # The CPU kernel executes attention sinks natively in bf16. If the
410
+ # sinks tensor is anything other than bf16, cast it to fp32 so it is
411
+ # executed in full float precision (done lazily here, after weights
412
+ # are loaded, rather than at __init__ time).
413
+ if self.sinks is not None and self.sinks.dtype not in [
414
+ torch.bfloat16,
415
+ torch.float32,
416
+ ]:
417
+ self.sinks = self.sinks.to(torch.float32)
418
ops.cpu_attention_with_kv_cache(
419
query=query[:num_actual_tokens],
420
key_cache=key_cache,