profile decode path and optimize tiny batch moe

worklogs/decode_analysis_report_2026-07-14.md

@@ -535,3 +535,135 @@ decode microbench 结果：
 4. 用 `ENGINEX_PROFILE_DECODE=1` 这类环境变量控制插桩，只在短压测时开启，避免污染正式结果。
 
 初步判断：custom all-reduce 不是第一大瓶颈；更可能的主战场是 xFormers decode attention、MoE 小 batch kernel、以及 TP 下大量小 kernel / 同步造成的低 GPU 利用率。
+
+## 2026-07-14：代码级 profiling 与第一轮 MoE 优化
+
+### profiling 插桩
+
+新增环境变量控制的 profiling：
+
+- `ENGINEX_PROFILE_DECODE=1`：开启 profiling。
+- `ENGINEX_PROFILE_EVERY=N`：每 N 次 model forward 打印一次累计统计。
+- `ENGINEX_PROFILE_SYNC=1`：每段计时前后 `torch.cuda.synchronize()`，用于定位 GPU 时间。
+
+插桩位置：
+
+- `vllm/worker/model_runner.py`
+  - `attn_state.begin_forward`
+  - `model_forward`
+  - `compute_logits`
+  - `sample`
+- `qwen3_6_scripts/qwen3_5.py`
+  - `full_attention`: qkv projection / norm+rope / paged attention / gate+o_proj
+  - `linear_attention`: GatedDeltaNet 整层
+  - `MoE`: gate / routing topk / routed experts / shared expert / TP all-reduce
+  - `DecoderLayer`: norm / attention / MLP
+
+注意：profiling 强制同步会显著拖慢请求，因此 profiling 结果只用于定位瓶颈，不作为真实性能分数。
+
+本地归档：
+
+- `worklogs/remote_results/2026-07-14-code-profile/server_profile_mode_eager_custom_ar_seq2_b8192.log`
+- `worklogs/remote_results/2026-07-14-code-profile/profile_mode_eager_custom_ar_short_c1_t24_r1.json`
+
+### 关键 profiling 结果
+
+短请求 `c1, max_tokens=24`，修正标签后，rank0 在 step=24 的主要累计耗时：
+
+| 模块 | 总耗时 | 平均单层/次 | 次数 | 判断 |
+| --- | ---: | ---: | ---: | --- |
+| `prefill.moe.routed_prefill_experts` | 5239.90ms | 65.50ms | 80 | 首 token 慢的最大来源 |
+| `prefill.layer.linear_attention` | 2944.76ms | 49.08ms | 60 | prefill 第二大来源 |
+| `decode.layer.mlp` | 1456.57ms | 1.66ms | 880 | decode 最大来源 |
+| `decode.layer.linear_attention` | 942.63ms | 1.43ms | 660 | decode 第二大来源 |
+| `decode.moe.routed_total` | 785.65ms | 0.89ms | 880 | MLP 中 routed expert 为主 |
+| `decode.moe.routed_decode_experts` | 540.24ms | 0.61ms | 880 | 单 token MoE expert 计算 |
+| `decode.layer.full_attention` | 322.06ms | 1.46ms | 220 | full attention 不是第一瓶颈 |
+| `decode.moe.tp_all_reduce` | 265.92ms | 0.30ms | 880 | 通信有成本，但不是最大项 |
+
+`ModelRunner` 粗粒度：
+
+```text
+decode.model_forward avg ~= 145.6ms
+decode.compute_logits avg ~= 1.17ms
+decode.sample avg ~= 1.01ms
+```
+
+结论：decode 慢主要发生在模型 forward 内部；logits 和 sampler 不是主瓶颈。
+
+### 为什么 GPU 算力打不满
+
+当前路径的 GPU 利用率低，不是因为单个大矩阵乘算不过来，而是因为每 token 被拆成大量小工作：
+
+1. **MoE 没有真正 fused kernel**
+   - 注释里已经说明 BI-V100 上缺少 `vllm_moe_topk_softmax / vllm_invoke_fused_moe_kernel`。
+   - 当前 routed expert 是纯 PyTorch 实现。
+   - prefill 通用路径按 expert 做 Python 循环，长 prompt 下会产生大量小 GEMM 和 CPU/GPU 同步点。
+
+2. **decode batch 太小**
+   - `max_num_seqs=2` 时每步只有 1-2 token。
+   - 小 batch 下矩阵乘规模小，kernel launch、Python 调度、TP 同步成本占比很高。
+
+3. **模型结构有大量 GatedDeltaNet linear_attention 层**
+   - profiling 显示 linear_attention 在 prefill 和 decode 都是大头之一。
+   - 这部分不是标准 paged attention，不能靠换 xFormers attention backend 直接解决。
+
+4. **TP all-reduce 不是第一瓶颈，但放大了小 kernel 问题**
+   - decode MoE all-reduce 单次约 0.30ms。
+   - 单次看不大，但每层一次、每 token 多次累积，且会让 rank 间等待更明显。
+
+### 第一轮针对性优化：MoE tiny-batch fast path
+
+发现：原代码只有 `T == 1` 的 MoE decode fast path。一旦并发 decode `T == 2`，会落入通用 `routed_prefill_experts` 路径：
+
+```python
+unique_eids = topk_ids.view(-1).unique().tolist()
+for eid in unique_eids:
+    ...
+```
+
+这条路径适合大 prefill，但不适合 `max_num_seqs=2` 的小批量 decode。
+
+本轮新增 `T <= 2` fast path：
+
+- 将 `T * top_k` 个选中 expert 展平。
+- 用两次 batched `torch.bmm` 计算 gate/up 和 down。
+- 避免 Python per-expert loop。
+- 保留 `T == 1` 原 fast path 不变。
+
+### 优化结果
+
+对比同口径 `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 MoE fast path | 100% | 4.42s | 5.95 | 10.60 |
+
+decode 聚合吞吐提升约 **31%**。TTFT 变差可能来自冷缓存/加载后首次请求抖动，后续需要用多轮 warmup 后再复测。
+
+本地归档：
+
+- `worklogs/remote_results/2026-07-14-code-profile/decode_tiny_batch_moe_short_c2_t128_r4.json`
+- `worklogs/remote_results/2026-07-14-code-profile/server_tiny_batch_moe_eager_custom_ar_seq2_b8192.log`
+
+### 下一步优化方向
+
+优先级从高到低：
+
+1. **MoE prefill 路径**
+   - 当前 `prefill.moe.routed_prefill_experts` 是 TTFT 最大来源。
+   - 需要把 Python per-expert loop 替换成更批量化的 grouped GEMM / batched GEMM。
+   - 官方负载长上下文输入占比极高，提升 prefill 会直接改善 TTFT 和 weighted throughput。
+
+2. **GatedDeltaNet linear_attention**
+   - decode 和 prefill 都是大头。
+   - 需要进一步拆分 projection、conv/update、state update、out_proj，确认是 recurrent update 还是投影占主。
+
+3. **更稳的并发策略**
+   - tiny-batch MoE 已证明 `T=2` 能受益。
+   - 后续可测试 `max_num_seqs=3/4`，但受 100K 上下文 KV cache 与 TTFT 影响，需要小心。
+
+4. **TP all-reduce 合并**
+   - 当前 MoE 每层 routed+shared 后做一次 all-reduce。
+   - 如果后续能将若干小通信或 residual 路径合并，可能进一步改善 decode 抖动，但优先级低于 MoE/linear_attention 计算本体。