tilelang量化内核拆解DeepSeek-V4.1-Flash-NVFP4的fp4_gemm与sparse_attn实现DeepSeek-V4.1-Flash-NVFP4 是 NVIDIA 用 Model Optimizer 将 552B 参数的 MoE 多模态大模型量化到 NVFP4 格式的开源权重仓库随附一套可直接运行的参考推理实现。它的核心正是 inference/kernel.py 里的一组 tilelang 量化内核其中fp4_gemmFP8×FP4 混合精度矩阵乘与sparse_attn稀疏注意力最关键。本文不用大段代码带你快速看懂这两个量化内核的设计思路与实现细节并了解它们如何被模型主流程驱动。 NVFP4 量化模型速览552B 参数如何瘦身先建立背景。DeepSeek-V4.1-Flash 原生是 MoE 架构CED 压缩稀疏注意力 CSA2支持 1M 上下文主干 552B 参数prefill 阶段仅激活 8B、decode 阶段激活 16B另带 196B 的 Engram 条件记忆。本仓库的量化要点项目说明量化格式NVFP4W4A4权重 FP4、激活 FP8组大小 16/32量化范围384 个路由专家 × 40 层的w1/w2/w3投影Engram 表从 FP8 转为 MXFP4E2M1 两值打包 每 32 元素一个 UE8M0 scale模型体积约 476 GiB → 约 400 GiB导出为 48 个 safetensors 分片目标硬件NVIDIA Blackwell实测 GB300完整模型卡见 README.md量化配置记录在 hf_quant_config.json 与 config.json 中。️ kernel.py 内核地图一个文件看懂 6 个 tilelang 内核整个推理内核只有约 600 行全部集中在 inference/kernel.py。每个内核都遵循同一套模式用tilelang.jit编译、通过T.symbolic让 M 维度动态化、内部 FP32 计算、shared memory fragment 两级缓冲。内核位置职责act_quantkernel.py#L40-L124按组 FP8 量化激活支持原地反量化回 BF16fp4_quantkernel.py#L127-L204按组 FP4 量化E8M0/E4M3 scale 可选fp8_gemmkernel.py#L207-L307FP8×FP8 分块缩放 GEMMfp4_gemmkernel.py#L477-L591FP8 激活 × FP4 权重混合精度 GEMMsparse_attnkernel.py#L310-L403top-k 索引收集 在线 softmax 稀疏注意力hc_split_sinkhornkernel.py#L406-L474Hyper-Connections 的双随机化混合矩阵一个值得注意的工程细节文件顶部通过pass_configs关闭了 warp specialization 与 TMA lower见 kernel.py#L9-L12保证内核在这类小规模块上稳定可编译。此外fast_log2_ceil/fast_pow2两个工具函数kernel.py#L22-L33直接改写 IEEE 754 位模式来计算 2 的幂避免慢速的 log/ceil 指令——这正是 E8M0 scale只能取 2 的幂快速取整的关键。⚡ fp4_gemm 实现拆解FP8 × FP4 混合精度 GEMM数据布局权重为什么占一半空间fp4_gemm计算C[M,N] A_fp8[M,K] B_fp4[N,K]^Tkernel.py#L477-L559A激活FP8 E4M3沿 K 方向每 32 或 128 元素共享一个 scaleact_block_size可配B权重FP4 E2M1每 2 个值打包进 1 字节存储形状是[N, K//2]每 32 元素配一个 E8M0 scalescales_b: [N, K//32]C输出BF16。这就是 NVFP4 的精髓权重每个数只花 4 bit靠更细的 scale 粒度1×32补偿低精度带来的动态范围损失。主循环两个累加器是关键设计内核切块为block_M32, block_N128, block_K32K 方向流水线 2 级。每轮迭代做四件事搬运把 A 的 FP8 块和 B 的 FP4 块拷入 shared memory升精度FP4 先转 FP32 再转 FP8注释写明是避免 C 重载歧义kernel.py#L538-L540这样就能复用 Tensor Core 的 FP8×FP8T.gemm缩放每行激活 scale、每列权重 scale 取出两级累加C_local_accum C_local * scale_a * scale_b。为什么要C_local和C_local_accum两个累加器因为 scale 是每 32 个 K 元素一个必须等一个 K 块的点积算完才能乘上对应 scale再并入总累加器——若混在一个累加器里就无法做分块缩放。同时T.use_swizzle(panel_size10)调整块调度顺序提升 L2 缓存命中率。激活侧的 FP8 量化由act_quant完成kernel.py#L40-L95每 32 个元素求 absmaxscale 取 2 的幂round_scale路径走位操作快速取整再 clamp 到 FP8 表示范围。inplaceTrue时做量化反量化融合用于 KV cache 原地写回——省一次显存往返。 sparse_attn 实现拆解索引收集 在线 softmaxDeepSeek-V4.1-Flash 的稀疏注意力CSA2思路是每个 query 只看滑动窗口内的原始 KVindexer 选出的 top-k 个压缩 KV 位置而不是全部历史。sparse_attn内核kernel.py#L310-L389把这个按索引取 KV 再做注意力一次完成并行粒度每个 (batch, 序列位置) 一个线程块256 线程。q 加载到 shared memory 后常驻top-k 的 KV 索引topk_idxsint32无效槽为 -1分块block64流水线收集索引收集kv_shared[i, j] idxs[i] ! -1 ? kv[by, idxs[i], j] : 0无效位置取 0随后打分时置为-infsoftmax 后自然归零在线 softmaxFlashAttention 式的 running max/sum。每来一批 KV 就更新scores_max旧的部分和acc_o乘上exp(旧max - 新max)修正、sum_exp同步缩放避免物化完整注意力矩阵attn_sink一个可学习的每头偏置attn_sink在最后并入sum_exp相当于给什么都不看的默认槽一个可学权重数值稳健running max 初值取-1e30而非-inf——若某行 top-k 全为 -1无有效索引exp(-inf - -inf)会产生 NaN用有限下界则输出自然全零与训练内核约定一致kernel.py#L352-L355。Python 封装sparse_attnkernel.py#L392-L403还会把不足 16 的注意力头数补齐到 16 再调用、结束后裁掉换取更规整的向量化宽度。 model.py 如何驱动这些量化内核内核不是孤立的inference/model.py 展示了它们在真实模型中的位置GEMM 分发linear()model.py#L181-L207按权重 dtype 自动选路——FP4 权重 →act_quantfp4_gemmFP8 权重 →act_quantfp8_gemm其余走 PyTorch 原生。MoE 的路由专家Expertmodel.py#L830-L851正是以 FP4 存储w1/w2/w3的注意力Attention.forwardmodel.py#L765-L789把滑动窗口 KVFP8 原地量化act_quant(..., True)与压缩 KV 的索引拼接后一次sparse_attn调完两路 KV压缩 KV 用 16 组 E4M3 scale 的fp4_act_quant写缓存model.py#L759-L761indexer 的 q/k 则用 32 组 E8M0——同一套量化内核两种配置Engram 条件记忆inference/engram.py 用压缩词表 素数桶哈希把 n-gram 查表结果门控进残差流查找表即上文提到的 MXFP4 Engram 权重Hyper-Connections残差流以hc_mult4份并行副本传递hc_split_sinkhorn内核负责把混合系数矩阵 Sinkhorn 迭代成双随机矩阵。 快速运行步骤编译内核并自测想亲手跑起来验证这些内核四步即可参考 inference/README.md# 1. 安装依赖tilelang 会按需 JIT 编译内核 python -m pip install -r requirements.txt # 2. 将 HuggingFace 权重转换为按 TP rank 切分的检查点专家权重转 FP4 python convert.py --hf-ckpt-path /path/to/ckpt --save-path /path/to/out \ --model-parallel 8 --expert-dtype fp4 --tokenizer-path /path/to/ckpt # 3. 运行示例TXT 或 Harmony JSON 输入 INPUT_FILEexamples/example.txt ./run.sh # 4. 形状与内核连通性自测构建小模型跑 prefill 22 步 decode # 真实调用 dense-fp8 / MoE-fp4 内核 python model.py注意python model.py的自测权重是未初始化的它验证的是张量形状与内核调用链路而非数值精度生产部署建议按 README.md 用 vLLM/SGLang 在 Blackwell GPU 上服务化。 小结这套内核设计教给我们的混合精度可以降权不降算fp4_gemm把权重存成 4 bit计算时升回 FP8 喂给 Tensor Core精度损失由更细的 1×32 E8M0 scale 兜底scale 粒度决定累加结构分块缩放天然要求块内点积 → 乘 scale → 块外累加的两级累加器模式稀疏注意力 索引收集 在线 softmaxtopk_idxs把看哪些位置变成纯数据内核内用 running max/sum 保证 O(1) 额外显存数值细节决定可用性-1e30替代-inf、FP4→FP8 必须过 FP32、位操作快速算 2 的幂——这些不起眼的代码正是参考实现最有价值的部分。如果你想继续深入建议按 inference/kernel.py → inference/model.py → inference/engram.py 的顺序通读配合python model.py自测即可完整复现这条 NVFP4 推理链路。创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考