执行摘要

• 一句话：修复FlexAttention掩码方法因绑定self导致跟踪崩溃
• 推荐动作：值得合入，属于低风险高质量修复。建议读者关注 @staticmethod 在避免意外闭包捕获方面的设计模式。

功能与动机

torch.nn.attention.flex_attention.create_block_mask 期望传入一个签名 (b, h, q_idx, kv_idx) 的纯可调用对象。当传入 self._causal_mask 这种绑定方法时，闭包会捕获整个 TorchFlexAttnBackend 实例，导致在特定 torch 版本下图形跟踪崩溃。PR body 明确描述了该问题。

实现拆解

在 python/sglang/srt/layers/attention/torch_flex_backend.py 中，

1. 定位 _causal_mask 和 _decode_mask 两个实例方法定义处。
2. 为两个方法添加 @staticmethod 装饰器，并移除参数列表中的 self。
3. 原有调用位置 self._causal_mask 和 self._decode_mask 无需修改——Python 会自动解析为底层函数。

关键文件：

• python/sglang/srt/layers/attention/torch_flex_backend.py（模块 注意力；类别 source；类型 core-logic；符号 _causal_mask, _decode_mask）: 核心修复文件，修改 _causal_mask 和 _decode_mask 为静态方法以消除绑定方法引起的闭包捕获问题。

关键符号：_causal_mask, _decode_mask

关键源码片段

python/sglang/srt/layers/attention/torch_flex_backend.py

核心修复文件，修改 _causal_mask 和 _decode_mask 为静态方法以消除绑定方法引起的闭包捕获问题。

# python/sglang/srt/layers/attention/torch_flex_backend.py
​
    @staticmethod
    def _causal_mask(b, h, q_idx, kv_idx):
        # 纯函数：判断 query 位置是否大于等于 key/value 位置，仅依赖参数
        return q_idx >= kv_idx
​
    @staticmethod
    def _decode_mask(b, h, q_idx, kv_idx):
        # 纯函数：判断 query 位置是否小于等于 key/value 位置，仅依赖参数
        return q_idx <= kv_idx

评论区精华

该 PR 无审核评论，但从 commit message 和 body 可以看出，确认了 create_block_mask 对纯可调用对象的依赖以及绑定方法导致闭包捕获的问题是唯一的动机。

• 暂无高价值评论线程

风险与影响

• 风险：风险极低：
  • 两个方法均为纯函数，不访问任何实例属性。
  • 所有外部调用均通过 self._causal_mask / self._decode_mask 方式，@staticmethod 兼容绑定方法调用语法。
  • 仅变更 2 行，不影响运行时性能。
  • 影响：直接解决使用最新 torch 时 TorchFlexAttnBackend 在 create_block_mask 阶段崩溃的问题，影响范围限定于使用 FlexAttention 的 attention 路径（EXTEND 和 DECODE 模式）。
  • 风险标记：暂无

关联脉络

• 暂无明显关联 PR