# How to Switch Attention Backend Between Flash-Attn and xFormers Using ATTN_BACKEND in TRELLIS.2

> Easily switch TRELLIS.2 attention backend between flash-attn and xformers by setting the ATTN_BACKEND environment variable. Control your implementation at runtime.

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

---

**Set the `ATTN_BACKEND` environment variable to `flash_attn` or `xformers` before importing TRELLIS.2 to control which attention implementation the library loads at runtime.**

The TRELLIS.2 repository by Microsoft provides flexible attention implementations for 3D generation workloads. You can switch attention backend between flash-attn and xformers using ATTN_BACKEND without modifying source code. This configuration variable determines whether the library uses FlashAttention, xFormers, or other supported backends for both dense and sparse attention operations.

## How TRELLIS.2 Selects the Attention Backend

The library determines its attention implementation at import time through two configuration modules. The `trellis2.modules.attention.config` module handles dense attention operations, while `trellis2.modules.sparse.config` manages sparse attention stacks. Both modules read the `ATTN_BACKEND` environment variable during initialization.

In [`trellis2/modules/attention/config.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/attention/config.py), the selection logic validates the environment variable against supported backends:

```python

# trellis2/modules/attention/config.py

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

```

Similarly, [`trellis2/modules/sparse/config.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/sparse/config.py) checks for `SPARSE_ATTN_BACKEND` first, then falls back to `ATTN_BACKEND`:

```python

# trellis2/modules/sparse/config.py

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 is not None and env_sparse_attn_backend in [
        'xformers', 'flash_attn', 'flash_attn_3']:
    ATTN = env_sparse_attn_backend

```

## Configuration Methods

You have three ways to set the attention backend, depending on your workflow requirements.

### Environment Variable (Recommended)

Setting `ATTN_BACKEND` before launching your script is the most reliable method. This ensures the backend is selected before any TRELLIS.2 modules are imported.

For Flash-Attention (default):

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

```

For xFormers:

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

```

### Runtime Setters (Interactive Use)

For Jupyter notebooks or interactive sessions where the environment variable was not preset, use the provided setter functions after importing the configuration modules:

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

set_backend('xformers')          # Sets global attention backend

set_attn_backend('xformers')     # Sets sparse attention backend

```

### In-Script Configuration (Import Order Critical)

You can set the environment variable programmatically, but you must do this **before** importing any TRELLIS.2 attention modules:

```python
import os

# Force xFormers for this process only

os.environ['ATTN_BACKEND'] = 'xformers'

# Import must happen AFTER setting the variable

import trellis2.modules.attention.config

# ... continue with model creation

```

## Supported Backend Values

The `ATTN_BACKEND` variable accepts different values depending on the attention type:

**Dense Attention** ([`trellis2/modules/attention/config.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/attention/config.py)):

- `flash_attn` — FlashAttention-2 (default)
- `flash_attn_3` — FlashAttention-3
- `xformers` — xFormers memory-efficient attention
- `sdpa` — PyTorch scaled dot-product attention
- `naive` — Pure PyTorch implementation

**Sparse Attention** ([`trellis2/modules/sparse/config.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/sparse/config.py)):

- `flash_attn` — FlashAttention-2
- `flash_attn_3` — FlashAttention-3
- `xformers` — xFormers implementation

## Separate Configuration for Sparse Attention

While `ATTN_BACKEND` controls both dense and sparse attention by default, you can override sparse attention independently using `SPARSE_ATTN_BACKEND`. This is useful when you want FlashAttention for dense layers but xFormers for sparse operations:

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

```

If `SPARSE_ATTN_BACKEND` is not set, the sparse config module automatically inherits the value from `ATTN_BACKEND`.

## Summary

- Set `ATTN_BACKEND` to `flash_attn` or `xformers` before importing TRELLIS.2 to control the attention implementation.
- Dense attention configuration resides in [`trellis2/modules/attention/config.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/attention/config.py), while sparse attention uses [`trellis2/modules/sparse/config.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/sparse/config.py).
- Use `set_backend()` and `set_attn_backend()` for runtime changes in interactive environments.
- For independent control, use `SPARSE_ATTN_BACKEND` to override sparse attention while keeping `ATTN_BACKEND` for dense attention.

## Frequently Asked Questions

### What happens if I set ATTN_BACKEND after importing TRELLIS.2?

The configuration is read at import time, so changing the environment variable after importing the module has no effect. To change backends mid-session, use the `set_backend()` and `set_attn_backend()` functions instead.

### Can I use FlashAttention-3 with TRELLIS.2?

Yes. Set `ATTN_BACKEND=flash_attn_3` to use FlashAttention-3. This value is supported in both the dense attention config ([`trellis2/modules/attention/config.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/attention/config.py)) and sparse attention config ([`trellis2/modules/sparse/config.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/sparse/config.py)).

### Why would I choose xFormers over FlashAttention?

xFormers provides memory-efficient attention implementations that may offer better compatibility with certain hardware configurations or PyTorch versions where FlashAttention is not available or causes compilation issues. According to the TRELLIS.2 source code, both backends are fully supported.

### Is there a difference between ATTN_BACKEND and ATTN_DEBUG?

Yes. `ATTN_DEBUG` is an alternative environment variable recognized by [`trellis2/modules/attention/config.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/attention/config.py) for debugging purposes, while `ATTN_BACKEND` is the primary configuration variable. For sparse attention, `SPARSE_ATTN_BACKEND` is the specific override that falls back to `ATTN_BACKEND` if not set.