chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 14 Packages Projects Releases Wiki Activity

[CPU] add mamba fla kernels for Qwen3-next (#12324)

Browse Source
This commit is contained in:
blzheng
2025-12-06 14:16:23 +08:00
committed by GitHub
parent 7b0c7ad163
commit d257bf87b9
4 changed files with 577 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
+28
View File
@@ -284,6 +284,25 @@ std::tuple<at::Tensor, at::Tensor> rotary_embedding_cpu(
// CPU and memory binding
std::string init_cpu_threads_env(const std::string& cpu_ids);
// fused_sigmoid_gating_delta_rule_update
at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
const at::Tensor& A_log,
const at::Tensor& dt_bias,
const at::Tensor& q,
const at::Tensor& k,
const at::Tensor& v,
const at::Tensor& a,
const at::Tensor& b,
at::Tensor& initial_state_source,
const at::Tensor& initial_state_indices,
const at::Tensor& cu_seqlens,
bool use_qk_l2norm_in_kernel,
double softplus_beta = 1.0,
double softplus_threshold = 20.0);
// fused_gdn_gating
std::tuple<at::Tensor, at::Tensor>
fused_gdn_gating_cpu(const at::Tensor& A_log, const at::Tensor& a, const at::Tensor& b, const at::Tensor& dt_bias);
// fused_qkvzba_split_reshape_cat_cpu
std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> fused_qkvzba_split_reshape_cat_cpu(
const at::Tensor& mixed_qkvz,
@@ -453,6 +472,15 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
// CPU and memory binding
m.def("init_cpu_threads_env(str cpu_ids) -> str");
// fused_sigmoid_gating_delta_rule_update
m.def(
"fused_sigmoid_gating_delta_rule_update_cpu(Tensor A_log, Tensor dt_bias, Tensor q, Tensor k, Tensor v, Tensor "
"a, Tensor b, Tensor(a!) initial_state_source, Tensor initial_state_indices, Tensor cu_seqlens, bool "
"use_qk_l2norm_in_kernel, float softplus_beta=1.0, float softplus_threshold=20.0) -> Tensor");
m.impl("fused_sigmoid_gating_delta_rule_update_cpu", torch::kCPU, &fused_sigmoid_gating_delta_rule_update_cpu);
// fused_gdn_gating
m.def("fused_gdn_gating_cpu(Tensor A_log, Tensor a, Tensor b, Tensor dt_bias) -> (Tensor, Tensor)");
m.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu);
// fused_qkvzba_split_reshape_cat_cpu
m.def(
"fused_qkvzba_split_reshape_cat_cpu(Tensor mixed_qkvz, Tensor mixed_ba, int num_heads_qk, int num_heads_v, int "