experiment moe router topk softmax

qwen3_6_scripts/qwen3_5.py

@@ -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: top-k logits -> softmax over selected experts.
-        routing_weights = torch.softmax(router_logits.float(), dim=-1)
+        # This is mathematically equivalent to softmax(all experts) -> top-k ->
-        topk_weights, topk_ids = torch.topk(
+        # renormalise, but avoids a full (T, num_experts) softmax in decode.
-            routing_weights, self.top_k, dim=-1)           # (T, top_k)
+        topk_logits, topk_ids = torch.topk(
-        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+            router_logits.float(), self.top_k, dim=-1)     # (T, top_k)
-        topk_weights = topk_weights.to(hidden_states.dtype)
+        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)

worklogs/decode_analysis_report_2026-07-14.md

@@ -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-15：MoE 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 路径做实验。

worklogs/remote_results/2026-07-15-official-like-baseline/eager_100k_router_topk_c1_r8_t256_cap40k.json

@@ -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
+    }
+  ]
+}