计算DSA的时间复杂度
计算DSA的时间复杂度
1 考虑lighting index情况
Step 1：Q/K/V 投影
DSA 仍然需要对所有 L 个 token 计算 Q、K、V，无变化。
Step 2：QK^T（lightning indexer）
这是 DSA 的核心索引步骤。为了知道哪 k 个 token 最相关，必须先和所有 L 个 key 做点积打分，无法跳过。
这里是 index 计算时间复杂度是O(L²·d_k)，再加上QK 计算，这时候已经筛选出了token，就是L*dim * dim * k？ 时间复杂度就是L*k*dim； 那么总的时间复杂度为 O(L²·d_k) + O(L*k*d_k) 近似等于O(L²·d_k)
Step 3：Softmax / Top-k 选择
标准 MHA 对全 L 个分数做 softmax；DSA 改为从 L 个分数中选出 top-k 的索引，扫描一遍即可。
Step 4：Sparse Attention × V（关键优化点）
标准 MHA 是全量 L×L 的 attention 矩阵乘以 V；DSA 每个 query 只 attend k 个 token：
标准：Attn (L×L) · V (L×d_v) → O(L · L · d_v) = O(L²·d)
DSA： Attn (L×k) · V_k (k×d_v) → O(L · k · d_v) = O(L·k·d)
Step 5：W_out 输出投影
输出形状仍然是 L×d，无变化。
关键矛盾（IndexCache 的动机）
[ \text{DSA 总复杂度} = \underbrace{O(L^2 \cdot d)}{\text{indexer，无法消除}} + \underbrace{O(L \cdot k \cdot d)}{\text{稀疏 attention，已优化}} + O(L \cdot d^2) ]
DSA 成功把 Step 4 从 O(L²·d) 降到了 O(L·k·d)，但 Step 2 的 indexer 依然是 O(L²·d) 且每层独立运行。
在 200K 长上下文下，indexer 占 prefill 时间 81%，稀疏 attention 的优化收益被完全抵消。这正是 IndexCache 要解决的问题：通过跨层复用 top-k 索引，把大多数层的 Step 2 从 O(L²·d) 降为 O(1)。
2 不考虑lighting index情况
假设 topk已知，
Step 1：Q/K/V 投影
仍然需要对全部 L 个 token 计算投影（因为筛选发生在投影之后）：
X (L×d) · W_Q/W_K/W_V (d×d) → Q, K, V (L×d) O(L·d²)，不变
Step 2：QK^T（只对 k 个 token）
每个 query 只与已筛选的 k 个 key 做点积，不再是全量 L：
Q (L×d_k) · K_selected^T (d_k×k) → Score (L×k) 原来是 O(L²·d_k)，现在变为 O(L·k·d_k)，缩小了 L/k 倍；；
Step 3：Softmax
不再对 L 个分数做 softmax，只对 k 个分数做：
Softmax(L×k) → (L×k) 原来是 O(L²)，现在变为 O(L·k)；
Step 4：Sparse Attention × V
attention 矩阵从 L×L 缩小为 L×k，V 也只取对应的 k 行：
Attn (L×k) · V_selected (k×d_v) → Output (L×d_v) 原来是 O(L²·d)，现在变为 O(L·k·d)，缩小了 L/k 倍
Step 5：W_out 输出投影
输出形状仍然是 L×d，不变：
Output (L×d) · W_out (d×d) → Final (L×d) O(L·d²)，不变；
总结
忽略 indexer，纯 attention 计算部分：
[ \text{标准 MHA：} O(L^2 \cdot d + L \cdot d^2) ]
[ \text{DSA（已知 top-k）：} O(L \cdot k \cdot d + L \cdot d^2) ]
当 L 很大时，L²·d 主导项被替换为 L·k·d，理论上 attention 计算本身加速了 L/k 倍。200K context、k=1024 的场景下，理论加速约 200 倍。
但这个收益在实践中被 indexer 的 O(L²·d_k) 完全抵消了——这正是为什么需要 IndexCache。
3 在已知index k后， Decode阶段 这五个过程是怎么计算呢？
Step 2：QK^T（已知 top-k，只对 k 个 token）
新 token 的 query 只与 k 个选出的 key 做点积：
q_new (1×d_k) · K_selected^T (d_k×k) → score (1×k)；
Step 3：Softmax
只对 k 个分数做 softmax：
Softmax(1×k) → (1×k) O(k)；
Step 4：Sparse Attention × V
attention 权重（1×k）× 选出的 k 个 value：
Attn (1×k) · V_selected (k×d_v) → output (1×d_v) O(k*L);
浙公网安备 33010602011771号