From 65a1a16594995bdc8b8feae2bf2ecf208d585be8 Mon Sep 17 00:00:00 2001 From: Tianmu Li Date: Wed, 29 Jul 2026 00:34:37 -0700 Subject: [PATCH] [CPU] Fix FP8 attention scratchpad sizing (#50194) Signed-off-by: Li, Tianmu Co-authored-by: Codex Co-authored-by: Li, Jiang --- csrc/cpu/cpu_attn.cpp | 9 ++++++--- csrc/cpu/torch_bindings.cpp | 6 ++++-- tests/kernels/attention/test_cpu_attn.py | 20 ++++++++++++++++++++ vllm/_custom_ops.py | 2 ++ vllm/v1/attention/backends/cpu_attn.py | 5 +++-- 5 files changed, 35 insertions(+), 7 deletions(-) diff --git a/csrc/cpu/cpu_attn.cpp b/csrc/cpu/cpu_attn.cpp index fa22861157..52c806f4fa 100644 --- a/csrc/cpu/cpu_attn.cpp +++ b/csrc/cpu/cpu_attn.cpp @@ -30,7 +30,8 @@ torch::Tensor get_scheduler_metadata( const torch::Tensor& query_start_loc, const bool causal, const int64_t window_size, const std::string& isa_hint, const bool enable_kv_split, - const std::optional& dynamic_causal) { + const std::optional& dynamic_causal, + const std::string& kv_cache_dtype) { cpu_attention::ISA isa; if (isa_hint == "amx") { isa = cpu_attention::ISA::AMX; @@ -65,9 +66,11 @@ torch::Tensor get_scheduler_metadata( input.dynamic_causal = dynamic_causal.has_value() ? dynamic_causal->data_ptr() : nullptr; + const int64_t kv_cache_idx = + static_cast(parse_fp8_kv_dtype(kv_cache_dtype)); VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() { - CPU_ATTN_DISPATCH(head_dim, isa, 0, [&]() { - input.elem_size = sizeof(scalar_t); + CPU_ATTN_DISPATCH(head_dim, isa, kv_cache_idx, [&]() { + input.elem_size = sizeof(attn_impl::kv_cache_t); input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t); input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t); input.output_buffer_elem_size = diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 80b9409d99..ee8e37110e 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -163,7 +163,8 @@ torch::Tensor get_scheduler_metadata( const torch::Tensor& query_start_loc, const bool casual, const int64_t window_size, const std::string& isa_hint, const bool enable_kv_split, - const std::optional& dynamic_causal); + const std::optional& dynamic_causal, + const std::string& kv_cache_dtype); void cpu_attn_reshape_and_cache(const torch::Tensor& key, const torch::Tensor& value, @@ -577,7 +578,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, " "int head_dim, Tensor seq_lens, ScalarType dtype, Tensor " "query_start_loc, bool casual, int window_size, str isa_hint, bool " - "enable_kv_split, Tensor? dynamic_causal) -> Tensor", + "enable_kv_split, Tensor? dynamic_causal, " + "str kv_cache_dtype=\"auto\") -> Tensor", &get_scheduler_metadata); ops.def( "cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) " diff --git a/tests/kernels/attention/test_cpu_attn.py b/tests/kernels/attention/test_cpu_attn.py index e296c226d7..3eb660aadb 100644 --- a/tests/kernels/attention/test_cpu_attn.py +++ b/tests/kernels/attention/test_cpu_attn.py @@ -534,6 +534,7 @@ def varlen_with_paged_kv( isa=isa, enable_kv_split=False, dynamic_causal=dynamic_causal_tensor, + kv_cache_dtype=kv_cache_dtype, ) out_without_split = torch.empty_like(query) @@ -569,6 +570,7 @@ def varlen_with_paged_kv( isa=isa, enable_kv_split=True, dynamic_causal=dynamic_causal_tensor, + kv_cache_dtype=kv_cache_dtype, ) out_with_split = torch.empty_like(query) @@ -803,6 +805,24 @@ def test_varlen_with_paged_kv_normal_amx( ) +@pytest.mark.skipif(not torch.cpu._is_amx_tile_supported(), reason="no AMX support.") +def test_varlen_with_paged_kv_fp8_large_prefill_amx() -> None: + varlen_with_paged_kv( + seq_lens=[(1024, 1024)] * 4, + num_heads=(16, 2), + head_size=256, + sliding_window=None, + dtype=torch.bfloat16, + block_size=2176, + soft_cap=None, + num_blocks=4, + use_alibi=False, + use_sink=False, + isa="amx", + kv_cache_dtype="fp8_e4m3", + ) + + @pytest.mark.parametrize("seq_lens", SEQ_LENS) @pytest.mark.parametrize("num_heads", NUM_HEADS) @pytest.mark.parametrize("head_size", HEAD_SIZES_VEC16) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index bc6d6d4afa..1e1e689083 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3814,6 +3814,7 @@ def cpu_attn_get_scheduler_metadata( isa: str, enable_kv_split: bool, dynamic_causal: torch.Tensor | None = None, + kv_cache_dtype: str = "auto", ) -> torch.Tensor: scheduler_metadata = torch.ops._C.get_scheduler_metadata( num_reqs, @@ -3828,6 +3829,7 @@ def cpu_attn_get_scheduler_metadata( isa, enable_kv_split, dynamic_causal, + kv_cache_dtype, ) return scheduler_metadata diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index e9c31f42f2..ea63f39522 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -155,12 +155,12 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata] if self.window_size is None: self.window_size = -1 self.block_size = vllm_config.cache_config.block_size - kv_cache_dtype_str = vllm_config.cache_config.cache_dtype + self.kv_cache_dtype = vllm_config.cache_config.cache_dtype self.isa = _get_attn_isa( self.dtype, self.block_size, self.head_dim, - kv_cache_dtype_str, + self.kv_cache_dtype, ) self.is_cross_attention = isinstance(kv_cache_spec, CrossAttentionSpec) self.is_encoder_only_attention = isinstance( @@ -234,6 +234,7 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata] isa=self.isa, enable_kv_split=envs.VLLM_CPU_ATTN_SPLIT_KV, dynamic_causal=dynamic_casual, + kv_cache_dtype=self.kv_cache_dtype, ) attn_metadata = CPUAttentionMetadata(