flagos-ai / flagos-ai/FlagTree
[BUG][mthreads] flagtune do_bench_musa_graph 在 MTT S5000 上 muCtxSynchronize 卡死,服务首请求死锁
- Dominant language
- Python
- Stars
- 350
- Forks
- 149
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 81
Description
## 根因
FlagTree rc2 新增的 flagtune replay benchmarker(`do_bench_musa_graph`)在 MTT S5000 上 `torch.musa.synchronize()` 卡死(muCtxSynchronize wedged),导致 flag_gems 任意算子的首次 autotune 死锁,sglang 服务首请求即 watchdog 超时 SIGQUIT。**根因不在算子、不在上层(flag_gems/sglang-plugin),在 FlagTree 的 flagtune graph benchmark 路径,应由 FlagTree 修。**
## 简单复现
只要 flagtree + flaggems + torch_musa 环境(**无需 sglang / 模型 / tp**):
```python
import time, torch, torch_musa, flag_gems
dev = flag_gems.device
M = N = K = 256
x = torch.randn(M, K, dtype=torch.float16, device=dev)
w = torch.randn(N, K, dtype=torch.float16, device=dev)
t0 = time.time()
with flag_gems.use_gems():
z = torch.nn.functional.linear(x, w)
torch_musa.synchronize()
print("linear first autotune: %.2f s" % (time.time() - t0))
```
**现象**:首次 autotune 要 **~16 秒**(本应 <1s)。卡在 `do_bench_musa_graph` 的 graph capture benchmark——`_calibrate_n_repeat` 先 capture 256 次 unroll 的 probe graph,再 capture `n_repeat` 次的正式 graph。tiny shape 下 kernel 极快,`n_repeat = rep/best` 被打到 `max_n_repeat=20000` 上限,20000 次 unroll 的 capture 在 MTT 上 wedged。
> 注:单算子单卡下能跑完(慢但不死锁)。死锁是同一路径在服务时序(tp 通信 + 多 graph capture 竞争 synchronize)下的重症,栈见下——**不用复现死锁也能定位**,慢 16s 已足够证明这条路径有问题。
## 细节
### 死锁栈(sglang 服务里 py-spy dump,卡在第 1 行)
```
__enter__ (torch_musa/musa_graph/graphs.py:201) # torch.musa.synchronize() ← 卡死在这
__enter__ (torchada/_patch.py:440)
do_bench_musa_graph (triton/flagtune/runtime/graph_benchmark.py:90)
replay_benchmark (triton/flagtune/runtime/benchmark_protocol.py:155)
_bench (autotuner.py:190)
bench (flag_gems/utils/libentry.py:1001)
...
linear (flag_gems/ops/linear.py:173)
```
`graphs.py:201` 正是 `__enter__` 第一行 `torch.musa.synchronize()`。
### 代码里已经承认的注释(graph_benchmark.py:62)
```python
# Cap the captured graph. A wedged muCtxSynchronize was observed on MTT
# S5000 while sweeping tiny shapes, and an unbounded unroll makes both the
# capture and every replay proportionally more exposed to it.
```
即 MTT S5000 的 muCtxSynchronize wedged 已知,现有做法只是 cap `max_n_repeat` 缓解,没根治。
### bisect:rc0 正常 / rc2 死锁
| 组合 | triton.flagtune | benchmark 路径 | 结果 |
| --- | --- | --- | --- |
| FlagTree rc0(2e62581) + FlagGems rc2 | 无 | fallback event timing | 正常(c64 压测 34min 不崩) |
| FlagTree rc2(5eda995) + FlagGems rc2 | 有 | do_bench_musa_graph | 首请求死锁 |
flag_gems 的 `libentry.py:52-68` 检测 `import triton.flagtune`:有则走 `triton.flagtune.runtime.benchmark_protocol.resolve_benchmarker`(死锁路径),无则 fallback `event timing`。该选择没有按 backend 判断,mthreads 上默认命中了有问题的 graph 路径。
### 版本
FlagTree `0.7.0-rc2.post1+triton3.6`(commit 5eda995;tag == 分支 HEAD,master 上同一文件也未修,`graph_benchmark.py` 历史仅 09-01 引入提交 cefac05)。
Contributor guide
Assessment
This issue has not been assessed yet.