flagos-ai / flagos-ai/FlagTree

[Kunlunxin] 昆仑芯使用FlageTree Triton3.0 在测试vllm Triton attention时也会报错

Open
#935 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
350
Forks
149
Avg merge
2d 4h
Merged PRs (30d)
81

Description

# Triton XPU 编译器在 Attention Pattern 下触发 UNREACHABLE 断言失败

## 环境

| 组件 | 版本 |
|------|------|
| PyTorch | 2.5.1+cu118 (昆仑芯定制版) |
| Triton | 3.0.0 (XPU backend, FlagTree) |
| Python | 3.10.15 |
| 硬件 | 昆仑芯 XPU |
| triton backends | `{'xpu': Backend(compiler=XPUBackend, driver=XPUDriver)}` |

## 问题描述

在同一个 triton kernel 中执行标准 attention 计算模式:

```
tl.dot (Q @ K^T) → softmax (max/exp/sum/div) → .to(fp16) → tl.dot (P @ V)
```

编译器在 `TritonSDNNToLLVM` 转换阶段触发断言失败,进程 abort。

## 错误信息

```
src1Type.getShape()[0] != 1
UNREACHABLE executed at /ssd2/dongjibin/flagtree/baidu/xpu/triton/third_party/xpu/lib/Conversion/TritonSDNNToLLVM/sdnn.h:2550!
```

另一个触发路径的错误信息:

```
UNREACHABLE executed at /data/zhengyang/cd-actions-runner_xpu/_work/FlagTree/FlagTree/third_party/xpu/lib/Dialect/TritonXPU/Transforms/CoreTiling.cpp:205!
```

## 最小复现

```python
import torch
import triton
import triton.language as tl

@triton.jit
def attention_kernel(
Q, K, V, Out,
N: tl.constexpr,
D: tl.constexpr,
sm_scale,
):
"""最小化的 attention kernel: Out = softmax(Q @ K^T * scale) @ V"""
offs_nd = tl.arange(0, N)[:, None] * D + tl.arange(0, D)[None, :]

q = tl.load(Q + offs_nd) # [N, D] fp16
k = tl.load(K + offs_nd) # [N, D] fp16
v = tl.load(V + offs_nd) # [N, D] fp16

# Step 1: S = Q @ K^T
s = tl.dot(q, tl.trans(k)) # [N, N] fp32
s = s * sm_scale

# Step 2: softmax
s_max = tl.max(s, axis=1)[:, None]
s = s - s_max
p = tl.exp(s)
p_sum = tl.sum(p, axis=1)[:, None]
p = p / p_sum # [N, N] fp32

# Step 3: cast + second dot ← 触发编译器 crash
p_fp16 = p.to(tl.float16)
o = tl.dot(p_fp16, v) # [N, D]

tl.store(Out + offs_nd, o)

N, D = 16, 16
sm_scale = 1.0 / (D ** 0.5)
Q = torch.randn(N, D, device='cuda', dtype=torch.float16)
K = torch.randn(N, D, device='cuda', dtype=torch.float16)
V = torch.randn(N, D, device='cuda', dtype=torch.float16)
Out = torch.empty(N, D, device='cuda', dtype=torch.float16)

attention_kernel[(1,)](Q, K, V, Out, N, D, sm_scale)
```

完整复现脚本(含对照测试):[`bug_reproduce.py`](https://github.com/user-attachments/files/30821529/bug_reproduce.py)

## 对照测试结果

以下操作**单独均可正常编译运行**:

| 测试 | 操作 | 结果 |
|------|------|------|
| tl.dot 单独 | `C = tl.dot(A, B)` | ✅ PASS |
| tl.dot + tl.trans | `S = tl.dot(Q, tl.trans(K))` | ✅ PASS |
| 2D softmax | `tl.max` + `tl.exp` + `tl.sum` + `/` | ✅ PASS |
| 两个 tl.dot + .to() cast(无 softmax) | `tl.dot(A.to(fp16), B)` | ✅ PASS |
| **dot → softmax → cast → dot(attention)** | 见上方 kernel | ❌ CRASH |

## 影响

此问题导致所有基于 triton 的 attention kernel 无法在昆仑芯上运行,包括但不限于:

- vLLM `triton_decode_attention` (TRITON_MLA)
- vLLM `triton_unified_attention` (TRITON_ATTN)
- vLLM `triton_prefill_attention`
- 任何包含标准 `QK^T → softmax → PV` 模式的 triton kernel

## 分析

根据错误位置 `sdnn.h:2550` 和断言 `src1Type.getShape()[0] != 1`,推测是 `TritonSDNNToLLVM` 在处理第二个 `tl.dot` 的操作数时,期望 `src1`(即 softmax 输出经 `.to()` cast 后的 tensor)的第一个维度为 1,但实际是 `[N, N]` 的 2D tensor。

可能的触发条件是:当 `tl.dot` 的输入来自一系列 elementwise 操作(exp、div)的结果而非直接从 memory load 时,编译器的 tiling/lowering 逻辑未正确处理这种 data flow。

## 期望行为

该 kernel 应正常编译并产生正确的 attention 计算结果,与 PyTorch reference 实现一致:

```python
s = (Q.float() @ K.float().T) * sm_scale
p = torch.softmax(s, dim=-1)
out = (p @ V.float()).half()
```

Contributor guide

Open the contributing guide

Research direction

Start by running the attached bug_reproduce.py on the stated Kunlunxin XPU environment and confirm the crash. Then inspect TritonSDNNToLLVM and CoreTiling.cpp at the reported assertion locations, tracing the second tl.dot after the softmax and cast. Done means the attention kernel compiles, runs, and matches the PyTorch reference without triggering either assertion.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.