[visual_gen] Optimize Wan2.2 Dataflow (Zhen Xie from VibeHPC)

[zhen-xie]

[visual_gen] Optimize Wan2.2 Dataflow (Zhen Xie from VibeHPC)

Description

This PR adds opt-in dataflow optimizations for Wan 2.2 inference in visual_gen, targeting single-GPU execution on NVIDIA B300.

The changes reduce QKV packing overhead, eliminate redundant FP8 activation quantization around the FFN, cache immutable FP8 weight metadata, and tune the fused Q/K normalization and RoPE kernel launch configuration.

All optimizations are disabled by default. This PR does not modify the model architecture, sampling procedure, MLPerf harness, dataset, request scheduler, or output format.

Changes

1. Preserve Packed QKV Storage Through Q/K Normalization and RoPE

The fused QKV projection already produces Q, K, and V in one shared allocation.

The original Q/K normalization and RoPE path creates separate output tensors. As a result, Transformer Engine repacks Q, K, and V before FP8 attention by launching a CatArrayBatchedCopy kernel.

This PR adds a Wan-specific CUDA kernel that:

• reads strided Q and K views directly from the fused QKV allocation;
• applies RMSNorm and RoPE in one kernel;
• optionally updates Q and K in place;
• preserves the shared [Q | K | V] storage layout.

Transformer Engine can then recognize the inputs as packed bsh3d storage and skip its per-attention CatArrayBatchedCopy.

The fused kernel preserves the BF16 rounding point between RMSNorm and RoPE to match the original Wan computation order.

Relevant files:

• visual_gen/csrc/DiTRMSNormRope/fused_qk_norm_rope_kernel.cu
• visual_gen/models/transformers/wan_transformer.py

2. Fuse Wan ApproximateGELU With FP8 Quantization

Wan uses the following FFN activation:

x * sigmoid(1.702 * x)

The original execution path is:

FP8 FFN-up GEMM
    ↓
ApproximateGELU
    ↓
Write BF16 activation
    ↓
Transformer Engine amax reduction
    ↓
Transformer Engine FP8 quantization
    ↓
FP8 FFN-down GEMM

This PR adds a Triton implementation that:

1. computes the exact Wan ApproximateGELU formula;
2. collects the activation amax;
3. calculates the E4M3 inverse scale;
4. quantizes the activation to FP8;
5. creates a Transformer Engine compatible Float8Tensor.

The optimized execution path is:

FP8 FFN-up GEMM
    ↓
Fused ApproximateGELU and amax collection
    ↓
FP8 quantization
    ↓
Transformer Engine compatible Float8Tensor
    ↓
FP8 FFN-down GEMM

The FFN-down projection consumes the generated Float8Tensor directly, avoiding a second activation amax and FP8 cast.

Unsupported tensor layouts and non-inference execution continue to use the original FeedForward path.

Relevant files:

• visual_gen/ops/wan_fused.py
• visual_gen/models/transformers/wan_transformer.py
• visual_gen/ops/linear.py

3. Cache Immutable FP8 Weight Wrappers

The FP8 linear path previously reconstructed the following objects during repeated inference calls:

• transposed FP8 weight representation;
• FP8 scale tensor;
• Transformer Engine Float8Tensor wrapper;
• related weight metadata.

Model weights remain unchanged during inference, so recreating these objects introduces unnecessary overhead.

This PR caches the prepared FP8 weight representation using:

weight.data_ptr()
weight._version
weight.shape
weight_scale.data_ptr()
weight_scale._version

The cache automatically refreshes if either the weight or its scale changes.

Pointer and version checks execute outside Torch Dynamo tracing to avoid graph breaks and DataPtrVariable errors.

Relevant file:

• visual_gen/ops/linear.py

4. Tune the Wan Q/K Norm and RoPE Launch Configuration

The fused Wan Q/K Norm and RoPE kernel supports the following thread-block sizes:

128 threads
256 threads
512 threads

This PR changes the default to 512 threads for the validated Wan 2.2 configuration:

Sequence length: 75,600
Attention heads: 40
Head dimension: 128
Hidden dimension: 5,120

Each warp continues to process one Q or K head row. The optimization only changes block scheduling and does not change the arithmetic performed within a warp.

The launch configuration can be overridden with:

export VISUAL_GEN_WAN_QK_ROPE_THREADS=128
export VISUAL_GEN_WAN_QK_ROPE_THREADS=256
export VISUAL_GEN_WAN_QK_ROPE_THREADS=512

Relevant file:

• visual_gen/csrc/DiTRMSNormRope/fused_qk_norm_rope_kernel.cu

Performance

Measurements were collected on a single NVIDIA B300 GPU using the Wan 2.2 A14B MLPerf Offline workload.

Test Configuration

GPU: NVIDIA B300 SXM6
Execution: Single GPU
Scenario: MLPerf Offline
Attention: Transformer Engine FP8
Linear layers: Transformer Engine FP8
Video token sequence: 75,600
Attention heads: 40
Head dimension: 128
Hidden dimension: 5,120
FFN dimension: 13,824

Isolated Operator Results

Optimization	Baseline	Optimized	Speedup
Packed QKV TE FP8 attention path	47.969 ms	46.528 ms	1.0310x
ApproximateGELU + FP8 quantization	1.356 ms	1.156 ms	1.1735x

Packed QKV Profiling Details

The real attention shape is:

Batch: 1
Sequence length: 75,600
Attention heads: 40
Head dimension: 128

The original path contains one QKV packing kernel per FP8 attention call:

CatArrayBatchedCopy average latency: 1.7514 ms

Measured results:

Separate QKV attention path: 47.969409 ms
Packed QKV attention path:   46.528168 ms
Saving per call:             1.441241 ms

Multiple Nsight Systems measurements indicated a local saving of 1.4 ms per attention invocation, depending on the run.

ApproximateGELU and FP8 Quantization Profiling Details

The real FFN activation shape is:

Rows: 75,600
Columns: 13,824
Data type: BF16
FP8 format: E4M3

Three independent long-run measurements produced:

Baseline mean:  1.356278 ms
Optimized mean: 1.155785 ms
Saving:         0.200493 ms per call
Latency change: -14.78%
Speedup:        1.1735x

The selected Triton quantization block size is:

4096 elements

End-to-End Performance Comparison

Configuration	Offline QPS	Change
Baseline	0.00814648	Reference
Packed QKV	0.00851063	+4.47%
Packed QKV + fused GELU/quantization	0.00864342	+6.10% vs. baseline
Packed QKV + weight cache	0.00852144	+0.127% vs. packed QKV

The weight-cache result is close to normal run-to-run variation. The optimization remains included because it removes redundant preparation work and did not introduce an observed memory or correctness issue.

Q/K Norm and RoPE Thread Sweep

Threads per block	Offline QPS	Single-sample latency
128	0.00815276	122.658 s
256	0.00814648	122.752 s
512	0.00816359	122.495 s

The 512-thread configuration measured 0.21% higher QPS than the 256-thread configuration.

Performance Result Scope

These measurements are development A/B results.

They are not MLPerf submission-valid performance results because the short exploratory runs do not satisfy the full MLPerf duration and query-count requirements.

Numerical Behavior

Packed QKV

Packed QKV and separate QKV produced identical outputs for the isolated attention comparison:

Cosine similarity:      1.0
Mean squared error:     0
Maximum absolute error: 0

The optimized path retains the same cuDNN FP8 SDPA kernel:

cudnn_generated_fort_native_sdpa_sm100_flash_fprop_f8_knob_7

The optimization changes the QKV storage layout presented to Transformer Engine but does not replace the attention computation.

Fused ApproximateGELU and FP8 Quantization

The fused kernel preserves Wan's activation formula:

x * sigmoid(1.702 * x)

Isolated numerical checks produced:

BF16 activation cosine similarity: 1.0
BF16 sampled MSE:                   0
Dequantized FP8 cosine similarity: 1.0
FP8 scale difference:              identical one FP32 ULP

Fused Q/K Norm and RoPE

The original path rounds the RMSNorm output to BF16 before applying the FP32 RoPE multiplication.

The fused CUDA kernel explicitly preserves this rounding boundary:

FP32 RMSNorm accumulation
    ↓
BF16 rounding
    ↓
FP32 RoPE multiplication
    ↓
BF16 output

This prevents the fusion from unintentionally carrying additional FP32 precision across the original BF16 boundary.

Accuracy Validation

The isolated numerical tests passed for the optimized operator paths.

Full MLPerf AccuracyOnly and VBench validation should still be completed before enabling these optimizations in a submission configuration.

Enabling the Optimizations

All behavior-changing paths are disabled by default.

Baseline

unset VISUAL_GEN_WAN_OPTIMIZATIONS

Enable the Optimized Wan Path

export VISUAL_GEN_WAN_OPTIMIZATIONS=1
export VISUAL_GEN_WAN_QK_ROPE_THREADS=512

When the master switch is enabled, the following optimizations are enabled by default:

Packed QKV
Fused ApproximateGELU and FP8 quantization
FP8 weight cache
512-thread Q/K Norm and RoPE scheduling

Disable Individual Components

Disable Packed QKV:

export VISUAL_GEN_WAN_PACKED_QKV=0

Disable fused ApproximateGELU and FP8 quantization:

export VISUAL_GEN_WAN_FUSED_GELU_QUANT=0

Disable the FP8 weight cache:

export VISUAL_GEN_WAN_FP8_WEIGHT_CACHE=0

Use the original 256-thread Q/K Norm and RoPE configuration:

export VISUAL_GEN_WAN_QK_ROPE_THREADS=256

The worker process must be restarted after changing these variables because the configuration is read when the Python modules are imported.

Fallback Behavior

The optimized FFN path checks the following conditions before using the fused ApproximateGELU and FP8 quantization implementation:

Inference mode
CUDA input
BF16 input
Contiguous input
Expected FeedForward module structure
Transformer Engine FP8 down projection
No active dropout

If any condition is not satisfied, execution falls back to the original FeedForward implementation.

The packed QKV path checks:

Self-attention
Fused QKV projection
BF16 Q and K
Matching Q, K, and V head dimensions
Available rotary embeddings
Compatible Q/K normalization weights

Unsupported cases use the original Q/K normalization and RoPE path.

Scope

This PR only changes the Wan 2.2 inference implementation in visual_gen.

It does not change:

• model weights;
• model architecture;
• denoising step count;
• scheduler behavior;
• prompt handling;
• negative-prompt handling;
• latent inputs;
• classifier-free guidance behavior;
• MLPerf LoadGen;
• MLPerf harness behavior;
• request scheduling;
• dataset handling;
• generated video dimensions;
• generated video frame count;
• output format.

Files Changed

visual_gen/csrc/DiTRMSNormRope/fused_qk_norm_rope_kernel.cu

Adds:

• Wan-specific fused Q/K RMSNorm and RoPE kernel;
• strided packed-QKV input support;
• in-place Q/K output support;
• configurable 128, 256, or 512-thread launch scheduling;
• PyTorch operator registration.

visual_gen/models/transformers/wan_transformer.py

Adds:

• master optimization switch;
• Packed QKV selection;
• fused Q/K Norm and RoPE dispatch;
• optimized FFN inference path;
• safe fallback to the original implementation.

visual_gen/ops/wan_fused.py

Adds:

• exact Wan ApproximateGELU kernel;
• fused amax collection;
• E4M3 scale calculation;
• FP8 quantization;
• Transformer Engine compatible Float8Tensor construction.

visual_gen/ops/linear.py

Adds:

• direct consumption of an existing Float8Tensor;
• cached FP8 weight wrapper and transposed storage;
• cache invalidation based on tensor pointer and version;
• Torch Dynamo safe cache lookup.

visual_gen/README-for-dataflow-optimizations.md

Documents:

• the master optimization switch;
• the Q/K Norm and RoPE thread configuration;
• the worker restart requirement.

Summary

This PR improves Wan 2.2 inference dataflow without changing the model architecture or benchmark framework.

The main improvements are:

• preserving packed QKV storage through Q/K normalization and RoPE;
• removing Transformer Engine QKV packing copies;
• fusing the exact Wan ApproximateGELU with FP8 quantization;
• allowing the FFN-down GEMM to consume an existing FP8 activation;
• caching immutable FP8 weight representations;
• tuning Q/K Norm and RoPE scheduling for the B300 workload.

[zhenhuaw-me]

Thank you Zhen Xie for the PR! I am very happy to see community contribution. Could you please follow https://github.com/NVIDIA/TensorRT-LLM/blob/main/CONTRIBUTING.md and bring this PR to product support shape?

[zhenhuaw-me]

Since the kernel is implemented for Wan, does it make sense to drop the env?

[luyiyun1021]

Hi Zhen, seems that you've contributed to the wrong branch feat/visual_gen, which is a deprecated branch and the newest commit stays at 2026-03-18. We are working on the main branch currently. Could you please rebase on main and see if this optimization still works?

[luyiyun1021]

Typically, we would both need kernel unittests and e2e lpips comparison to ensure there is no accuracy regression. check scripts/visualgen_eval/visual_gen_lpips_score_eval.py

[luyiyun1021]

Wan's norm_q/norm_k normalize over the full hidden dimension (num_heads * head_dim, i.e. 5120 for A14B), before splitting into heads. Here each warp only accumulates one head's 128 elements, which changes the operation to per-head RMSNorm. Could we reduce across all heads of each token and add a parity test with different input magnitudes across heads?

[luyiyun1021]

dim is local to the current head, but Wan's RMSNorm weights span all num_heads * head_dim dimensions. This currently reuses the first head's weights for every head. Should this index be weight[headIdx * head_dim + dim]? A parity test with nonuniform norm weights across heads would catch this independently of the reduction-domain issue.

[luyiyun1021]

Wan's FFN uses activation_fn="gelu-approximate", which resolves to GELU(approximate="tanh") in the pinned diffusers 0.36.0. This helper computes x * sigmoid(1.702 * x), a different activation. The .proj guard also accepts the original GELU module, so the optimized path replaces Wan's activation. Could we fuse the tanh GELU formula and validate parity against the actual Wan FFN? The dispatch should also check the activation type and approximation setting.

[o-stoner]

Thanks for the writeup! For change 1, this looks like it may already be implemented on main; see Attention.forward() / apply_packed_qk_norm_rope() in tensorrt_llm/_torch/visual_gen/modules/attention.py (wired into Wan via fuse_qk_norm_rope in transformer_wan.py). Could you check whether your changes should be perhaps used to extend the existing kernel instead?