| 1 | #include "cache.h" |
| 2 | #include "ops.h" |
| 3 | #include "core/registration.h" |
| 4 | |
| 5 | #include <torch/library.h> |
| 6 | |
| 7 | // Note: overwrite the external definition for sharing same name between |
| 8 | // libraries use different ISAs. |
| 9 | #define TORCH_EXTENSION_NAME _C |
| 10 | |
| 11 | void release_dnnl_matmul_handler(int64_t handler); |
| 12 | |
| 13 | int64_t create_onednn_scaled_mm_handler(const torch::Tensor& b, |
| 14 | const torch::Tensor& b_scales, |
| 15 | at::ScalarType output_type, |
| 16 | bool dynamic_act_quant, bool use_azp, |
| 17 | int64_t primitive_cache_size); |
| 18 | |
| 19 | void onednn_scaled_mm(torch::Tensor& c, const torch::Tensor& a, |
| 20 | const torch::Tensor& a_scales, |
| 21 | const std::optional<torch::Tensor>& azp, |
| 22 | const std::optional<torch::Tensor>& azp_adj, |
| 23 | const std::optional<torch::Tensor>& bias, |
| 24 | const torch::Tensor& handler_tensor); |
| 25 | |
| 26 | int64_t create_onednn_mm_handler(const torch::Tensor& b, |
| 27 | int64_t primitive_cache_size); |
| 28 | |
| 29 | void onednn_mm(torch::Tensor& c, const torch::Tensor& a, |
| 30 | const std::optional<torch::Tensor>& bias, |
| 31 | const torch::Tensor& handler_tensor); |
| 32 | |
| 33 | bool is_onednn_acl_supported(); |
| 34 | |
| 35 | void mla_decode_kvcache(torch::Tensor& out, torch::Tensor& query, |
| 36 | torch::Tensor& kv_cache, double scale, |
| 37 | torch::Tensor& block_tables, torch::Tensor& seq_lens); |
| 38 | |
| 39 | int64_t init_shm_manager(const std::string& name, const int64_t group_size, |
| 40 | const int64_t rank, const int64_t thread_num); |
| 41 | |
| 42 | std::string join_shm_manager(int64_t handle, const std::string& name); |
| 43 | |
| 44 | void shm_allreduce(int64_t handle, torch::Tensor& data); |
| 45 | |
| 46 | void shm_gather(int64_t handle, torch::Tensor& data, |
| 47 | const std::optional<std::vector<torch::Tensor>>& outputs, |
| 48 | int64_t dst); |
| 49 | |
| 50 | void shm_all_gather(int64_t handle, const torch::Tensor& data, |
| 51 | torch::Tensor& output); |
| 52 | |
| 53 | void shm_send_tensor_list(int64_t handle, |
| 54 | const std::vector<torch::Tensor>& tensor_list, |
| 55 | int64_t dst); |
| 56 | |
| 57 | std::vector<torch::Tensor> shm_recv_tensor_list(int64_t handle, int64_t src); |
| 58 | |
| 59 | // SGL CPU kernels |
| 60 | |
| 61 | void weight_packed_linear(at::Tensor& out, const at::Tensor& mat1, |
| 62 | const at::Tensor& mat2, |
| 63 | const std::optional<at::Tensor>& bias, bool is_vnni); |
| 64 | |
| 65 | at::Tensor convert_weight_packed(at::Tensor& weight); |
| 66 | |
| 67 | at::Tensor convert_scale_packed(at::Tensor& scale); |
| 68 | |
| 69 | void fused_experts_cpu(at::Tensor& out, at::Tensor& hidden_states, |
| 70 | at::Tensor& w1, at::Tensor& w2, at::Tensor& topk_weights, |
| 71 | at::Tensor& topk_ids, int64_t moe_comp_method, |
| 72 | const std::optional<at::Tensor>& w1_scale, |
| 73 | const std::optional<at::Tensor>& w2_scale, |
| 74 | const std::optional<at::Tensor>& w1_zero, |
| 75 | const std::optional<at::Tensor>& w2_zero, |
| 76 | const std::optional<at::Tensor>& a1_scale, |
| 77 | const std::optional<std::vector<int64_t>> block_size, |
| 78 | const std::optional<at::Tensor>& w1_bias, |
| 79 | const std::optional<at::Tensor>& w2_bias, |
| 80 | const std::optional<double>& alpha, |
| 81 | const std::optional<double>& limit, bool is_vnni); |
| 82 | |
| 83 | at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2, |
| 84 | at::Tensor& scales2, |
| 85 | const std::optional<at::Tensor>& bias, |
| 86 | at::ScalarType out_dtype, bool is_vnni); |
| 87 | |
| 88 | // Adapted from sglang: FP8 W8A16 kernel |
| 89 | at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2, |
| 90 | at::Tensor& scales2, |
| 91 | std::vector<int64_t> block_size, |
| 92 | const std::optional<at::Tensor>& bias, |
| 93 | at::ScalarType out_dtype, bool is_vnni); |
| 94 | |
| 95 | // Adapted from sglang: FP8 W8A8 kernels |
| 96 | std::tuple<at::Tensor, at::Tensor> float8_linear_prepack_impl( |
| 97 | const at::Tensor& weight, const at::Tensor& scales); |
| 98 | at::Tensor fp8_scaled_mm_with_quant(const at::Tensor& act, |
| 99 | const std::optional<at::Tensor>& act_scales, |
| 100 | bool channelwise, const at::Tensor& weight, |
| 101 | const at::Tensor& weight_scales, |
| 102 | const std::optional<at::Tensor>& bias, |
| 103 | at::ScalarType output_dtype); |
| 104 | std::tuple<at::Tensor, at::Tensor> quantize_fp8e4m3_vec( |
| 105 | const at::Tensor& t, bool channelwise, c10::optional<at::Tensor> scale_opt); |
| 106 | |
| 107 | // Adapted from sglang: MLA CPU kernels (AMX-only) |
| 108 | void decode_attention_cpu(at::Tensor& query, at::Tensor& k_buffer, |
| 109 | at::Tensor& v_buffer, at::Tensor& output, |
| 110 | const std::optional<at::Tensor>& key, |
| 111 | const std::optional<at::Tensor>& value, |
| 112 | const std::optional<at::Tensor>& loc, |
| 113 | at::Tensor& attn_logits, at::Tensor& req_to_token, |
| 114 | at::Tensor& req_pool_indices, at::Tensor& seq_lens, |
| 115 | double sm_scale, double logit_cap, bool is_cross_attn, |
| 116 | int64_t sliding_window_size, |
| 117 | std::optional<at::Tensor> encoder_lens, |
| 118 | std::optional<at::Tensor> sinks); |
| 119 | |
| 120 | void extend_attention_cpu( |
| 121 | at::Tensor& q_extend, const std::optional<at::Tensor>& k_extend, |
| 122 | const std::optional<at::Tensor>& v_extend, at::Tensor& o_extend, |
| 123 | at::Tensor& k_buffer, at::Tensor& v_buffer, at::Tensor& req_to_token, |
| 124 | at::Tensor& req_pool_indices, at::Tensor& seq_lens, |
| 125 | at::Tensor& extend_seq_lens, at::Tensor& extend_start_loc, |
| 126 | int64_t max_len_extend, double sm_scale, double logit_cap, |
| 127 | bool is_cross_attn, int64_t sliding_window_size, |
| 128 | std::optional<at::Tensor> encoder_lens, std::optional<at::Tensor> sinks, |
| 129 | std::optional<at::Tensor> tree_mask); |
| 130 | |
| 131 | void bmm_cpu(at::Tensor& out, at::Tensor& mat1, at::Tensor& mat2, bool is_vnni, |
| 132 | const std::optional<at::Tensor>& scale); |
| 133 | |
| 134 | // vLLM-native: CPU cache-write op for MLA's single-latent-buffer KV cache |
| 135 | // (the CUDA-only `concat_and_cache_mla` has no CPU dispatch). |
| 136 | void concat_and_cache_mla_cpu(const at::Tensor& kv_c_normed, |
| 137 | const at::Tensor& k_pe, at::Tensor& kv_cache, |
| 138 | const at::Tensor& slot_mapping); |
| 139 | |
| 140 | // Adapted from sglang: DeepSeek-V4 mHC gating kernels |
| 141 | std::tuple<at::Tensor, at::Tensor, at::Tensor> hc_pre_fused_cpu( |
| 142 | at::Tensor& x, at::Tensor& hc_fn, at::Tensor& hc_scale, at::Tensor& hc_base, |
| 143 | int64_t hc_mult, int64_t sinkhorn_iters, double rms_eps, double hc_eps); |
| 144 | |
| 145 | at::Tensor hc_post_fused_cpu(at::Tensor& x, at::Tensor& residual, |
| 146 | at::Tensor& post, at::Tensor& comb); |
| 147 | |
| 148 | at::Tensor hc_head_fused_cpu(at::Tensor& x, at::Tensor& hc_fn, |
| 149 | at::Tensor& hc_scale, at::Tensor& hc_base, |
| 150 | double hc_eps, double norm_eps); |
| 151 | |
| 152 | // Adapted from sglang: DeepSeek-V4 fp8_ds_mla cache-write kernel |
| 153 | at::Tensor fused_qnorm_rope_kv_insert_cpu(at::Tensor& q, at::Tensor& kv, |
| 154 | at::Tensor& positions, |
| 155 | at::Tensor& swa_kv_cache_2d, |
| 156 | at::Tensor& slot_mapping, |
| 157 | at::Tensor& cos_sin_cache, |
| 158 | int64_t q_head_padded, double eps, |
| 159 | int64_t cache_block_size); |
| 160 | |
| 161 | // vLLM-native: DeepSeek-V4 sparse MQA attention kernel for the fp8_ds_mla |
| 162 | // cache layout (SWA window + optional compressed top-k index sets). |
| 163 | void flash_mla_with_kvcache_cpu( |
| 164 | at::Tensor& out, at::Tensor& q, at::Tensor& window_cache_2d, |
| 165 | at::Tensor& window_slots, int64_t window_block_size, |
| 166 | at::Tensor& compressed_cache_2d, at::Tensor& compressed_slots, |
| 167 | int64_t compressed_block_size, at::Tensor& attn_sink, double scale); |
| 168 | |
| 169 | // Adapted from sglang: DeepSeek-V4 compressor state-cache write + fused |
| 170 | // compress+RMSNorm+quant+RoPE+store kernels. |
| 171 | void save_partial_states_cpu(at::Tensor& kv, at::Tensor& score, at::Tensor& ape, |
| 172 | at::Tensor& positions, at::Tensor& state_cache, |
| 173 | at::Tensor& slot_mapping); |
| 174 | |
| 175 | void compress_norm_rope_store_cpu( |
| 176 | at::Tensor& state_cache, at::Tensor& gather_slots, at::Tensor& positions, |
| 177 | at::Tensor& kv_slot_mapping, at::Tensor& rms_norm_weight, |
| 178 | double rms_norm_eps, at::Tensor& cos_sin_cache, at::Tensor& kv_cache_2d, |
| 179 | int64_t kv_cache_block_size, int64_t compress_ratio); |
| 180 | |
| 181 | // head_dim=128 indexer-compressor variant of the above (single fp8 quant |
| 182 | // block + raw fp32 scale -- matches the paged indexer K-cache layout |
| 183 | // fp8_paged_mqa_logits_cpu/topk_transform_512_cpu read). |
| 184 | void compress_norm_rope_store_indexer_cpu( |
| 185 | at::Tensor& state_cache, at::Tensor& gather_slots, at::Tensor& positions, |
| 186 | at::Tensor& kv_slot_mapping, at::Tensor& rms_norm_weight, |
| 187 | double rms_norm_eps, at::Tensor& cos_sin_cache, at::Tensor& kv_cache_2d, |
| 188 | int64_t kv_cache_block_size, int64_t compress_ratio); |
| 189 | |
| 190 | // Adapted from sglang: DeepSeek-V4 sparse indexer paged MQA-logits + |
| 191 | // top-512 transform kernels (DECODE path -- reads the paged K-cache |
| 192 | // directly via `page_table`, whole batch in one call). |
| 193 | at::Tensor fp8_paged_mqa_logits_cpu(at::Tensor& q_fp8, at::Tensor& kvcache_fp8, |
| 194 | at::Tensor& weight, at::Tensor& seq_lens, |
| 195 | at::Tensor& page_table, int64_t block_size, |
| 196 | int64_t max_seq_len); |
| 197 | |
| 198 | void topk_transform_512_cpu(at::Tensor& scores, at::Tensor& seq_lens, |
| 199 | at::Tensor& page_tables, |
| 200 | at::Tensor& out_page_indices, int64_t page_size, |
| 201 | const std::optional<at::Tensor>& out_raw_indices); |
| 202 | |
| 203 | // Adapted from sglang: DeepSeek-V4 MoE routing kernels (flat biased top-k and |
| 204 | // hash-routed-layer top-k). |
| 205 | std::tuple<at::Tensor, at::Tensor> biased_topk_cpu( |
| 206 | at::Tensor& hidden_states, at::Tensor& gating_output, |
| 207 | at::Tensor& correction_bias, int64_t topk, bool renormalize, |
| 208 | std::string scoring_func, int64_t num_fused_shared_experts, |
| 209 | std::optional<double> routed_scaling_factor, |
| 210 | bool apply_routed_scaling_factor_on_output); |
| 211 | |
| 212 | std::tuple<at::Tensor, at::Tensor> hash_topk_cpu( |
| 213 | at::Tensor& gating_output, at::Tensor& tid2eid, int64_t topk, |
| 214 | std::string scoring_func, int64_t num_fused_shared_experts, |
| 215 | int64_t num_experts, double routed_scaling_factor); |
| 216 | |
| 217 | // vLLM-native: DeepSeek-V4 sparse indexer Q-side RoPE + FP8 quant kernel |
| 218 | // (fp8 path only; MXFP4 stays on triton-cpu). |
| 219 | void fused_indexer_q_rope_quant_cpu(at::Tensor& positions, at::Tensor& index_q, |
| 220 | at::Tensor& index_q_cos_sin_cache, |
| 221 | at::Tensor& index_q_fp8, |
| 222 | at::Tensor& index_weights, |
| 223 | double index_weights_softmax_scale, |
| 224 | double index_weights_head_scale, |
| 225 | at::Tensor& index_weights_out); |
| 226 | |
| 227 | // vLLM-native: inverse GPT-J RoPE for DeepSeek-V4 CPU attention's `_o_proj` |
| 228 | // output de-rotation. |
| 229 | at::Tensor inverse_gptj_rope_o_proj_cpu(at::Tensor& o, at::Tensor& positions, |
| 230 | at::Tensor& cos_sin_cache, |
| 231 | int64_t rope_dim); |
| 232 | |
| 233 | // Adapted from sglang: INT4 W4A8 kernels |
| 234 | std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp( |
| 235 | at::Tensor qweight, // awq: (*, K, N / 8) || gptq: (*, K / 8, N) , int32 |
| 236 | at::Tensor qzeros, // awq: (*, K / group_size, N / 8) || gptq: (*, K / |
| 237 | // group_size, N / 8) , int32 |
| 238 | at::Tensor scales, // awq: (*, K / group_size, N) || gptq: (*, K / |
| 239 | // group_size, N) , bfloat16 |
| 240 | int64_t quant_method_4bit); |
| 241 | |
| 242 | at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros, |
| 243 | at::Tensor& w_scales, |
| 244 | std::optional<at::Tensor> bias); |
| 245 | |
| 246 | // Adapted from sglang: GDN |
| 247 | std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu( |
| 248 | const at::Tensor& query, const at::Tensor& key, const at::Tensor& value, |
| 249 | const at::Tensor& g, const at::Tensor& beta, |
| 250 | const at::Tensor& initial_state, bool output_final_state, |
| 251 | const at::Tensor& cu_seqlens, bool head_first, bool use_qk_l2norm_in_kernel, |
| 252 | const at::Tensor& initial_state_indices, double eps = 1e-5); |
| 253 | |
| 254 | at::Tensor fused_sigmoid_gating_delta_rule_update_cpu( |
| 255 | const at::Tensor& A_log, const at::Tensor& dt_bias, const at::Tensor& q, |
| 256 | const at::Tensor& k, const at::Tensor& v, const at::Tensor& a, |
| 257 | const at::Tensor& b, at::Tensor& initial_state_source, |
| 258 | const at::Tensor& initial_state_indices, const at::Tensor& cu_seqlens, |
| 259 | bool use_qk_l2norm_in_kernel, double softplus_beta = 1.0, |
| 260 | double softplus_threshold = 20.0); |
| 261 | |
| 262 | at::Tensor fused_sigmoid_gating_delta_rule_update_spec_cpu( |
| 263 | const at::Tensor& A_log, const at::Tensor& dt_bias, const at::Tensor& q, |
| 264 | const at::Tensor& k, const at::Tensor& v, const at::Tensor& a, |
| 265 | const at::Tensor& b, at::Tensor& initial_state_source, |
| 266 | const at::Tensor& spec_state_indices, const at::Tensor& num_accepted_tokens, |
| 267 | const at::Tensor& cu_seqlens, bool use_qk_l2norm_in_kernel, |
| 268 | double softplus_beta = 1.0, double softplus_threshold = 20.0); |
| 269 | |
| 270 | std::tuple<at::Tensor, at::Tensor> fused_gdn_gating_cpu( |
| 271 | const at::Tensor& A_log, const at::Tensor& a, const at::Tensor& b, |
| 272 | const at::Tensor& dt_bias); |
| 273 | |
| 274 | // Adapted from sglang: casual_conv1d kernels |
| 275 | at::Tensor causal_conv1d_weight_pack(const at::Tensor& weight); |
| 276 | |
| 277 | at::Tensor causal_conv1d_fwd_cpu( |
| 278 | const at::Tensor& x, const at::Tensor& weight, |
| 279 | const std::optional<at::Tensor>& bias, |
| 280 | const std::optional<at::Tensor>& conv_states, |
| 281 | const std::optional<at::Tensor>& query_start_loc, |
| 282 | const std::optional<at::Tensor>& cache_indices, |
| 283 | const std::optional<at::Tensor>& has_initial_state, bool silu_activation, |
| 284 | int64_t pad_slot_id, bool is_weight_packed); |
| 285 | |
| 286 | at::Tensor causal_conv1d_update_cpu( |
| 287 | const at::Tensor& x, const at::Tensor& conv_states, |
| 288 | const at::Tensor& weight, const std::optional<at::Tensor>& bias, |
| 289 | bool silu_activation, const std::optional<at::Tensor>& num_accepted_tokens, |
| 290 | const std::optional<at::Tensor>& conv_state_indices, int64_t pad_slot_id, |
| 291 | bool is_weight_packed); |
| 292 | |
| 293 | void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input, |
| 294 | const std::string& activation); |
| 295 | |
| 296 | bool cpu_attn_has_isa(const std::string& isa); |
| 297 | bool cpu_has_amx_fp8(); |
| 298 | |
| 299 | torch::Tensor get_scheduler_metadata( |
| 300 | const int64_t num_req, const int64_t num_heads_q, |
| 301 | const int64_t num_heads_kv, const int64_t head_dim, |
| 302 | const torch::Tensor& seq_lens, at::ScalarType dtype, |
| 303 | const torch::Tensor& query_start_loc, const bool casual, |
| 304 | const int64_t window_size, const std::string& isa_hint, |
| 305 | const bool enable_kv_split, |
| 306 | const std::optional<torch::Tensor>& dynamic_causal, |
| 307 | const std::string& kv_cache_dtype); |
| 308 | |
| 309 | void cpu_attn_reshape_and_cache(const torch::Tensor& key, |
| 310 | const torch::Tensor& value, |
| 311 | torch::Tensor& key_cache, |
| 312 | torch::Tensor& value_cache, |
| 313 | const torch::Tensor& slot_mapping, |
| 314 | const std::string& isa, const double k_scale, |
| 315 | const double v_scale, |
| 316 | const std::string& kv_cache_dtype); |
| 317 | |
| 318 | void cpu_attention_with_kv_cache( |
| 319 | const torch::Tensor& query, const torch::Tensor& key_cache, |
| 320 | const torch::Tensor& value_cache, torch::Tensor& output, |
| 321 | const torch::Tensor& query_start_loc, const torch::Tensor& seq_lens, |
| 322 | const double scale, const bool causal, |
| 323 | const std::optional<torch::Tensor>& alibi_slopes, |
| 324 | const int64_t sliding_window_left, const torch::Tensor& block_table, |
| 325 | const double softcap, const torch::Tensor& scheduler_metadata, |
| 326 | const std::optional<torch::Tensor>& s_aux, |
| 327 | const std::optional<torch::Tensor>& dynamic_causal, const double k_scale, |
| 328 | const double v_scale, const std::string& kv_cache_dtype); |
| 329 | |
| 330 | // Note: just for avoiding importing errors |
| 331 | void placeholder_op() { TORCH_CHECK(false, "Unimplemented"); } |
| 332 | |
| 333 | void cpu_gemm_wna16(const torch::Tensor& input, const torch::Tensor& q_weight, |
| 334 | torch::Tensor& output, const torch::Tensor& scales, |
| 335 | const std::optional<torch::Tensor>& zeros, |
| 336 | const std::optional<torch::Tensor>& bias, |
| 337 | const int64_t pack_factor, const std::string& isa_hint); |
| 338 | |
| 339 | void prepack_moe_weight(const torch::Tensor& weight, |
| 340 | torch::Tensor& packed_weight, const std::string& isa); |
| 341 | |
| 342 | void cpu_fused_moe(torch::Tensor& output, const torch::Tensor& input, |
| 343 | const torch::Tensor& w13, const torch::Tensor& w2, |
| 344 | const std::optional<torch::Tensor>& w13_bias, |
| 345 | const std::optional<torch::Tensor>& w2_bias, |
| 346 | const torch::Tensor& topk_weights, |
| 347 | const torch::Tensor& topk_id, const bool skip_weighted, |
| 348 | const std::string& act, const std::string& isa); |
| 349 | |
| 350 | void prepack_moe_weight_int8(const torch::Tensor& weight, |
| 351 | torch::Tensor& packed_weight, |
| 352 | const std::string& isa); |
| 353 | |
| 354 | void cpu_fused_moe_int8(torch::Tensor& output, const torch::Tensor& input, |
| 355 | const torch::Tensor& w13, const torch::Tensor& w2, |
| 356 | const torch::Tensor& w13_scale, |
| 357 | const torch::Tensor& w2_scale, |
| 358 | const std::optional<torch::Tensor>& w13_bias, |
| 359 | const std::optional<torch::Tensor>& w2_bias, |
| 360 | const torch::Tensor& topk_weights, |
| 361 | const torch::Tensor& topk_id, const bool skip_weighted, |
| 362 | const std::string& act, const std::string& isa); |
| 363 | |
| 364 | void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc, |
| 365 | const torch::Tensor positions, |
| 366 | const torch::Tensor block_table, |
| 367 | torch::Tensor slot_mapping, |
| 368 | const int64_t block_size); |
| 369 | |
| 370 | at::Tensor causal_conv1d_update_cpu_impl( |
| 371 | at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight, |
| 372 | const c10::optional<at::Tensor>& bias, |
| 373 | const c10::optional<std::string>& activation, |
| 374 | const c10::optional<at::Tensor>& conv_state_indices, |
| 375 | const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id); |
| 376 | |
| 377 | void selective_state_update_cpu_impl( |
| 378 | at::Tensor& state, const at::Tensor& x, const at::Tensor& dt, |
| 379 | const at::Tensor& A, const at::Tensor& B, const at::Tensor& C, |
| 380 | const c10::optional<at::Tensor>& D, const c10::optional<at::Tensor>& z, |
| 381 | const c10::optional<at::Tensor>& dt_bias, bool dt_softplus, |
| 382 | const c10::optional<at::Tensor>& state_batch_indices, |
| 383 | const c10::optional<at::Tensor>& dst_state_batch_indices, |
| 384 | int64_t null_block_id, at::Tensor& out, |
| 385 | const c10::optional<at::Tensor>& num_accepted_tokens, |
| 386 | const c10::optional<at::Tensor>& cu_seqlens); |
| 387 | |
| 388 | void mamba_chunk_scan_fwd_cpu_impl(at::Tensor& out, at::Tensor& final_states, |
| 389 | const at::Tensor& x, const at::Tensor& dt, |
| 390 | const at::Tensor& A, const at::Tensor& B, |
| 391 | const at::Tensor& C, |
| 392 | const c10::optional<at::Tensor>& D, |
| 393 | const c10::optional<at::Tensor>& z, |
| 394 | const at::Tensor& cu_seqlens); |
| 395 | |
| 396 | torch::Tensor fused_gumbel_argmax(const torch::Tensor& logits, |
| 397 | const torch::Tensor& seeds); |
| 398 | |
| 399 | torch::Tensor greedy_argmax(const torch::Tensor& logits); |
| 400 | |
| 401 | void init_cpu_memory_env(std::vector<int64_t> node_ids); |
| 402 | |
| 403 | namespace cpu_utils { |
| 404 | void eagle_prepare_inputs_padded_kernel_impl( |
| 405 | const torch::Tensor& cu_num_draft_tokens, |
| 406 | const torch::Tensor& valid_sampled_tokens_count, |
| 407 | const torch::Tensor& query_start_loc_gpu, |
| 408 | torch::Tensor& token_indices_to_sample, |
| 409 | torch::Tensor& num_rejected_tokens_gpu, const int64_t num_reqs); |
| 410 | void eagle_prepare_next_token_padded_kernel_impl( |
| 411 | const torch::Tensor& sampled_token_ids, |
| 412 | const torch::Tensor& discard_request_mask, |
| 413 | const torch::Tensor& backup_next_token_ids, torch::Tensor& next_token_ids, |
| 414 | torch::Tensor& valid_sampled_tokens_count, const int64_t vocab_size, |
| 415 | const int64_t num_sampled_tokens_per_req, const int64_t num_reqs); |
| 416 | void eagle_step_slot_mapping_metadata_kernel_impl( |
| 417 | const torch::Tensor& positions, const torch::Tensor& block_table, |
| 418 | torch::Tensor& seq_lens, torch::Tensor& out_clamped_positions, |
| 419 | torch::Tensor& out_slot_mapping, const int64_t block_size, |
| 420 | const int64_t max_model_len, const int64_t PAD_ID); |
| 421 | void copy_and_expand_eagle_inputs_kernel_impl( |
| 422 | const torch::Tensor& target_token_ids, |
| 423 | const torch::Tensor& target_positions, const torch::Tensor& next_token_ids, |
| 424 | torch::Tensor& out_input_ids, torch::Tensor& out_positions, |
| 425 | torch::Tensor& out_is_rejected_token_mask, |
| 426 | torch::Tensor& out_is_masked_token_mask, |
| 427 | torch::Tensor& out_new_token_indices, |
| 428 | torch::Tensor& out_hidden_state_mapping, |
| 429 | const torch::Tensor& query_start_loc, const torch::Tensor& query_end_loc, |
| 430 | const int64_t padding_token_id, const int64_t parallel_drafting_token_id, |
| 431 | const int64_t total_input_tokens, |
| 432 | const int64_t num_padding_slots_per_request, const bool shift_input_ids); |
| 433 | void copy_and_expand_dflash_inputs_kernel_impl( |
| 434 | const torch::Tensor& next_token_ids, const torch::Tensor& target_positions, |
| 435 | torch::Tensor& out_input_ids, torch::Tensor& out_context_positions, |
| 436 | torch::Tensor& out_query_positions, torch::Tensor& out_context_slot_mapping, |
| 437 | torch::Tensor& out_query_slot_mapping, torch::Tensor& out_token_indices, |
| 438 | const torch::Tensor& block_table, const torch::Tensor& query_start_loc, |
| 439 | const std::optional<torch::Tensor>& num_rejected_tokens, |
| 440 | const int64_t parallel_drafting_token_id, const int64_t block_size, |
| 441 | const int64_t num_query_per_req, const int64_t num_speculative_tokens, |
| 442 | const int64_t total_input_tokens, const bool has_num_rejected); |
| 443 | void rejection_greedy_sample_kernel_impl( |
| 444 | torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, |
| 445 | const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax, |
| 446 | const torch::Tensor& bonus_token_ids, |
| 447 | const std::optional<torch::Tensor>& is_greedy, const int64_t max_spec_len); |
| 448 | void rejection_random_sample_kernel_impl( |
| 449 | torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, |
| 450 | const torch::Tensor& draft_token_ids, |
| 451 | const std::optional<torch::Tensor>& draft_probs, |
| 452 | const torch::Tensor& target_probs, const torch::Tensor& bonus_token_ids, |
| 453 | const torch::Tensor& recovered_token_ids, |
| 454 | const torch::Tensor& uniform_probs, |
| 455 | const std::optional<torch::Tensor>& is_greedy, const int64_t max_spec_len, |
| 456 | const int64_t vocab_size, const bool no_draft_probs); |
| 457 | void expand_kernel_impl(torch::Tensor& output, const torch::Tensor& input, |
| 458 | const torch::Tensor& cu_num_tokens, |
| 459 | const int64_t replace_from, const int64_t replace_to); |
| 460 | void sample_recovered_tokens_kernel_impl( |
| 461 | torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, |
| 462 | const torch::Tensor& draft_token_ids, |
| 463 | const std::optional<torch::Tensor>& draft_probs, |
| 464 | const torch::Tensor& target_probs, const torch::Tensor& inv_q, |
| 465 | const int64_t vocab_size, const bool no_draft_probs); |
| 466 | } // namespace cpu_utils |
| 467 | |
| 468 | TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { |
| 469 | // vLLM custom ops |
| 470 | |
| 471 | ops.def( |
| 472 | "dynamic_4bit_int_moe(" |
| 473 | "Tensor x, Tensor topk_ids, Tensor topk_weights," |
| 474 | "Tensor w13_packed, Tensor w2_packed," |
| 475 | "int hidden_size, int intermediate_size," |
| 476 | "int group_size, bool apply_router_weight_on_input, int activation_kind" |
| 477 | ") -> Tensor"); |
| 478 | |
| 479 | ops.impl("dynamic_4bit_int_moe", torch::kCPU, &dynamic_4bit_int_moe_cpu); |
| 480 | |
| 481 | // Activation ops |
| 482 | |
| 483 | // Activation function used in SwiGLU. |
| 484 | ops.def("silu_and_mul(Tensor! out, Tensor input) -> ()"); |
| 485 | ops.impl("silu_and_mul", torch::kCPU, &silu_and_mul); |
| 486 | |
| 487 | // Activation function used in GeGLU with `none` approximation. |
| 488 | ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()"); |
| 489 | ops.impl("gelu_and_mul", torch::kCPU, &gelu_and_mul); |
| 490 | |
| 491 | // Activation function used in GeGLU with `tanh` approximation. |
| 492 | ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()"); |
| 493 | ops.impl("gelu_tanh_and_mul", torch::kCPU, &gelu_tanh_and_mul); |
| 494 | |
| 495 | // GELU tanh implementation. |
| 496 | ops.def("gelu_tanh(Tensor! out, Tensor input) -> ()"); |
| 497 | ops.impl("gelu_tanh", torch::kCPU, &gelu_tanh); |
| 498 | |
| 499 | // GELU implementation used in GPT-2. |
| 500 | ops.def("gelu_new(Tensor! out, Tensor input) -> ()"); |
| 501 | ops.impl("gelu_new", torch::kCPU, &gelu_new); |
| 502 | |
| 503 | // Approximate GELU implementation. |
| 504 | ops.def("gelu_fast(Tensor! out, Tensor input) -> ()"); |
| 505 | ops.impl("gelu_fast", torch::kCPU, &gelu_fast); |
| 506 | |
| 507 | // Quick GELU implementation. |
| 508 | ops.def("gelu_quick(Tensor! out, Tensor input) -> ()"); |
| 509 | ops.impl("gelu_quick", torch::kCPU, &gelu_quick); |
| 510 | |
| 511 | #if (defined(__aarch64__) && !defined(__APPLE__)) |
| 512 | |
| 513 | ops.def( |
| 514 | "activation_lut_bf16(Tensor! out, Tensor input, str activation)" |
| 515 | " -> ()"); |
| 516 | ops.impl("activation_lut_bf16", torch::kCPU, &activation_lut_bf16); |
| 517 | |
| 518 | #endif // (defined(__aarch64__) && !defined(__APPLE__)) |
| 519 | |
| 520 | #if (defined(__aarch64__) && defined(ARM_BF16_SUPPORT) && \ |
| 521 | !defined(__APPLE__)) || \ |
| 522 | (defined(__AVX512BF16__) && defined(__AVX512F__) && \ |
| 523 | defined(__AVX512VNNI__)) |
| 524 | |
| 525 | // Adapted from sglang: casual_conv1d kernels |
| 526 | ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor"); |
| 527 | ops.impl("causal_conv1d_weight_pack", torch::kCPU, |
| 528 | &causal_conv1d_weight_pack); |
| 529 | |
| 530 | ops.def( |
| 531 | "causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? " |
| 532 | "conv_states, Tensor? query_start_loc," |
| 533 | "Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, " |
| 534 | "int pad_slot_id, bool is_weight_packed) -> " |
| 535 | "Tensor"); |
| 536 | ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu); |
| 537 | |
| 538 | ops.def( |
| 539 | "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor " |
| 540 | "weight, Tensor? bias, bool silu_activation," |
| 541 | "Tensor? num_accepted_tokens, Tensor? conv_state_indices, int " |
| 542 | "pad_slot_id, " |
| 543 | "bool is_weight_packed) -> Tensor"); |
| 544 | ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu); |
| 545 | #endif |
| 546 | |
| 547 | // Layernorm |
| 548 | // Apply Root Mean Square (RMS) Normalization to the input tensor. |
| 549 | ops.def( |
| 550 | "rms_norm(Tensor! out, Tensor input, Tensor? weight, float epsilon) -> " |
| 551 | "()"); |
| 552 | ops.impl("rms_norm", torch::kCPU, &rms_norm); |
| 553 | |
| 554 | // In-place fused Add and RMS Normalization. |
| 555 | ops.def( |
| 556 | "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, " |
| 557 | "float epsilon) -> ()"); |
| 558 | ops.impl("fused_add_rms_norm", torch::kCPU, &fused_add_rms_norm); |
| 559 | |
| 560 | // Rotary embedding |
| 561 | // Apply GPT-NeoX or GPT-J style rotary embedding to query and key. |
| 562 | ops.def( |
| 563 | "rotary_embedding(Tensor positions, Tensor! query," |
| 564 | " Tensor!? key, int head_size," |
| 565 | " Tensor cos_sin_cache, bool is_neox, int " |
| 566 | "rope_dim_offset=0, bool inverse=False) -> ()"); |
| 567 | ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding); |
| 568 | |
| 569 | // Quantization |
| 570 | #if defined(__AVX512F__) || defined(__AVX2__) || \ |
| 571 | (defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__) || \ |
| 572 | defined(__riscv_v) || defined(__s390x__) |
| 573 | // Helper function to release oneDNN handlers |
| 574 | ops.def("release_dnnl_matmul_handler(int handler) -> ()", |
| 575 | &release_dnnl_matmul_handler); |
| 576 | |
| 577 | // Create oneDNN GEMM handler |
| 578 | ops.def( |
| 579 | "create_onednn_mm_handler(Tensor b, int " |
| 580 | "primitive_cache_size) -> int", |
| 581 | &create_onednn_mm_handler); |
| 582 | |
| 583 | // oneDNN GEMM |
| 584 | ops.def( |
| 585 | "onednn_mm(Tensor! c, Tensor a, Tensor? bias, " |
| 586 | "Tensor handler_tensor) -> ()"); |
| 587 | ops.impl("onednn_mm", torch::kCPU, &onednn_mm); |
| 588 | |
| 589 | // Check if oneDNN was built with ACL backend |
| 590 | ops.def("is_onednn_acl_supported() -> bool", &is_onednn_acl_supported); |
| 591 | |
| 592 | // Create oneDNN W8A8 handler |
| 593 | ops.def( |
| 594 | "create_onednn_scaled_mm_handler(Tensor b, Tensor b_scales, ScalarType " |
| 595 | "output_type, bool dynamic_act_quant, bool use_azp, int " |
| 596 | "primitive_cache_size) -> int", |
| 597 | &create_onednn_scaled_mm_handler); |
| 598 | |
| 599 | // oneDNN scaled_mm for W8A8 with static per-tensor activation quantization |
| 600 | ops.def( |
| 601 | "onednn_scaled_mm(Tensor! c, Tensor a, Tensor a_scales, Tensor? azp, " |
| 602 | "Tensor? azp_adj, Tensor? bias, Tensor handler_tensor) -> ()"); |
| 603 | ops.impl("onednn_scaled_mm", torch::kCPU, &onednn_scaled_mm); |
| 604 | |
| 605 | // Compute int8 quantized tensor for given scaling factor. |
| 606 | ops.def( |
| 607 | "static_scaled_int8_quant(Tensor! out, Tensor input, Tensor scale," |
| 608 | "Tensor? azp) -> ()"); |
| 609 | ops.impl("static_scaled_int8_quant", torch::kCPU, &static_scaled_int8_quant); |
| 610 | |
| 611 | // Compute int8 quantized tensor and scaling factor |
| 612 | ops.def( |
| 613 | "dynamic_scaled_int8_quant(Tensor! out, Tensor input, Tensor! scale, " |
| 614 | "Tensor!? azp) -> ()"); |
| 615 | ops.impl("dynamic_scaled_int8_quant", torch::kCPU, |
| 616 | &dynamic_scaled_int8_quant); |
| 617 | #endif |
| 618 | |
| 619 | // SHM CCL |
| 620 | #if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__)) || \ |
| 621 | defined(__powerpc64__) |
| 622 | ops.def( |
| 623 | "init_shm_manager(str name, int group_size, int rank, int thread_num) -> " |
| 624 | "int", |
| 625 | &init_shm_manager); |
| 626 | ops.def("join_shm_manager(int handle, str name) -> str", &join_shm_manager); |
| 627 | ops.def("shm_allreduce(int handle, Tensor! data) -> ()"); |
| 628 | ops.impl("shm_allreduce", torch::kCPU, &shm_allreduce); |
| 629 | ops.def( |
| 630 | "shm_gather(int handle, Tensor data, Tensor[](a!)? outputs, int dst) -> " |
| 631 | "()"); |
| 632 | ops.impl("shm_gather", torch::kCPU, &shm_gather); |
| 633 | ops.def( |
| 634 | "shm_all_gather(int handle, Tensor data, Tensor! output) -> " |
| 635 | "()"); |
| 636 | ops.impl("shm_all_gather", torch::kCPU, &shm_all_gather); |
| 637 | ops.def( |
| 638 | "shm_send_tensor_list(int handle, Tensor[](a) tensor_list, int dst) -> " |
| 639 | "()"); |
| 640 | ops.impl("shm_send_tensor_list", torch::kCPU, &shm_send_tensor_list); |
| 641 | ops.def("shm_recv_tensor_list(int handle, int src) -> Tensor[](a)", |
| 642 | &shm_recv_tensor_list); |
| 643 | #endif // #if defined(__AVX512F__) || defined(__aarch64__) |
| 644 | |
| 645 | // sgl-kernels |
| 646 | #if defined(__AVX512BF16__) && defined(__AVX512F__) && defined(__AVX512VNNI__) |
| 647 | ops.def( |
| 648 | "weight_packed_linear(Tensor(a0!) out, Tensor(a1) mat1, Tensor(a2) mat2, " |
| 649 | "Tensor(a3)? " |
| 650 | "bias, bool is_vnni) -> ()"); |
| 651 | ops.impl("weight_packed_linear", torch::kCPU, &weight_packed_linear); |
| 652 | ops.def("convert_weight_packed(Tensor! weight) -> Tensor"); |
| 653 | ops.impl("convert_weight_packed", torch::kCPU, &convert_weight_packed); |
| 654 | ops.def("convert_scale_packed(Tensor! scale) -> Tensor"); |
| 655 | ops.impl("convert_scale_packed", torch::kCPU, &convert_scale_packed); |
| 656 | ops.def( |
| 657 | "fused_experts_cpu(Tensor(a0!) out, Tensor hidden_states, Tensor w1, " |
| 658 | "Tensor w2, Tensor topk_weights, Tensor topk_ids, " |
| 659 | "int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, " |
| 660 | "Tensor? w1_zero, Tensor? w2_zero, Tensor? a1_scale, int[]? block_size, " |
| 661 | "Tensor? w1_bias, Tensor? w2_bias, float? alpha, float? limit, " |
| 662 | "bool is_vnni) -> " |
| 663 | "()"); |
| 664 | ops.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu); |
| 665 | ops.def( |
| 666 | "int8_scaled_mm_with_quant(Tensor mat1, Tensor mat2, Tensor scales2, " |
| 667 | "Tensor? bias, ScalarType out_dtype, bool is_vnni) -> Tensor"); |
| 668 | ops.impl("int8_scaled_mm_with_quant", torch::kCPU, |
| 669 | &int8_scaled_mm_with_quant); |
| 670 | |
| 671 | // Adapted from sglang: FP8 W8A16 kernel |
| 672 | ops.def( |
| 673 | "fp8_scaled_mm_cpu(Tensor(a0!) mat1, Tensor(a1!) mat2, Tensor(a2!) " |
| 674 | "scales2, SymInt[] block_size, Tensor? bias, ScalarType out_dtype, " |
| 675 | "bool is_vnni) -> Tensor"); |
| 676 | ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu); |
| 677 | |
| 678 | // Adapted from sglang: FP8 W8A8 kernels |
| 679 | ops.def( |
| 680 | "float8_linear_prepack_cpu(Tensor weight, Tensor scales) -> (Tensor, " |
| 681 | "Tensor)"); |
| 682 | ops.impl("float8_linear_prepack_cpu", torch::kCPU, |
| 683 | &float8_linear_prepack_impl); |
| 684 | ops.def( |
| 685 | "fp8_scaled_mm_with_quant(Tensor act, Tensor? act_scales, bool " |
| 686 | "channelwise, " |
| 687 | "Tensor weight, Tensor weight_scales, Tensor? bias, ScalarType " |
| 688 | "output_dtype) -> Tensor"); |
| 689 | ops.impl("fp8_scaled_mm_with_quant", torch::kCPU, &fp8_scaled_mm_with_quant); |
| 690 | ops.def( |
| 691 | "quantize_fp8e4m3_vec(Tensor input, bool channelwise, Tensor? scale_opt) " |
| 692 | "-> (Tensor, Tensor)"); |
| 693 | ops.impl("quantize_fp8e4m3_vec", torch::kCPU, &quantize_fp8e4m3_vec); |
| 694 | |
| 695 | // Adapted from sglang: MLA CPU kernels (AMX-only, DeepSeek V2/V3/R1) |
| 696 | ops.def( |
| 697 | "decode_attention_cpu(Tensor query, Tensor k_buffer, Tensor v_buffer, " |
| 698 | "Tensor(a!) output, Tensor? key, Tensor? value, Tensor? loc, Tensor " |
| 699 | "attn_logits, Tensor req_to_token, Tensor req_pool_indices, Tensor " |
| 700 | "seq_lens, float sm_scale, float logit_cap, bool is_cross_attn, int " |
| 701 | "sliding_window_size, Tensor? encoder_lens, Tensor? sinks) -> ()"); |
| 702 | ops.impl("decode_attention_cpu", torch::kCPU, &decode_attention_cpu); |
| 703 | |
| 704 | ops.def( |
| 705 | "extend_attention_cpu(Tensor q_extend, Tensor? k_extend, Tensor? " |
| 706 | "v_extend, Tensor(a!) o_extend, Tensor k_buffer, Tensor v_buffer, " |
| 707 | "Tensor req_to_token, Tensor req_pool_indices, Tensor seq_lens, Tensor " |
| 708 | "extend_seq_lens, Tensor extend_start_loc, int max_len_extend, float " |
| 709 | "sm_scale, float logit_cap, bool is_cross_attn, int " |
| 710 | "sliding_window_size, Tensor? encoder_lens, Tensor? sinks, Tensor? " |
| 711 | "tree_mask=None) -> ()"); |
| 712 | ops.impl("extend_attention_cpu", torch::kCPU, &extend_attention_cpu); |
| 713 | |
| 714 | ops.def( |
| 715 | "bmm_cpu(Tensor(a!) out, Tensor mat1, Tensor mat2, bool is_vnni, " |
| 716 | "Tensor? scale) -> ()"); |
| 717 | ops.impl("bmm_cpu", torch::kCPU, &bmm_cpu); |
| 718 | |
| 719 | ops.def( |
| 720 | "concat_and_cache_mla_cpu(Tensor kv_c_normed, Tensor k_pe, " |
| 721 | "Tensor(a!) kv_cache, Tensor slot_mapping) -> ()"); |
| 722 | ops.impl("concat_and_cache_mla_cpu", torch::kCPU, &concat_and_cache_mla_cpu); |
| 723 | |
| 724 | // Adapted from sglang: DeepSeek-V4 mHC gating kernels |
| 725 | ops.def( |
| 726 | "hc_pre_fused_cpu(Tensor x, Tensor hc_fn, Tensor hc_scale, " |
| 727 | "Tensor hc_base, int hc_mult, int sinkhorn_iters, float rms_eps, " |
| 728 | "float hc_eps) -> (Tensor, Tensor, Tensor)"); |
| 729 | ops.impl("hc_pre_fused_cpu", torch::kCPU, &hc_pre_fused_cpu); |
| 730 | |
| 731 | ops.def( |
| 732 | "hc_post_fused_cpu(Tensor x, Tensor residual, Tensor post, " |
| 733 | "Tensor comb) -> Tensor"); |
| 734 | ops.impl("hc_post_fused_cpu", torch::kCPU, &hc_post_fused_cpu); |
| 735 | |
| 736 | ops.def( |
| 737 | "hc_head_fused_cpu(Tensor x, Tensor hc_fn, Tensor " |
| 738 | "hc_scale, Tensor hc_base, float hc_eps, float norm_eps) -> " |
| 739 | "Tensor"); |
| 740 | ops.impl("hc_head_fused_cpu", torch::kCPU, &hc_head_fused_cpu); |
| 741 | |
| 742 | // Adapted from sglang: DeepSeek-V4 fp8_ds_mla cache-write kernel |
| 743 | ops.def( |
| 744 | "fused_qnorm_rope_kv_insert_cpu(Tensor q, Tensor kv, Tensor positions, " |
| 745 | "Tensor(a!) swa_kv_cache_2d, Tensor slot_mapping, Tensor " |
| 746 | "cos_sin_cache, int q_head_padded, float eps, int cache_block_size) " |
| 747 | "-> Tensor"); |
| 748 | ops.impl("fused_qnorm_rope_kv_insert_cpu", torch::kCPU, |
| 749 | &fused_qnorm_rope_kv_insert_cpu); |
| 750 | |
| 751 | // vLLM-native: DeepSeek-V4 sparse MQA attention kernel |
| 752 | ops.def( |
| 753 | "flash_mla_with_kvcache_cpu(Tensor(a!) out, Tensor q, Tensor " |
| 754 | "window_cache_2d, Tensor window_slots, int window_block_size, Tensor " |
| 755 | "compressed_cache_2d, Tensor compressed_slots, int " |
| 756 | "compressed_block_size, Tensor attn_sink, float scale) -> ()"); |
| 757 | ops.impl("flash_mla_with_kvcache_cpu", torch::kCPU, |
| 758 | &flash_mla_with_kvcache_cpu); |
| 759 | |
| 760 | // Adapted from sglang: DeepSeek-V4 compressor kernels |
| 761 | ops.def( |
| 762 | "save_partial_states_cpu(Tensor kv, Tensor score, Tensor ape, " |
| 763 | "Tensor positions, Tensor(a!) state_cache, Tensor slot_mapping) " |
| 764 | "-> ()"); |
| 765 | ops.impl("save_partial_states_cpu", torch::kCPU, &save_partial_states_cpu); |
| 766 | |
| 767 | ops.def( |
| 768 | "compress_norm_rope_store_cpu(Tensor state_cache, Tensor " |
| 769 | "gather_slots, Tensor positions, Tensor kv_slot_mapping, Tensor " |
| 770 | "rms_norm_weight, float rms_norm_eps, Tensor cos_sin_cache, " |
| 771 | "Tensor(a!) kv_cache_2d, int kv_cache_block_size, int compress_ratio) " |
| 772 | "-> ()"); |
| 773 | ops.impl("compress_norm_rope_store_cpu", torch::kCPU, |
| 774 | &compress_norm_rope_store_cpu); |
| 775 | |
| 776 | ops.def( |
| 777 | "compress_norm_rope_store_indexer_cpu(Tensor state_cache, Tensor " |
| 778 | "gather_slots, Tensor positions, Tensor kv_slot_mapping, Tensor " |
| 779 | "rms_norm_weight, float rms_norm_eps, Tensor cos_sin_cache, " |
| 780 | "Tensor(a!) kv_cache_2d, int kv_cache_block_size, int compress_ratio) " |
| 781 | "-> ()"); |
| 782 | ops.impl("compress_norm_rope_store_indexer_cpu", torch::kCPU, |
| 783 | &compress_norm_rope_store_indexer_cpu); |
| 784 | |
| 785 | // Adapted from sglang: DeepSeek-V4 sparse indexer paged MQA-logits kernel |
| 786 | // (DECODE path) |
| 787 | ops.def( |
| 788 | "fp8_paged_mqa_logits_cpu(Tensor q_fp8, Tensor kvcache_fp8, Tensor " |
| 789 | "weight, Tensor seq_lens, Tensor page_table, int block_size, int " |
| 790 | "max_seq_len) -> Tensor"); |
| 791 | ops.impl("fp8_paged_mqa_logits_cpu", torch::kCPU, &fp8_paged_mqa_logits_cpu); |
| 792 | |
| 793 | ops.def( |
| 794 | "topk_transform_512_cpu(Tensor scores, Tensor seq_lens, Tensor " |
| 795 | "page_tables, Tensor(a!) out_page_indices, int page_size, Tensor(b!)? " |
| 796 | "out_raw_indices) -> ()"); |
| 797 | ops.impl("topk_transform_512_cpu", torch::kCPU, &topk_transform_512_cpu); |
| 798 | |
| 799 | ops.def( |
| 800 | "fused_indexer_q_rope_quant_cpu(Tensor positions, Tensor index_q, " |
| 801 | "Tensor index_q_cos_sin_cache, Tensor(a!) index_q_fp8, Tensor " |
| 802 | "index_weights, float index_weights_softmax_scale, float " |
| 803 | "index_weights_head_scale, Tensor(b!) index_weights_out) -> ()"); |
| 804 | ops.impl("fused_indexer_q_rope_quant_cpu", torch::kCPU, |
| 805 | &fused_indexer_q_rope_quant_cpu); |
| 806 | |
| 807 | // Adapted from sglang: DeepSeek-V4 MoE routing kernels |
| 808 | ops.def( |
| 809 | "biased_topk_cpu(Tensor hidden_states, Tensor gating_output, Tensor " |
| 810 | "correction_bias, int topk, bool renormalize, str scoring_func, int " |
| 811 | "num_fused_shared_experts, float? routed_scaling_factor, bool " |
| 812 | "apply_routed_scaling_factor_on_output) -> (Tensor, Tensor)"); |
| 813 | ops.impl("biased_topk_cpu", torch::kCPU, &biased_topk_cpu); |
| 814 | |
| 815 | ops.def( |
| 816 | "hash_topk_cpu(Tensor gating_output, Tensor tid2eid, int topk, str " |
| 817 | "scoring_func, int num_fused_shared_experts, int num_experts, float " |
| 818 | "routed_scaling_factor) -> (Tensor, Tensor)"); |
| 819 | ops.impl("hash_topk_cpu", torch::kCPU, &hash_topk_cpu); |
| 820 | |
| 821 | // vLLM-native: inverse GPT-J RoPE for _o_proj's output de-rotation |
| 822 | ops.def( |
| 823 | "inverse_gptj_rope_o_proj_cpu(Tensor o, Tensor positions, Tensor " |
| 824 | "cos_sin_cache, int rope_dim) -> Tensor"); |
| 825 | ops.impl("inverse_gptj_rope_o_proj_cpu", torch::kCPU, |
| 826 | &inverse_gptj_rope_o_proj_cpu); |
| 827 | #endif |
| 828 | |
| 829 | #if (defined(__AVX512BF16__) && defined(__AVX512F__) && \ |
| 830 | defined(__AVX512VNNI__)) || \ |
| 831 | defined(__riscv) |
| 832 | // Adapted from sglang: INT4 W4A8 kernels |
| 833 | ops.def( |
| 834 | "convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor " |
| 835 | "scales, int quant_method_4bit) -> (Tensor, " |
| 836 | "Tensor, Tensor)"); |
| 837 | ops.impl("convert_weight_packed_scale_zp", torch::kCPU, |
| 838 | &convert_weight_packed_scale_zp); |
| 839 | |
| 840 | ops.def( |
| 841 | "int4_scaled_mm_cpu(Tensor(a0!) x, Tensor(a1!) w, Tensor(a2!) w_zeros, " |
| 842 | "Tensor(a3!) w_scales, Tensor? bias) -> Tensor"); |
| 843 | ops.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu); |
| 844 | #endif |
| 845 | |
| 846 | // Adapted from sglang: GDN kernels |
| 847 | ops.def( |
| 848 | "chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, " |
| 849 | "Tensor g, Tensor beta, " |
| 850 | "Tensor initial_state, bool output_final_state, Tensor cu_seqlens, bool " |
| 851 | "head_first, " |
| 852 | "bool use_qk_l2norm_in_kernel, Tensor initial_state_indices, float " |
| 853 | "eps=1e-5) -> (Tensor, Tensor)"); |
| 854 | ops.impl("chunk_gated_delta_rule_cpu", torch::kCPU, |
| 855 | &chunk_gated_delta_rule_cpu); |
| 856 | ops.def( |
| 857 | "fused_sigmoid_gating_delta_rule_update_cpu(Tensor A_log, Tensor " |
| 858 | "dt_bias, Tensor q, Tensor k, Tensor v, Tensor " |
| 859 | "a, Tensor b, Tensor(a!) initial_state_source, Tensor " |
| 860 | "initial_state_indices, Tensor cu_seqlens, bool " |
| 861 | "use_qk_l2norm_in_kernel, float softplus_beta=1.0, float " |
| 862 | "softplus_threshold=20.0) -> Tensor"); |
| 863 | ops.impl("fused_sigmoid_gating_delta_rule_update_cpu", torch::kCPU, |
| 864 | &fused_sigmoid_gating_delta_rule_update_cpu); |
| 865 | ops.def( |
| 866 | "fused_sigmoid_gating_delta_rule_update_spec_cpu(Tensor A_log, Tensor " |
| 867 | "dt_bias, Tensor q, Tensor k, Tensor v, Tensor a, Tensor b, " |
| 868 | "Tensor(a!) initial_state_source, Tensor spec_state_indices, " |
| 869 | "Tensor num_accepted_tokens, Tensor cu_seqlens, bool " |
| 870 | "use_qk_l2norm_in_kernel, float softplus_beta=1.0, float " |
| 871 | "softplus_threshold=20.0) -> Tensor"); |
| 872 | ops.impl("fused_sigmoid_gating_delta_rule_update_spec_cpu", torch::kCPU, |
| 873 | &fused_sigmoid_gating_delta_rule_update_spec_cpu); |
| 874 | ops.def( |
| 875 | "fused_gdn_gating_cpu(Tensor A_log, Tensor a, Tensor b, Tensor dt_bias) " |
| 876 | "-> (Tensor, Tensor)"); |
| 877 | ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu); |
| 878 | |
| 879 | // Pure runtime CPU ISA capability check. |
| 880 | ops.def("cpu_has_amx_fp8() -> bool", &cpu_has_amx_fp8); |
| 881 | |
| 882 | // CPU attention kernels |
| 883 | ops.def("cpu_attn_has_isa(str isa) -> bool", &cpu_attn_has_isa); |
| 884 | ops.def( |
| 885 | "get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, " |
| 886 | "int head_dim, Tensor seq_lens, ScalarType dtype, Tensor " |
| 887 | "query_start_loc, bool casual, int window_size, str isa_hint, bool " |
| 888 | "enable_kv_split, Tensor? dynamic_causal, " |
| 889 | "str kv_cache_dtype=\"auto\") -> Tensor", |
| 890 | &get_scheduler_metadata); |
| 891 | ops.def( |
| 892 | "cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) " |
| 893 | "key_cache, Tensor(a3!) value_cache, Tensor slot_mapping, str isa, " |
| 894 | "float k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> " |
| 895 | "()", |
| 896 | &cpu_attn_reshape_and_cache); |
| 897 | ops.def( |
| 898 | "cpu_attention_with_kv_cache(Tensor query, Tensor key_cache, Tensor " |
| 899 | "value_cache, Tensor(a3!) output, Tensor query_start_loc, Tensor " |
| 900 | "seq_lens, float scale, bool causal, Tensor? alibi_slopes, SymInt " |
| 901 | "sliding_window_size, Tensor block_table, " |
| 902 | "float softcap, Tensor scheduler_metadata, Tensor? s_aux, Tensor? " |
| 903 | "dynamic_causal, " |
| 904 | "float k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> " |
| 905 | "()", |
| 906 | &cpu_attention_with_kv_cache); |
| 907 | |
| 908 | // placeholders |
| 909 | ops.def("static_scaled_fp8_quant() -> ()", placeholder_op); |
| 910 | ops.def("dynamic_scaled_fp8_quant() -> ()", placeholder_op); |
| 911 | ops.def("dynamic_per_token_scaled_fp8_quant() -> ()", placeholder_op); |
| 912 | |
| 913 | // WNA16 |
| 914 | #if defined(__AVX512F__) || defined(__riscv_v) || defined(__s390x__) || \ |
| 915 | defined(__powerpc__) |
| 916 | ops.def( |
| 917 | "cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, " |
| 918 | "Tensor scales, Tensor? zeros, Tensor? bias, SymInt " |
| 919 | "pack_factor, str isa_hint) -> ()"); |
| 920 | ops.impl("cpu_gemm_wna16", torch::kCPU, &cpu_gemm_wna16); |
| 921 | #endif |
| 922 | |
| 923 | // fused moe |
| 924 | ops.def( |
| 925 | "prepack_moe_weight(Tensor weight, Tensor(a1!) packed_weight, str isa) " |
| 926 | "-> ()"); |
| 927 | ops.impl("prepack_moe_weight", torch::kCPU, &prepack_moe_weight); |
| 928 | ops.def( |
| 929 | "cpu_fused_moe(Tensor(a0!) output, Tensor input, Tensor w13, Tensor w2, " |
| 930 | "Tensor? w13_bias, Tensor? w2_bias, Tensor topk_weights, Tensor topk_id, " |
| 931 | "bool skip_weighted, " |
| 932 | "str act, str isa) -> ()"); |
| 933 | ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe); |
| 934 | |
| 935 | #if (defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && \ |
| 936 | !defined(__APPLE__)) || \ |
| 937 | defined(__powerpc64__) |
| 938 | ops.def( |
| 939 | "prepack_moe_weight_int8(Tensor weight, Tensor(a1!) packed_weight, " |
| 940 | "str isa) -> ()"); |
| 941 | ops.impl("prepack_moe_weight_int8", torch::kCPU, &prepack_moe_weight_int8); |
| 942 | ops.def( |
| 943 | "cpu_fused_moe_int8(Tensor(a0!) output, Tensor input, Tensor w13, " |
| 944 | "Tensor w2, Tensor w13_scale, Tensor w2_scale, Tensor? w13_bias, " |
| 945 | "Tensor? w2_bias, Tensor topk_weights, Tensor topk_id, bool " |
| 946 | "skip_weighted, str act, str isa) -> ()"); |
| 947 | ops.impl("cpu_fused_moe_int8", torch::kCPU, &cpu_fused_moe_int8); |
| 948 | #endif // ARM_I8MM+BF16 or __powerpc64__ |
| 949 | ops.def( |
| 950 | "mla_decode_kvcache(" |
| 951 | " Tensor! out, Tensor query, Tensor kv_cache," |
| 952 | " float scale, Tensor block_tables, Tensor seq_lens) -> ()"); |
| 953 | ops.impl("mla_decode_kvcache", torch::kCPU, &mla_decode_kvcache); |
| 954 | |
| 955 | ops.def( |
| 956 | "compute_slot_mapping_kernel_impl(Tensor query_start_loc, Tensor " |
| 957 | "positions, Tensor block_table, Tensor(a3!) slot_mapping, SymInt " |
| 958 | "block_size) -> ()", |
| 959 | &compute_slot_mapping_kernel_impl); |
| 960 | |
| 961 | // Mamba CPU kernels |
| 962 | ops.def( |
| 963 | "causal_conv1d_update_cpu_vec(" |
| 964 | "Tensor(a0!) x, Tensor(a1!) conv_state, Tensor weight, " |
| 965 | "Tensor? bias, str? activation, Tensor? conv_state_indices, " |
| 966 | "Tensor? query_start_loc, SymInt pad_slot_id) -> Tensor", |
| 967 | &causal_conv1d_update_cpu_impl); |
| 968 | |
| 969 | ops.def( |
| 970 | "selective_state_update_cpu(" |
| 971 | "Tensor(a0!) state, Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, " |
| 972 | "Tensor? D, Tensor? z, Tensor? dt_bias, bool dt_softplus, " |
| 973 | "Tensor? state_batch_indices, Tensor? dst_state_batch_indices, " |
| 974 | "SymInt null_block_id, Tensor(a13!) out, " |
| 975 | "Tensor? num_accepted_tokens, Tensor? cu_seqlens) -> ()", |
| 976 | &selective_state_update_cpu_impl); |
| 977 | |
| 978 | ops.def( |
| 979 | "mamba_chunk_scan_fwd_cpu(" |
| 980 | "Tensor(a0!) out, Tensor(a1!) final_states, " |
| 981 | "Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, " |
| 982 | "Tensor? D, Tensor? z, Tensor cu_seqlens) -> ()", |
| 983 | &mamba_chunk_scan_fwd_cpu_impl); |
| 984 | |
| 985 | ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env); |
| 986 | |
| 987 | // Fused sampling kernels |
| 988 | ops.def("fused_gumbel_argmax(Tensor logits, Tensor seeds) -> Tensor"); |
| 989 | ops.impl("fused_gumbel_argmax", torch::kCPU, &fused_gumbel_argmax); |
| 990 | |
| 991 | ops.def("greedy_argmax(Tensor logits) -> Tensor"); |
| 992 | ops.impl("greedy_argmax", torch::kCPU, &greedy_argmax); |
| 993 | |
| 994 | // Speculative decoding kernels |
| 995 | ops.def( |
| 996 | "eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, " |
| 997 | "Tensor valid_sampled_tokens_count, Tensor query_start_loc_gpu, " |
| 998 | "Tensor(a3!) token_indices_to_sample, " |
| 999 | "Tensor(a4!) num_rejected_tokens_gpu, " |
| 1000 | "SymInt num_reqs) -> ()", |
| 1001 | &cpu_utils::eagle_prepare_inputs_padded_kernel_impl); |
| 1002 | ops.def( |
| 1003 | "eagle_prepare_next_token_padded_kernel_impl(" |
| 1004 | "Tensor sampled_token_ids, Tensor discard_request_mask, " |
| 1005 | "Tensor backup_next_token_ids, Tensor(a3!) next_token_ids, " |
| 1006 | "Tensor(a4!) valid_sampled_tokens_count, SymInt vocab_size, " |
| 1007 | "SymInt num_sampled_tokens_per_req, SymInt num_reqs) -> ()", |
| 1008 | &cpu_utils::eagle_prepare_next_token_padded_kernel_impl); |
| 1009 | ops.def( |
| 1010 | "eagle_step_slot_mapping_metadata_kernel_impl(" |
| 1011 | "Tensor positions, Tensor block_table, Tensor(a2!) seq_lens, " |
| 1012 | "Tensor(a3!) out_clamped_positions, Tensor(a4!) out_slot_mapping, " |
| 1013 | "SymInt block_size, SymInt max_model_len, SymInt PAD_ID) -> ()", |
| 1014 | &cpu_utils::eagle_step_slot_mapping_metadata_kernel_impl); |
| 1015 | ops.def( |
| 1016 | "copy_and_expand_eagle_inputs_kernel_impl(" |
| 1017 | "Tensor target_token_ids, Tensor target_positions, " |
| 1018 | "Tensor next_token_ids, Tensor(a3!) out_input_ids, " |
| 1019 | "Tensor(a4!) out_positions, " |
| 1020 | "Tensor(a5!) out_is_rejected_token_mask, " |
| 1021 | "Tensor(a6!) out_is_masked_token_mask, " |
| 1022 | "Tensor(a7!) out_new_token_indices, " |
| 1023 | "Tensor(a8!) out_hidden_state_mapping, " |
| 1024 | "Tensor query_start_loc, Tensor query_end_loc, " |
| 1025 | "SymInt padding_token_id, SymInt parallel_drafting_token_id, " |
| 1026 | "SymInt total_input_tokens, SymInt num_padding_slots_per_request, " |
| 1027 | "bool shift_input_ids) -> ()", |
| 1028 | &cpu_utils::copy_and_expand_eagle_inputs_kernel_impl); |
| 1029 | ops.def( |
| 1030 | "copy_and_expand_dflash_inputs_kernel_impl(" |
| 1031 | "Tensor next_token_ids, Tensor target_positions, " |
| 1032 | "Tensor(a2!) out_input_ids, Tensor(a3!) out_context_positions, " |
| 1033 | "Tensor(a4!) out_query_positions, " |
| 1034 | "Tensor(a5!) out_context_slot_mapping, " |
| 1035 | "Tensor(a6!) out_query_slot_mapping, " |
| 1036 | "Tensor(a7!) out_token_indices, Tensor block_table, " |
| 1037 | "Tensor query_start_loc, Tensor? num_rejected_tokens, " |
| 1038 | "SymInt parallel_drafting_token_id, SymInt block_size, " |
| 1039 | "SymInt num_query_per_req, SymInt num_speculative_tokens, " |
| 1040 | "SymInt total_input_tokens, bool has_num_rejected) -> ()", |
| 1041 | &cpu_utils::copy_and_expand_dflash_inputs_kernel_impl); |
| 1042 | ops.def( |
| 1043 | "rejection_greedy_sample_kernel_impl(" |
| 1044 | "Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, " |
| 1045 | "Tensor draft_token_ids, Tensor target_argmax, " |
| 1046 | "Tensor bonus_token_ids, Tensor? is_greedy, " |
| 1047 | "SymInt max_spec_len) -> ()", |
| 1048 | &cpu_utils::rejection_greedy_sample_kernel_impl); |
| 1049 | ops.def( |
| 1050 | "rejection_random_sample_kernel_impl(" |
| 1051 | "Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, " |
| 1052 | "Tensor draft_token_ids, Tensor? draft_probs, " |
| 1053 | "Tensor target_probs, Tensor bonus_token_ids, " |
| 1054 | "Tensor recovered_token_ids, Tensor uniform_probs, " |
| 1055 | "Tensor? is_greedy, SymInt max_spec_len, SymInt vocab_size, " |
| 1056 | "bool no_draft_probs) -> ()", |
| 1057 | &cpu_utils::rejection_random_sample_kernel_impl); |
| 1058 | ops.def( |
| 1059 | "expand_kernel_impl(Tensor(a0!) output, Tensor input, " |
| 1060 | "Tensor cu_num_tokens, SymInt replace_from, " |
| 1061 | "SymInt replace_to) -> ()", |
| 1062 | &cpu_utils::expand_kernel_impl); |
| 1063 | ops.def( |
| 1064 | "sample_recovered_tokens_kernel_impl(" |
| 1065 | "Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, " |
| 1066 | "Tensor draft_token_ids, Tensor? draft_probs, " |
| 1067 | "Tensor target_probs, Tensor inv_q, SymInt vocab_size, " |
| 1068 | "bool no_draft_probs) -> ()", |
| 1069 | &cpu_utils::sample_recovered_tokens_kernel_impl); |
| 1070 | } |
| 1071 | |
| 1072 | REGISTER_EXTENSION(TORCH_EXTENSION_NAME) |