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

[diffusion] fix: avoid flashattention forward context lookup (#27153)

Browse Source
This commit is contained in:
Mick
2026-06-04 08:11:50 +08:00
committed by GitHub
parent cfb7fb4fad
commit 3790173b3b
1 changed files with 5 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py
+3 -3
View File
@@ -7,7 +7,6 @@ import torch
from sglang.jit_kernel.flash_attention import flash_attn_varlen_func from sglang.jit_kernel.flash_attention import flash_attn_varlen_func
from sglang.multimodal_gen.runtime.layers.utils import register_custom_op from sglang.multimodal_gen.runtime.layers.utils import register_custom_op
from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
from sglang.multimodal_gen.runtime.platforms import ( from sglang.multimodal_gen.runtime.platforms import (
AttentionBackendEnum, AttentionBackendEnum,
) )
@@ -381,9 +380,10 @@ class FlashAttentionImpl(AttentionImpl):
*, *,
return_softmax_lse: bool = False, return_softmax_lse: bool = False,
): ):
attn_metadata: FlashAttentionMetadata = get_forward_context().attn_metadata if attn_metadata is not None:
if attn_metadata is not None and attn_metadata.max_seqlen_q is None: if attn_metadata.max_seqlen_q is None:
attn_metadata.max_seqlen_q = query.shape[1] attn_metadata.max_seqlen_q = query.shape[1]
if attn_metadata.max_seqlen_k is None:
attn_metadata.max_seqlen_k = key.shape[1] attn_metadata.max_seqlen_k = key.shape[1]
max_seqlen_q = attn_metadata.max_seqlen_q max_seqlen_q = attn_metadata.max_seqlen_q
max_seqlen_k = attn_metadata.max_seqlen_k max_seqlen_k = attn_metadata.max_seqlen_k