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 写法)
|