experiment moe router topk softmax

This commit is contained in:
2026-07-15 04:19:56 +08:00
parent 16d52b469c
commit eea0653ace
3 changed files with 179 additions and 6 deletions

View File

@@ -849,12 +849,13 @@ class Qwen3_5MoeSparseBlock(nn.Module):
Output is partial (pre-all-reduce), same contract as FusedMoE
with reduce_results=False.
"""
# Routing: softmax → topk → renormalise
routing_weights = torch.softmax(router_logits.float(), dim=-1)
topk_weights, topk_ids = torch.topk(
routing_weights, self.top_k, dim=-1) # (T, top_k)
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
topk_weights = topk_weights.to(hidden_states.dtype)
# Routing: top-k logits -> softmax over selected experts.
# This is mathematically equivalent to softmax(all experts) -> top-k ->
# renormalise, but avoids a full (T, num_experts) softmax in decode.
topk_logits, topk_ids = torch.topk(
router_logits.float(), self.top_k, dim=-1) # (T, top_k)
topk_weights = torch.softmax(
topk_logits, dim=-1).to(hidden_states.dtype) # (T, top_k)
w13 = self.experts.w13_weight # (E, 2*I, H)
w2 = self.experts.w2_weight # (E, H, I)

View File

@@ -1014,3 +1014,51 @@ KeyError: 'invalid tool call parser: qwen3_coder (chose from { hermes,internlm,l
2. 快速性能迭代优先用 100K/cap40K 官方相似数据,减少单轮成本。
3. 确定有效优化后,再回到 245K 配置做确认。
4. 当前离目标最大的差距仍是 prefill/cache 吞吐和 decode MoE 吞吐,而不是缓存命中率或请求成功率。
## 2026-07-15MoE router top-k softmax 实验
### 改动
尝试把 MoE router 从:
```python
routing_weights = torch.softmax(router_logits.float(), dim=-1)
topk_weights, topk_ids = torch.topk(routing_weights, self.top_k, dim=-1)
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
```
改为:
```python
topk_logits, topk_ids = torch.topk(router_logits.float(), self.top_k, dim=-1)
topk_weights = torch.softmax(topk_logits, dim=-1)
```
数学上这等价于 “full softmax -> topk -> renormalize”但避免了完整 `(T, num_experts)` softmax。预期收益主要来自减少每层 decode 路由开销。
### 结果
本地归档:
- `worklogs/remote_results/2026-07-15-official-like-baseline/eager_100k_router_topk_c1_r8_t256_cap40k.json`
| 指标 | 100K 基线 | router top-k softmax | 目标 |
| --- | ---: | ---: | ---: |
| Success Rate | 100.00% | 100.00% | >= 99% |
| Cache Hit Rate | 74.17% | 74.17% | >= 50% |
| TTFT P90 | 9.803s | 13.048s | <= 5s |
| Output TPS P10 | 5.825 | 5.664 | >= 20 |
| Aggregate Output TPS | 5.789 | 5.379 | - |
| Uncached Input TPS | 178.330 | 172.454 | - |
| Cache TPS | 512.015 | 495.142 | - |
| Weighted Throughput | 883.109 | 850.328 | >= 8000 |
### 结论
该改动没有带来收益,反而略降。可能原因:
- full softmax 相对 expert GEMM 和 TP all-reduce 占比很小;
- topk logits 后 softmax 改变了 kernel 调度形态,但没有减少主瓶颈;
- 小样本中输出 token 数不同,仍需看趋势而不是单条。
操作结论:不保留该代码改动,回退到 full-softmax router。下一步继续围绕 MoE expert GEMM、all-reduce、prefill/cache 路径做实验。

View File

@@ -0,0 +1,124 @@
{
"created_at": "2026-07-14T20:11:28",
"url": "http://127.0.0.1:1111",
"model": "llm",
"dataset": "/root/work/logs/synthetic_cumulative_8_s3_cap40k.jsonl",
"concurrency": 1,
"max_requests": 8,
"max_tokens": 256,
"wall_sec": 337.77776083163917,
"success_rate": 1.0,
"ttft_p90_sec": 13.047750872373578,
"output_tps_p10_per_request": 5.664154516073774,
"aggregate_output_tps": 5.379276585665033,
"aggregate_input_tps_uncached": 172.45362707296303,
"aggregate_cache_tps": 495.1421301041857,
"cache_hit_rate": 0.7416795639891973,
"weighted_token_throughput": 850.3276245683974,
"totals": {
"requests": 8,
"success": 8,
"prompt_tokens": 225499,
"input_tokens_uncached": 58251,
"cached_tokens": 167248,
"completion_tokens": 1817,
"reasoning_tokens": 1685
},
"targets": {
"output_tps_p10_per_request_gte_20": false,
"ttft_p90_lte_5": false,
"cache_hit_rate_gte_50pct": true,
"success_rate_gte_99pct": true,
"weighted_token_throughput_gte_8000": false
},
"results": [
{
"ok": true,
"elapsed_sec": 26.321524418890476,
"ttft_sec": 7.245203051716089,
"prompt_tokens": 7498,
"cached_tokens": 0,
"completion_tokens": 160,
"reasoning_tokens": 74,
"output_tps": 8.387361321943352,
"error": null
},
{
"ok": true,
"elapsed_sec": 41.994782442227006,
"ttft_sec": 21.734486425295472,
"prompt_tokens": 34177,
"cached_tokens": 7488,
"completion_tokens": 121,
"reasoning_tokens": 75,
"output_tps": 5.972272068427839,
"error": null
},
{
"ok": true,
"elapsed_sec": 46.809104984626174,
"ttft_sec": 3.47783805988729,
"prompt_tokens": 35384,
"cached_tokens": 34160,
"completion_tokens": 256,
"reasoning_tokens": 256,
"output_tps": 5.907974037422925,
"error": null
},
{
"ok": true,
"elapsed_sec": 41.270918272435665,
"ttft_sec": 9.324864206835628,
"prompt_tokens": 19274,
"cached_tokens": 7472,
"completion_tokens": 256,
"reasoning_tokens": 256,
"output_tps": 8.013509257647705,
"error": null
},
{
"ok": true,
"elapsed_sec": 47.37227803096175,
"ttft_sec": 3.5642303358763456,
"prompt_tokens": 37099,
"cached_tokens": 35376,
"completion_tokens": 256,
"reasoning_tokens": 256,
"output_tps": 5.843675157172532,
"error": null
},
{
"ok": true,
"elapsed_sec": 36.956629026681185,
"ttft_sec": 5.4115986954420805,
"prompt_tokens": 14661,
"cached_tokens": 7472,
"completion_tokens": 256,
"reasoning_tokens": 256,
"output_tps": 8.115382908555416,
"error": null
},
{
"ok": true,
"elapsed_sec": 48.93905151821673,
"ttft_sec": 3.3529286701232195,
"prompt_tokens": 38199,
"cached_tokens": 37088,
"completion_tokens": 256,
"reasoning_tokens": 256,
"output_tps": 5.615744090653815,
"error": null
},
{
"ok": true,
"elapsed_sec": 48.10511343367398,
"ttft_sec": 3.073553144931793,
"prompt_tokens": 39207,
"cached_tokens": 38192,
"completion_tokens": 256,
"reasoning_tokens": 256,
"output_tps": 5.684901841253756,
"error": null
}
]
}