DeepSeek-V4原生稀疏注意力：CSA/HCA内核与TileLang实现解析

1. 项目概述：为什么DeepSeek-V4的CSA/HCA代码值得逐行精读？

你有没有过这种体验：翻开源码，看到一个叫sparse_attn的函数，名字很酷，但点进去全是T.gemm、T.Parallel、acc_s、sum_exp这类符号，像在读天书？我第一次打开DeepSeek-V4的kernel.py时就是这感觉。但后来发现，这恰恰是当前大模型工程最硬核、也最值得深挖的一环——它不是在“模拟”高效，而是在CUDA底层用TileLang原语，把稀疏注意力的数学本质，一针一线地织进了GPU的warps和shared memory里。这篇解读，不讲空泛的“稀疏注意力有多好”，而是带你真正看懂：当topk_idxs = [3, 17, 42, 105]这个数组传进kernel时，GPU上256个线程是如何协作，在不到10微秒内，把离散索引变成连续gemm、完成online softmax、并把结果精准写回输出张量的。

核心关键词“LLM, DeepSeek-V4, 原生稀疏注意力”在这里不是标签，而是三个锚点：LLM代表问题域——我们处理的是超长序列下的KV cache爆炸；DeepSeek-V4是具体载体——它没有走FlashAttention-2那种“优化现有dense attn”的老路，而是从头设计了CSA（Compressed Sparse Attention）和HCA（Hierarchical Compressed Attention）两种原生稀疏范式；而“原生稀疏注意力”则是灵魂——它的稀疏性不是靠剪枝或mask后验生成，而是由Indexer模块在前向传播中主动、可学习地选择top-k位置，再由sparse_attn_kernel在硬件层面无损执行。这意味着，你看到的每一行TileLang代码，都是对“如何用最少的计算，换取最长的有效上下文”这一终极命题的直接回答。适合谁？如果你正在做模型推理优化、想理解现代LLM的KV cache管理本质、或是准备自己动手写CUDA kernel，那么这篇就是你的实操手册。它不假设你精通CUDA，但要求你愿意跟着代码指针，一行行看懂数据在内存中的流转路径。

2. 整体架构与设计哲学：三层注意力的协同逻辑

2.1 为什么必须放弃“全量KV attention”？——显存与计算的双重暴政

在深入代码前，得先直面那个被所有论文轻描淡写带过的残酷现实：一个7B参数的模型，head_dim=128，n_heads=32，当序列长度达到32K时，仅存储KV cache就需要多少显存？简单算笔账：KV cache形状是[batch, seq_len, n_heads, head_dim]，但DeepSeek-V4的MLA（Multi-head Latent Attention）有个关键简化——所有head共享同一组KV，所以实际是[batch, seq_len, head_dim]。单精度FP32下，一个token的KV占128×4=512字节。32K tokens就是32768×512≈16MB。看起来不多？错。这是单个layer的cache。一个32层的模型，就是32×16MB≈512MB。这还只是KV，没算中间激活、梯度、优化器状态。更致命的是计算量：标准attention的复杂度是O(seq²)，32K序列意味着单次attention要计算约10亿个QK点积。即使A100满血跑，也要毫秒级延迟。这就是为什么DeepSeek-V4的代码里，你看不到torch.einsum('bshd,bthd->bhst', q, k)这种“教科书式”写法——它在工程上已经死了。CSA/HCA的设计，本质上是一场精密的“资源配给制”：用滑动窗口保局部精度，用KV压缩减缓长程衰减，再用稀疏索引确保只计算真正重要的交互。三者不是并列选项，而是嵌套的层级结构。

2.2 CSA与HCA：两种原生稀疏范式的定位差异

代码里反复出现的compress_ratio参数，是理解整个架构的钥匙。它不是一个魔法数字，而是定义了模型“记忆粒度”的开关。当compress_ratio == 0时，模型退化为纯滑动窗口attention，这是最省资源的模式，但长程依赖完全丢失。当compress_ratio == 4时，CSA（Compressed Sparse Attention）被激活：每4个历史token被Compressor模块压缩成1个“摘要token”，这些摘要被存入KV cache的后半段，形成一个“压缩区”。此时，topk_idxs会同时包含两部分：滑动窗口内的原始索引（如[100,101,102,103]）和压缩区内的摘要索引（如[32000,32001,32002,32003]）。而HCA（Hierarchical Compressed Attention）则更进一步，它出现在compress_ratio > 4的场景，比如ratio=8或16，此时Compressor会进行多级压缩——第一级把8个token压成1个，第二级再把8个一级摘要压成1个二级摘要，形成树状KV cache。代码中Compressor类的overlap机制，正是为HCA服务的：当ratio=4且overlap=True时，相邻的4-token窗口会重叠1个位置（即窗口0: tokens[0-3], 窗口1: tokens[3-6]），这样压缩后的摘要就能平滑过渡，避免在窗口边界处出现信息断层。这种设计哲学，和人类阅读时“扫视标题+精读段首+跳读细节”的认知模式惊人地一致——模型不是在“记住一切”，而是在不同粒度上构建“记忆索引”。

2.3 数据流全景图：从x到o的七步旅程

把model.py中Attention.forward的543行代码，浓缩成一条清晰的数据流，能帮你建立全局视角。这不是简单的输入输出映射，而是一场在GPU内存中精心编排的“数据交响乐”：

1. 输入注入：x（shape[b, s, dim]）进入，这是上一层的输出，也是本层的全部信息源。
2. Q的低秩锻造：x先经wq_a降维到q_lora_rank（通常远小于n_heads*head_dim），再q_norm归一化，最后wq_b升维。这步的妙处在于，qr = q_norm(wq_a(x))这个低秩Q，不仅用于后续attention计算，还被直接喂给Indexer模块做top-k选择——用更小的计算量，换来同样有效的索引决策。
3. KV的单头统一：x经wkv投影，得到[b, s, head_dim]的单头KV。注意，这里没有n_heads维度，这是MLA的核心，大幅削减参数量和显存。
4. RoPE与量化双处理：对KV的后rd维（rope_head_dim）加旋转位置编码，对前head_dim-rd维做FP8 QAT（Quantization-Aware Training）模拟，这是为后续kernel的高效计算铺路。
5. 索引的动态编织：get_window_topk_idxs生成滑动窗口索引；若compress_ratio>0，self.indexer(qr)或get_compress_topk_idxs生成压缩区索引；两者cat拼接，形成最终的topk_idxs。这个数组，就是整个稀疏attention的“作战地图”。
6. KV Cache的时空管理：Prefill阶段，原始KV写入cache前window_size位置（环形缓冲），Compressor(x)生成的压缩KV则拼接到cache末尾；Decode阶段，则是环形写入滑动窗口，并触发增量压缩。
7. 稀疏Kernel的终极执行：sparse_attn(q, kv_cache, topk_idxs, ...)调用TileLang kernel，将离散索引转化为连续计算，产出o，再经反向RoPE和分组低秩投影wo_a/wo_b，得到最终输出。

这七步，每一步都对应着代码中一个关键函数或一段核心逻辑。理解它们之间的因果和时序，比死记硬背某行代码更重要。

3. Attention类核心解析：从__init__到forward的工程密码

3.1 __init__里的参数玄机：为什么Q/KV/O要用三种不同的低秩策略？

Attention.__init__（L439-482）看似只是参数初始化，实则埋藏着DeepSeek-V4对计算-参数-精度三角关系的深刻权衡。我们逐个拆解：

• Q的双阶段低秩（wq_a+wq_b）：wq_a将dim维输入压缩到q_lora_rank（例如64），wq_b再将其映射到n_heads*head_dim（例如32×128=4096）。表面看是“先压再扩”，但q_norm插在中间，是点睛之笔。它让低秩Q的分布更稳定，避免了wq_b在升维时因数值范围过大而引入梯度爆炸。更重要的是，qr = q_norm(wq_a(x))这个中间产物，被复用为Indexer的输入。试想，如果Indexer直接用全尺寸Q，计算量会陡增一个数量级；而用qr，维度从4096降到64，索引选择的开销几乎可以忽略。这是一种典型的“计算复用”设计，代码里没有注释，但逻辑上严丝合缝。

• KV的单头共享（wkv）：wkv的权重形状是[dim, head_dim]，而非[dim, n_heads, head_dim]。这意味着所有32个attention head，共用同一组KV向量。这直接砍掉了n_heads-1倍的KV参数量和显存。但代价是什么？是每个head失去了独立的“视角”。DeepSeek-V4的应对之道，是把“视角差异”转移到Q上——通过wq_b的升维权重，让不同head的Q在同一个KV空间上“投射”出不同的注意力模式。这就像一群人用同一台望远镜看风景，但每个人调整焦距和角度的方式不同。代码中apply_rotary_emb(kv[..., -rd:], freqs_cis)只对后rd维加RoPE，也是为了在有限的head_dim里，把位置信息的表达效率最大化。

• O的分组低秩（wo_a+wo_b）：wo_a的权重是[n_groups, head_dim, group_dim]，wo_b是[n_groups * group_dim, dim]。n_groups通常设为n_heads // 2或类似值，是一种GQA（Grouped-Query Attention）的变体。它的工程价值在于，wo_a的计算可以按组并行，einsum("bsgd,grd->bsgr", o, wo_a.weight)这行代码，把原本需要n_heads次独立计算的wo_a投影，压缩成了n_groups次。wo_b则负责把分组结果重新融合。这种设计，在保证输出表达力的同时，显著降低了wo层的FLOPs。我在实测一个13B模型时发现，仅wo_a/wo_b的分组策略，就让单次decode的延迟降低了约12%。

提示：q_lora_rank、n_groups、compress_ratio这三个参数，是模型工程师调优的黄金三角。增大q_lora_rank能提升Q的表达力，但会增加Indexer负担；增大n_groups能降低wo计算，但可能削弱head间的独立性；增大compress_ratio能延长上下文，但会增加Compressor的计算开销。没有银弹，只有根据你的硬件和任务需求做的精细平衡。

3.2 forward中的QK-norm与RoPE：稳定训练与位置编码的底层实现

forward函数（L484-543）的前三步，是整个attention计算的基石。其中QK-norm和RoPE的实现细节，直接决定了模型能否稳定收敛。

• QK-norm的数值稳定性（L498）：q *= rsqrt(mean(q^2) + eps)这行代码，是QK-norm的PyTorch实现。rsqrt是1/sqrt(x)的快速近似。它的作用，是让Q向量的L2范数趋近于1。为什么重要？因为在attention的Q @ K.T计算中，如果Q或K的范数过大，点积结果会急剧膨胀，导致softmax的输入值极大，exp()运算极易溢出，产生inf或nan。QK-norm通过强制归一化，在计算前就把数值范围“框定”在一个安全区间。eps=1e-6是防止mean(q^2)为0时除零。这个操作在wq_b之后、unflatten之前执行，确保归一化作用于低秩升维后的完整Q，而不是中间的qr。这是一个非常务实的工程技巧，比在loss里加梯度裁剪更治本。

• RoPE的精确切片（L499 & L504）：apply_rotary_emb(q[..., -rd:], freqs_cis)和apply_rotary_emb(kv[..., -rd:], freqs_cis)这两行，揭示了RoPE应用的精确位置。...表示前面所有维度，-rd:表示取最后rd个维度。rd（rope_head_dim）通常设为head_dim // 2，即只对Q/KV的一半维度应用旋转。这是DeepSeek-V4的一个关键优化。完整的RoPE会作用于全部head_dim维，但实验表明，对一半维度应用，既能保留足够的位置信息表达能力，又能显著减少freqs_cis的内存占用和apply_rotary_emb的计算量。freqs_cis是一个预计算好的复数频率表，形状为[max_seq_len, rd//2]，apply_rotary_emb函数内部会将其展开并与Q/KV的最后rd维做复数乘法。代码中q[..., -rd:]的切片，确保了RoPE只影响被设计的那部分维度，其他维度保持原样，为模型保留了更大的自由度去学习非位置相关的特征。

3.3 KV Cache布局与管理：环形缓冲与惰性压缩的实战智慧

kv_cache的布局（[batch, window_size + max_seq_len/ratio, head_dim]）和管理逻辑，是CSA/HCA能落地的物理基础。它不是一个静态数组，而是一个动态演化的“记忆器官”。

• 环形缓冲（Sliding Window）的实现（L520-521）：self.kv_cache[:bsz, cutoff:win], self.kv_cache[:bsz, :cutoff] = kv[:, -win:].split(...)这行代码，是环形写入的经典手法。win = window_size，cutoff = start_pos % win。kv[:, -win:]取输入KV的最后win个token。split操作将其切成两段：第一段长度为win - cutoff，写入cache的[cutoff:win]位置；第二段长度为cutoff，写入cache的[:cutoff]位置。这样，无论start_pos是多少，新token总能被写入start_pos % win这个索引，而旧token则被自然覆盖。这就像一个无限长的胶卷，你只关心最近的window_size帧，后面的帧自动被新帧覆盖。代码中cutoff的计算，是环形逻辑的核心。

• 惰性压缩（Lazy Compression）的触发机制（L524 & L530）：kv_compress = self.compressor(x, start_pos)这行在Prefill和Decode中都存在，但行为迥异。Prefill时，x是整个序列，Compressor一次性处理所有token，生成全部压缩KV。Decode时，x是单个token，Compressor的forward方法（L343-359）会先将这个token的KV和score存入kv_state/score_state缓冲区，然后检查(start_pos+1) % ratio == 0。只有当累积满ratio个token时，才真正触发一次压缩计算。这种“攒够了再干”的惰性策略，完美匹配了自回归decode的逐token特性，避免了为单个token就启动一次完整压缩的浪费。Compressor类中kv_state和score_state的register_buffer声明，确保了这些状态能在训练/推理中跨step持久化，这是实现惰性压缩的基础设施。

注意：kv_cache的两个区域——滑动窗口区和压缩区——在内存中是连续的，但逻辑上是隔离的。topk_idxs的拼接（cat([window_idxs, compress_idxs])）之所以能工作，是因为sparse_attn_kernel接收的是整个kv_cache张量，它内部的gather操作（L324-325）会根据topk_idxs中的绝对索引，直接从这个连续内存块中抓取数据。这种“逻辑分区、物理连续”的设计，既保证了访问效率，又维持了概念清晰。

4. sparse_attn_kernel深度剖析：TileLang如何驾驭GPU的野性

4.1 TileLang的革命性：从“写CUDA”到“描述计算”

kernel.py:276-352的sparse_attn_kernel，是整篇代码的皇冠。它用TileLang写的事实，本身就宣告了一种新的GPU编程范式。传统CUDA需要你手动管理block、thread、shared memory、warp shuffle，而TileLang让你专注于“我要做什么”，把“怎么做”交给编译器。with T.Kernel(m, b, threads=threads) as (bx, by):这行，就是TileLang的宣言：我声明一个二维grid，第一维m（query positions）对应bx，第二维b（batch）对应by，每个block有256个线程。剩下的，T.copy、T.gemm、T.Parallel这些原语，都是对计算意图的声明，而非对硬件的指令。这就像你告诉建筑师“我要一个客厅、一个厨房、一个卧室”，而不是告诉他“砖要怎么砌、电线要怎么拉”。

• Grid与Block的语义映射：bx（block x）直接映射到m（query sequence length），意味着每个block负责处理一个query position的所有计算。by（block y）映射到b（batch size），意味着每个block也负责一个batch sample。所以，总共m x b个block，每个block处理1 query x 1 batch的完整稀疏attention。这种映射，让kernel的并行逻辑一目了然：query之间天然独立，是最理想的并行维度。

• Shared Memory的智能调度：kv_shared[i, j]（L324-325）是一个[block, d]大小的shared memory数组。block在这里是topk（通常是64），d是head_dim。T.Parallel(block, d)声明了对这个数组的并行加载。TileLang编译器会自动将这个循环映射到256个线程上，让每个线程负责一个(i,j)对。kv_shared的存在，是性能的关键——它把从global memory中离散gather来的KV，变成了shared memory中连续的矩阵，后续的Q @ KV.T就可以用高效的T.gemm完成，避免了global memory的高延迟随机访问。这是TileLang“描述计算”带来的最大红利：你只需说“我要把离散索引的KV放到shared memory里”，编译器就为你生成最优的内存搬运代码。

4.2 Gather操作：离散索引到连续gemm的魔法转换

Gather（L322-327）是sparse_attn_kernel的第一个核心环节，它解决了稀疏attention的根本难题：如何把topk_idxs = [3, 17, 42, 105]这样的离散地址，变成GPU可以高效计算的连续数据块。

• 索引加载与越界处理（L322-323）：idxs[i] = T.if_then_else(t * block + i < topk, topk_idxs[by, bx, t * block + i], -1)。t是block内处理的第t个topk块（因为topk可能大于block，需要分块处理），i是线程ID。这行代码的意思是：对于当前block要处理的第i个索引，如果它在topk范围内，就从topk_idxs中加载；否则，填入-1作为无效标记。-1是一个哨兵值，后续所有基于它的操作都会被屏蔽。这比用0或max_int做哨兵更安全，因为-1在数组索引中天然非法。

• 离散Gather到Shared Memory（L324-325）：kv_shared[i, j] = T.if_then_else(idxs[i] != -1, kv[by, idxs[i], j], 0)。这才是魔法所在。kv[by, idxs[i], j]是global memory中的随机访问，idxs[i]是刚才加载的离散索引。T.if_then_else确保了只有有效索引才进行读取，无效索引则写入0。最终，kv_shared这个shared memory数组，就包含了topk个被选中的KV向量，按topk_idxs的顺序排列，且内存连续。这为下一步的T.gemm铺平了道路。你可以把它想象成一个“数据整理员”，把散落在各处的快递（KV），按照一张订单（topk_idxs）的地址，整齐地摆放在一个临时货架（kv_shared）上，等待打包（gemm）。

4.3 Online Softmax：FlashAttention风格的数值稳定引擎

sparse_attn_kernel的Online Softmax（L328-343）是其技术含量最高的部分。它没有像PyTorch那样先算出完整的[m, topk]attention score矩阵，而是采用流式计算，边算边归一化，彻底规避了O(m×topk)的内存占用。

• Running Max与Sum的维护：scores_max和sum_exp是两个核心状态变量，它们在每个topk块的处理中被迭代更新。scores_max = max(scores_max, acc_s)不断更新当前遇到的最大score值。sum_exp = sum_exp * scores_scale + sum(acc_s)则用scores_scale = exp(scores_max_prev - scores_max)来修正历史sum_exp，使其与当前acc_s的scale对齐。这个exp(old_max - new_max)因子，是数值稳定的精髓。假设旧sum_exp是基于old_max=10计算的，现在new_max=12，那么exp(10-12)=exp(-2)≈0.135，就把旧sum_exp缩小了约7.4倍，使其与exp(score-12)的新值在同一量级上相加。这保证了无论acc_s的值有多大，sum_exp都不会溢出。

• Accumulator的累加逻辑（L340-342）：acc_o = acc_o * scores_scale先修正历史输出，acc_o += acc_s @ KV再累加当前块的贡献。acc_s @ KV是当前块的[block, d]score矩阵与kv_shared的[block, d]KV矩阵的gemm，产出[block, d]的output片段。acc_o是一个running accumulator，最终acc_o /= sum_exp完成归一化。整个过程，acc_o和sum_exp始终只占用O(d)和O(1)的内存，与topk大小无关。这是FlashAttention思想的完美复刻，也是sparse_attn_kernel能处理超长序列而不爆显存的根基。

4.4 Attention Sink：可学习的“不关注”机制

Attention Sink（L345-348）是sparse_attn_kernel中最具思想性的设计。它引入了一个可学习的标量attn_sink[i]（每个head一个），并将其融入softmax的分母计算：sum_exp[i] += T.exp(attn_sink[i] - scores_max[i])。

• Sink的物理意义：attn_sink不是一个真实的KV token，而是一个虚拟的、不提供任何value的“空槽位”。它只参与softmax的归一化分母sum_exp的计算。T.exp(attn_sink[i] - scores_max[i])这个项，相当于给分母增加了一个常数偏置。当所有真实KV的score都很低时（例如，当前query与所有top-k KV都不相关），scores_max会很小，exp(attn_sink - scores_max)就会很大，从而显著增大sum_exp，稀释掉所有真实KV的注意力权重，使模型的输出更接近于一个“默认”状态。这就像人类在面对完全陌生的信息时，会本能地“不置可否”，而不是强行关联。

• Sink的训练价值：attn_sink是nn.Parameter，在训练中会被梯度更新。模型会学会为每个head设置一个合适的attn_sink值。在代码中，它被初始化为一个较小的负数（如-2.0），鼓励模型在初期更倾向于关注真实KV。随着训练深入，模型会根据任务需要，动态调整这个值。例如，在问答任务中，当query是“请总结上文”，模型可能会学到一个较大的attn_sink，以抑制对细节token的关注，转而聚焦于摘要token。这是一种内生于attention机制的、可学习的“不确定性建模”，比外部加一个额外的sink token更优雅、更高效。

5. Compressor类详解：学习型门控池化的实现艺术

5.1 Compressor的核心公式与工程实现

Compressor（L279-377）是CSA/HCA的“记忆压缩器”，其核心公式compressed_kv = Σ (kv_i * softmax(score_i))（i = 0..ratio-1）看似简单，但其实现充满了工程巧思。

• Score的生成与位置编码（L328-330）：score = self.wgate(x)生成原始门控分数，score = score.unflatten(1, (-1, ratio))将其reshape为[b, cutoff/ratio, ratio, head_dim]。关键的一步是score = score + self.ape，其中self.ape是[ratio, coff*head_dim]的可学习绝对位置编码。coff = 1 + overlap，当overlap=True时，coff=2，ape的shape变为[ratio, 2*head_dim]。这意味着，ape的前head_dim维用于当前窗口，后head_dim维用于重叠窗口。score += ape的操作，让模型不仅能区分“第几个token”，还能区分“这个token属于哪个窗口”，为后续的softmax提供了丰富的上下文感知能力。

• Overlap Transform的错位艺术（L307-314）：overlap_transform是处理overlap=True时的核心函数。它接收一个[b, s, ratio, 2d]的张量（s是窗口数），并将其重组为[b, s, 2*ratio, d]。关键的错位赋值：

  • new_tensor[:, :, ratio:] ← current_data：将当前窗口的数据（原张量的后半d维）放到新张量的后ratio列。
  • new_tensor[:, 1:, :ratio] ← previous_data：将上一个窗口的数据（原张量的前半d维）放到新张量的前ratio列，但起始位置是第1行（[:, 1:, ...]），即跳过第0行。 这样，第i个新窗口的前ratio列，来自第i-1个旧窗口；后ratio列，来自第i个旧窗口。窗口0的前ratio列，由于没有i-1，会被初始化为0（kv_state的初始值），softmax后权重为0，不影响结果。这种错位，实现了窗口间的无缝衔接，让压缩后的摘要token能同时“看到”前一组和当前组的token，极大地缓解了窗口边界效应。

5.2 Prefill与Decode的双路径设计

Compressor.forward（L316-377）严格区分了Prefill和Decode两种模式，体现了对不同推理场景的深刻理解。

• Prefill路径（L325-342）：这是“批量处理”模式。x是整个输入序列，cutoff是序列长度。x被unflatten(1, (-1, ratio))按ratio分组，score加上ape后做softmax，沿dim=2（即ratio维度）加权求和，得到压缩KV。尾部不足ratio的余数，被存入kv_state/score_state，为后续Decode做准备。这个路径的计算是密集的、一次性的，充分利用了GPU的并行能力。

• Decode路径（L343-359）：这是“流式处理”模式。每次只来一个token，x的shape是[b, 1, dim]。它被立即投影为kv和score，并存入kv_state/score_state。if (start_pos+1) % ratio == 0:是触发条件，只有当累积满ratio个token时，才调用_compress_step进行一次压缩。overlap模式下，_compress_step会联合kv_state（当前窗口）和kv_state_prev（上一窗口）进行压缩，然后将kv_state滑动到kv_state_prev的位置，为下一个窗口做准备。这种设计，让Compressor在Decode时的计算开销变得极其平滑，不会出现Prefill时那种“爆发式”的计算峰值，对实时性要求高的应用至关重要。

5.3 数据流与量化：从FP32到FP4的端到端链路

Compressor的后处理（L360-377）完成了从数学计算到硬件部署的最后一步。

• RMSNorm与RoPE的叠加（L362-364）：压缩后的kv先经过self.norm（RMSNorm），再对最后rope_head_dim维施加apply_rotary_emb。RMSNorm的作用是稳定压缩后KV的数值分布，避免因加权求和导致的方差坍缩。RoPE的再次应用，则确保了压缩后的摘要token，依然携带了精确的位置信息，这对于长程依赖建模不可或缺。

• 量化策略的选择（L366-367）：if self.rotate: quantize_fp4(kv, ...) else: act_quant(kv, ...)。self.rotate是一个布尔标志，控制量化方式。quantize_fp4是更激进的4-bit量化，适用于对精度要求不高的压缩区KV；act_quant则是常规的FP8 QAT模拟，精度更高。这种混合量化策略，是模型压缩的高级技巧：对“摘要”用更低精度，对“原始”用更高精度，在整体精度损失可控的前提下，最大化显存节省。act_quant函数内部，会对KV的非RoPE部分进行FP8量化，而RoPE部分保持FP16/FP32，确保位置编码的精度不被破坏。

6. 实操心得与常见问题排查

6.1 我踩过的坑：TileLang编译失败的三大元凶

在本地复现sparse_attn_kernel时，我被编译错误折磨了整整两天。分享三个最痛的教训，帮你绕开这些深坑：

• 坑一：topk不是常量，TileLang无法推导。topk在代码中是一个Python变量，但在T.Kernel的for循环里，T.Parallel(block, d)中的block必须是编译期可知的常量。解决方案：在T.Kernel定义时，把topk作为参数传入，with T.Kernel(m, b, topk, threads=threads) as (bx, by, tz):，然后在T.Parallel中使用tz。否则，编译器会报Cannot infer shape。

• 坑二：kv_cache的shape未对齐。kv_cache的第二个维度是window_size + max_seq_len/ratio，但如果max_seq_len不能被ratio整除，max_seq_len/ratio会是浮点数，导致kv_cache的shape非法。解决方案：在Attention.__init__中，max_seq_len必须是ratio的整数倍，或者在计算kv_cachesize时，用math.ceil(max_seq_len / ratio)向上取整。

• 坑三：attn_sink的梯度未注册。attn_sink是nn.Parameter(torch.empty(h))，但如果你在forward中直接用了attn_sink[i]，而没有在__init__中用self.register_parameter('attn_sink', ...)显式注册，训练时attn_sink的梯度就不会被收集。解决方案：务必在__init__中用self.register_parameter，并在forward中用self.attn_sink[i]访问。

6.2 性能