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

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). The system defaults to Flash-Attention but reads the ATTN_BACKEND environment variable at import time to override this.

# 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), which lazily imports the selected library on first use.

Switching via Environment Variable

Set ATTN_BACKEND before importing any TRELLIS.2 modules:

# Use XFormers for all dense attention operations

export ATTN_BACKEND=xformers
python train.py

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

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). This allows fine-grained control over sparse vs. dense attention implementations.

# 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) and follows the same lazy import pattern as the dense implementation.

Environment Variable Control

# Dedicated sparse backend control

export SPARSE_ATTN_BACKEND=xformers
python train_sparse_model.py

# 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

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:

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:

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

Runtime Switching in Notebooks

# 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 — Global BACKEND variable and set_backend() function
• trellis2/modules/attention/full_attn.py — Dense attention kernel dispatch
• trellis2/modules/sparse/config.py — Sparse-specific ATTN variable and set_attn_backend() function
• 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 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:

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}")

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}