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

[Diffusion] Revert CPU AMX optimizations (#30716)

Browse Source
This commit is contained in:
Mick
2026-07-10 09:09:38 +08:00
committed by GitHub
parent 1e75ba236e
commit 5ce5e1ee3e
12 changed files with 13 additions and 166 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/flash_attn.cpp
+4 -5
View File
@@ -443,8 +443,7 @@ at::Tensor flash_attn_varlen_func(
const at::Tensor& cu_seqlens_k,
int64_t max_seqlen_q,
int64_t max_seqlen_k,
bool causal,
const std::optional<double>& sm_scale) {
bool causal) {
CHECK_LAST_DIM_CONTIGUOUS_INPUT(q);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(k);
CHECK_LAST_DIM_CONTIGUOUS_INPUT(v);
@@ -481,7 +480,7 @@ at::Tensor flash_attn_varlen_func(
TORCH_CHECK(head_size_v % 2 == 0, "invalid head_size_v ", head_size_v);
// softmax scale
double _sm_scale = sm_scale.has_value() ? sm_scale.value() : 1.0 / std::sqrt(static_cast<double>(head_size));
double sm_scale = 1.0 / std::sqrt(static_cast<double>(head_size));
// check whether the batch has variant lengths
const bool is_varlen =
@@ -523,7 +522,7 @@ at::Tensor flash_attn_varlen_func(
k_strideH,
v_strideN,
v_strideH,
_sm_scale,
sm_scale,
sz,
causal);
} else {
@@ -546,7 +545,7 @@ at::Tensor flash_attn_varlen_func(
k_strideH,
v_strideN,
v_strideH,
_sm_scale,
sm_scale,
sz,
causal);
}
sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
+2 -3
View File
@@ -235,8 +235,7 @@ at::Tensor flash_attn_varlen_func(
const at::Tensor& cu_seqlens_k,
int64_t max_seqlen_q,
int64_t max_seqlen_k,
bool causal,
const std::optional<double>& sm_scale);
bool causal);
// linear attention
std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
@@ -677,7 +676,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
// 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, float? sm_scale) -> Tensor");
"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