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

[Al2O3Al2O3]

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:

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

# 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:

try:
    pm.run(mod, 'make_ttxir')
except Exception as e:
    from triton import OutOfResources
    raise OutOfResources(0, 0, f"uni_sram {e}")

to:

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:

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.

[wbavon]

@Al2O3Al2O3 Kunlunxin has released a new version, 0.6.1a1+xpu3.6. According to R&D feedback, the issue should be fixed, pls helop to test against

[zhzhcookie]

@Al2O3Al2O3 Kunlunxin has released a new version, 0.6.1a1+xpu3.6. According to R&D feedback, the issue should be fixed, pls helop to test against

@Al2O3Al2O3 Please try 0.7.0rc2+xpu3.6