Use flashdreams.accelerated in Wan family models

[fangjunzhou-nv]

Second part for #486

[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]

Greptile Summary

The PR introduces accelerated attention, FP8 quantization, Triton kernels, and supporting benchmarks, then makes those implementations selectable across Wan-family models and integrations.

• Adds shared accelerated multi-head-attention, FlashAttention, FP8, RoPE, and KV-cache implementations.
• Wires accelerated attention into Wan, Lingbot, and OmniDreams transformer configurations.
• Adds correctness tests, GPU benchmarks, benchmark scripts, and dependency metadata.
• Adds NVTX instrumentation around pipeline and diffusion stages.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains within the eligible follow-up-review scope.

No blocking failure remains.

Important Files Changed

Filename	Overview
flashdreams/flashdreams/accelerated/multi_head_attention_triton.py	Adds the primary accelerated attention implementation, including fused projection options, backend selection, cache management, and FP8 support.
flashdreams/flashdreams/accelerated/triton/rms_rope_kv_cache.py	Adds fused Triton kernels for normalization, rotary embeddings, and KV-cache writes.
flashdreams/flashdreams/accelerated/triton/flash_attention.py	Adds the Triton FlashAttention implementation used by the accelerated backend.
flashdreams/flashdreams/recipes/wan/transformer/impl/modules.py	Makes Wan transformer blocks select between the existing and accelerated attention implementations.
integrations/omnidreams/omnidreams/transformer/impl/modules.py	Integrates selectable accelerated attention into OmniDreams transformer modules.
flashdreams/flashdreams/infra/pipeline/base.py	Adds NVTX annotations around pipeline lifecycle stages without changing their state-transition semantics.
uv.lock	Updates the lockfile for benchmark, instrumentation, environment, and lint dependencies; the flagged aiohttp version is unchanged from base.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  Config["Wan-family model config"] --> Backend{"Attention backend"}
  Backend --> Native["Existing Wan attention"]
  Backend --> Accelerated["flashdreams.accelerated"]
  Accelerated --> Projection["Fused Q/K/V projections"]
  Projection --> RoPE["RoPE + KV-cache update"]
  RoPE --> SDPA{"SDPA backend"}
  SDPA --> Triton["Triton FlashAttention"]
  SDPA --> CuDNN["cuDNN SDPA"]
  Triton --> Output["Output projection"]
  CuDNN --> Output

Loading

_{Reviews (3): Last reviewed commit: "Removed Torch MHA and simplified Triton ..." | Re-trigger Greptile}