| 1 | #include "cpu_attn_dispatch_generated.h" |
| 2 | |
| 3 | // Runtime check for AMX-FP8, implemented in cpu_isa.cpp. |
| 4 | extern bool runtime_has_amx_fp8(); |
| 5 | |
| 6 | // Maps kv_cache_dtype string to Fp8KVCacheDataType enum. |
| 7 | // "auto" -> kAuto(0); "fp8"/"fp8_e4m3" -> kFp8E4M3; "fp8_e5m2" -> kFp8E5M2. |
| 8 | static inline cpu_attention::Fp8KVCacheDataType parse_fp8_kv_dtype( |
| 9 | const std::string& kv_cache_dtype) { |
| 10 | if (kv_cache_dtype == "fp8_e5m2") |
| 11 | return cpu_attention::Fp8KVCacheDataType::kFp8E5M2; |
| 12 | if (kv_cache_dtype == "fp8_e4m3" || kv_cache_dtype == "fp8") |
| 13 | return cpu_attention::Fp8KVCacheDataType::kFp8E4M3; |
| 14 | return cpu_attention::Fp8KVCacheDataType::kAuto; |
| 15 | } |
| 16 | |
| 17 | bool cpu_attn_has_isa(const std::string& isa) { |
| 18 | if (isa == "rvv") { |
| 19 | #if defined(__riscv) && defined(__riscv_v_min_vlen) && \ |
| 20 | (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256) |
| 21 | return true; |
| 22 | #else |
| 23 | return false; |
| 24 | #endif |
| 25 | } |
| 26 | if (isa == "amx_fp8") { |
| 27 | #ifdef CPU_CAPABILITY_AMXFP8 |
| 28 | // Guard with runtime CPUID check: the binary may be compiled with |
| 29 | // -mamx-fp8 (CPU_CAPABILITY_AMXFP8 defined) but run on a CPU that |
| 30 | // has AMX-BF16 without AMX-FP8 (e.g. Sapphire Rapids, Emerald Rapids). |
| 31 | return runtime_has_amx_fp8(); |
| 32 | #else |
| 33 | return false; |
| 34 | #endif |
| 35 | } |
| 36 | return false; |
| 37 | } |
| 38 | |
| 39 | torch::Tensor get_scheduler_metadata( |
| 40 | const int64_t num_req, const int64_t num_heads_q, |
| 41 | const int64_t num_heads_kv, const int64_t head_dim, |
| 42 | const torch::Tensor& seq_lens, at::ScalarType dtype, |
| 43 | const torch::Tensor& query_start_loc, const bool causal, |
| 44 | const int64_t window_size, const std::string& isa_hint, |
| 45 | const bool enable_kv_split, |
| 46 | const std::optional<torch::Tensor>& dynamic_causal, |
| 47 | const std::string& kv_cache_dtype) { |
| 48 | cpu_attention::ISA isa; |
| 49 | if (isa_hint == "amx") { |
| 50 | isa = cpu_attention::ISA::AMX; |
| 51 | } else if (isa_hint == "amx_fp8") { |
| 52 | isa = cpu_attention::ISA::AMX_FP8; |
| 53 | } else if (isa_hint == "vec") { |
| 54 | isa = cpu_attention::ISA::VEC; |
| 55 | } else if (isa_hint == "vec16") { |
| 56 | isa = cpu_attention::ISA::VEC16; |
| 57 | } else if (isa_hint == "neon") { |
| 58 | isa = cpu_attention::ISA::NEON; |
| 59 | } else if (isa_hint == "vxe") { |
| 60 | isa = cpu_attention::ISA::VXE; |
| 61 | } else if (isa_hint == "rvv") { |
| 62 | isa = cpu_attention::ISA::RVV; |
| 63 | } else if (isa_hint == "vsx") { |
| 64 | isa = cpu_attention::ISA::VSX; |
| 65 | } else { |
| 66 | TORCH_CHECK(false, "Unsupported CPU attention ISA hint: " + isa_hint); |
| 67 | } |
| 68 | |
| 69 | cpu_attention::AttentionScheduler::ScheduleInput input; |
| 70 | input.num_reqs = num_req; |
| 71 | input.num_heads_q = num_heads_q; |
| 72 | input.num_heads_kv = num_heads_kv; |
| 73 | input.head_dim = head_dim; |
| 74 | input.query_start_loc = query_start_loc.data_ptr<int32_t>(); |
| 75 | input.seq_lens = seq_lens.data_ptr<int32_t>(); |
| 76 | |
| 77 | input.sliding_window_size = window_size; |
| 78 | input.causal = causal; |
| 79 | input.isa = isa; |
| 80 | input.enable_kv_split = enable_kv_split; |
| 81 | input.dynamic_causal = |
| 82 | dynamic_causal.has_value() ? dynamic_causal->data_ptr<bool>() : nullptr; |
| 83 | |
| 84 | const int64_t kv_cache_idx = |
| 85 | static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype)); |
| 86 | VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() { |
| 87 | CPU_ATTN_DISPATCH(head_dim, isa, kv_cache_idx, [&]() { |
| 88 | input.elem_size = sizeof(attn_impl::kv_cache_t); |
| 89 | input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t); |
| 90 | input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t); |
| 91 | input.output_buffer_elem_size = |
| 92 | sizeof(attn_impl::partial_output_buffer_t); |
| 93 | input.max_num_q_per_iter = attn_impl::MaxQHeadNumPerIteration; |
| 94 | input.kv_block_alignment = attn_impl::BlockSizeAlignment; |
| 95 | }); |
| 96 | }); |
| 97 | |
| 98 | cpu_attention::AttentionScheduler scheduler; |
| 99 | torch::Tensor metadata = scheduler.schedule(input); |
| 100 | return metadata; |
| 101 | } |
| 102 | |
| 103 | void cpu_attn_reshape_and_cache( |
| 104 | const torch::Tensor& key, // [token_num, head_num, head_size] |
| 105 | const torch::Tensor& value, // [token_num, head_num, head_size] |
| 106 | torch::Tensor& |
| 107 | key_cache, // [num_blocks, num_kv_heads, block_size, head_size] |
| 108 | torch::Tensor& |
| 109 | value_cache, // [num_blocks, num_kv_heads, block_size, head_size] |
| 110 | const torch::Tensor& slot_mapping, const std::string& isa, |
| 111 | const double k_scale = 1.0, const double v_scale = 1.0, |
| 112 | const std::string& kv_cache_dtype = "auto") { |
| 113 | TORCH_CHECK_EQ(key.dim(), 3); |
| 114 | TORCH_CHECK_EQ(value.dim(), 3); |
| 115 | TORCH_CHECK_EQ(key_cache.dim(), 4); |
| 116 | TORCH_CHECK_EQ(value_cache.dim(), 4); |
| 117 | TORCH_CHECK_EQ(key.stride(2), 1); |
| 118 | TORCH_CHECK_EQ(value.stride(2), 1); |
| 119 | |
| 120 | const int64_t kv_cache_idx = |
| 121 | static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype)); |
| 122 | const bool is_fp8 = (kv_cache_idx != 0); |
| 123 | |
| 124 | if (is_fp8) { |
| 125 | TORCH_CHECK(key_cache.scalar_type() == at::ScalarType::Byte, |
| 126 | "key_cache must be uint8 for FP8 path"); |
| 127 | TORCH_CHECK(value_cache.scalar_type() == at::ScalarType::Byte, |
| 128 | "value_cache must be uint8 for FP8 path"); |
| 129 | TORCH_CHECK(k_scale > 0, "k_scale must be positive for FP8 path"); |
| 130 | TORCH_CHECK(v_scale > 0, "v_scale must be positive for FP8 path"); |
| 131 | } |
| 132 | |
| 133 | const float k_inv = is_fp8 ? 1.0f / static_cast<float>(k_scale) : 0.0f; |
| 134 | const float v_inv = is_fp8 ? 1.0f / static_cast<float>(v_scale) : 0.0f; |
| 135 | |
| 136 | const int64_t token_num = key.size(0); |
| 137 | const int64_t head_num = key.size(1); |
| 138 | const int64_t head_dim = key.size(2); |
| 139 | const int64_t num_blocks = key_cache.size(0); |
| 140 | const int64_t num_blocks_stride = key_cache.stride(0); |
| 141 | const int64_t cache_head_num_stride = key_cache.stride(1); |
| 142 | const int64_t block_size = key_cache.size(2); |
| 143 | const int64_t block_size_stride = key_cache.stride(2); |
| 144 | |
| 145 | cpu_attention::ISA isa_tag = [&]() { |
| 146 | if (isa == "amx") { |
| 147 | return cpu_attention::ISA::AMX; |
| 148 | } else if (isa == "amx_fp8") { |
| 149 | return cpu_attention::ISA::AMX_FP8; |
| 150 | } else if (isa == "vec") { |
| 151 | return cpu_attention::ISA::VEC; |
| 152 | } else if (isa == "vec16") { |
| 153 | return cpu_attention::ISA::VEC16; |
| 154 | } else if (isa == "neon") { |
| 155 | return cpu_attention::ISA::NEON; |
| 156 | } else if (isa == "vxe") { |
| 157 | return cpu_attention::ISA::VXE; |
| 158 | } else if (isa == "rvv") { |
| 159 | return cpu_attention::ISA::RVV; |
| 160 | } else if (isa == "vsx") { |
| 161 | return cpu_attention::ISA::VSX; |
| 162 | } else { |
| 163 | TORCH_CHECK(false, "Invalid ISA type: " + isa); |
| 164 | } |
| 165 | }(); |
| 166 | |
| 167 | if (is_fp8) { |
| 168 | TORCH_CHECK(isa_tag == cpu_attention::ISA::AMX || |
| 169 | isa_tag == cpu_attention::ISA::AMX_FP8 || |
| 170 | isa_tag == cpu_attention::ISA::VEC, |
| 171 | "FP8 KV cache is only supported on x86 (AMX_FP8/AMX/VEC) ISA"); |
| 172 | } |
| 173 | |
| 174 | VLLM_DISPATCH_FLOATING_TYPES( |
| 175 | key.scalar_type(), "cpu_attn_reshape_and_cache", [&]() { |
| 176 | CPU_ATTN_DISPATCH(head_dim, isa_tag, kv_cache_idx, [&]() { |
| 177 | using kv_t = typename attn_impl::kv_cache_t; |
| 178 | attn_impl::reshape_and_cache( |
| 179 | key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(), |
| 180 | reinterpret_cast<kv_t*>(key_cache.data_ptr()), |
| 181 | reinterpret_cast<kv_t*>(value_cache.data_ptr()), |
| 182 | slot_mapping.data_ptr<int64_t>(), token_num, key.stride(0), |
| 183 | value.stride(0), head_num, key.stride(1), value.stride(1), |
| 184 | num_blocks, num_blocks_stride, cache_head_num_stride, block_size, |
| 185 | block_size_stride, k_inv, v_inv); |
| 186 | }); |
| 187 | }); |
| 188 | } |
| 189 | |
| 190 | void cpu_attention_with_kv_cache( |
| 191 | const torch::Tensor& query, // [num_tokens, num_heads, head_size] |
| 192 | const torch::Tensor& |
| 193 | key_cache, // [num_blocks, num_kv_heads, block_size, head_size] |
| 194 | const torch::Tensor& |
| 195 | value_cache, // [num_blocks, num_kv_heads, block_size, head_size] |
| 196 | torch::Tensor& output, // [num_tokens, num_heads, head_size] |
| 197 | const torch::Tensor& query_start_loc, // [num_tokens + 1] |
| 198 | const torch::Tensor& seq_lens, // [num_tokens] |
| 199 | const double scale, const bool causal, |
| 200 | const std::optional<torch::Tensor>& alibi_slopes, // [num_heads] |
| 201 | const int64_t sliding_window, |
| 202 | const torch::Tensor& block_table, // [num_tokens, max_block_num] |
| 203 | const double softcap, const torch::Tensor& scheduler_metadata, |
| 204 | const std::optional<torch::Tensor>& s_aux, // [num_heads] |
| 205 | const std::optional<torch::Tensor>& dynamic_causal, // [num_reqs] |
| 206 | const double k_scale = 1.0, const double v_scale = 1.0, |
| 207 | const std::string& kv_cache_dtype = "auto") { |
| 208 | TORCH_CHECK_EQ(query.dim(), 3); |
| 209 | TORCH_CHECK_EQ(query.stride(2), 1); |
| 210 | TORCH_CHECK_EQ(key_cache.dim(), 4); |
| 211 | TORCH_CHECK_EQ(value_cache.dim(), 4); |
| 212 | |
| 213 | const int64_t kv_cache_idx = |
| 214 | static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype)); |
| 215 | const bool is_fp8 = (kv_cache_idx != 0); |
| 216 | if (is_fp8) { |
| 217 | TORCH_CHECK(key_cache.scalar_type() == at::ScalarType::Byte, |
| 218 | "key_cache must be uint8 for FP8 path"); |
| 219 | TORCH_CHECK(value_cache.scalar_type() == at::ScalarType::Byte, |
| 220 | "value_cache must be uint8 for FP8 path"); |
| 221 | TORCH_CHECK(k_scale > 0, "k_scale must be positive for FP8 path"); |
| 222 | TORCH_CHECK(v_scale > 0, "v_scale must be positive for FP8 path"); |
| 223 | } |
| 224 | |
| 225 | cpu_attention::AttentionInput input; |
| 226 | input.metadata = reinterpret_cast<cpu_attention::AttentionMetadata*>( |
| 227 | scheduler_metadata.data_ptr()); |
| 228 | input.num_tokens = query.size(0); |
| 229 | input.num_heads = query.size(1); |
| 230 | input.num_kv_heads = key_cache.size(1); |
| 231 | input.block_size = key_cache.size(2); |
| 232 | input.query = query.data_ptr(); |
| 233 | input.query_num_tokens_stride = query.stride(0); |
| 234 | input.query_num_heads_stride = query.stride(1); |
| 235 | input.cache_num_blocks_stride = key_cache.stride(0); |
| 236 | input.cache_num_kv_heads_stride = key_cache.stride(1); |
| 237 | input.blt_num_tokens_stride = block_table.stride(0); |
| 238 | input.key_cache = key_cache.data_ptr(); |
| 239 | input.value_cache = value_cache.data_ptr(); |
| 240 | input.output = output.data_ptr(); |
| 241 | input.query_start_loc = query_start_loc.data_ptr<int32_t>(); |
| 242 | input.seq_lens = seq_lens.data_ptr<int32_t>(); |
| 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 | // 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; |
| 262 | input.causal = causal; |
| 263 | input.sliding_window_size = sliding_window; |
| 264 | input.softcap = static_cast<float>(softcap); |
| 265 | |
| 266 | if (is_fp8) { |
| 267 | input.k_scale_fp8 = static_cast<float>(k_scale); |
| 268 | input.v_scale_fp8 = static_cast<float>(v_scale); |
| 269 | TORCH_CHECK(input.metadata->isa == cpu_attention::ISA::AMX || |
| 270 | input.metadata->isa == cpu_attention::ISA::AMX_FP8 || |
| 271 | input.metadata->isa == cpu_attention::ISA::VEC, |
| 272 | "FP8 KV cache is only supported on x86 (AMX_FP8/AMX/VEC) ISA"); |
| 273 | } |
| 274 | |
| 275 | VLLM_DISPATCH_FLOATING_TYPES( |
| 276 | query.scalar_type(), "cpu_attention_with_kv_cache", [&]() { |
| 277 | CPU_ATTN_DISPATCH( |
| 278 | query.size(2), input.metadata->isa, kv_cache_idx, [&]() { |
| 279 | TORCH_CHECK_EQ(input.block_size % attn_impl::BlockSizeAlignment, |
| 280 | 0); |
| 281 | cpu_attention::AttentionMainLoop<attn_impl> mainloop; |
| 282 | mainloop(&input); |
| 283 | }); |
| 284 | }); |
| 285 | } |