标准 Attention 的可拆分性

标准 attention 对一个 query token q 的计算是：

o = \text{softmax}\left(\frac{q K^T}{\sqrt{d}}\right) V

关键性质：如果 KV 可以分成两段不相交的子集 K = [K_A; K_B]，则：

\text{LSE}_A = \log \sum_{j \in A} \exp\left(\frac{q k_j}{\sqrt{d}}\right), \quad o_A = \text{softmax}(q K_A^T) V_A

\text{LSE}_{total} = \log\left(e^{\text{LSE}_A} + e^{\text{LSE}_B}\right)

o = e^{\text{LSE}_A - \text{LSE}_{total}} \cdot o_A + e^{\text{LSE}_B - \text{LSE}_{total}} \cdot o_B

这就是 merge_state_v2 做的事，两段分开算再合并，数学上完全等价于一次完整 attention。

Topk = 1 时没有问题

设前缀 KV 长度为 L，draft tokens 长度为 S：

Request i:  [prefix_KV (L tokens)] + [draft_KV (S tokens)]

一次 flash_attn_with_kvcache 搞定，cache_seqlens = L + S。

Topk > 1 引入了什么问题

以 topk=4、S=3 为例，Target Verify 阶段：

Request i 有 4 条候选树，每条 3 个 draft tokens：
  Branch 0: [d_00, d_01, d_02]  各自需要 attend: prefix_KV_i + 本分支祖先
  Branch 1: [d_10, d_11, d_12]  各自需要 attend: prefix_KV_i + 本分支祖先
  Branch 2: [d_20, d_21, d_22]  ──────────────────────────────────────────
  Branch 3: [d_30, d_31, d_32]        共享同一段 prefix_KV_i !!

朴素展开（Naive Expand）方案：将 batch 展成 B × topk 个虚拟请求：

prefix_KV_i 被重复读取 topk 次，总读 KV 量：

\text{Naive} = B \times topk \times (L + S) = B \times 4 \times (L + 3)

当 L = 1000 时，主要开销都在反复读那 1000 token 的 prefix。

Cascade 方案的数学结构

将 KV 拆成两段，利用 prefix 被所有分支共享：

第一段 attention（共享 prefix，metadata.page_table）：

B \text{ 个请求，每个请求的 Q} = \text{所有 topk 分支的全部 query tokens} \text{KV} = \text{prefix\_KV}_i \quad (\text{每个请求只读一次！})

输出：o^{(A)}，\text{LSE}^{(A)}

第二段 attention（各自独立的 draft KV，forward_metadata_spec_decode_expand.page_table）：

B \times topk \text{ 个虚拟请求，每个请求的 Q} = \text{本分支 query tokens} \text{KV} = \text{各自分支的 draft tokens (很短，只有 } S \text{ 个)}

输出：o^{(B)}，\text{LSE}^{(B)}

合并：merge_state_v2(o_A, LSE_A, o_B, LSE_B) → 精确等价于原始完整 attention。

总读 KV 量：

\text{Cascade} = \underbrace{B \times L}_{\text{prefix 只读一次}} + \underbrace{B \times topk \times S}_{\text{draft 各自读}}

多条候选 token（topk）影响的核心

当 topk = 4，L = 2000，S = 5： - Naive 每请求读 8020 个 KV token - Cascade 每请求读 2020 个 KV token - 节省约 4 倍，且 L 越长、topk 越大，优势越明显

这就是为什么 use_cascade_attn 的触发条件是 topk > 1——topk = 1 时没有共享 prefix 可利用，cascade 无意义；topk 越大，prefix KV 的重复读取浪费越严重，cascade 收益越大。