OpenAI Privacy Filter源码解析128专家MoE路由与Triton融合算子的实现细节【免费下载链接】privacy-filterOpenAI Privacy Filter项目地址: https://gitcode.com/gh_mirrors/pr/privacy-filterOpenAI Privacy Filter 是一款可在本地运行的 PII个人身份信息检测与脱敏模型一次前向传播即可为每个 token 打上隐私标签实现高速文本脱敏。它的核心亮点是 1.5B 总参数、每 token 只激活 4 个专家约 50M 活跃参数的 MoE 架构以及用 Triton 手写算子把专家矩阵乘与 SwiGLU 激活融合成单次 kernel显著减少显存往返。本文带你从零读懂它的 128 专家 MoE 路由与 Triton 融合算子的实现细节。一、模型定位一次前向传播完成脱敏 与逐 token 生成的语言模型不同Privacy Filter 是双向 token 分类模型输入一段文本模型对每个 token 输出 33 维 logit背景类O 8 类隐私标签 × BIOES 四种边界标签再经约束 Viterbi 解码得到连贯的 span。一次前向 全部打标这是它高吞吐的关键。模型支持 128K 上下文、8 类隐私标签邮箱、电话、地址、人名、账号、URL、日期、密钥等架构细节见 README.md8 层 Transformer block残差宽度d_model 640分组查询注意力GQA14 个 query 头 / 2 个 KV 头稀疏 MoE 前馈共 128 个专家每 token 只路由 top-4双向带状注意力带宽 128有效窗口 257 token安装后可直接用 CLI 体验pip install -e . opf Alice was born on 1990-01-02.二、128 专家 MoE 路由从 Gate 到 Token 打包 MoE 前馈块MLPBlock位于 opf/_model/model.py。它的路由流程分四步1️⃣ Gate 打分与 Top-4 选择每个 token 先经 RMSNorm再过一个hidden_size → 128的线性 Gatemodel.py#L743-L750对 128 个专家分数做topk(k4)只保留分数最高的 4 个专家对选中分数做 softmax 后再除以 4保证专家权重总和为 12️⃣ Token 按专家排序打包这是理解 Triton 路径的核心model.py#L763-L782。由于每个 token 要去 4 个不同专家代码先把(token_id, expert_id)对展平然后argsort(expert_ids)把 token 按专家 ID 排序同一专家的 token 在显存中连续存放bincount统计每个专家分到的行数cumsum生成offsets起始偏移这一步把每个专家处理不同长度的一批 token变成了一块连续紧凑的矩阵——后续 grouped GEMM 才能高效执行。3️⃣ 分专家矩阵乘打包后的 token 矩阵乘第一层专家权重结果按路由权重加权用index_add_累加回原 token 位置model.py#L820-L829最后乘回experts_per_token修正缩放。4️⃣ CPU 回退路径无 Triton 时走_moe_chunk的 batched bmm 回退路径model.py#L831-L867逐 token 扩展 4 份、fp32 计算、按OPF_TORCH_OPS_BATCH分块控显存结果与 Triton 路径数值对齐_batched_linear_with_parity。三、Triton 融合算子省掉中间激活的显存往返 真正的性能核心在 opf/_model/triton_moe.py。朴素实现需要 3 次大张量落盘mlp1 输出 → SwiGLU 输出 → mlp2 输出。Triton 版本把它们压成2 个 kernel。3.1 分组 GEMM 内核_grouped_matmul_kernel的 grid 是三维的(行块, 列块, 专家ID)每个专家的程序实例先读取自己的offsets/lengths行数不足的块直接返回天然处理各专家负载不均沿 K 维分块循环累加累加器用 fp32 保证精度块大小按 N 自适应N ≥ 4096用 128×128×32、8 warps否则 128×64triton_moe.py#L107-L1183.2 SwiGLU 第二层投影融合内核_grouped_swiglu_w2_kernel是点睛之笔它在同一次 GEMM 循环内完成h_glu clamp(h_glu, max7.0) h_lin clamp(h_lin, min-7.0, max7.0) act h_glu * sigmoid(1.702 * h_glu) * (h_lin 1.0) # SwiGLU acc tl.dot(act, w) # 第二层投影 acc acc b # 偏置激活值act只在寄存器/共享内存中生成直接参与tl.dot从不写回全局显存LIMIT7.0限幅与ALPHA1.702作为编译期常量tl.constexpr烘焙进 kernel与纯 PyTorch 版 swiglu() 完全一致grouped_swiglu_w2封装默认启用OPF_MOE_FUSED_SWIGLU_W2true设OPF_MOE_FUSED_SWIGLU_W20可切回分离激活 第二次 grouped GEMM路径对比验证。3.3 环境变量开关一览变量默认作用OPF_MOE_TRITONCUDA 上开启使用/禁用 Triton MoE 内核OPF_MOE_FUSED_SWIGLU_W2true启用 SwiGLUmlp2 融合内核OPF_EXPERTS_PER_TOKEN检查点配置值4运行时覆盖每 token 专家数OPF_ALLOW_TF32关闭允许 TF32 加速矩阵乘OPF_ATTN_LOW_PRECISION关闭注意力用 bf16 低精度路径环境变量的解析逻辑在 opf/_common/env.py支持1/true/yes/on四种真值。四、权重加载MXFP4 反量化与多卡分片 专家权重在检查点中以MXFP4 块浮点格式存储Checkpoint.get()负责按需反量化opf/_model/weights.py每个字节含两个 4-bit nibble分别查 16 值FP4_VALUES查找表±0/0.5/1/1.5/2/3/4/6每行一个 int8 scale-127偏置还原用torch.ldexp一次完成指数展开分块循环每块约 32768×1024 行控制峰值显存参数名映射表PARAM_NAME_MAPweights.py#L32-L49把逻辑名mlp1_weight对应到磁盘上的swiglu.weight系列张量并优先取 bf16 融合张量。多 GPU 并行时按专家维度分片model.py#L1031-L1046mlp1_weight按输出维2×intermediate切分——对应 SwiGLU 的两半mlp2_weight按输入维intermediate切分两卡各算一半all_reduce(SUM)汇总因此 Triton 路径内嵌了dist.all_reducemodel.py#L818-L819五、配套机制带状注意力与 Viterbi 解码 MoE 之外还有两个与性能/精度强相关的细节双向带状注意力sdpa()用F.pad unfold为每个 token 构造 257 宽的窗口model.py#L431-L473RoPE 采用 YaRN 风格缩放rope_scaling_factor32把 4096 初始长度外推到 128K并额外拼接一组可学习的 sink logit 提升长文本稳定性。约束 Viterbi 解码ViterbiCRFDecoderopf/_core/decoding.py预计算 BIOES 合法转移表并暴露 6 个转移偏置背景保持、进入 span、延续、闭合等运行时可调出偏召回或偏精确的工作点避免逐 token argmax 导致的碎片化边界。六、源码导读清单 按依赖顺序阅读两天即可吃透全链路opf/_model/model.py —— 架构主体RMSNorm、RoPE、GQA 注意力、MLPBlockMoE 路由opf/_model/triton_moe.py —— 两个 Triton 内核分组 GEMM 与 SwiGLU 融合投影opf/_model/weights.py —— 检查点索引、MXFP4 反量化、参数名映射opf/_core/decoding.py —— 约束 Viterbi 解码与转移偏置校准opf/_core/runtime.py —— 推理运行时装配分词、窗口切分、span 后处理七、总结它做对了什么✨打包排序 grouped GEMM把变长专家负载变成连续矩阵分块避免 128 次独立小 matmul 的启动开销寄存器级融合SwiGLU 激活不落显存直接喂给第二层 GEMM省掉一次完整往返fp32 累加 常数烘焙tl.constexpr把限幅/alpha 编译进 kernel精度与纯 PyTorch 参考实现对齐优雅回退CPU、无 Triton、多卡场景都有对应路径且数值一致这套实现让一个能在浏览器里跑的 1.5B 模型同时拥有企业级数据清洗所需的高吞吐——这正是开源 OpenAI Privacy Filter 最值得借鉴的工程范本。【免费下载链接】privacy-filterOpenAI Privacy Filter项目地址: https://gitcode.com/gh_mirrors/pr/privacy-filter创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
OpenAI Privacy Filter源码解析：128专家MoE路由与Triton融合算子的实现细节
OpenAI Privacy Filter源码解析128专家MoE路由与Triton融合算子的实现细节【免费下载链接】privacy-filterOpenAI Privacy Filter项目地址: https://gitcode.com/gh_mirrors/pr/privacy-filterOpenAI Privacy Filter 是一款可在本地运行的 PII个人身份信息检测与脱敏模型一次前向传播即可为每个 token 打上隐私标签实现高速文本脱敏。它的核心亮点是 1.5B 总参数、每 token 只激活 4 个专家约 50M 活跃参数的 MoE 架构以及用 Triton 手写算子把专家矩阵乘与 SwiGLU 激活融合成单次 kernel显著减少显存往返。本文带你从零读懂它的 128 专家 MoE 路由与 Triton 融合算子的实现细节。一、模型定位一次前向传播完成脱敏 与逐 token 生成的语言模型不同Privacy Filter 是双向 token 分类模型输入一段文本模型对每个 token 输出 33 维 logit背景类O 8 类隐私标签 × BIOES 四种边界标签再经约束 Viterbi 解码得到连贯的 span。一次前向 全部打标这是它高吞吐的关键。模型支持 128K 上下文、8 类隐私标签邮箱、电话、地址、人名、账号、URL、日期、密钥等架构细节见 README.md8 层 Transformer block残差宽度d_model 640分组查询注意力GQA14 个 query 头 / 2 个 KV 头稀疏 MoE 前馈共 128 个专家每 token 只路由 top-4双向带状注意力带宽 128有效窗口 257 token安装后可直接用 CLI 体验pip install -e . opf Alice was born on 1990-01-02.二、128 专家 MoE 路由从 Gate 到 Token 打包 MoE 前馈块MLPBlock位于 opf/_model/model.py。它的路由流程分四步1️⃣ Gate 打分与 Top-4 选择每个 token 先经 RMSNorm再过一个hidden_size → 128的线性 Gatemodel.py#L743-L750对 128 个专家分数做topk(k4)只保留分数最高的 4 个专家对选中分数做 softmax 后再除以 4保证专家权重总和为 12️⃣ Token 按专家排序打包这是理解 Triton 路径的核心model.py#L763-L782。由于每个 token 要去 4 个不同专家代码先把(token_id, expert_id)对展平然后argsort(expert_ids)把 token 按专家 ID 排序同一专家的 token 在显存中连续存放bincount统计每个专家分到的行数cumsum生成offsets起始偏移这一步把每个专家处理不同长度的一批 token变成了一块连续紧凑的矩阵——后续 grouped GEMM 才能高效执行。3️⃣ 分专家矩阵乘打包后的 token 矩阵乘第一层专家权重结果按路由权重加权用index_add_累加回原 token 位置model.py#L820-L829最后乘回experts_per_token修正缩放。4️⃣ CPU 回退路径无 Triton 时走_moe_chunk的 batched bmm 回退路径model.py#L831-L867逐 token 扩展 4 份、fp32 计算、按OPF_TORCH_OPS_BATCH分块控显存结果与 Triton 路径数值对齐_batched_linear_with_parity。三、Triton 融合算子省掉中间激活的显存往返 真正的性能核心在 opf/_model/triton_moe.py。朴素实现需要 3 次大张量落盘mlp1 输出 → SwiGLU 输出 → mlp2 输出。Triton 版本把它们压成2 个 kernel。3.1 分组 GEMM 内核_grouped_matmul_kernel的 grid 是三维的(行块, 列块, 专家ID)每个专家的程序实例先读取自己的offsets/lengths行数不足的块直接返回天然处理各专家负载不均沿 K 维分块循环累加累加器用 fp32 保证精度块大小按 N 自适应N ≥ 4096用 128×128×32、8 warps否则 128×64triton_moe.py#L107-L1183.2 SwiGLU 第二层投影融合内核_grouped_swiglu_w2_kernel是点睛之笔它在同一次 GEMM 循环内完成h_glu clamp(h_glu, max7.0) h_lin clamp(h_lin, min-7.0, max7.0) act h_glu * sigmoid(1.702 * h_glu) * (h_lin 1.0) # SwiGLU acc tl.dot(act, w) # 第二层投影 acc acc b # 偏置激活值act只在寄存器/共享内存中生成直接参与tl.dot从不写回全局显存LIMIT7.0限幅与ALPHA1.702作为编译期常量tl.constexpr烘焙进 kernel与纯 PyTorch 版 swiglu() 完全一致grouped_swiglu_w2封装默认启用OPF_MOE_FUSED_SWIGLU_W2true设OPF_MOE_FUSED_SWIGLU_W20可切回分离激活 第二次 grouped GEMM路径对比验证。3.3 环境变量开关一览变量默认作用OPF_MOE_TRITONCUDA 上开启使用/禁用 Triton MoE 内核OPF_MOE_FUSED_SWIGLU_W2true启用 SwiGLUmlp2 融合内核OPF_EXPERTS_PER_TOKEN检查点配置值4运行时覆盖每 token 专家数OPF_ALLOW_TF32关闭允许 TF32 加速矩阵乘OPF_ATTN_LOW_PRECISION关闭注意力用 bf16 低精度路径环境变量的解析逻辑在 opf/_common/env.py支持1/true/yes/on四种真值。四、权重加载MXFP4 反量化与多卡分片 专家权重在检查点中以MXFP4 块浮点格式存储Checkpoint.get()负责按需反量化opf/_model/weights.py每个字节含两个 4-bit nibble分别查 16 值FP4_VALUES查找表±0/0.5/1/1.5/2/3/4/6每行一个 int8 scale-127偏置还原用torch.ldexp一次完成指数展开分块循环每块约 32768×1024 行控制峰值显存参数名映射表PARAM_NAME_MAPweights.py#L32-L49把逻辑名mlp1_weight对应到磁盘上的swiglu.weight系列张量并优先取 bf16 融合张量。多 GPU 并行时按专家维度分片model.py#L1031-L1046mlp1_weight按输出维2×intermediate切分——对应 SwiGLU 的两半mlp2_weight按输入维intermediate切分两卡各算一半all_reduce(SUM)汇总因此 Triton 路径内嵌了dist.all_reducemodel.py#L818-L819五、配套机制带状注意力与 Viterbi 解码 MoE 之外还有两个与性能/精度强相关的细节双向带状注意力sdpa()用F.pad unfold为每个 token 构造 257 宽的窗口model.py#L431-L473RoPE 采用 YaRN 风格缩放rope_scaling_factor32把 4096 初始长度外推到 128K并额外拼接一组可学习的 sink logit 提升长文本稳定性。约束 Viterbi 解码ViterbiCRFDecoderopf/_core/decoding.py预计算 BIOES 合法转移表并暴露 6 个转移偏置背景保持、进入 span、延续、闭合等运行时可调出偏召回或偏精确的工作点避免逐 token argmax 导致的碎片化边界。六、源码导读清单 按依赖顺序阅读两天即可吃透全链路opf/_model/model.py —— 架构主体RMSNorm、RoPE、GQA 注意力、MLPBlockMoE 路由opf/_model/triton_moe.py —— 两个 Triton 内核分组 GEMM 与 SwiGLU 融合投影opf/_model/weights.py —— 检查点索引、MXFP4 反量化、参数名映射opf/_core/decoding.py —— 约束 Viterbi 解码与转移偏置校准opf/_core/runtime.py —— 推理运行时装配分词、窗口切分、span 后处理七、总结它做对了什么✨打包排序 grouped GEMM把变长专家负载变成连续矩阵分块避免 128 次独立小 matmul 的启动开销寄存器级融合SwiGLU 激活不落显存直接喂给第二层 GEMM省掉一次完整往返fp32 累加 常数烘焙tl.constexpr把限幅/alpha 编译进 kernel精度与纯 PyTorch 参考实现对齐优雅回退CPU、无 Triton、多卡场景都有对应路径且数值一致这套实现让一个能在浏览器里跑的 1.5B 模型同时拥有企业级数据清洗所需的高吞吐——这正是开源 OpenAI Privacy Filter 最值得借鉴的工程范本。【免费下载链接】privacy-filterOpenAI Privacy Filter项目地址: https://gitcode.com/gh_mirrors/pr/privacy-filter创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考