6deca6aab1
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
93 lines
4.3 KiB
Markdown
93 lines
4.3 KiB
Markdown
# 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 ✓
|
||
```
|
||
|
||
## 风险点
|
||
|
||
1. **K2 `tl.argmax` 在国产 NPU 上可能不支持** — 降级方案:把 K2 换成 `torch.topk` 1 行
|
||
2. **K1 / K3 中间 buffer 大小** — `sel` 占 `T*N*4 bytes`(fp32),DeepSeek-V3 风格 (T=128, N=256) = 128KB, 8 芯片评估时 buffer 分配可能有限制
|
||
3. **`tl.where` + 迭代 K 次的循环** — Triton 编译复杂度,某些 NPU 可能 register spilling
|
||
|
||
## 预期结果 (我的估计)
|
||
|
||
| 指标 | 估计 | 信心 |
|
||
|------|------|------|
|
||
| 跑通芯片数 | 6-7/8 | 中 (tl.argmax 是最大变数) |
|
||
| 平均加速比 | 3-8× | 中 (3 kernel 限制上限) |
|
||
| 排名 | 跟 sitraliqui(9.05x) 接近或略低 | 低 |
|
||
|
||
## 之之的提交步骤
|
||
|
||
1. 打开平台 flagos.net → 第 3 批 → Task 38 → "提交代码"
|
||
2. 上传 `/Users/zhizhi/Desktop/sigmoid_gate_topk_renorm.zip`
|
||
3. 提交 → 等结果(通常 1-3 分钟)
|
||
4. 看到结果告诉我,如果有芯片 Failed 我会分析
|
||
|
||
## 之后怎么调 (如果第一次不理想)
|
||
|
||
- **跑通但分低**(6-8×): 把 K1 和 K3 合并成 1 个 kernel, 减少 launch overhead
|
||
- **某芯片 Failed**: 看错误, 大概率是 tl.argmax, 改用 torch.topk
|
||
- **跨芯片不稳**: 加 7 芯片特化版 (跟 K30 一样 1+7 写法)
|