Attention Backend¶

Available backend¶

Cache-DiT supports multiple Attention backends for better performance. The supported list is as follows:

backend	details	parallelism	attn_mask
native	Native SDPA Attention, w/ cache-dit optimized	✅	✅
_sdpa_cudnn	CUDNN Attention via SDPA API, w/ cache-dit optimized	✅	✅
_native_cudnn	CUDNN Attention via SDPA API, w/o cache-dit optimized	✅	✖️
flash	official FlashAttention-2	✅	✖️
_flash_3	official FlashAttention-3	✅	✖️
sage	FP8 SageAttention	✅	✖️
_native_npu	Optimized Ascend NPU Attention	✅	✅
_npu_fia	NPU Attention for Ring Parallelism	✅	✅

Single GPU Inference¶

Users can specify Attention backend by setting the attention_backend parameter of enable_cache API or use set_attn_backend interface directly.

import cache_dit

# Setting the `attention_backend` parameter of `enable_cache` API
cache_dit.enable_cache(pipe_or_adapter, ..., attention_backend="_sdpa_cudnn")
# Or, use `set_attn_backend` interface directly.  
cache_dit.set_attn_backend(pipe_or_adapter, attention_backend="_sdpa_cudnn")

Distributed inference¶

Users also can specify Attention backend by setting the attention_backend parameter of parallelism_config in the cases of distributed inference:

from cache_dit import ParallelismConfig

cache_dit.enable_cache(
  pipe_or_adapter, 
  cache_config=DBCacheConfig(...),
  parallelism_config=ParallelismConfig(
    ulysses_size=2, # or, tp_size=2
    # flash, native(sdpa), _native_cudnn, _sdpa_cudnn, sage
    attention_backend="_sdpa_cudnn",
  ),
)

FP8 Attention¶

For FP8 Attention, users must install sage-attention. Then, pass the sage attention backend to the parallelism_config as an extra parameter. Please note that attention mask is not currently supported for FP8 sage attention.

# pip3 install git+https://github.com/thu-ml/SageAttention.git 
from cache_dit import ParallelismConfig

cache_dit.enable_cache(
  pipe_or_adapter, 
  cache_config=DBCacheConfig(...),
  parallelism_config=ParallelismConfig(
    ulysses_size=2, # or, tp_size=2
    attention_backend="sage",
  ),
)