Skip to content

v2.3.6 - Enhanced SageAttention3 (SA3) Integration

Choose a tag to compare

View all tags
@ussoewwin ussoewwin released this 23 Jan 18:42
· 91 commits to main since this release
v2.3.6
3872a80

Enhanced SageAttention3 (SA3) Integration

This release enhances the Patch Sage Attention DM node with improved SageAttention3 (SA3) support, including version detection, constraint handling, and automatic fallback mechanisms.

Key Features

1. SA3 Version Detection

  • Added dedicated get_sage_attention3_info() function for SA3 version detection
  • Detects Blackwell GPU support (FP4 kernel availability)
  • Returns version information, availability status, and Blackwell support flag

2. Improved Logging

  • SA2 version logs are now skipped when SA3 modes are selected
  • SA3-specific logging with version information (e.g., "SageAttention3 3.0.0.b1 (Blackwell FP4)")
  • Clear distinction between SA2 and SA3 mode logging

3. Constraint Handling and Fallback

  • Automatic fallback to PyTorch SDPA when SA3 constraints are not met:
    • headdim >= 256: SA3 FP4 kernel does not support head dimensions >= 256
    • attn_mask != None: SA3 does not support attention masks
  • Seamless fallback ensures compatibility with all model configurations

4. Tensor Layout Conversion

  • Automatic conversion from ComfyUI's default NHD layout [batch, seq_len, heads, dim] to SA3's expected HND layout [batch, heads, seq_len, dim]
  • Proper layout restoration after SA3 processing

5. Per-Block Mean Support

  • Support for both sageattn3 (standard mode, accuracy-focused) and sageattn3_per_block_mean (fast mode, Triton-optimized)
  • Per-block mean processing uses 128-token blocks with Triton kernels for improved performance

Technical Details

Modified File: nodes/sa.py

Changes:

  • Lines 102-125: Added get_sage_attention3_info() function
  • Lines 151-152: Added SA2 log skip condition for SA3 modes
  • Lines 184-188: Added SA3-specific version detection and logging
  • Lines 195-200: Clarified tensor layout conversion
  • Lines 202-210: Added fallback detection logic
  • Lines 211-218: Improved fallback processing and SA3 call branching

Usage

  1. Start ComfyUI with --use-sage-attention flag:

    python main.py --use-sage-attention

  2. Add the Patch Sage Attention DM node to your workflow

  3. Select SA3 mode from the sage_attention dropdown:

    • sageattn3: Standard SA3 mode (accuracy-focused)
    • sageattn3_per_block_mean: Fast SA3 mode (Triton-optimized)

Expected Log Output

When SA2 is selected:

Patching comfy attention to use SageAttention 2.1.1+cu128torch2.7

When SA3 is selected:

Patching comfy attention to use SageAttention3 3.0.0.b1 (Blackwell FP4)

Compatibility

  • Requires sageattn3 package (v3.0.0.b1 or higher recommended)
  • Optimized for Blackwell architecture GPUs (RTX 5060 Ti 16GB and similar)
  • Automatic fallback ensures compatibility with all model configurations
  • Compatible with ComfyUI's attention function format via wrap_attn decorator

Notes

  • SA3 constraints (headdim < 256, no attention mask) are automatically handled with fallback to PyTorch SDPA
  • The node dynamically patches attention on each model execution and automatically cleans up afterward
  • Version information is logged on every generation for debugging and verification
Assets 2
Loading