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

[CPU] optimize flash_attn_varlen_func (#15708)

Browse Source
This commit is contained in:
Ma Mingfei
2026-01-29 22:07:05 -08:00
committed by GitHub
parent 336dc4579e
commit 88f7759402
10 changed files with 1189 additions and 339 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
+17
View File
@@ -109,6 +109,17 @@ void extend_attention_cpu(
double sm_scale,
double logit_cap);
// flash attention
at::Tensor flash_attn_varlen_func(
const at::Tensor& q,
const at::Tensor& k,
const at::Tensor& v,
const at::Tensor& cu_seqlens_q,
const at::Tensor& cu_seqlens_k,
int64_t max_seqlen_q,
int64_t max_seqlen_k,
bool causal);
// linear attention
std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
const at::Tensor& query,
@@ -382,6 +393,12 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
"extend_start_loc, int max_len_extend, float sm_scale, float logit_cap) -> ()");
m.impl("extend_attention_cpu", torch::kCPU, &extend_attention_cpu);
// flash attn
m.def(
"flash_attn_varlen_func(Tensor q, Tensor k, Tensor v, Tensor cu_seqlens_q, Tensor cu_seqlens_k, "
"int max_seqlen_q, int max_seqlen_k, bool causal) -> Tensor");
m.impl("flash_attn_varlen_func", torch::kCPU, &flash_attn_varlen_func);
// linear attn
m.def(
"chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, Tensor g, Tensor beta, "