flagos-ai / flagos-ai/FlagTree
XPU Compiler: OutOfResources on GDN fused_post_conv_kernel (Qwen3.6)
- Dominant language
- Python
- Stars
- 350
- Forks
- 149
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 81
Description
# FlagTree XPU Compiler: OutOfResources on GDN fused_post_conv_kernel
## Environment
- **FlagTree version**: 0.6.1+xpu3.6
- **Triton version**: 3.6.0
- **Hardware**: KunlunXin P800 OAM (XPU)
- **Driver**: xre5.37.1
- **vLLM**: 0.20.2+flagos
- **Model**: Qwen3.6-27B (uses GatedDeltaNet attention)
## Issue Description
FlagTree XPU backend fails to compile the `_fused_post_conv_kernel` from vLLM's GatedDeltaNet (GDN) implementation with `OutOfResources: uni_sram PassManager::run failed`.
## Error Stack Trace
```
File "/opt/flagtree/triton/backends/xpu/compiler.py", line 367, in make_ttxir
raise OutOfResources(0, 0, f"uni_sram {e}")
triton.runtime.errors.OutOfResources: out of resource: uni_sram PassManager::run failed, Required: 0, Hardware limit: 0. Reducing block sizes or `num_stages` may help.
```
Full stack:
```
File "/flagos/lib/python3.10/site-packages/vllm/model_executor/layers/fla/ops/fused_gdn_prefill_post_conv.py", line 215, in fused_post_conv_prep
_fused_post_conv_kernel[grid](...)
File "/opt/flagtree/triton/runtime/jit.py", line 733, in run
kernel = self._do_compile(key, signature, device, constexprs, options, attrs, warmup)
File "/opt/flagtree/triton/compiler/compiler.py", line 378, in compile
cur_module = run_stage(ext, compile_ir, cur_module)
File "/opt/flagtree/triton/backends/xpu/compiler.py", line 367, in make_ttxir
raise OutOfResources(0, 0, f"uni_sram {e}")
```
## Kernel Characteristics
The failing kernel is a complex fused operation that:
- Splits conv1d output into Q/K/V/gating components
- Applies L2 normalization to Q/K
- Computes gating (softplus/exp transforms)
- Writes to 5 separate output tensors
**Launch config:**
- `num_warps=4`
- `num_stages=2` ← likely the problem
- `BLOCK_T=16`, `BK=triton.next_power_of_2(K)`, `BV=triton.next_power_of_2(V)`
- For Qwen3.6-27B: H=128, HV=16, K=128, V=128
**Kernel source:** https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/fla/ops/fused_gdn_prefill_post_conv.py#L20
## Root Cause
The error is raised in `make_ttxir()` after running the PassManager pipeline:
```python
try:
pm.run(mod, 'make_ttxir')
except Exception as e:
from triton import OutOfResources
raise OutOfResources(0, 0, f"uni_sram {e}")
```
**Problem 1: The catch-all exception handler masks the real compilation error.** The actual failure could be in any of the XPU passes (legalize, vectorize, alloca, etc.), but we only see "uni_sram PassManager::run failed".
**Problem 2: The uni_sram allocation fails for this complex kernel with `num_stages=2`.** The kernel has:
- Multiple control-flow branches (Q/K path vs V/gating path)
- Multiple memory loads/stores per block
- Intermediate accumulations (L2 norm, softplus)
- 2-stage pipelining requires buffering for both stages
## Reproducibility
```bash
# Start vLLM serve with Qwen3.6-27B
PYTHONPATH=/opt/flagtree CUDA_VISIBLE_DEVICES=0 \
python3 -m vllm.entrypoints.openai.api_server \
--model Qwen3.6-27B \
--tensor-parallel-size 1 \
--trust-remote-code
```
The error occurs during model initialization in `_warmup_prefill_kernels()`.
## Suggested Fixes
### Option 1: Improve Error Reporting (High Priority)
In `/opt/flagtree/triton/backends/xpu/compiler.py:367`, change:
```python
try:
pm.run(mod, 'make_ttxir')
except Exception as e:
from triton import OutOfResources
raise OutOfResources(0, 0, f"uni_sram {e}")
```
to:
```python
try:
pm.run(mod, 'make_ttxir')
except Exception as e:
from triton import OutOfResources
import traceback
# Include the full exception chain so users can see which pass failed
raise OutOfResources(0, 0, f"uni_sram {e}\n{traceback.format_exc()}")
```
This would help diagnose which specific pass is failing.
### Option 2: Add Fallback for num_stages
The XPU compiler could automatically retry with `num_stages=1` when `num_stages=2` fails:
```python
try:
pm.run(mod, 'make_ttxir')
except Exception as e:
if metadata.get("num_stages", 1) > 1:
warnings.warn(f"Compilation failed with num_stages={metadata['num_stages']}, retrying with num_stages=1")
metadata["num_stages"] = 1
# rebuild PM and retry
else:
from triton import OutOfResources
raise OutOfResources(0, 0, f"uni_sram {e}")
```
### Option 3: Improve uni_sram Allocation
The long-term fix is to improve the resource allocation in the relevant passes (likely `tritonxpu_legalize_pass` or `tritonxpu_alloca_pass`) to handle complex multi-output kernels with pipelining.
## Workarounds
For now, users can:
1. Patch vLLM's kernel to use `num_stages=1` for Kunlunxin backend
2. Use models without GatedDeltaNet (e.g., standard Qwen2 with regular attention)
3. Wait for compiler improvements
## Additional Context
- The same kernel compiles successfully on NVIDIA GPUs with triton/CUDA backend
- Other simpler triton kernels in vLLM work fine on Kunlunxin
- This appears to be a resource limit specific to complex kernels with multi-stage pipelining on XPU
---
Would appreciate any guidance on which pass is likely failing and whether there are tuning knobs (env vars like `TRITONXPU_BUFFER_SIZE`, etc.) that could help allocate more uni_sram.
Contributor guide
Assessment
This issue has not been assessed yet.