执行摘要

• 一句话：修复 FlexAttention 编码器块掩码编译爆炸
• 推荐动作：该 PR 解决了实际的编译性能问题，改动简洁且附带完善的测试覆盖。虽然 FlexAttention 未来可能被废弃，但在当前阶段合并此 PR 是值得的，特别是对使用 AMD GPU 进行编码器推理的用户。建议在合并后关注是否有用户报告编码器注意力下的精度问题。

功能与动机

FlexAttention 在编码器（encoder-only）场景下默认使用 16 令牌的小块大小，导致 torch.compile/Inductor 在构建块掩码时需要处理大量物理到逻辑映射，造成编译爆炸，首次请求冷缓存耗时高达 300 秒（在 AMD MI355 上）。将默认块大小提升至 128 可大幅降低编译开销，冷缓存测试降至 38 秒。

实现拆解

1. 区分 encoder 和 decoder：在 FlexAttentionMetadataBuilder.__init__（vllm/v1/attention/backends/flex_attention.py）中增加 uses_paged_kv 变量，通过 isinstance(kv_cache_spec, EncoderOnlyAttentionSpec) 判断是否为编码器模式，并传递给 _get_block_sizes。

2. 调整默认块大小逻辑：在静态方法 _get_block_sizes 中新增 uses_paged_kv 参数。原先仅依赖 supports_small_blocks（PyTorch 版本）决定默认块大小；现在改为 supports_small_blocks and uses_paged_kv，即只有 paged KV 注意力且 PyTorch >= 2.9 时才使用小块（16 / cache_block_size），否则统一使用 128。这确保了编码器注意力始终使用 128 块大小，避免编译爆炸。

3. 更新配置文档：在 vllm/config/attention.py 中更新 flex_attn_q_block_size 和 flex_attn_kv_block_size 的 docstring，明确默认值区分场景：paged KV 注意力在 PyTorch >= 2.9 时使用小块，编码器注意力或旧 PyTorch 时使用 128。

4. 新增参数化单元测试：在 tests/kernels/test_flex_attention.py 中添加 test_flex_attention_default_block_sizes（参数化覆盖四种组合）以及 test_flex_attention_explicit_block_sizes_override_encoder_defaults（验证显式覆盖）。

5. 增强编码器正确性测试：将 test_encoder_flex_attention_vs_default_backend 的 prompts 替换为跨越 128 token 块边界的重复文本（120/130/254 词），并将 max_model_len 从 100 提升至 384，确保注意力跨越多个块且覆盖边界。

6. 确定性化长文本测试：在 tests/entrypoints/pooling/embed/test_online_long_text.py 中，将 _generate_random_text 函数改为使用 random.Random(word_count) 局部种子随机数生成器，替换全局 random，确保 token 计数和编译器形状在多次运行间完全可重现。

关键文件：

• vllm/v1/attention/backends/flex_attention.py（模块 注意力后端；类别 source；类型 core-logic；符号 FlexAttentionMetadataBuilder.init, FlexAttentionMetadataBuilder._get_block_sizes）: 核心逻辑变更：在 init 中判断是否为 paged KV 并传递标志；_get_block_sizes 方法根据新策略决定默认块尺寸
• tests/kernels/test_flex_attention.py（模块 注意力测试；类别 test；类型 test-coverage；符号 test_flex_attention_default_block_sizes, test_flex_attention_explicit_block_sizes_override_encoder_defaults）: 新增参数化测试覆盖默认块大小组合；修改 encoder 正确性测试以跨越 128 块边界
• vllm/config/attention.py（模块 配置；类别 source；类型 core-logic）: 更新 flex_attn_q/kv_block_size 的文档注释，反映 encoder/paged KV 不同的默认行为
• tests/entrypoints/pooling/embed/test_online_long_text.py（模块 长文本测试；类别 test；类型 test-coverage；符号 _generate_random_text）: 将 _generate_random_text 改为确定性生成，确保 token 计数和编译器形状可重现

关键符号：FlexAttentionMetadataBuilder._get_block_sizes, FlexAttentionMetadataBuilder.init, test_flex_attention_default_block_sizes, test_flex_attention_explicit_block_sizes_override_encoder_defaults, _generate_random_text

关键源码片段

vllm/v1/attention/backends/flex_attention.py

核心逻辑变更：在 init 中判断是否为 paged KV 并传递标志；_get_block_sizes 方法根据新策略决定默认块尺寸

    @staticmethod
    def _get_block_sizes(
        attn_cfg,
        supports_small_blocks: bool,
        cache_block_size: int,
        uses_paged_kv: bool, # 新增：是否为 paged KV 注意力（decoder）
    ) -> tuple[int, int]:
        # 只有 paged KV 且 PyTorch >= 2.9 才使用小块（16），否则统一 128
        use_small_blocks = supports_small_blocks and uses_paged_kv
        q_block_size = 16 if use_small_blocks else 128
        kv_block_size = cache_block_size if use_small_blocks else 128
​
        # 如果用户显式设置了块大小，则覆盖默认值
        q_block_size = attn_cfg.flex_attn_q_block_size or q_block_size
        kv_block_size = attn_cfg.flex_attn_kv_block_size or kv_block_size
        # 校验块大小必须为 2 的幂且不超过最大值（与原来一致）
        ...
        return q_block_size, kv_block_size

评论区精华

核心讨论围绕 FlexAttention 的废弃前景展开。hmellor 评论说：“Potentially unneeded if FlexAttention is soon to be deprecated?” 作者 AndreasKaratzas 回应：“True but until it is deprecated, we will need to render CI green on some of the platforms (like MI355).” 审批人 MatthewBonanni 赞同：“LGTM, doesn't hurt to merge while we decide on deprecation”，认为在决定废弃前合并无害，最终批准 PR。

• FlexAttention 废弃前景下的 PR 必要性 (design): MatthewBonanni 审批，认为合并无害（“doesn't hurt to merge while we decide on deprecation”），批准 PR。

风险与影响

• 风险：主要风险包括：
  • 回归风险：如果未来有编码器模型依赖于小块默认值（如某些定制 mask mod），强制 128 块可能导致精度或性能退化。但测试已覆盖跨越块边界的正确性，且显式覆盖机制存在，回归可能性低。
  • 废弃计划不确定性：FlexAttention 后端可能在未来版本中被废弃，届时此项改动将失去价值，但不会带来损害。
  • 平台差异：此修正是针对 AMD MI355 平台测试的，其他 GPU 平台的 behavior 可能略有不同，但逻辑是通用的。
  • 影响：对用户的影响：编码器注意力模型的首次请求编译时间从 300 秒降至约 38 秒（在 MI355 上），极大提升冷启动体验。对 decoder 模型无影响。对系统：无显著额外开销。对团队：维护了 FlexAttention 在编码器场景下的可用性，并为可能的废弃提供了兼容过渡。
  • 风险标记：encoder-only 默认变更, FlexAttention 废弃风险

关联脉络

• 暂无明显关联 PR