Mavis (阿念) 6deca6aab1 ✨ ZZ-FLAGOS-S2-2026-D246-001 · Task 38 sigmoid_gate_topk_renorm 准备就绪

3-kernel 写法 (K1 sigmoid+bias / K2 自写 topk 迭代 argmax / K3 gather+cat+sigmoid+renorm+scale+split)
+ 国产 NPU 套路全套 (指针 cast int16 / int64 stride / enable_fp_fusion=False / 1/(1+exp(-x))
+ 算法 8 case 测试全过 (T=1~128, N=8~256, S=1~2, k=2~8, 含 DeepSeek-V3 风格)
+ 之之 D246 00:12 决策:用 Triton 自写 topk, 走性能路线

提交次数 1/30 已用 (K35), 剩 29 次 · 截止 2026-09-03 19:59 (剩 ~19h41m)

阿念 (Mavis, ICE-GL-AN-001) · 之之的家 · 2026-09-03 D246 00:18

2026-09-03 00:18:02 +08:00

4.3 KiB

Raw Blame History

Task 38 · sigmoid_gate_topk_renorm 提交清单

生成时间: 2026-09-03 00:17 (D246) 提交次数: 今日 1/30 (用了 1 次在 K35),还剩 29 次截止: 2026-09-03 19:59 (还剩 19h41m)

准备好的文件

/Users/zhizhi/Desktop/sigmoid_gate_topk_renorm.zip  (2.6 KB)
└── sigmoid_gate_topk_renorm.py  (7.6 KB · 通用版 · 无 7 芯片特化)

主入口: reference = sigmoid_gate_topk_renorm 函数签名: def reference(logits, k, n_shared_experts, route_scale, global_scale, bias): 返回: (routed_w[T,k], indices[T,k] int32, shared_w[T,S])

实现方案 · 3-kernel 写法

┌─ K1: _sigmoid_bias_kernel ─────────────────┐
│   sel = sigmoid(routed) + bias              │
│   sel 存为 fp32 中间 buffer [T, N]          │
└────────────────────┬────────────────────────┘
                     ↓
┌─ K2: _topk_kernel (自写) ──────────────────┐
│   迭代 tl.argmax K 次, 每次选最大索引       │
│   把已选位置 mask 为 -inf, 防重复选         │
│   写 indices [T, k] int32                   │
└────────────────────┬────────────────────────┘
                     ↓
┌─ K3: _gate_finalize_kernel ────────────────┐
│   gather(routed, indices) → routed_vals[k]  │
│   load shared_logits → shared_vals[S]       │
│   cat → sigmoid → /sum → *route_scale*gs    │
│   split → routed_w [T,k] / shared_w [T,S]   │
└─────────────────────────────────────────────┘

国产 NPU 套路 (D245 验证过)

套路	作用
`with torch.get_device_module(x.device).device(x.device):`	切设备 (K30/K35 跑通的必要条件)
`enable_fp_fusion=False, num_warps=4`	国产 NPU 关 fp fusion
`pid.to(int64) + tl.arange(0, BLOCK).to(int64)`	防 stride overflow
`input_ptr.to(tl.pointer_type(tl.int16))`	防 NaN 在 load 时被吃
`1.0 / (1.0 + tl.exp(-x))` (不用 `tl.sigmoid`)	tl.sigmoid 国产 NPU 不支持

本地测试 (算法层, 8/8 pass)

跑 test_algorithm.py,对比平台 reference vs 算法模拟(纯 torch 拆分 K1/K2/K3):

T=1   N=64  S=1 k=8 : indices match ✓  routed close ✓  shared close ✓
T=1   N=64  S=1 k=4 : indices match ✓  routed close ✓  shared close ✓
T=8   N=64  S=1 k=8 : indices match ✓  routed close ✓  shared close ✓
T=32  N=128 S=2 k=8 : indices match ✓  routed close ✓  shared close ✓
T=64  N=256 S=1 k=8 : indices match ✓  routed close ✓  shared close ✓  (DeepSeek-V3 风格)
T=128 N=256 S=1 k=6 : indices match ✓  routed close ✓  shared close ✓
T=1   N=8   S=1 k=2 : indices match ✓  routed close ✓  shared close ✓
T=16  N=32  S=1 k=4 : indices match ✓  routed close ✓  shared close ✓

ALL PASS ✓

风险点

K2 tl.argmax 在国产 NPU 上可能不支持 — 降级方案:把 K2 换成 torch.topk 1 行
K1 / K3 中间 buffer 大小 — sel 占 T*N*4 bytes(fp32),DeepSeek-V3 风格 (T=128, N=256) = 128KB, 8 芯片评估时 buffer 分配可能有限制
tl.where + 迭代 K 次的循环 — Triton 编译复杂度,某些 NPU 可能 register spilling

预期结果 (我的估计)

指标	估计	信心
跑通芯片数	6-7/8	中 (tl.argmax 是最大变数)
平均加速比	3-8×	中 (3 kernel 限制上限)
排名	跟 sitraliqui(9.05x) 接近或略低	低

之之的提交步骤

打开平台 flagos.net → 第 3 批 → Task 38 → "提交代码"
上传 /Users/zhizhi/Desktop/sigmoid_gate_topk_renorm.zip
提交 → 等结果(通常 1-3 分钟)
看到结果告诉我,如果有芯片 Failed 我会分析

之后怎么调 (如果第一次不理想)

跑通但分低(6-8×): 把 K1 和 K3 合并成 1 个 kernel, 减少 launch overhead
某芯片 Failed: 看错误, 大概率是 tl.argmax, 改用 torch.topk
跨芯片不稳: 加 7 芯片特化版 (跟 K30 一样 1+7 写法)

4.3 KiB Raw Blame History Unescape Escape