flashdreams.accelerated API v0

[fangjunzhou-nv]

Add FlashDreams accelerated primitives and integrate them with Omnidreams

Summary

Introduce reusable quantization and optimized multi-head attention primitives under flashdreams.accelerated, integrate them with Omnidreams, and add reproducible performance benchmarks and plotting tools.

1. flashdreams.accelerated

1.1 Quantization toolkit

• Add tensor- and slice-granularity quantization for:
  • INT8
  • FP8 E4M3
  • FP8 E5M2
• Provide Torch reference and Triton CUDA implementations for quantization and dequantization.
• Add quantized non-persistent linear layers backed by CUDA integer/scaled GEMMs.
• Support dynamically quantized and prequantized activations.
• Keep derived quantized weights and scales out of state_dict, preserving source checkpoint compatibility.
• Cover scale validation, empty inputs, output dtypes, round trips, and quantized GEMMs with CPU and GPU tests.

1.2 Optimized MHA

• Add a shared multi-head attention interface for streaming self-attention and static cross-attention.
• Add a Torch reference implementation and an optimized CUDA implementation.
• Support:
  • cuDNN SDPA
  • Triton FlashAttention 2
  • optional TMA FlashAttention 2 kernels
  • full QKV fusion, fused KV, or unfused projections
  • optional INT8/FP8 projection quantization
  • optional FP8 SDPA and KV caches
  • head- or inner-scoped Q/K normalization
  • interleaved or split RoPE before or after KV-cache storage
• Reuse BlockKVCache and preserve caller-managed cache lifecycles.
• Preserve checkpoint-native projection names by keeping fused and quantized derived weights non-persistent.
• Require CUDA FP16/BF16 inputs and compute capability 9.0 or newer for the optimized path.
• Add numerical parity tests across attention types, SDPA backends, fusion policies, RoPE policies, quantization modes, and TMA/non-TMA kernels.

2. Integration: Omnidreams

• Add independently configurable self- and cross-attention backends while retaining the existing Omnidreams implementation as the default.
• Adapt Omnidreams attention modules to the shared optimized MHA interface without changing checkpoint keys or cache behavior.
• Thread optimized attention policies through the block and network configurations.
• Register new runner presets:
  • omnidreams-triton-fa2
  • omnidreams-cuda-cudnn
  • omnidreams-cuda-sparge
  • omnidreams-cuda-sage3fp8
• Skip final KV-cache advancement in steady-state performance presets.
• Fix native CUDA extension builds on GB300:
  • stop forcing all devices to compile for 12.0a
  • enable SageAttention 3 sources only for validated SM120a devices
  • defer other architecture selection to PyTorch or explicit environment overrides
  • isolate extension names and caches by CUDA architecture
• Add tests for backend selection, optimized-policy propagation, cache lifecycles, checkpoint compatibility, runner registration, and native architecture detection.

3. Benchmark and plot scripts

• Add pytest-benchmark suites for:
  • quantization and dequantization
  • quantized linear layers and GEMMs
  • self- and cross-attention
  • Omnidreams attention modules and DiT blocks
  • complete Omnidreams network and pipeline execution
• Compare Torch, optimized cuDNN, optimized FA2, and native CUDA configurations using matched production-shaped workloads.
• Include warmup rounds to exclude compilation and autotuning overhead.
• Record implementation, device, tensor geometry, precision, and backend metadata in benchmark results.
• Add scripts that run benchmark suites, export JSON results, and generate comparison plots.
• Add aggregate scripts for running or plotting all benchmark groups.
• Document the Omnidreams test and benchmark workflows.
• Add pytest-benchmark configuration and dependencies.

Testing

The change includes CPU, GPU, and manual benchmark coverage. GPU validation requires a supported NVIDIA GPU; the optimized MHA path requires compute capability 9.0 or newer.

Generated benchmark JSON and figures are intentionally not committed.

[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 reusable accelerated quantization and multi-head-attention primitives, integrates selectable optimized attention backends into OmniDreams, and adds benchmark tooling.

• Adds Torch, Triton, cuDNN, FP8, and quantized projection paths under flashdreams.accelerated.
• Adds configurable optimized self- and cross-attention implementations while preserving checkpoint-native projection parameters.
• Refreshes non-persistent fused and quantized projection weights after checkpoint loading and device or dtype transformations.
• Adds runner presets, architecture-aware native extension configuration, tests, benchmarks, and plotting scripts.

Confidence Score: 5/5

The PR appears safe to merge.

The previously reported stale derived-weight issue is resolved because optimized attention now registers a post-load hook that rebuilds fused and quantized projections from the loaded canonical parameters, and no blocking failure remains.

Important Files Changed

Filename	Overview
flashdreams/flashdreams/accelerated/multi_head_attention/optimized.py	Adds optimized attention policy, projection fusion, quantization, cache validation, and post-load/device-transform refresh of derived execution weights.
integrations/omnidreams/omnidreams/transformer/impl/modules.py	Adds selectable optimized OmniDreams attention adapters while retaining canonical checkpoint projection names and cache contracts.
flashdreams/flashdreams/accelerated/quantization/quantizer.py	Implements reference and accelerated tensor- and slice-granularity quantization interfaces.
flashdreams/flashdreams/accelerated/quantization/linear.py	Adds non-persistent quantized linear execution for dynamic and prequantized activations.
integrations/omnidreams/omnidreams/transformer/impl/network.py	Propagates independently configurable self- and cross-attention policies through the OmniDreams network.
integrations/omnidreams/omnidreams/config.py	Registers optimized attention policies in new OmniDreams runner presets.
flashdreams/flashdreams/accelerated/multi_head_attention/cudnn/native_fp8.py	Adds cached cuDNN Frontend graphs for FP8 scaled-dot-product attention.
uv.lock	Records the new benchmark, development, and lint dependencies without changing the pre-existing flagged Torch or Pillow resolutions.

Sequence Diagram

sequenceDiagram
  participant Network as OmniDreams Network
  participant Attention as Optimized Attention
  participant Loader as load_state_dict
  participant Derived as Derived Projections
  participant Forward as Inference
  Network->>Attention: Construct canonical Q/K/V projections
  Attention->>Derived: Build initial non-persistent projections
  Loader->>Attention: Load checkpoint parameters
  Attention->>Derived: Post-load hook refreshes fused/quantized weights
  Network->>Attention: Move or cast module
  Attention->>Derived: _apply refreshes weights on final device/dtype
  Forward->>Derived: Execute accelerated projections

Loading

_{Reviews (7): Last reviewed commit: "Add benchmark and plot scripts for flash..." | Re-trigger Greptile}

[ArielG-NV]

Question on inclusion of file