[Shn0930]

Purpose

Add SageAttention 2 as an opt-in backend to FlashDreams' shared attention
abstractions so generation models can use it without implementing a
model-specific SageAttention wrapper. The existing cuDNN path remains the
default and is unchanged.

Scope

• Add backend="sage2" support to NativeAttention and
  ContextParallelAttention, including optional dependency loading, input/output
  validation, LSE handling, and ring/Ulysses dispatch. This is framework-level
  support available to models built on these shared attention classes.
• Reuse the existing HND attention path: BSHD inputs are transposed before the
  backend call and transposed back afterward. This avoids layout-specific CP
  branches; direct NHD improved a complete steady WAN AR block by only 0.155%.
• Expose an opt-in self-attention backend in the WAN, Cosmos, Template, and
  OmniDreams framework/Python DiT configurations. Cross-attention remains on
  cuDNN because the measured short cross-attention shape did not benefit from
  Sage2.
• Reject sage2 with use_cuda_graph=True. CUDA-graph replay produced
  incorrect video after capture; Sage2 remains compatible with
  compile_network=True when explicit CUDA graphs are disabled.
• Keep OmniDreams native optimized DiT separate: Sage2 selects the validated
  framework/Python DiT path and rejects native acceleration instead of silently
  ignoring the requested backend.
• Add CPU contract and model-configuration tests.
• Do not add SageAttention as a required project dependency and do not change
  default behavior. Model-specific native attention stacks remain independent
  from the shared framework adapter.

Example opt-in configuration:

--pipeline.diffusion-model.transformer.network.self-attention-backend sage2 \
--pipeline.diffusion-model.transformer.use-cuda-graph False

Cosmos and Template expose the same nested network.self_attention_backend
setting. OmniDreams additionally requires native_dit_acceleration="disabled"
to select its framework/Python DiT path.

Test Plan

1. Run the Sage2 attention and model-configuration CPU tests:

   pytest -q \
     flashdreams/tests/test_sage2_attention.py \
     flashdreams/tests/test_wan_context_parallel.py \
     integrations_v2/omnidreams/tests/test_transformer_attention_backend.py

2. Run formatting, lint, and type checks on the changed attention files and
   tests:

   pre-commit run --files \
     flashdreams/flashdreams/core/attention/native.py \
     flashdreams/flashdreams/core/attention/cp.py \
     flashdreams/tests/test_sage2_attention.py

3. Compare cuDNN and Sage2 in fresh deterministic GPU processes using identical
   checkpoints, prompts, seeds, and generation settings:

   • Bidirectional WAN 2.1 T2V 1.3B: 480x832, 81 frames, 50 steps.
   • Causal Self-Forcing WAN 2.1 T2V 1.3B: 480x832, 12 AR blocks,
     141 frames, four denoise steps per block.
   • Two-rank H20 ring and Ulysses smoke tests.
   • CUDA-graph replay correctness probe.

4. Compare HND and direct-NHD layouts using two crossed H20 runs, and verify
   generated-video hashes.

5. As supplemental shared-backend quality evidence, compare cuDNN and Sage2 on
   1,048 paired OmniDreams clips. OmniDreams integration itself is not changed
   by this PR.

Test Result

• Sage2/model routing tests: 38 passed.
• Template and OmniDreams preset regression tests: 17 passed, 8 GPU/manual
  tests deselected.
• Full CPU suite: 1,989 passed, 2 skipped, 1 manual, 1,586 deselected.
• Ruff formatting/lint and type checks: passed.
• Bidirectional WAN hot-cache end-to-end time: 223.072 s -> 151.412 s
  (1.473x).
• Causal Self-Forcing steady full AR block: 2177.93 ms -> 1749.41 ms
  (1.245x); steady diffusion stage: 1253.44 ms -> 911.29 ms
  (1.375x).
• Two-rank H20 ring and Ulysses smoke tests: passed.
• Direct NHD versus HND improved the steady full AR block by only 0.155%;
  all four generated MP4 hashes matched. The PR therefore keeps the simpler
  shared HND path.
• Sage2 CUDA-graph replay produced incorrect video from the capture point. The
  unsupported configuration is now rejected before model construction.
• Supplemental 1,048-clip OmniDreams A/B results:
  • Full-pipeline median: 20.306 s -> 16.496 s (1.231x).
  • Paired frame CLIP cosine: 0.990432.
  • Clean-FID, cuDNN/Sage2: 7.1127 / 7.1509.
  • FVD, cuDNN/Sage2: 40.2894 / 40.8775.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; the previous Windows test failure is addressed.

Summary

Adds optional SageAttention 2 support to shared attention and model settings, including an experimental native OmniDreams path.

• Shared attention can run SageAttention 2 for model self-attention.
• OmniDreams can route SageAttention 2 through its framework or eligible native path.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Attention setting] --> B{Sage2 requested?}
  B -->|No| C[Existing backend]
  B -->|Yes| D{CUDA graphs enabled?}
  D -->|Yes| E[Reject or use allowed framework fallback]
  D -->|No| F{OmniDreams native requested?}
  F -->|No| G[Shared Sage2 adapter]
  F -->|Yes| H[Check model target device]
  H --> I[Native Sage2 or configured framework fallback]

Loading

_{Reviews (7) · Last reviewed commit: "test(omnidreams): isolate Sage2 platform..." · Reviewed by Greptile}

[jmccaffrey-nv]

Thanks @Shn0930 ! we'll review

[jmccaffrey-nv]

/ok to test c9c9fca

[jarcherNV]

Looks good to me overall:

• SageAttention 2 remains opt-in, with cuDNN unchanged as the default.
• The fallback, CUDA graph, and target-device handling look sound.
• Coverage and performance/quality evidence are thorough.

Approved from my side. Please resolve the current merge conflicts and then we can merge this.