[Lunemere]

Summary

This PR significantly accelerates DreamZero inference on NVIDIA Thor (sm_110a) by introducing:

1. FlashAttention 4 (FA4) integration for self-attention.
2. A Triton-fused RMSNorm kernel matching PyTorch 2.10 semantics.
3. High-performance native BF16 FFN and QKV projection kernels (CUTLASS + cuBLASLt).

In end-to-end benchmarks on DreamZero-Droid, these changes achieve:

• -24.11% full-chunk latency (17.53s → 13.30s)
• -28.62% denoising-loop latency (14.23s → 10.15s)
• 13%–34% speedup across targeted GEMM operators

---

Key Changes

1. Attention & Normalization

• FlashAttention 4 Integration: Added DreamZero-Droid configuration support leveraging FA4 for self-attention.
• Fused RMSNorm (Triton): Fused self-attention (Q/K) and cross-attention (Q) RMSNorm using a custom Triton kernel. Strictly matches PyTorch 2.10 reduction precision and BF16 rounding behavior.

2. Native GEMM Optimizations

• FFN Up-Projection (CUTLASS): Fused GEMM + bias + tanh GELU. Implemented thread-local execution plan caching to eliminate plan recreation overhead without pinning model weights or intermediate activation buffers.
• FFN Down-Projection (cuBLASLt): Fused GEMM + bias with a pre-pinned heuristic algorithm, eliminating runtime algorithm search and workspace allocation jitter.
• Self-Attention QKV Projections (CUTLASS): Optimized all three Q/K/V projections with fused bias. Specialized for contiguous BF16 inputs of shape [1785, 5120].

3. Architecture Dispatch & Safety Fallbacks

• Targeted Hardware Dispatch: Automatically routes to native FFN and QKV kernels upon detecting NVIDIA Thor devices.
• Graceful Fallback: Automatically falls back to standard PyTorch projections for non-contiguous inputs, unsupported tensor shapes, or when LoRA adapters are active.

---

Build & Environment Notes

• Target Architecture: Compiles specifically for sm_110a (NVIDIA Thor).
• Prerequisites: CUTLASS extensions resolve headers via the CUTLASS_PATH environment variable (e.g., export CUTLASS_PATH=/path/to/cutlass before compilation).

---

Performance Evaluation

Benchmark Configuration

• Model: DreamZero-Droid (Image-to-Video & Action)
• Inputs & Output: 3 camera views, 352 × 640 resolution @ 5 FPS
• Denoising Setup: 15 chunks, 16 video/action denoising steps
• DiT Scheduling: Executed at steps [1, 2, 3, 7, 11, 14, 15, 16]
• Generation Settings: CFG enabled (guidance scale: 5.0), DiT quantization: OFF, torch.compile: OFF
• Memory/Offloading: T5 CPU offload ON; Model & VAE CPU offload OFF

Cumulative End-to-End Latency

Each row includes all optimizations from previous stages.

Cumulative Stage	Full Chunk (s)	Denoising Loop (s)	Per-Step Latency (ms)	Chunk Speedup
Baseline (Original)	17.529	14.226	1774.85	—
+ FA4	16.837	13.541	1688.32	+3.95%
+ Fused RMSNorm	14.490	11.288	1407.56	+17.34%
+ CUTLASS FFN Up	14.065	10.880	1357.76	+19.76%
+ cuBLASLt FFN Down	13.360	10.313	1285.33	+23.78%
+ Native QKV Projections	13.302	10.154	1265.27	+24.11%

Kernel Microbenchmarks

Operator	Baseline (ms)	Optimized (ms)	Latency Reduction
FFN Up (+ Bias + GELU)	224.75	148.97	-33.72%
FFN Down (+ Bias)	203.74	177.31	-12.97%
QKV Projection (+ Bias)	226.88	157.59	-30.54%