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

This commit is contained in:
Mick
2026-06-04 08:11:50 +08:00
committed by GitHub
parent cfb7fb4fad
commit 3790173b3b
@@ -7,7 +7,6 @@ import torch
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.managers.forward_context import get_forward_context
from sglang.multimodal_gen.runtime.platforms import (
AttentionBackendEnum,
)
@@ -381,10 +380,11 @@ class FlashAttentionImpl(AttentionImpl):
*,
return_softmax_lse: bool = False,
):
attn_metadata: FlashAttentionMetadata = get_forward_context().attn_metadata
if attn_metadata is not None and attn_metadata.max_seqlen_q is None:
attn_metadata.max_seqlen_q = query.shape[1]
attn_metadata.max_seqlen_k = key.shape[1]
if attn_metadata is not None:
if attn_metadata.max_seqlen_q is None:
attn_metadata.max_seqlen_q = query.shape[1]
if attn_metadata.max_seqlen_k is None:
attn_metadata.max_seqlen_k = key.shape[1]
max_seqlen_q = attn_metadata.max_seqlen_q
max_seqlen_k = attn_metadata.max_seqlen_k
else: