[RFC]: Support Model FLOPs Utilization (MFU) Metrics for Diffusion Models (DiTs)

1. Background

vLLM currently exposes MFU and memory throughput metrics (e.g., vllm:estimated_flops_per_gpu_total, vllm:estimated_read_bytes_per_gpu_total) when started with the --enable-mfu-metrics flag. These metrics are critical for identifying compute bottlenecks and evaluating hardware efficiency.

However, vllm-omni's diffusion models (such as Wan2.1, HunyuanVideo, Flux, etc.) use Diffusion Transformers (DiTs). DiTs possess vastly different FLOP profiles compared to standard auto-regressive LLMs, mainly due to spatial-temporal attention, cross-attention with context latents, and highly variable sequence lengths (e.g., $100K+$ tokens for long videos). Currently, vllm-omni does not accurately estimate or report the MFU for these diffusion workloads.

2. Motivation

For DiT architectures in vllm-omni, inference—especially for long video generation and large batch image generation—is heavily compute-bound (bottlenecked by Tensor Cores) rather than memory-bound.

Adding accurate MFU tracking for diffusion models will allow us to:

1. Evaluate Parallelism Efficiency: Track MFU when scaling sequence length across multiple GPUs (SP/USP, Ring Attention). A drop in MFU will clearly expose communication overhead (e.g., All-Gather/Reduce-Scatter) masking compute.
2. Benchmark Attention Backends: With sequence lengths easily exceeding 100K, $O(N^2)$ self-attention dominates FLOPs. MFU is needed to directly compare FlashAttention-3, SageAttention, and custom kernels.
3. Assess Kernel & Launch Overheads: Help identify un-fused operations or CPU overhead when scaling to large batches.

3. Proposed Implementation

We propose extending the existing vLLM metrics pipeline to correctly calculate inference FLOPs for DiT models.

Taking inspiration from FastVideo's MFU calculation, the theoretical inference FLOPs (forward pass only) for a single DiT layer can be defined as:

# Variables from batch and model config:
# seq_len, context_len, hidden_dim, ffn_dim

# 1. QKV + Out Projections
qkv_out_flops = 8 * hidden_dim * hidden_dim * seq_len

# 2. Cross-Attention Projections
cross_attn_proj_flops = (4 * hidden_dim * hidden_dim * seq_len) + \
                        (4 * hidden_dim * hidden_dim * context_len)

# 3. MLP
mlp_flops = 4 * hidden_dim * ffn_dim * seq_len

# 4. Self-Attention Matmuls
self_attn_flops = 4 * seq_len * seq_len * hidden_dim

# 5. Cross-Attention Matmuls
cross_attn_flops = 4 * seq_len * context_len * hidden_dim

# Total inference FLOPs per layer (Forward only)
flops_per_layer = qkv_out_flops + cross_attn_proj_flops + mlp_flops + self_attn_flops + cross_attn_flops

# Total FLOPs per GPU
total_flops = batch_size * flops_per_layer * num_layers

Integration Points:

1. Model Runner / Profiler: Implement a compute_estimated_flops method in the DiT model runner (e.g., OmniNPUModelRunner, DiffusionWorker, etc.) that factors in the precise seq_len and context_len of the current batch.
2. Metrics Exporter: Ensure vllm:estimated_flops_per_gpu_total consumes this calculated value properly when --enable-mfu-metrics is active.

4. Open Questions & Edge Cases

1. The "Fake Work" (Padding) Problem: In large batch inference with mixed resolutions/durations, padding tokens are introduced. Including padding FLOPs artificially inflates MFU while "Goodput" (useful work) remains low. Question: Should we calculate FLOPs strictly based on unpadded seq_len per request to reflect true MFU?
2. Varied Architectures: Different models (Wan2.1 vs. HunyuanVideo) may have slightly different MLPs or attention mechanics (e.g., grouped query attention, varying 3D RoPE calculations). Should we maintain a model-specific FLOP calculation interface, or is a generic DiT estimation sufficient?

5. Next Steps

• Gather feedback from core contributors on the FLOP formulas for our currently supported diffusion models.
• Implement a draft PR introducing the DiT FLOPs calculator.
• Validate MFU outputs on H100s for a Wan2.2 video generation workload.