执行摘要

• 一句话：避免 FlashAttention 中全局 forward context 查找
• 推荐动作：值得合并。这是一个小而精巧的优化：既消除了不必要的全局查找，又修复了潜在的元数据覆盖 bug。代码审查中的建议已被采纳，逻辑正确。推荐阅读以了解如何安全移除全局依赖。

功能与动机

避免在每次 forward 调用时通过全局 forward context 重新读取 attention 元数据，改用调用方显式传递的 attn_metadata 参数，减少间接查找开销并简化数据流。PR 描述中明确提到 "stop rereading the global forward context" 并使用 "the supplied attention metadata when present; otherwise fall back to the dense Q/K shapes"。

实现拆解

1. 移除全局上下文导入和调用：删掉 import 语句中的 get_forward_context，并在 forward 方法中移除 attn_metadata = get_forward_context().attn_metadata 这一行。
2. 修复条件逻辑：将原本的一层条件 if attn_metadata is not None and attn_metadata.max_seqlen_q is None 拆分为两层：先判断 attn_metadata is not None，再分别判断 max_seqlen_q 和 max_seqlen_k 是否为 None，从而只在需要时填充，避免覆盖已有的元数据值。
3. 保持 fallback 路径：当 attn_metadata 为 None 时，仍然使用 query.shape[1] 和 key.shape[1] 作为回退。
   无测试、配置或部署配套改动。

关键文件：

• python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py（模块 扩散模块；类别 source；类型 dependency-wiring）: 核心变更文件，修改了 import 和 FlashAttentionImpl.forward 逻辑，移除全局 forward context 依赖并修复条件分支。

关键符号：FlashAttentionImpl.forward

关键源码片段

python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py

核心变更文件，修改了 import 和 FlashAttentionImpl.forward 逻辑，移除全局 forward context 依赖并修复条件分支。

# python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py
​
class FlashAttentionImpl(AttentionImpl):
    # ... __init__ 省略 ...
​
    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: AttentionMetadata = None, # 现在由调用方传入
        *,
        return_softmax_lse: bool = False,
    ):
        # 移除：attn_metadata = get_forward_context().attn_metadata
        # 改为直接使用传入的 attn_metadata 参数
        if attn_metadata is not None:
            # 仅在元数据中对应字段为 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:
            # 没有元数据时，直接使用 query/key 长度作为回退
            max_seqlen_q = query.shape[1]
            max_seqlen_k = key.shape[1]
        # 后续 FA 调用使用 max_seqlen_q 和 max_seqlen_k ...

评论区精华

AI 审查（gemini-code-assist[bot]）指出原有条件逻辑存在 bug：当 attn_metadata 非空且 max_seqlen_q 已存在时，原条件为 false，会进入 else 分支并用 query.shape[1] / key.shape[1] 覆盖元数据中的值。审查建议采用嵌套 if 结构修复。该建议已被作者采纳并体现在最终代码中。

• 逻辑 bug：条件分支导致已有元数据被覆盖 (correctness): 作者采纳建议，将条件拆分为两层嵌套 if，分别检查 max_seqlen_q 和 max_seqlen_k 是否为 None。

风险与影响

• 风险：风险极低。变更仅涉及单个函数内约 10 行代码，逻辑清晰，且已有 A/B 验证（Cosmos3 Nano T2V 模型输出 SHA 匹配、性能持平）。唯一潜在风险是调用方可能未正确传递 attn_metadata，但 fallback 路径保证了向后兼容。
• 影响：影响范围限于 sglang/multimodal_gen 模块的 FlashAttention 后端。对于使用该后端的扩散模型推理，消除了每次 forward 调用中对全局上下文的查找，可能带来轻微性能提升（PR 中 A/B 测试显示候选时间 51.0760s vs 基线 51.2327s）。不改变 API 接口或行为。
• 风险标记：无测试覆盖

关联脉络

• PR #27143 PR Stack base (推测): PR body 中提到 A/B 测试基于 #27143 stack，表明该 PR 属于同一功能线。