执行摘要

• 一句话：FlashInfer 接受 KV-sharing 省略的 K/V，修复启动崩溃
• 推荐动作：值得精读：源码改动只有约 20 行，但 review 讨论覆盖了 CPU 开销约束下的代码放置、契约穷尽性、以及未初始化张量静默算错三类问题，是「小 diff、高信息密度」的典型样本。建议重点关注三个决策点：一是 K/V 切片下沉到 DCP 调用点而非在 padding 块加 None 判断；二是 key is None 必须配合 gemma4_mtp 改传 None 才「correct by construction」；三是删除单测的取舍——若要复用该方案，建议补一个不依赖内部实现细节的端到端用例，或在 DCP 组合上保留一个 canary。

功能与动机

54917 为避免在 KV-shared 层计算无用的 K/V 投影，把 Gemma4/Gemma3n 的共享层改为 self.attn(q, None, None)，但 FlashInfer 后端仍无条件执行 key[:num_actual_tokens] / value[:num_actual_tokens]，导致 nightly 构建 #87584 的 :nvidia: (B200) Spec Decode Speculators + MTP Nightly 中 test_mtp_correctness[gemma4-e4b] 在引擎初始化阶段抛 TypeError: 'NoneType' object is not subscriptable。#55789 提出整体回滚 #54917，其 issue body 明确写道 A narrow fix is preferable to this revert，并指出 #54917 修的是真实问题（save_pretrained/fine-tuned checkpoint 会省略共享层的 k_proj/v_proj/k_norm）。本 PR 因此选择前向修复：让 FlashInfer 直接读取共享 KV cache 并接受 None，而不是回滚（#55789）或让 Gemma4 继续传占位张量（#55871，两者均未合并）。

实现拆解

1. 放宽后端函数契约（vllm/v1/attention/backends/flashinfer.py）：FlashInferImpl.forward 的 key / value 形参类型由 torch.Tensor 改为 torch.Tensor | None，docstring 补充「KV-sharing decoder layer 时为 None」。调用侧（Attention.forward → unified_attention_with_output）早已允许传 None，只有后端未实现该契约，因此这一步是把隐含假设显式化，不引入新分支。

2. K/V 的 padding 切片从公共块下沉到 DCP prefill 调用点：删除 query = query[:num_actual_tokens] 旁边的两行 K/V 切片，改为在 use_dcp 分支调用 prefill_wrapper.run(...) 时传入 key[num_decode_tokens:num_actual_tokens] / value[num_decode_tokens:num_actual_tokens]。原因是该方法运行在 piece-wise CUDA Graph 的 eager 区域，函数内 NOTE(woosuk) 明确警告 view/slice 的 CPU 开销不可忽视；K/V 全文只在该调用点被消费，下沉后非 DCP 的常见路径（含 KV-sharing decode）完全不触碰 K/V，既去掉 None 判断也去掉两次张量操作。

3. DCP × KV-sharing 组合显式拒绝：在 use_dcp 分支入口加 if key is None or value is None: raise NotImplementedError("FlashInfer DCP prefill does not support KV-sharing layers")。DCP prefill 需要当前层的 new-token K/V 参与 A2A 交换，共享层无法提供，因此把「静默算错」改为「明确报错」。注释说明这是本 PR 有意保留的已知限制。

4. 移除 gemma4_mtp 的占位 K/V（vllm/model_executor/models/gemma4_mtp.py）：Gemma4MTPAttention.forward 删掉 num_tokens / kv_dummy = torch.empty(...) 及注释，改为 self.attn(q, None, None)。收益有三：省掉每个 draft 层 forward 的一次分配；让后端 key is None 判断真正穷尽（否则占位张量会绕过第 3 步的检查，把未初始化显存当 new-token K/V 送进 BatchDCPPrefillWrapper.run）；与目标模型 gemma4.py:548 及既有 Attention 层契约保持一致。共享的抽象 Attention 接口未改动。

5. 测试与验证配套：初版在 tests/v1/attention/test_attention_backends.py 增加了 cache-only decode 与 DCP prefill 拒绝两个单测，但 LucasWilkinson 认为测试过度依赖内部实现细节，作者在 f767a7e4 全部删除，最终 PR 不含测试文件改动；回归改由既有 Gemma4 nightly 与 PR 描述中的 GSM8K 评测兜底（E2B 混合后端 strict-match 0.4829 / flexible-extract 0.6550，参考值 0.485；加 MTP 助理后 strict-match 0.4829）。

关键文件：

• vllm/v1/attention/backends/flashinfer.py（模块 注意力后端；类别 source；类型 core-logic；符号 FlashInferImpl.forward）: FlashInfer 后端唯一入口，本 PR 的核心修复点：放宽 key/value 为可选、把 K/V padding 切片下沉到 DCP prefill 调用点、为 DCP × KV-sharing 增加显式 NotImplementedError。
• vllm/model_executor/models/gemma4_mtp.py（模块 投机解码；类别 source；类型 data-contract；符号 Gemma4MTPAttention.forward）: 移除每次 draft 层 forward 都会分配的未初始化占位 K/V，改为传 None，使 FlashInfer 的 key is None 判断真正穷尽，并避免 DCP prefill 消费脏显存。

关键符号：FlashInferImpl.forward, Gemma4MTPAttention.forward

关键源码片段

vllm/v1/attention/backends/flashinfer.py

FlashInfer 后端唯一入口，本 PR 的核心修复点：放宽 key/value 为可选、把 K/V padding 切片下沉到 DCP prefill 调用点、为 DCP × KV-sharing 增加显式 NotImplementedError。

# vllm/v1/attention/backends/flashinfer.py —— FlashInferImpl.forward（整理后的关键片段）
def forward(
    self,
    layer: torch.nn.Module,
    query: torch.Tensor,
    # KV-sharing 层不再计算 K/V 投影，Attention 层会直接传 None，
    # 因此后端契约必须允许 None：K/V 由共享目标层的 KV cache 提供
    key: torch.Tensor | None,
    value: torch.Tensor | None,
    kv_cache: torch.Tensor,
    attn_metadata: FlashInferMetadata,
    output: torch.Tensor,
    output_scale: torch.Tensor | None = None,
    output_block_scale: torch.Tensor | None = None,
) -> torch.Tensor:
    if attn_metadata is None:
        # Profiling run（预热阶段），直接返回零输出
        return output.fill_(0)
​
    num_actual_tokens = attn_metadata.num_actual_tokens
​
    # Inputs and outputs may be padded for CUDA graphs
    query = query[:num_actual_tokens]
    output_padded = output
    output = output[:num_actual_tokens]
    # 这里不再有 key = key[:num_actual_tokens] / value = value[:num_actual_tokens]：
    # 本方法运行在 piece-wise CUDA Graph 的 eager 区域，
    # NOTE(woosuk) 要求尽量减少 view / slice 这类 CPU 开销偏高的操作；
    # K/V 全文只在下面 DCP prefill 一处被消费，切片因此下沉到该调用点，
    # 非 DCP 的常见路径（含 KV-sharing decode）完全不触碰 K/V。
​
    # ...（cascade 分支、decode / prefill 公共准备逻辑保持不变）
​
    use_dcp = self.dcp_world_size > 1
​
    if num_prefill_tokens > 0:
        # ...（prefill_query 切分、maybe_quant_query 等预处理）
        if not prefill_use_trtllm:
            assert isinstance(attn_metadata.prefill, FIPrefill)
            prefill_wrapper = attn_metadata.prefill.wrapper
            assert prefill_wrapper is not None
            if use_dcp:
                # DCP prefill 需要当前层的 new-token K/V 参与 A2A 交换，
                # KV-sharing 层拿不到真实张量，必须显式报错；
                # 否则占位 / 未初始化张量会被静默算成错误的 draft attention
                if key is None or value is None:
                    raise NotImplementedError(
                        "FlashInfer DCP prefill does not support KV-sharing layers"
                    )
                assert isinstance(prefill_wrapper, BatchDCPPrefillWrapper)
                # ...（window_left / logits_soft_cap / sm_scale 一致性断言）
​
                prefill_wrapper.run(
                    layer,
                    prefill_query,
                    kv_cache_tuple,
                    # 边界从公共块搬到这里，语义不变：仍只取 [num_decode_tokens, num_actual_tokens)
                    key[num_decode_tokens:num_actual_tokens],
                    value[num_decode_tokens:num_actual_tokens],
                    out=output[num_decode_tokens:],
                )
            else:
                # 非 DCP 分支只消费 paged KV cache，不读 K/V，
                # 因此 KV-sharing 层传 None 也能正常跑——本次修复的核心
                kv_cache_for_fi = nvfp4_kv_data if self.is_kvcache_nvfp4 else kv_cache_tuple
                # ...（BatchPrefillWithPagedKVCacheWrapper.run 等）

vllm/model_executor/models/gemma4_mtp.py

移除每次 draft 层 forward 都会分配的未初始化占位 K/V，改为传 None，使 FlashInfer 的 key is None 判断真正穷尽，并避免 DCP prefill 消费脏显存。

# vllm/model_executor/models/gemma4_mtp.py —— Gemma4MTPAttention.forward（整理后的完整实现）
class Gemma4MTPAttention(nn.Module):
    # __init__ 中：kv_sharing_target_layer_name 由 Gemma4Proposer._setup_gemma4_kv_sharing()
    # 在模型构造完成后回填，本层因此不创建 k_proj / v_proj / k_norm，
    # 并设置 self.is_kv_shared_layer = True
​
    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        **kwargs,
    ) -> torch.Tensor:
        q, _ = self.q_proj(hidden_states)
​
        q = q.unflatten(-1, (self.num_heads, self.head_dim))
        q = self.q_norm(q)
        q = q.flatten(-2, -1)
​
        q, _ = self.rotary_emb(positions, q, None)
​
        # Attention 从共享目标层的 KV cache 读 K/V，直接传 None 表示本层不产出 K/V。
        # 旧实现构造 kv_dummy = torch.empty(...) 冒充 K/V，有两个问题：
        # 1) 每个 draft 层 forward 都多一次显存分配；
        # 2) 占位张量非 None，会绕过 FlashInfer DCP prefill 的 key is None 检查，
        # 把未初始化显存当 new-token K/V 送进 BatchDCPPrefillWrapper.run，
        # 静默算出错误的 draft attention。
        attn_output = self.attn(q, None, None)
        output, _ = self.o_proj(attn_output)
        return output

评论区精华

1）gemma4_mtp 占位 K/V 必须随本 PR 一起清掉（hmellor → 已解决）

key is None isn't exhaustive: gemma4_mtp.py:249-258 allocates kv_dummy = torch.empty(...) and passes it as both K and V, so that layer walks past this check and hands uninitialized memory to BatchDCPPrefillWrapper.run as new-token K/V, silently computing wrong draft attention instead of raising here.

hmellor 的依据很硬：attention.py:533-544 在 K/V 为 None 时已会跳过 .view() 与 unified_kv_cache_update，gemma4.py:548 的目标模型在同一后端与 compile 路径上就是传 None，占位张量只因为本 PR 要修的 bug 才存在。作者在 356f859e 照做，并补充 Gemma4 MTP 评测。

2）K/V 切片放在哪一层（性能约束下的位置选择，已采纳）

Sinking the slice to the DCP use site (the pattern in rocm_attn.py:461-462) drops both the None handling and two tensor ops from every non-DCP call, which is the common path.

hmellor 指出在 padding 块里加 if key is not None 会给最常见的非 DCP 路径加两个 Python 分支，而函数上方 25 行的注释恰恰在警告 slice 的 CPU 开销。作者随后把切片移到 DCP 调用点，并保留 num_actual_tokens 上界。

3）DCP 组合的拒绝时机（hmellor 提建议，作者明确留作后续）

this rejects the unsupported combination lazily, once per forward, so a DCP Gemma4 deployment starts up clean ... and then raises from inside a compiled region on the first prefill.

hmellor 建议在 __init__ 加校验（dcp_world_size 此时已可用），但也承认 gemma4.py:360 是在构造后回填 kv_sharing_target_layer_name，init 看不到该字段，因此属于可选增强。作者回复「leaving the initialization-time validation out of scope and keeping the runtime error」，该线程以遗留项结束。

4）删除「过度细化」的单测（LucasWilkinson → 已执行）

nit: please remove the over-specified test

LucasWilkinson 对 tests/v1/attention/test_attention_backends.py 的两条相同 nit 被作者在 f767a7e4 执行，理由是该回归已被 Gemma4 nightly 覆盖；代价是本 PR 最终不含任何自动化测试。

5）抽象接口的 None 注解缺口（遗留）

Attention.forward, unified_attention_with_output, and AttentionImpl.forward all still annotate key: torch.Tensor despite None already flowing through them; the custom-op schema doesn't catch it because eager dispatch does no boxed type check.

hmellor 认为这是既有问题、可以不改，但本 PR 是「start closing that gap」的正确起点。

• gemma4_mtp 的占位 K/V 必须随本 PR 一起清掉 (correctness): 已采纳：gemma4_mtp 改传 None，删除 kv_dummy 分配，E2B MTP strict-match 0.4829 ± 0.0138（参考 0.485）验证无精度回退。
• K/V 切片下沉到 DCP 调用点，避免拖慢常见非 DCP 路径 (performance): 已采纳：padding 块不再处理 K/V，切片移入 prefill_wrapper.run 调用的实参。
• DCP × KV-sharing 的拒绝时机：运行时抛错 vs 初始化校验 (design): 未实施，明确留作后续：本 PR 仅保留运行时 NotImplementedError，DCP + KV-sharing 仍是已知不支持组合。
• 删除过度细化的 FlashInfer 单元测试 (testing): 已删除：最终 PR 不含测试改动，回归保障依赖 :nvidia: (B200) Spec Decode Speculators + MTP Nightly 与 PR 中的 GSM8K 评测。
• 抽象 Attention 接口的类型注解仍未体现 None 契约 (design): 遗留未修：仅在 review 中记录，本 PR 只改了具体后端的签名与 docstring。

风险与影响

• 风险：1）核心注意力路径被触碰：FlashInferImpl.forward 是 FlashInfer 后端唯一入口，cascade、decode（XQA / TRTLLM_GEN / FlashInfer native）、prefill 各分支共享同一段 padding 与 K/V 处理代码。本次虽只改尾部 DCP 分支与 padding 块，但切片位置变化属于所有 backend 组合共用的路径，回归面比 diff 行数看起来更大。

2）类型保护减弱：key / value 从 Tensor 放宽为 Tensor | None，而抽象层 Attention.forward、unified_attention_with_output、AttentionImpl.forward 的注解仍是 torch.Tensor，custom-op schema 又因 eager dispatch 不做 boxed 类型检查而无法拦截。后续若在非 DCP 路径误用 K/V，只会在运行时才暴露。

3）测试覆盖缺失：最终版本删除了 tests/v1/attention/test_attention_backends.py 中的两个单测，NotImplementedError 这条新异常路径没有任何自动化覆盖，全靠 Gemma4 nightly 与人工 GSM8K 评测；若后续有人打开 DCP × KV-sharing 组合，缺少红灯信号。

4）DCP 部署限制：FlashInfer + DCP 下 KV-sharing 模型无法 prefill（运行时抛 NotImplementedError，且在首次 prefill 的编译区域内抛出），作者已声明为 out of scope，属于需要写进使用文档的已知限制。

5）CI 结论不干净：作者报告 5 个失败 step group（MI355/H200/CPU Multimodal Processor、H200 Extra Initialization 指向既有 HyperCLOVAX hyperclovax_vision_v2 不识别问题，L4 Distributed Models 是 Llama-4 Maverick 初始化），均未进入 flashinfer.py；但合并时并没有全绿 CI 作为证据。

• 影响：用户侧：Gemma4 E2B/E4B、Gemma3n 等 num_kv_shared_layers > 0 的模型在 FlashInfer（含 sliding_window=FLASH_ATTN 的混合后端配置）下恢复可用，MTP 助理模型不再在引擎初始化阶段崩溃；同时消除 gemma4_mtp 在 DCP prefill 下静默使用未初始化显存算错 draft attention 的隐患，正确性收益高于修复本身。系统侧：非 DCP 常见路径去掉了两次张量切片，CPU 侧开销不增反降，符合该文件对 piece-wise CUDA Graph eager 区域的性能纪律。团队侧：确立「KV-sharing 层传 None」为跨后端统一契约，flash_attn / triton_attn 仅在 encoder 分支解引用 K/V、rocm_attn 已处理 None，因此其他后端无需跟随修改，但抽象层类型注解缺口被显式记录。
• 风险标记：核心注意力后端路径变更, 无新增自动化测试覆盖, DCP×KV-sharing 仍不支持, 抽象层 None 类型契约未同步

关联脉络

• PR #54917 [Bugfix][Gemma] Conditionally create KV projections/norms on KV-shared layers: 引入本 PR 所修回归的 PR：它让 KV-shared 层不再创建 k_proj/v_proj/k_norm 并调用 self.attn(q, None, None)，触发 FlashInfer 无条件切片 None 而崩溃。本 PR 是对它的前向修复，保留其收益。
• PR #55789 Revert "[Bugfix][Gemma] Conditionally create KV projections/norms on KV-shared layers" (#54917): 与本次修复同源的 issue/兜底回滚 PR，其 body 记录 nightly #87584 的崩溃栈并明确希望 a narrow fix is preferable to this revert；该 PR 已关闭未合并，由本 PR 取代。
• PR #55871 向 FlashInfer 传占位 K/V 张量的替代修复（原始标题未在材料中给出，据 PR#55864 描述整理）: 重复工作检查中提到的另一方案：从 Gemma4 侧继续传占位 K/V，已被关闭未合并；本 PR 改为让 FlashInfer 接受省略 K/V，并顺带删除 gemma4_mtp 的占位张量。
• PR #50439 [Attention] Extend XQA decode support on SM90: 近期同样修改 vllm/v1/attention/backends/flashinfer.py 的 PR，说明 FlashInferImpl.forward 的 decode 分支（XQA/TRTLLM）正在持续演进，本 PR 触及其共享的 padding 与 K/V 处理约定。