optimize moe tiny batch tokenwise path

worklogs/decode_analysis_report_2026-07-14.md

@@ -667,3 +667,67 @@ decode 聚合吞吐提升约 **31%**。TTFT 变差可能来自冷缓存/加载
 4. **TP all-reduce 合并**
    - 当前 MoE 每层 routed+shared 后做一次 all-reduce。
    - 如果后续能将若干小通信或 residual 路径合并，可能进一步改善 decode 抖动，但优先级低于 MoE/linear_attention 计算本体。
+
+## 2026-07-14：MoE decode fast path v2（tokenwise）
+
+### 背景
+
+上一轮 `T <= 2` tiny-batch fast path 使用 batched `torch.bmm`，将 `T * top_k` 个 expert 选择展平后批量计算。它已经将 c2 aggregate Output TPS 从 8.07 提升到 10.60。
+
+继续分析后发现，对于极小 batch（`T=2`）：
+
+- 第一段 gate/up projection 用 `torch.bmm` 实际是 `T * top_k` 个 `1 x H` 小矩阵乘。
+- 原 `T == 1` 路径用的是 `F.linear(hidden, w13_sel.reshape(-1, H))`，会形成一个更大的 GEMM，通常更适合 GPU。
+
+因此新增 v2：`tokenwise` tiny-batch 实现。
+
+### 实现
+
+新增环境变量：
+
+- `ENGINEX_MOE_TINY_IMPL=tokenwise|bmm`
+- `ENGINEX_MOE_TINY_MAX=4`
+
+默认：
+
+```text
+ENGINEX_MOE_TINY_IMPL=tokenwise
+ENGINEX_MOE_TINY_MAX=4
+```
+
+核心逻辑：
+
+- 对 `T <= 4` 的小批量 decode，逐 token 走已经验证过的 `T == 1` 大 `F.linear` 路径。
+- 每个 token 内仍然把 top-k expert 的 `w13` 拼成一个大权重，减少第一段 projection 的小 GEMM。
+- 保留上一轮 batched bmm 实现，可通过 `ENGINEX_MOE_TINY_IMPL=bmm` 回退做 A/B 测试。
+
+代码位置：
+
+- `qwen3_6_scripts/qwen3_5.py`
+- `Qwen3_5MoeSparseBlock._pure_pytorch_experts`
+
+### 结果
+
+同口径 `short c2, max_tokens=128, requests=4`：
+
+| 版本 | 成功率 | TTFT P90 | per-request Output TPS P10 | Aggregate Output TPS |
+| --- | ---: | ---: | ---: | ---: |
+| custom all-reduce on，优化前 | 100% | 1.46s | 4.22 | 8.07 |
+| tiny-batch bmm | 100% | 4.42s | 5.95 | 10.60 |
+| tiny-batch tokenwise v2 | 100% | 4.49s | 6.43 | 11.28 |
+
+相对上一版 bmm，aggregate Output TPS 又提升约 **6.4%**；相对优化前提升约 **39.8%**。
+
+本地归档：
+
+- `worklogs/remote_results/2026-07-14-code-profile/decode_moe_tokenwise_short_c2_t128_r4.json`
+- `worklogs/remote_results/2026-07-14-code-profile/server_moe_tokenwise_eager_custom_ar_seq2_b8192.log`
+
+### 结论
+
+这个结果说明：对当前 BI-V100 + PyTorch fallback MoE 路径，**将小 GEMM 尽量合并为较大的 per-token GEMM** 比把所有 token/expert 都塞进 batched bmm 更好。瓶颈确实集中在 decode MoE routed expert 的小矩阵计算与调度开销。
+
+下一步可继续沿两个方向走：
+
+1. 测试 `max_num_seqs=4`，利用 `T <= 4` tokenwise fast path，看 aggregate Output TPS 是否继续上涨。
+2. 继续优化 shared expert / linear_attention，因为 routed expert 已经明显改善，decode 下一个大头会逐渐转向 `linear_attention` 和 shared expert。