[diffusion] fix: avoid flashattention forward context lookup (#27153)
This commit is contained in:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user