[diffusion] kernel fusion: ScaleResidualNormScaleShift

[yingluosanqian]

Related to #12799 (Roadmap)
Related to #14437 (apply fused kernel to diffustion model, this pr only impl kernel for sgl-kernel repo)

Motivation

This PR fuses the residual-add, normalization, and scale-shift steps to eliminate redundant memory accesses and reduce GPU communication overhead in diffusion models. i.e.

1. residual_out = residual + gate * x
2. normalized = norm(residual_out)
3. modulated = (1 + scale) * normalized + shift

Modifications

1. Add a new fused kernel ScaleResidualNormScaleShift to sgl-kernel
2. Integrated the kernel into models such as Hunyuan and Wan

Accuracy Tests

Use the native python implementation as the baseline:

• For FP32, we require 1e-5 absolute and 1e-5 relative error.
• For FP16, we use 1e-2 absolute and 1e-2 relative error.

(These tolerances may be further refined based on broader model-level evaluation.)

Benchmarking and Profiling

Single Kernel Test: 4.99x speedup

Benchmarking on H200 across multiple (seq_len, hidden_dim) settings shows that fused kernel achieves an average 4.99× speedup over the unfused baseline (details here), with smaller shapes yielding larger gains. .

E2E Test

Performance Comparison Report: Wan-AI/Wan2.1-T2V-1.3B-Diffusers (1.07x speedup)

1. High-level Summary

Metric	Baseline	New	Diff	Status
E2E Latency	87712.82 ms	81929.56 ms	-5783.26 ms (-6.6%)	✅
Throughput	0.01 req/s	0.01 req/s	-	-

2. Stage Breakdown

Stage Name	Baseline (ms)	New (ms)	Diff (ms)	Diff (%)	Status
InputValidationStage	0.04	0.04	+0.00	+0.5%	⚪️
TextEncodingStage	2649.54	2546.71	-102.83	-3.9%	⚪️
ConditioningStage	0.05	0.02	-0.03	-60.0%	⚪️
TimestepPreparationStage	3.02	2.51	-0.51	-16.9%	⚪️
LatentPreparationStage	16.25	12.18	-4.07	-25.1%	⚪️
DenoisingStage	81457.85	75774.33	-5683.52	-7.0%	⚪️
DecodingStage	3568.70	3584.35	+15.65	+0.4%	⚪️

Metadata

• Baseline Commit: e6ced88040d0d35622e3ee46166cd2d41bd57091
• New Commit: e6ced88040d0d35622e3ee46166cd2d41bd57091
• Timestamp: 2025-12-04T14:12:54.485444

Performance Comparison Report: hunyuanvideo-community/HunyuanVideo

run with gpu-nums=4

No clear speedup is seen on Hunyuan, likely because its much longer sequence length (115,200 in hunyuan vs. 32,760 in wan) makes attention kernel the dominant cost.

1. High-level Summary

Metric	Baseline	New	Diff	Status
E2E Latency	395672.21 ms	394468.82 ms	-1203.40 ms (-0.3%)
Throughput	0.00 req/s	0.00 req/s	-	-

2. Stage Breakdown

Stage Name	Baseline (ms)	New (ms)	Diff (ms)	Diff (%)	Status
InputValidationStage	0.05	0.05	+0.01	+18.8%	⚪️
TextEncodingStage	1254.86	1416.92	+162.06	+12.9%	🔴
ConditioningStage	0.02	0.02	+0.00	+9.5%	⚪️
TimestepPreparationStage	57.30	58.78	+1.48	+2.6%	⚪️
LatentPreparationStage	324.41	323.24	-1.18	-0.4%	⚪️
DenoisingStage	376379.62	375239.32	-1140.30	-0.3%	⚪️
DecodingStage	17619.50	17395.27	-224.22	-1.3%	⚪️

Metadata

• Baseline Commit: d276cc35c6e5c5ca9453c7a6e369a383cf206263
• New Commit: d276cc35c6e5c5ca9453c7a6e369a383cf206263
• Timestamp: 2025-12-04T15:27:43.426476

Profiling: 78% SOL bandwidth

Using a real case from the wan model (batch=1, seq=32760, hidden_dim=1536), the fused kernel processes approximately:

• ~400MB of reads
  • [1 * 32760 * 1536 * 2 (residual, x) + 1 * 1 * 1536 (gate) + 1536 * 2 (scale, shift) + 1536 * 2 (weight, bias)] * 4B
• ~400MB of writes
  • [1 * 32760 * 1536 * 2 (modulated, residual_out)] * 4B

The kernel completes in 213 µs, yielding an effective bandwidth of: 800 MB / 213 µs ≈ 3.76 TB/s. This reaches ~78% of the theoretical H200 HBM bandwidth (4.8 TB/s), consistent with ncu report.

Checklist

• cuda kernel impl
• unit test
  • acc test
  • perf test
• e2e test
• clean code

Others

Kernel Impl

1. The kernel follows below 3-step：

• elementwise: residual_out = residual + gate * x;
• reduce: normalized = norm(residual_out);
• elementwise: modulated = (1 + scale) * normalized + shift.

2. Elementwise ops use ldg.128 (ldg.64 for fp16) to load/store 4 elements at a time.
3. Reduction uses the Welford algorithm for numerical stability:

• hidden_dim ≤ 1024: one warp per hidden_dim (thread → warp reduce).
• hidden_dim > 1024: one CTA per hidden_dim (thread → warp → CTA reduce).

4. residual_out is not stored in shared memory, relying on L1 cache instead.
5. covers all data types and shapes supported by the original kernel.

[mickqian]

https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/docs/contributing.md

[BBuf]

For kernels related to sglang diffusion, you can create a new folder named sgl_diffusion under csrc, and then move the kernel folder from here into it.

[BBuf]

Please write separate test and benchmark cases for this kernel, referring to other kernels under sgl-kernel.

[yingluosanqian]

For kernels related to sglang diffusion, you can create a new folder named sgl_diffusion under csrc, and then move the kernel folder from here into it.

ok.

Please write separate test and benchmark cases for this kernel, referring to other kernels under sgl-kernel.

added: sgl-kernel/tests/sgl_diffusion/test_fused_scale_residual_norm_scale_shift.py

Due to the large number of possible combinations, I did not generate all of them but ensured each parameter appears at least once.

[BBuf]

Should we add a pointer align check? Such as: assert(reinterpret_cast<uintptr_t>(ptr) % 16 == 0);

[yingluosanqian]

thanks. I fixed the typo, improved the alignment checks, and added tests for unaligned addresses.

Previously, I only used D % 4 == 0 to decide whether vectorized loads were allowed. Now I use a stricter check based on the actual data pointer alignment:
reinterpret_cast<uintptr_t>(t.data_ptr()) % 16 == 0 for FP32, and
reinterpret_cast<uintptr_t>(t.data_ptr()) % 8 == 0 for FP16/BF16.