flagos-ai / flagos-ai/FlagTree
【FR】Bug: FlagTree Ascend backend fails to lower vLLM apply_top_k_top_p_triton (TritonIR→LinalgIR, () → tensor<1xf32> cannot be legalized)
- Dominant language
- Python
- Stars
- 350
- Forks
- 149
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 81
Description
## Bug: FlagTree Ascend backend fails to lower vLLM `apply_top_k_top_p_triton` kernel
FlagTree Ascend 后端在 **TritonIR → LinalgIR** lowering 阶段无法 legalize vLLM 内置采样 kernel `apply_top_k_top_p_triton` 中的一处 `() → tensor<1xf32>` 材化(scalar ↔ 1-element tensor),导致 `MLIRCompilationError`,EngineCore 直接 `EngineDeadError` 崩溃。
**影响面**:凡使用 top-p / top-k 采样的模型,一进入采样即崩。本批次已确认阻塞 **openbmb/AgentCPM-Report** 与 **openbmb/MiniCPM4.1-8B** 两个模型的精度评测(评测 0 分、无法产出结果)。
### 关键结论(责任归属)
- **崩溃发生在 FlagTree 的 Ascend 编译后端**:`triton/spec/ascend/compiler/compiler.py:346` → `raise MLIRCompilationError`(此环境的 `triton` 包即 FlagTree,含 `spec/ascend`、`experimental/tle` 特有布局)。
- **与 FlagGems 无关**:V1(不注入 FlagGems)与 V2(注入 FlagGems)**报同一个错**。
- **与华为 vllm-ascend 无关**:基础镜像 `flagrelease_ascend_vllm020plugin_base:no_vllm_ascend`,环境未安装华为 vllm-ascend 插件。
- kernel **源码**是 vLLM 的(`vllm/v1/sample/ops/topk_topp_triton.py`),但**编译它的是 FlagTree**;失败发生在 FlagTree 自身的 MLIR 转换 pass(`ConvertTritonIRToLinalgIR`),尚未进入 CANN/bishengir 最终 codegen。
### Environment
| Item | Value |
|------|-------|
| Hardware | 华为 Ascend910B |
| FlagTree | 0.6.0 |
| Triton backend | `triton/backends/ascend`(FlagTree Ascend backend) |
| vLLM | 0.20.2 |
| torch_npu | 2.10.0 |
| CANN | 9.0.0 |
| FlagGems | 5.3.4(开/不开均复现) |
| plugin | vllm-plugin-FL 0.2.0 |
| Base image | `flagrelease_ascend_vllm020plugin_base:no_vllm_ascend`(无 vllm-ascend) |
### Steps to Reproduce
1. 在 Ascend910B + 上述 FlagTree 0.6.0 环境启动 vLLM 0.20.2 服务(**无需开启 FlagGems**)。
2. 用任一开启 top-p / top-k 采样的模型(如 `generation_config.json` 含 `temperature`/`top_p`)发起推理,触发 batch 采样路径。
3. vLLM 调用 `apply_top_k_top_p_triton` → FlagTree 编译 `_topk_topp_kernel` → MLIR lowering 失败 → EngineCore 崩溃。
> 已验证:V1(native,不开 FlagGems)与 V2(开 FlagGems)均 100% 复现,报错完全相同。
### Root Cause(精确定位)
出错的 Triton 源码行(vLLM `topk_topp_triton.py:677`):
```python
idx = tl.cast(p * 200, tl.int32)
```
FlagTree 的 `ConvertTritonIRToLinalgIR` pass 报:
```
topk_topp_triton.py:677:34: error: failed to legalize unresolved materialization
from () to ('tensor<1xf32>') that remained live after conversion
idx = tl.cast(p * 200, tl.int32)
^
topk_topp_triton.py:677:34: note: see existing live user here:
%extracted_67 = tensor.extract %88[%c0] : tensor<1xf32>
topk_topp_triton.py:94:0: error: failed to apply Conversion Patterns
```
即:一个标量值 `p`(来自 `tensor.extract %88[%c0] : tensor<1xf32>`)在 Triton→Linalg 转换后,其 `() ↔ tensor<1xf32>` 的 unrealized_conversion_cast 无法被 legalize,`PassManager::run failed`。这属于 **FlagTree Ascend 后端 scalar/tensor 材化处理的缺陷**。
### Error Logs(节选,两个模型一致)
```
/vllm-workspace/vllm/vllm/v1/sample/ops/topk_topp_triton.py:677:34: error: failed to legalize unresolved materialization from () to ('tensor<1xf32>') that remained live after conversion
idx = tl.cast(p * 200, tl.int32)
^
/vllm-workspace/vllm/vllm/v1/sample/ops/topk_topp_triton.py:677:34: note: see existing live user here: %extracted_67 = tensor.extract %88[%c0] : tensor<1xf32>
/vllm-workspace/vllm/vllm/v1/sample/ops/topk_topp_triton.py:94:0: error: failed to apply Conversion Patterns
Traceback (most recent call last):
File ".../vllm/v1/sample/ops/topk_topp_sampler.py", line 106, in forward_native
logits = apply_top_k_top_p(logits, k, p)
File ".../vllm/v1/sample/ops/topk_topp_sampler.py", line 252, in apply_top_k_top_p
return apply_top_k_top_p_triton(logits, k, p)
File ".../vllm/v1/sample/ops/topk_topp_triton.py", line 1035, in apply_top_k_top_p_triton
_topk_topp_kernel[(NUM_PROGRAMS,)](
File ".../site-packages/triton/spec/ascend/runtime/jit.py", line 733, in run
kernel = self._do_compile(...)
File ".../site-packages/triton/experimental/tle/__init__.py", line 61, in tle_compile
compiled_kernel = original_compile_fn(src, target, options)
File ".../site-packages/triton/spec/ascend/compiler/compiler.py", line 346, in compile
raise MLIRCompilationError(stage_name, error_detail) from e
triton.compiler.errors.MLIRCompilationError:
///------------------[ERROR][Triton][BEG]------------------
[ConvertTritonIRToLinalgIR] encounters error:
PassManager::run failed
///------------------[ERROR][Triton][END]------------------
vllm.v1.engine.exceptions.EngineDeadError: EngineCore encountered an issue.
```
### Expected Behavior
FlagTree Ascend 后端应能正确 lower 该 kernel(正确处理 scalar → 1-element tensor 的材化),使 top-p/top-k 采样在 Ascend 上正常编译执行。
### Possible Directions
- 修复 `ConvertTritonIRToLinalgIR` 对 `() ↔ tensor<1xf32>` unrealized_conversion_cast 的 legalization(补充 conversion pattern,或在 lowering 时消除残留的 scalar↔tensor 材化)。
- 复现最小用例:对一个从 `tensor.extract` 得到的标量做 `tl.cast(scalar * const, tl.int32)`,走 Ascend 后端编译。
- 临时规避(非修复):压低并发避开 batch 采样编译路径(实测 `FLAGOS_MAX_CONCURRENCY=4` 仍无法产出有效结果,不作为解决方案)。
### Affected Models (this batch)
- openbmb/AgentCPM-Report — 评测采样崩溃,0 分
- openbmb/MiniCPM4.1-8B — 评测采样崩溃,0 分
Contributor guide
Research direction
Start with triton/spec/ascend/compiler/compiler.py:346 and trace the ConvertTritonIRToLinalgIR pass. Reproduce the failure through vLLM's vllm/v1/sample/ops/topk_topp_triton.py, especially line 677 and the tensor.extract path, then reduce it to the scalar-to-one-element-tensor case described in the issue. Done means the Ascend backend lowers the kernel and top-p/top-k sampling runs without MLIRCompilationError or EngineDeadError.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100