执行摘要

• 一句话：支持 FlashInfer 非因果注意力用于 DFlash 推测解码
• 推荐动作：本 PR 变更集中、逻辑清晰，值得需要理解 vLLM 注意力后端选择机制和 DFlash 实现的开发者精读。其中关于非因果路径的门控设计可作为类似功能扩展的参考。

功能与动机

RTX 4090 上使用 FP8 KV 缓存需要 DFlash + FlashInfer 组合（PR body）。DFlash 草案模型需要非因果注意力后端；此前 FlashInfer 后端不支持非因果模式，导致 DFlash 无法与 FlashInfer 配合使用。

实现拆解

实现分为三步：

1. 在 FlashInferBackend 中添加 supports_non_causal 类方法返回 True，使其能通过注意力选择器的非因果检测。
2. 在 FlashInferMetadata 中新增 causal 字段，并在 FlashInferMetadataBuilder.build 中根据 causal 标志决定批次拆分逻辑：因果模式下保持 prefill/decode 拆分；非因果模式下将所有请求作为 prefill 处理，跳过 TRTLLM decode 路径。
3. 在 _get_prefill_wrapper 中新增非因果分支，当 causal=False 时创建独立的 BatchPrefillWithPagedKVCacheWrapper 用于非因果 prefill，并对 DCP/NVFP4 等不兼容配置显式抛出 NotImplementedError。
   测试文件 test_attention_selector.py 中调整 FlashInfer 非因果测试预期从 False 改为 True。
   文档 attention_backends.md 中 FlashInfer 的 Non-Causal 列由 ❌ 更新为 ✅。

关键文件：

• vllm/v1/attention/backends/flashinfer.py（模块 注意力后端；类别 source；类型 core-logic；符号 supports_non_causal, FlashInferMetadata.causal, _get_prefill_wrapper, build）: 核心实现，添加非因果支持，包括supports_non_causal方法、causal字段、非因果prefill包装器、构建逻辑中非因果分支。
• tests/kernels/attention/test_attention_selector.py（模块 测试；类别 test；类型 test-coverage）: 更新测试预期，验证FlashInfer支持非因果
• docs/design/attention_backends.md（模块 文档；类别 docs；类型 documentation）: 更新 FlashInfer 非因果支持状态

关键符号：supports_non_causal, _get_prefill_wrapper, build

关键源码片段

vllm/v1/attention/backends/flashinfer.py

核心实现，添加非因果支持，包括supports_non_causal方法、causal字段、非因果prefill包装器、构建逻辑中非因果分支。

@classmethod
def supports_non_causal(cls) -> bool:
    # 声明 FlashInfer 后端支持非因果注意力（用于 DFlash 草案模型）
    return True
​
def _get_prefill_wrapper(
    self,
    causal: bool = True,
) -> BatchPrefillWithPagedKVCacheWrapper | BatchDCPPrefillWrapper:
    if not causal:
        # 非因果模式：禁用 DCP（尚未支持）和 NVFP4 KV 缓存（不兼容）
        if self.use_dcp:
            raise NotImplementedError(
                "FlashInfer non-causal prefill is not supported with DCP yet."
            )
        if self.is_kvcache_nvfp4:
            raise NotImplementedError(
                "FlashInfer non-causal attention is not supported with "
                "NVFP4 KV cache."
            )
        # 创建或复用独立的非因果 prefill 包装器，
        # 避免与因果 prefill 包装器共用导致状态混乱
        if self._noncausal_prefill_wrapper is None:
            self._noncausal_prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
                self._get_workspace_buffer(),
                get_kv_cache_layout(),
                backend="auto",
            )
        return self._noncausal_prefill_wrapper
    # 因果模式：原有逻辑
    if self._prefill_wrapper is None:
        ...
    return self._prefill_wrapper

评论区精华

reviewer benchislett 指出：当 FlashInfer 宣称 supports_non_causal 时，不应自动选择 TRTLLM 后端然后失败；应通过门控 prefill_use_trtllm 和 all_uses_trtllm 来避免选择 TRTLLM。另外要求在非因果且使用了 TRTLLM decode 的情况下打印一次性 fallback 警告，并注意日志简洁。

• 非因果模式下避免自动选择 TRTLLM 后端 (design): PR 作者实现了门控，在非因果时设置 prefill_use_trtllm=False 并限制 decode_use_trtllm。
• 非因果 + TRTLLM decode 时打印一次性 fallback 警告 (design): PR 作者添加了一条一次性警告日志。

风险与影响

• 风险：核心风险在于非因果路径的 prefill 是否会被错误地选择到不支持的 TRTLLM、DCP 或 NVFP4 配置。当前通过显式检查确保在非因果时禁用这些路径，但若未来添加新后端或 KV 缓存 dtype 可能遗漏检查。另外，非因果模式下将所有 decode 请求作为 prefill 处理可能带来性能影响，但当前设计仅用于 DFlash 草案模型，影响可控。
• 影响：对使用 DFlash 推测解码且需要 FlashInfer 后端的用户（如 RTX 4090 FP8 用户）是必要支持。对不使用 DFlash 的用户无影响。系统层面减少了 FlashInfer 和 DFlash 的组合限制，拓展了推测解码硬件覆盖范围。
• 风险标记：核心后端变更, 新功能测试覆盖有限, 不支持的配置显式抛异常

关联脉络

• PR #39995 Original large PR for DFlash + FlashInfer support: 本PR从此PR拆分而来，聚焦FlashInfer后端非因果支持。