minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[CPU] add fused input proj for qwen3.5 (#31171)

Browse Source
This commit is contained in:
Ma Mingfei
2026-07-15 15:06:24 +08:00
committed by GitHub
parent a649b5a9db
commit 41e0b4b369
5 changed files with 238 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
+9
View File
@@ -502,6 +502,10 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> fused_qkvzba_split_re
int64_t head_qk,
int64_t head_v);
// fused_input_proj_cpu
std::tuple<at::Tensor, at::Tensor>
fused_input_proj_cpu(at::Tensor& hidden_states, at::Tensor& qkvz_weight, at::Tensor& ba_weight, bool is_vnni);
// image preprocessor
std::tuple<at::Tensor, at::Tensor> image_preprocess_cpu(
at::TensorList images,
@@ -844,6 +848,11 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
"num_heads_v, int "
"head_qk, int head_v) -> (Tensor, Tensor, Tensor, Tensor)");
m.impl("fused_qkvzba_split_reshape_cat_contiguous_cpu", torch::kCPU, &fused_qkvzba_split_reshape_cat_contiguous_cpu);
// fused_input_proj_cpu
m.def(
"fused_input_proj_cpu(Tensor hidden_states, Tensor qkvz_weight, Tensor ba_weight, bool is_vnni) -> (Tensor, "
"Tensor)");
m.impl("fused_input_proj_cpu", torch::kCPU, &fused_input_proj_cpu);
// image preprocessor
m.def(