DeepSeek-V4性能优化PR #14428已合入 main作者 ivyilike · 2026-09-03 · 改动 1 文件 +20/-7
DeepSeek-V4 的 DSA 稀疏注意力需要按 token 位置取 RoPE 的 cos/sin 表
（get_cos_and_sin_dsa，ops/rope_dsv4.py）。优化前的写法是经典的
「先索引、再拷贝」两步式：
curr_cos = full_rope_cos[pos_tensor] —— 高级索引，分配一个新张量（host 侧 malloc）buf_cos[:num_tokens].copy_(curr_cos) —— 再拷进常驻缓冲（又一次 kernel 分发）cos/sin 各来一遍 = 每对 4 个算子。decode 路径每步、每层都要做， host 侧的分发次数和临时对象分配都翻倍——而真正想要的只是「把表里若干行写进缓冲」这一个动作。
改动后：用 torch.gather(..., out=...) 把「选行 + 写入」合成一个算子、
直接写进常驻缓冲——每对 4 ops → 2 ops，省 2 次分发和 2 次中间分配（消除 host 侧分配/GC 压力）。
gather 索引的构造技巧：把 pos_tensor 变形为 (num_tokens,1,1,1) 再沿旋转维 expand，
让 dim=0 的 gather 恰好选中 pos_tensor[i] 那一行。
单文件改动（ops/rope_dsv4.py，+20/-7）。缓存路径（use_cache）换成 gather 直写；
无缓存路径保留原索引写法（本就要分配新张量，无收益）；语义完全等价，无 API 变化。
curr_cos = full_rope_cos[pos_tensor] # 高级索引：选行并分配临时张量（host malloc + 一次分发）
curr_sin = full_rope_sin[pos_tensor]
if use_cache: # 常驻缓冲路径（decode 主路径）
group_buffers = (
_ROPE_STATE.runtime_buffer.get(config_key, {}).get(group_name)
if draft_index is None
else _ROPE_STATE.spec_runtime_buffer.get(config_key, {}).get(group_name)
)
if group_buffers is None:
continue
buf_cos, buf_sin = group_buffers # 常驻缓冲（本意是避免每步重复分配）
num_tokens = pos_tensor.size(0)
if draft_index is None:
buf_cos[:num_tokens].copy_(curr_cos) # 第二个算子：拷进缓冲（cos/sin 各一遍 = 每对 4 ops）
buf_sin[:num_tokens].copy_(curr_sin) # sin 同理
batch_result[config_key][group_name] = (buf_cos[:num_tokens], buf_sin[:num_tokens])
else:
buf_cos[draft_index - 1][:num_tokens].copy_(curr_cos)
buf_sin[draft_index - 1][:num_tokens].copy_(curr_sin)
batch_result[config_key][group_name] = (
buf_cos[draft_index - 1][:num_tokens], gather_idx = ( # ⚡构造 gather 索引：token 位置沿旋转维广播（等价原选行语义）
pos_tensor.to(torch.long).reshape(-1, 1, 1, 1).expand(num_tokens, 1, 1, full_rope_cos.size(-1))
)
if draft_index is None:
torch.gather(full_rope_cos, 0, gather_idx, out=buf_cos[:num_tokens]) # ⚡选行+直写一步完成，写入常驻缓冲
torch.gather(full_rope_sin, 0, gather_idx, out=buf_sin[:num_tokens]) # sin 同理；每对 4 ops → 2 ops
batch_result[config_key][group_name] = (buf_cos[:num_tokens], buf_sin[:num_tokens])
else:
torch.gather(full_rope_cos, 0, gather_idx, out=buf_cos[draft_index - 1][:num_tokens])
torch.gather(full_rope_sin, 0, gather_idx, out=buf_sin[draft_index - 1][:num_tokens])
batch_result[config_key][group_name] = (
buf_cos[draft_index - 1][:num_tokens],
buf_sin[draft_index - 1][:num_tokens],
)
else:
curr_cos = full_rope_cos[pos_tensor] # ↳无缓存路径保留原写法（本就要分配新张量，无收益）
curr_sin = full_rope_sin[pos_tensor]
batch_result[config_key][group_name] = (curr_cos, curr_sin)
return RopeDataProxy(batch_result, is_cos=True), RopeDataProxy(batch_result, is_cos=False)效果（PR 实测）：get_cos_and_sin_dsa 单次调用
814μs → 701μs（约 -14%）。decode 路径每步都调用，稳定缩小每步的 host 开销。