[CPU] Fix FP8 attention scratchpad sizing (#50194)
Signed-off-by: Li, Tianmu <tianmu.li@intel.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Tianmu Li committed
Jul 29, 2026 at 00:34 UTC
65a1a16594995bdc8b8feae2bf2ecf208d585be8
5 files changed
+35
-7
csrc/cpu/cpu_attn.cpp
+6
-3
@@ -30,7 +30,8 @@ torch::Tensor get_scheduler_metadata(
30
const torch::Tensor& query_start_loc, const bool causal,
31
const int64_t window_size, const std::string& isa_hint,
32
const bool enable_kv_split,
33
- const std::optional<torch::Tensor>& dynamic_causal) {
33
+ const std::optional<torch::Tensor>& dynamic_causal,
34
+ const std::string& kv_cache_dtype) {
35
cpu_attention::ISA isa;
36
if (isa_hint == "amx") {
37
isa = cpu_attention::ISA::AMX;
@@ -65,9 +66,11 @@ torch::Tensor get_scheduler_metadata(
66
input.dynamic_causal =
67
dynamic_causal.has_value() ? dynamic_causal->data_ptr<bool>() : nullptr;
68
69
+ const int64_t kv_cache_idx =
70
+ static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype));
71
VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() {
69
- CPU_ATTN_DISPATCH(head_dim, isa, 0, [&]() {
70
- input.elem_size = sizeof(scalar_t);
72
+ CPU_ATTN_DISPATCH(head_dim, isa, kv_cache_idx, [&]() {
73
+ input.elem_size = sizeof(attn_impl::kv_cache_t);
74
input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t);
75
input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t);
76
input.output_buffer_elem_size =
csrc/cpu/torch_bindings.cpp
+4
-2
@@ -163,7 +163,8 @@ torch::Tensor get_scheduler_metadata(
163
const torch::Tensor& query_start_loc, const bool casual,
164
const int64_t window_size, const std::string& isa_hint,
165
const bool enable_kv_split,
166
- const std::optional<torch::Tensor>& dynamic_causal);
166
+ const std::optional<torch::Tensor>& dynamic_causal,
167
+ const std::string& kv_cache_dtype);
168
169
void cpu_attn_reshape_and_cache(const torch::Tensor& key,
170
const torch::Tensor& value,
@@ -577,7 +578,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
578
"get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, "
579
"int head_dim, Tensor seq_lens, ScalarType dtype, Tensor "
580
"query_start_loc, bool casual, int window_size, str isa_hint, bool "
580
- "enable_kv_split, Tensor? dynamic_causal) -> Tensor",
581
+ "enable_kv_split, Tensor? dynamic_causal, "
582
+ "str kv_cache_dtype=\"auto\") -> Tensor",
583
&get_scheduler_metadata);
584
ops.def(
585
"cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) "
tests/kernels/attention/test_cpu_attn.py
+20
@@ -534,6 +534,7 @@ def varlen_with_paged_kv(
534
isa=isa,
535
enable_kv_split=False,
536
dynamic_causal=dynamic_causal_tensor,
537
+ kv_cache_dtype=kv_cache_dtype,
538
)
539
540
out_without_split = torch.empty_like(query)
@@ -569,6 +570,7 @@ def varlen_with_paged_kv(
570
isa=isa,
571
enable_kv_split=True,
572
dynamic_causal=dynamic_causal_tensor,
573
+ kv_cache_dtype=kv_cache_dtype,
574
)
575
576
out_with_split = torch.empty_like(query)
@@ -803,6 +805,24 @@ def test_varlen_with_paged_kv_normal_amx(
805
)
806
807
808
+@pytest.mark.skipif(not torch.cpu._is_amx_tile_supported(), reason="no AMX support.")
809
+def test_varlen_with_paged_kv_fp8_large_prefill_amx() -> None:
810
+ varlen_with_paged_kv(
811
+ seq_lens=[(1024, 1024)] * 4,
812
+ num_heads=(16, 2),
813
+ head_size=256,
814
+ sliding_window=None,
815
+ dtype=torch.bfloat16,
816
+ block_size=2176,
817
+ soft_cap=None,
818
+ num_blocks=4,
819
+ use_alibi=False,
820
+ use_sink=False,
821
+ isa="amx",
822
+ kv_cache_dtype="fp8_e4m3",
823
+ )
824
+
825
+
826
@pytest.mark.parametrize("seq_lens", SEQ_LENS)
827
@pytest.mark.parametrize("num_heads", NUM_HEADS)
828
@pytest.mark.parametrize("head_size", HEAD_SIZES_VEC16)
vllm/_custom_ops.py
+2
@@ -3814,6 +3814,7 @@ def cpu_attn_get_scheduler_metadata(
3814
isa: str,
3815
enable_kv_split: bool,
3816
dynamic_causal: torch.Tensor | None = None,
3817
+ kv_cache_dtype: str = "auto",
3818
) -> torch.Tensor:
3819
scheduler_metadata = torch.ops._C.get_scheduler_metadata(
3820
num_reqs,
@@ -3828,6 +3829,7 @@ def cpu_attn_get_scheduler_metadata(
3829
isa,
3830
enable_kv_split,
3831
dynamic_causal,
3832
+ kv_cache_dtype,
3833
)
3834
return scheduler_metadata
3835
vllm/v1/attention/backends/cpu_attn.py
+3
-2
@@ -155,12 +155,12 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
155
if self.window_size is None:
156
self.window_size = -1
157
self.block_size = vllm_config.cache_config.block_size
158
- kv_cache_dtype_str = vllm_config.cache_config.cache_dtype
158
+ self.kv_cache_dtype = vllm_config.cache_config.cache_dtype
159
self.isa = _get_attn_isa(
160
self.dtype,
161
self.block_size,
162
self.head_dim,
163
- kv_cache_dtype_str,
163
+ self.kv_cache_dtype,
164
)
165
self.is_cross_attention = isinstance(kv_cache_spec, CrossAttentionSpec)
166
self.is_encoder_only_attention = isinstance(
@@ -234,6 +234,7 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
234
isa=self.isa,
235
enable_kv_split=envs.VLLM_CPU_ATTN_SPLIT_KV,
236
dynamic_causal=dynamic_casual,
237
+ kv_cache_dtype=self.kv_cache_dtype,
238
)
239
240
attn_metadata = CPUAttentionMetadata(