flagos-ai / flagos-ai/FlagTree

XPU Compiler: OutOfResources on GDN fused_post_conv_kernel (Qwen3.6)

Open
#1,126 1 comment 0 reactions 2 assignees Claimed by @zhzhcookie View on GitHub
bug flagos2.2-rc0 P0 xpu
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.