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

[diffusion] fix: fix using the sub-optimal fa kernel (#16382)

Browse Source
This commit is contained in:
Qi Yuhang
2026-01-05 16:25:41 +08:00
committed by GitHub
parent bebd625ba1
commit dcfb92ddbc
1 changed files with 5 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py
+5 -8
View File
@@ -201,18 +201,15 @@ class FlashAttentionImpl(AttentionImpl):
if use_upstream: if use_upstream:
bsz, seqlen, nheads_q, d = q_shape bsz, seqlen, nheads_q, d = q_shape
bsz_k, seqlen_k, nheads_k, d_k = k_shape q_ = query.contiguous()
bsz_v, seqlen_v, nheads_v, d_v = v_shape k_ = key.contiguous()
q_ = query.contiguous().reshape(bsz * seqlen, nheads_q, d) v_ = value.contiguous()
k_ = key.contiguous().reshape(bsz * seqlen, nheads_k, d)
v_ = value.contiguous().reshape(bsz * seqlen, nheads_v, d)
cu = _get_cu_seqlens(q_.device.index, bsz, seqlen)
out = flash_attn_varlen_func_upstream( out = flash_attn_varlen_func_upstream(
q_, q_,
k_, k_,
v_, v_,
cu, None,
cu, None,
seqlen, seqlen,
seqlen, seqlen,
softmax_scale=self.softmax_scale, softmax_scale=self.softmax_scale,