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, and CUDA-oriented quantization and attention implementations.
• Preserves canonical checkpoint parameters while rebuilding non-persistent fused and quantized execution weights after checkpoint loading.
• Adds Omnidreams backend policies, runner presets, compatibility tests, and performance benchmarks.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains.

The previously reported stale-derived-weight defect is addressed by rebuilding fused and quantized projections in a registered post-load hook after canonical checkpoint tensors are loaded.

Important Files Changed

Filename	Overview
flashdreams/flashdreams/accelerated/multi_head_attention/optimized.py	Defines optimized attention and now rebuilds non-persistent derived projections through a module post-load hook and after device or dtype transformations.
integrations/omnidreams/omnidreams/transformer/impl/modules.py	Adds optimized self- and cross-attention adapters that initialize checkpoint-compatible canonical projections before registering derived-weight refresh behavior.
integrations/omnidreams/omnidreams/transformer/init.py	Loads checkpoints through the parent network, allowing child post-load hooks to refresh optimized execution weights before inference.
integrations/omnidreams/tests/test_transformer_attention_backend.py	Covers strict checkpoint-key compatibility and optimized backend dispatch, with no remaining blocking issue identified.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Omnidreams checkpoint] --> B[Canonical Q/K/V parameters]
  B --> C[Post-load refresh hook]
  C --> D[Non-persistent fused or quantized weights]
  D --> E[Optimized attention backend]
  E --> F[Autoregressive rollout]

Loading

_{Reviews (25): Last reviewed commit: "Fixed rope latex format" | Re-trigger Greptile}

[ArielG-NV]

Question on inclusion of file

[jarcherNV]

I guess we should restrict the license header stuff to just the first 2 lines, no need for the rest.

[fangjunzhou-nv]

As all the source code in this project are putting full license header. You may want to submit a new issue to change the AGENT.md for this over the project. This is outside the scope of this PR.

[jarcherNV]

I am not sure this is true. We originally used the longer style although switched to the shorter style based on some internal guidance and that is mainly the style that newer changes have used. The split is probably around 50/50 overall for flashdreams at this point. I guess I will do a general license cleanup later to capture stuff like this. I think my preference would be if newer files did not use the longer style if we know we will definitely change it. The skill file probably needs to be updated.

[fangjunzhou-nv]

Could we add a document summarizing the performance findings from this work so far?

I'll compose a formal documentation on flashdreams.accelerated and put final benchmark numbers on RTX PRO 6000 and GB 300 in the doc.

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.

[fangjunzhou-nv]

/ok to test b032f75

[fangjunzhou-nv]

I've just realized CUDA graph really don't like the fp8 cudnn in flashdreams/flashdreams/accelerated/multi_head_attention/cudnn/native_fp8.py I need to fix it...

[fangjunzhou-nv]

/ok to test 9e809fb

[fangjunzhou-nv]

/ok to test f88c93c

[fangjunzhou-nv]

/ok to test 9e44fb0

[jarcherNV]

Would it make sense to add a mermaid diagram or two to this file? A single left-to-right figure from projection through Q/K norm, RoPE, cache write, and SDPA, with each OptimizedImplConfig knob annotated at the stage it controls, would show how these stages relate in a way the per-knob bullets can't, and would make BEFORE_KV_CACHE vs AFTER_KV_CACHE obvious at a glance. The derived-weight lifecycle in the integration section might be worth one too.