From 3790173b3b9ae79a909c9e495723f276093aa685 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 4 Jun 2026 08:11:50 +0800 Subject: [PATCH] [diffusion] fix: avoid flashattention forward context lookup (#27153) --- .../runtime/layers/attention/backends/flash_attn.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py index 31372e2e1..8f8ff57c9 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py @@ -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: