From 88afb9dcac7672f8e4a16590e67c2f09fe173dff Mon Sep 17 00:00:00 2001 From: Ankit-Jaiswal-AMD Date: Wed, 23 Sep 2026 19:25:09 +0530 Subject: [PATCH] [CPU] Adds support for fp32 attention sinks (#56252) Signed-off-by: Ankit Jaiswal Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> --- csrc/cpu/cpu_attn.cpp | 14 ++++++- csrc/cpu/cpu_attn_impl.hpp | 33 +++++++++++---- tests/kernels/attention/test_cpu_attn.py | 52 ++++++++++++++++++++++-- vllm/v1/attention/backends/cpu_attn.py | 9 ++++ 4 files changed, 95 insertions(+), 13 deletions(-) diff --git a/csrc/cpu/cpu_attn.cpp b/csrc/cpu/cpu_attn.cpp index ea0b4e49b2..13e430fb41 100644 --- a/csrc/cpu/cpu_attn.cpp +++ b/csrc/cpu/cpu_attn.cpp @@ -243,7 +243,19 @@ void cpu_attention_with_kv_cache( input.block_table = block_table.data_ptr(); input.alibi_slopes = alibi_slopes.has_value() ? alibi_slopes->data_ptr() : nullptr; - input.s_aux = s_aux.has_value() ? s_aux->data_ptr() : nullptr; + // Attention sinks may be bf16 (native path) or fp32. Anything that is not + // bf16 must be provided as fp32 and is executed in full float precision. + if (s_aux.has_value()) { + TORCH_CHECK(s_aux->scalar_type() == at::ScalarType::BFloat16 || + s_aux->scalar_type() == at::ScalarType::Float, + "cpu_attention_with_kv_cache: s_aux (attention sinks) dtype ", + s_aux->scalar_type(), " must be bfloat16 or float32"); + input.s_aux = s_aux->data_ptr(); + input.s_aux_is_bf16 = s_aux->scalar_type() == at::ScalarType::BFloat16; + } else { + input.s_aux = nullptr; + input.s_aux_is_bf16 = false; + } input.dynamic_causal = dynamic_causal.has_value() ? dynamic_causal->data_ptr() : nullptr; input.scale = scale; diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp index 4ade47eff3..d914e15063 100644 --- a/csrc/cpu/cpu_attn_impl.hpp +++ b/csrc/cpu/cpu_attn_impl.hpp @@ -879,7 +879,11 @@ struct AttentionInput { int32_t* seq_lens; int32_t* block_table; float* alibi_slopes; - c10::BFloat16* s_aux; + // Attention sinks pointer. May reference bf16 or fp32 data depending on + // `s_aux_is_bf16`. bf16 sinks keep the native bf16 path; anything else is + // provided as fp32 and executed in full float precision. + const void* s_aux; + bool s_aux_is_bf16; bool* dynamic_causal; float scale; bool causal; @@ -1565,7 +1569,8 @@ class AttentionMainLoop { const float scale = input->scale; const float softcap_scale = input->softcap; const float* alibi_slopes = input->alibi_slopes; - const c10::BFloat16* s_aux = input->s_aux; + const void* s_aux = input->s_aux; + const bool s_aux_is_bf16 = input->s_aux_is_bf16; const bool* dynamic_causal = input->dynamic_causal; const bool is_dynamic_causal = dynamic_causal != nullptr; @@ -1773,9 +1778,6 @@ class AttentionMainLoop { const float* curr_alibi_slopes = (alibi_slopes != nullptr ? alibi_slopes + q_head_start_idx : nullptr); - const c10::BFloat16* curr_s_aux = - (s_aux != nullptr ? s_aux + q_head_start_idx : nullptr); - // copy the Q tile to q_buffer, the logical layout of q_buffer is // [actual_q_token_num, curr_q_heads_per_kv, head_dim] { @@ -1787,10 +1789,23 @@ class AttentionMainLoop { if (use_sink) { alignas(64) float s_aux_fp32[16]; - // All other platforms have BF16Vec16 available - vec_op::BF16Vec16 vec_bf16(curr_s_aux); - vec_op::FP32Vec16 vec_fp32(vec_bf16); - vec_fp32.save(s_aux_fp32); + // Sinks may be provided as bf16 or fp32. Either way they are + // loaded into fp32: bf16 sinks are up-converted (native + // BF16Vec16 path, available on all platforms), while fp32 + // sinks are loaded directly and executed in full precision. + if (s_aux_is_bf16) { + const c10::BFloat16* curr_s_aux = + static_cast(s_aux) + + q_head_start_idx; + vec_op::BF16Vec16 vec_bf16(curr_s_aux); + vec_op::FP32Vec16 vec_fp32(vec_bf16); + vec_fp32.save(s_aux_fp32); + } else { + const float* curr_s_aux = + static_cast(s_aux) + q_head_start_idx; + vec_op::FP32Vec16 vec_fp32(curr_s_aux); + vec_fp32.save(s_aux_fp32); + } float* __restrict__ curr_sum_buffer = sum_buffer; float* __restrict__ curr_max_buffer = max_buffer; diff --git a/tests/kernels/attention/test_cpu_attn.py b/tests/kernels/attention/test_cpu_attn.py index 7c4b7811c7..3046e2d2ac 100644 --- a/tests/kernels/attention/test_cpu_attn.py +++ b/tests/kernels/attention/test_cpu_attn.py @@ -428,6 +428,7 @@ def varlen_with_paged_kv( k_scale: float = 1.0, v_scale: float = 1.0, dynamic_causal: list[bool] | None = None, + s_aux_dtype: torch.dtype = torch.bfloat16, ) -> None: set_random_seed(0) num_seqs = len(seq_lens) @@ -449,9 +450,7 @@ def varlen_with_paged_kv( # 2^(-8/n) alibi_slopes = _get_alibi_slopes(num_query_heads) if use_alibi else None - s_aux = ( - 15 * torch.rand((num_query_heads,), dtype=torch.bfloat16) if use_sink else None - ) + s_aux = 15 * torch.rand((num_query_heads,), dtype=s_aux_dtype) if use_sink else None is_fp8 = kv_cache_dtype != "auto" if is_fp8 and current_platform.get_cpu_architecture() != CpuArchEnum.X86: @@ -538,6 +537,8 @@ def varlen_with_paged_kv( ) out_without_split = torch.empty_like(query) + if s_aux is not None and s_aux.dtype != torch.bfloat16: + s_aux = s_aux.to(torch.float32) cpu_attention_with_kv_cache( query=query, key_cache=packed_key_cache, @@ -1128,6 +1129,51 @@ def test_varlen_with_paged_kv_sink( ) +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3"]) +@pytest.mark.parametrize("seq_lens", SEQ_LENS) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +@pytest.mark.parametrize("head_size", [96]) +@pytest.mark.parametrize("block_size", [128]) +@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("soft_cap", [None]) +@pytest.mark.parametrize("num_blocks", NUM_BLOCKS) +@pytest.mark.parametrize("use_alibi", [False]) +@pytest.mark.parametrize("use_sink", [True]) +@pytest.mark.parametrize("isa", [get_attn_isa()]) +@pytest.mark.parametrize("s_aux_dtype", [torch.float16]) +def test_varlen_with_paged_kv_sink_fp16( + seq_lens: list[tuple[int, int]], + num_heads: tuple[int, int], + head_size: int, + sliding_window: int | None, + dtype: torch.dtype, + block_size: int, + soft_cap: float | None, + num_blocks: int, + use_alibi: bool, + use_sink: bool, + isa: str, + kv_cache_dtype: str, + s_aux_dtype: torch.dtype, +) -> None: + varlen_with_paged_kv( + seq_lens=seq_lens, + num_heads=num_heads, + head_size=head_size, + sliding_window=sliding_window, + dtype=dtype, + block_size=block_size, + soft_cap=soft_cap, + num_blocks=num_blocks, + use_alibi=use_alibi, + use_sink=use_sink, + isa=isa, + kv_cache_dtype=kv_cache_dtype, + s_aux_dtype=s_aux_dtype, + ) + + @pytest.mark.parametrize( "kv_cache_dtype", [ diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index 24be5b8ec9..36df340967 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -406,6 +406,15 @@ class CPUAttentionBackendImpl(AttentionImpl): kv_cache_dtype=self.kv_cache_dtype, ) + # The CPU kernel executes attention sinks natively in bf16. If the + # sinks tensor is anything other than bf16, cast it to fp32 so it is + # executed in full float precision (done lazily here, after weights + # are loaded, rather than at __init__ time). + if self.sinks is not None and self.sinks.dtype not in [ + torch.bfloat16, + torch.float32, + ]: + self.sinks = self.sinks.to(torch.float32) ops.cpu_attention_with_kv_cache( query=query[:num_actual_tokens], key_cache=key_cache,