Optimize DSA backward SM100 with enhanced topk_length loop and TMA fast path

[xxyux]

[Perf] DSA backward SM100: per-query topk_length 循环上界 + 连续 KV tile 的 TMA 批量加载

背景 / 动机

DeepSeek Sparse Attention 的 backward kernel(SM100/Blackwell)在加载 KV 时逐行发 cp.async,并对每个 query 固定遍历 topk 个 tile。两点可优化:

1. 很多 query 的有效 topk 条目远少于 topk,尾部整段空洞仍被遍历;
2. 部分 tile 选中的 KV 行索引恰好连续,逐行 gather 浪费了 TMA 批量搬运能力。

本 PR 在 kernel 侧支持按 topk_length 收缩循环上界,并对连续 tile 走 TMA 快路径,同时保证稀疏场景完全正确。(调用侧传入 topk_length 的改动在 PaddleFleet 主库,不在本 PR 范围。)

改动内容(仅 dsa_bwd_sm100.py,+205/-33)

• per-query 循环上界:消费传入的 mTopkLength,kernel 只遍历到该上界为止,跳过尾部全空 tile。

• 连续 KV 的 TMA 快路径:新增 tma_atom_KV + 带运行时行偏移的 gmem 视图,连续 tile 用单条 TMA 搬 64 行。是否连续由 4 个 load warp 在 shared kv_vote 上投票决定;投票槽按 tile_index & 1 双缓冲,配合每迭代一次 load_KV_sync_barrier,把 warp 间偏移限制在一个迭代内。TMA 事务直接记在管线自身 full mbarrier 上(128 次 arrive + expect_tx),与慢路径 producer_commit 等价。

• 慢路径无分支化:负索引用位运算 clamp 到合法地址无条件拷贝,cp_async_wait 后按 OR-reduce 符号位判断整块是否含空洞,再对空洞行补零;指令流不分叉,利于 ptxas 调度。

• compute 路径重排:tmem fence + consumer_release 后移到 STSM+commit 之后,删除主循环冗余的 compute_sync_barrier;T2R 由 Repetition(8) 拆成两次 Repetition(4) 以重叠计算与访存延迟。

• reduce 路径提前释放:T2R 取到寄存器 + fence 后即 consumer_release,再发 fire-and-forget 的 atomic_add;not same_hdim_kv 分支的 dKV2/3 从 store_dKV 改为寄存器版 reduce_dKV_from_reg。

正确性说明

• 稀疏安全:TMA 仅在 tile 的 64 个 KV 索引严格连续且无空洞时启用;任何跳跃/空洞/warp 间不对齐都会让投票落空,自动回退逐行 cp.async gather,行为与优化前一致。
• 无越界:能通过投票 ⇒ 64 行全有效 ⇒ 起始行 B ≤ KV行数 - block_tile,TMA 坐标恒在界内。
• 同步收支平衡:load_KV_sync_barrier 因 kv_fast 在全体 producer 线程间取值一致,快/慢分支不会相位错位;reduce_sync_barrier 仅存在于 same_hdim_kv 分支,两分支均为 const_expr 编译期统一。