# How to Switch Between Flash-Attention and XFormers Backends in TRELLIS.2

> Learn to switch between Flash-Attention and XFormers backends in TRELLIS.2 using environment variables or the Python API for faster model performance.

- Repository: [Microsoft/TRELLIS.2](https://github.com/microsoft/TRELLIS.2)
- Tags: how-to-guide
- Published: 2026-08-04

---

**TRELLIS.2 supports multiple attention backends—Flash-Attention, XFormers, and others—controlled via environment variables or Python API calls before model initialization.**

The Microsoft TRELLIS.2 repository provides a flexible attention configuration system that lets you toggle between high-performance CUDA kernels and cross-platform alternatives. Whether you're optimizing for maximum throughput on A100/H100 GPUs or need fallback compatibility, understanding how to switch between **Flash-Attention** and **XFormers** backends is essential for efficient model deployment.

## Global Attention Backend Configuration

The primary attention backend for dense transformer layers is managed in [[`trellis2/modules/attention/config.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/attention/config.py)](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/attention/config.py). The system defaults to **Flash-Attention** but reads the `ATTN_BACKEND` environment variable at import time to override this.

```python

# From trellis2/modules/attention/config.py

BACKEND = 'flash_attn'

env_attn_backend = os.environ.get('ATTN_BACKEND')
if env_attn_backend in ['xformers', 'flash_attn', 'flash_attn_3', 'sdpa', 'naive']:
    BACKEND = env_attn_backend

```

Available backend options include:

- `flash_attn` — Default, highly optimized for modern NVIDIA GPUs
- `flash_attn_3` — Experimental next-generation Flash-Attention
- `xformers` — Flexible PyTorch-based implementation
- `sdpa` — PyTorch native scaled dot-product attention
- `naive` — Pure PyTorch reference implementation

The actual kernel dispatch occurs in [[`trellis2/modules/attention/full_attn.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/attention/full_attn.py)](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/attention/full_attn.py), which lazily imports the selected library on first use.

### Switching via Environment Variable

Set `ATTN_BACKEND` **before** importing any TRELLIS.2 modules:

```bash

# Use XFormers for all dense attention operations

export ATTN_BACKEND=xformers
python train.py

```

```bash

# Restore Flash-Attention (default behavior)

export ATTN_BACKEND=flash_attn
python inference.py

```

### Switching via Python API

For Jupyter notebooks or runtime configuration, use `set_backend()`:

```python
from trellis2.modules.attention.config import set_backend

# Switch to XFormers mid-session

set_backend('xformers')

# Verify the change

from trellis2.modules.attention import config
print(config.BACKEND)  # 'xformers'

```

**Critical timing requirement:** Call `set_backend()` **before** importing any model classes that instantiate attention layers. The backend value is read once at module initialization and cached for subsequent operations.

## Sparse Attention Backend Configuration

Sparse transformer layers in TRELLIS.2 maintain a separate backend configuration in [[`trellis2/modules/sparse/config.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/sparse/config.py)](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/sparse/config.py). This allows fine-grained control over sparse vs. dense attention implementations.

```python

# From trellis2/modules/sparse/config.py

ATTN = 'flash_attn'

# Hierarchical fallback: SPARSE_ATTN_BACKEND → ATTN_BACKEND → default

env_sparse_attn_backend = os.environ.get('SPARSE_ATTN_BACKEND')
if env_sparse_attn_backend is None:
    env_sparse_attn_backend = os.environ.get('ATTN_BACKEND')

if env_sparse_attn_backend in ['xformers', 'flash_attn', 'flash_attn_3']:
    ATTN = env_sparse_attn_backend

```

The sparse attention dispatcher lives in [[`trellis2/modules/sparse/attention/full_attn.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/sparse/attention/full_attn.py)](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/sparse/attention/full_attn.py) and follows the same lazy import pattern as the dense implementation.

### Environment Variable Control

```bash

# Dedicated sparse backend control

export SPARSE_ATTN_BACKEND=xformers
python train_sparse_model.py

```

```bash

# Use Flash-Attention for sparse layers while keeping dense layers configurable

export SPARSE_ATTN_BACKEND=flash_attn
export ATTN_BACKEND=xformers  # Dense layers use XFormers

python train.py

```

### Python API Control

```python
from trellis2.modules.sparse.config import set_attn_backend

# Configure sparse attention independently

set_attn_backend('flash_attn')

```

## Practical Configuration Patterns

### Unified Backend (Simplest)

When you want both dense and sparse layers using the same implementation:

```bash
export ATTN_BACKEND=xformers
python train.py  # Both dense and sparse use XFormers

```

### Split Backend (Performance-Tuned)

Optimize for hardware characteristics where Flash-Attention excels on dense matrices but XFormers handles sparse patterns efficiently:

```bash
export ATTN_BACKEND=flash_attn
export SPARSE_ATTN_BACKEND=xformers
python train.py

```

### Runtime Switching in Notebooks

```python

# Configure backends before any model imports

from trellis2.modules.attention.config import set_backend
from trellis2.modules.sparse.config import set_attn_backend

set_backend('flash_attn')
set_attn_backend('xformers')

# Now safe to import models—backends are locked in

from trellis2.models.sparse_structure_vae import SparseStructureVAE

model = SparseStructureVAE(config_path="configs/vae.yaml")

```

## Backend Selection Reference

| Scenario | Recommended Backend | Environment Setup |
|----------|---------------------|-------------------|
| Maximum throughput on Ampere/Hopper | `flash_attn` | `export ATTN_BACKEND=flash_attn` |
| Older GPUs (pre-A100) | `xformers` | `export ATTN_BACKEND=xformers` |
| Debugging attention behavior | `naive` | `export ATTN_BACKEND=naive` |
| PyTorch-native only deployment | `sdpa` | `export ATTN_BACKEND=sdpa` |
| Mixed sparse/dense optimization | Split configuration | Set both `ATTN_BACKEND` and `SPARSE_ATTN_BACKEND` |

## Key Implementation Files

Understanding these source locations helps with debugging and custom modifications:

- **[`trellis2/modules/attention/config.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/attention/config.py)** — Global `BACKEND` variable and `set_backend()` function
- **[`trellis2/modules/attention/full_attn.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/attention/full_attn.py)** — Dense attention kernel dispatch
- **[`trellis2/modules/sparse/config.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/sparse/config.py)** — Sparse-specific `ATTN` variable and `set_attn_backend()` function  
- **[`trellis2/modules/sparse/attention/full_attn.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/sparse/attention/full_attn.py)** — Sparse attention kernel dispatch

Both attention implementations use deferred imports to minimize startup overhead and avoid hard dependencies on unused libraries.

## Summary

- **Two configuration domains** exist: global (`ATTN_BACKEND`) and sparse (`SPARSE_ATTN_BACKEND`), with hierarchical fallback behavior
- **Environment variables** must be set before Python imports TRELLIS.2 modules to take effect
- **Python API** (`set_backend()`, `set_attn_backend()`) enables runtime switching when called before model instantiation
- **Five backends** are supported: `flash_attn`, `flash_attn_3`, `xformers`, `sdpa`, and `naive`
- **Kernel dispatch** occurs in [`full_attn.py`](https://github.com/microsoft/TRELLIS.2/blob/main/full_attn.py) files for both dense and sparse paths, with lazy library importing

## Frequently Asked Questions

### What happens if I set the backend after importing TRELLIS.2 models?

The change has no effect on already-instantiated attention layers. The `BACKEND` and `ATTN` variables are read once during module initialization and cached. To apply a new backend, you must restart the Python process or use a fresh interpreter.

### Can I use different backends for different models in the same script?

Not simultaneously within the same process. The backend configuration is global and module-scoped. For heterogeneous backend requirements, use subprocess isolation or sequentially load models with different environment variable presets.

### What if Flash-Attention is installed but fails at runtime?

TRELLIS.2 falls back gracefully only for the `naive` backend. For production stability with uncertain GPU compatibility, explicitly set `ATTN_BACKEND=xformers` to avoid CUDA kernel errors during attention operations.

### How do I verify which backend is actually active?

Inspect the config module after model initialization:

```python
from trellis2.modules.attention import config
print(f"Dense attention: {config.BACKEND}")

from trellis2.modules.sparse import config as sparse_config
print(f"Sparse attention: {sparse_config.ATTN}")

```