[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>

csrc/cpu/cpu_attn.cpp

@@ -243,7 +243,19 @@ void cpu_attention_with_kv_cache(
   input.block_table = block_table.data_ptr<int32_t>();
   input.alibi_slopes =
       alibi_slopes.has_value() ? alibi_slopes->data_ptr<float>() : nullptr;
-  input.s_aux = s_aux.has_value() ? s_aux->data_ptr<c10::BFloat16>() : 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<bool>() : nullptr;
   input.scale = scale;

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<const c10::BFloat16*>(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<const float*>(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;

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",
     [

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,