[diffusion] Fuse ERNIE AdaLN residual path by BBuf · Pull Request #30170 · sgl-project/sglang

BBuf · 2026-07-05T15:46:06Z

Motivation

Use existing diffusion fused kernels in the ERNIE-Image Turbo shared AdaLN block to reduce repeated RMSNorm, scale/shift, residual, and gate-add overhead in the denoising transformer path.

The optimized block keeps the checkpoint parameter names unchanged while replacing this hot eager pattern with guarded fused calls:

hidden_states = norm(hidden_states) * scale + shift
hidden_states = residual + gate * branch_out

Modifications

Add private ERNIE fast-path helpers around existing fused_norm_scale_shift, fused_scale_residual_norm_scale_shift, and residual_gate_add_cuda.
Fuse the first RMSNorm + scale/shift before attention.
Fuse attention residual + gate add + MLP-input RMSNorm + scale/shift.
Reuse the residual-gate CUDA op for the final MLP residual update.
Keep adaLN_sa_ln and adaLN_mlp_ln checkpoint parameter names unchanged.
Guard CUDA/dtype/layout/runtime support and fall back to the original PyTorch expressions on unsupported inputs or one-time runtime failure.

Accuracy Tests

python3 -m py_compile   python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py   test/registered/jit/diffusion/test_ernie_fused_adaln.py
python3 -m ruff format --check ...
python3 -m ruff check ...
git diff --check

Result: local syntax/format/lint/diff checks passed.

H200 unit test:

CUDA_VISIBLE_DEVICES=6 PYTHONPATH=python FLASHINFER_DISABLE_VERSION_CHECK=1 pytest -q test/registered/jit/diffusion/test_ernie_fused_adaln.py -q

Result: 6 passed.

Speed Tests and Profiling

H200 native-backend A/B on ernie-image-turbo with the benchmark preset. Native backend gate passed: no Falling back to diffusers backend, Using diffusers backend, or Loaded diffusers pipeline in the benchmark log.

E2E command, run once at the PR parent baseline and once at this PR head. Only the checkout and --perf-dump-path label changed between the two runs.

cd /tmp/sglang_local_validate
CUDA_VISIBLE_DEVICES=7 PYTHONPATH=python sglang generate \
  --backend=sglang \
  --model-path=baidu/ERNIE-Image-Turbo \
  --prompt="A futuristic cyberpunk city at night, neon lights reflecting on wet streets" \
  --seed=42 \
  --width=1024 \
  --height=1024 \
  --save-output \
  --warmup \
  --enable-torch-compile \
  --perf-dump-path /tmp/sglang_local_validate/outputs/pr_candidate/ernie-image-turbo_fused_adaln.json

Compare the baseline/head perf dumps with:

PYTHONPATH=python python python/sglang/multimodal_gen/benchmarks/compare_perf.py \
  /tmp/sglang_local_validate/outputs/pr_baseline/ernie-image-turbo_baseline.json \
  /tmp/sglang_local_validate/outputs/pr_candidate/ernie-image-turbo_fused_adaln.json

Implementation	E2E ms	DenoisingStage ms	Delta E2E	Delta Denoise
Baseline	22154.41	14652.40	-	-
Fused AdaLN	20071.36	13334.39	-9.4%	-9.0%

Perf dump: /tmp/sglang_local_validate/outputs/pr_candidate/ernie-image-turbo_fused_adaln.json on the H200 validation box.

H200 Fused-vs-Unfused Accuracy

The generated-image A/B compares the PR parent baseline against this PR head:

Side	Commit
Baseline unfused	`3ea875fef48f6f01fa3bddd9e2197ad190cef29d`
Fused AdaLN	`8d2b755a167be19a7a45971243c47e8f5c1e9946`

To isolate the denoising transformer change, the ERNIE prompt enhancement output was fixed once and then reused for both runs with use_pe=false. This avoids PromptEnhancementStage nondeterminism; a baseline-vs-baseline repeat with the same fixed prompt is bitwise identical.

Check	Prompt / Seed / Shape	Comparison	Result
Native backend gate	fixed ERNIE enhanced prompt, `seed=42`, `1024x1024`	logs for baseline and fused generation	no diffusers fallback signals
Baseline repeat determinism	same fixed prompt, same seed	baseline vs baseline repeat	`SSIM=1.000000`, `PSNR=inf`, `max_abs=0`
Fused-vs-unfused output	same fixed prompt, same seed	fused image vs unfused image	`SSIM=0.720635`, `PSNR=18.06 dB`, `max_abs=255`

Quality conclusion: the fused full-image sample is not pixel-identical to the unfused path, but the visual layout/content is preserved in the side-by-side sample below. The difference is from the fused residual/norm path changing floating-point accumulation in the denoising trajectory, not from prompt rewriting.

Result Image Comparison

H200 ERNIE-Image Turbo A/B output sample (seed=42, 1024x1024). The figure compares the original unfused path against the fused AdaLN residual path and includes an amplified absolute-difference view.

Image-level comparison: SSIM=0.720635, PSNR=18.06 dB, max_abs=255. Baseline repeat check: SSIM=1.000000, PSNR=inf, max_abs=0.

Checklist

Default ERNIE path tries the fused CUDA/Triton helpers first and falls back to the existing PyTorch implementation when unsupported.
No checkpoint parameter names changed.
H200 unit test included for fused helper behavior.
H200 native-backend benchmark evidence included.
H200 same-prompt generated-image comparison and SSIM/PSNR included.

CI States

Latest PR Test (Base): ❌ Run #28763999706
Latest PR Test (Extra): ❌ Run #28763999617

gemini-code-assist

Code Review

This pull request introduces fused CUDA kernel optimizations for ERNIE DiT models, adding fast-path helpers for AdaLN, residual scaling, and gate additions, along with corresponding unit tests. The review feedback suggests making the gate tensor contiguous to prevent silent fallbacks to the slow path, robustly validating residual and gate tensors to avoid permanently disabling the fast path due to transient exceptions, and updating the unit tests to simulate non-contiguous chunked tensors.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-07-05T15:48:05Z

+            if can_use_residual_gate_add_cuda(residual, update, gate):
+                return residual_gate_add_cuda(residual, update, gate)


In production, gate_mlp is chunked from mod_params and is therefore non-contiguous. Since can_use_residual_gate_add_cuda strictly requires contiguous inputs, it will return False and silently fall back to the slow PyTorch path. Making gate contiguous before checking ensures the CUDA fast path is successfully utilized.

Suggested change

if can_use_residual_gate_add_cuda(residual, update, gate):

return residual_gate_add_cuda(residual, update, gate)

gate_contig = gate.contiguous()

if can_use_residual_gate_add_cuda(residual, update, gate_contig):

return residual_gate_add_cuda(residual, update, gate_contig)

gemini-code-assist · 2026-07-05T15:48:05Z

+def _ernie_can_use_fused_adaln(
+    x: torch.Tensor,
+    norm: RMSNorm,
+    shift: torch.Tensor,
+    scale: torch.Tensor,
+) -> bool:
+    return (
+        x.is_cuda
+        and x.dim() == 3
+        and x.dtype in (torch.float16, torch.bfloat16, torch.float32)
+        and norm.weight is not None
+        and norm.weight.is_cuda
+        and norm.weight.dtype == x.dtype
+        and shift.is_cuda
+        and scale.is_cuda
+        and shift.dtype == x.dtype
+        and scale.dtype == x.dtype
+        and x.shape[-1] % 256 == 0
+        and x.shape[-1] <= 8192
+    )


To prevent transient incompatible inputs (e.g., mismatched device or dtype for residual or gate) from raising exceptions and permanently disabling the fused AdaLN fast path globally, we should validate residual and gate in _ernie_can_use_fused_adaln.

def _ernie_can_use_fused_adaln( x: torch.Tensor, norm: RMSNorm, shift: torch.Tensor, scale: torch.Tensor, residual: Optional[torch.Tensor] = None, gate: Optional[torch.Tensor] = None, ) -> bool: return ( x.is_cuda and x.dim() == 3 and x.dtype in (torch.float16, torch.bfloat16, torch.float32) and norm.weight is not None and norm.weight.is_cuda and norm.weight.dtype == x.dtype and shift.is_cuda and scale.is_cuda and shift.dtype == x.dtype and scale.dtype == x.dtype and x.shape[-1] % 256 == 0 and x.shape[-1] <= 8192 and (residual is None or (residual.is_cuda and residual.dtype == x.dtype and residual.shape == x.shape)) and (gate is None or (gate.is_cuda and gate.dtype == x.dtype)) )

gemini-code-assist · 2026-07-05T15:48:05Z

+    if not _ERNIE_FUSED_ADALN_DISABLED and _ernie_can_use_fused_adaln(
+        x, norm, shift, scale
+    ):


Pass residual and gate to _ernie_can_use_fused_adaln to perform robust validation before executing the fused kernel.

Suggested change

if not _ERNIE_FUSED_ADALN_DISABLED and _ernie_can_use_fused_adaln(

x, norm, shift, scale

):

if not _ERNIE_FUSED_ADALN_DISABLED and _ernie_can_use_fused_adaln(

x, norm, shift, scale, residual, gate

):

gemini-code-assist · 2026-07-05T15:48:05Z

+    shift = torch.randn((1, 1, 256), device="cuda", dtype=dtype)
+    scale = torch.randn((1, 1, 256), device="cuda", dtype=dtype)


Simulate non-contiguous chunked tensors in the test to match production conditions and ensure the fast path handles them correctly.

Suggested change

shift = torch.randn((1, 1, 256), device="cuda", dtype=dtype)

scale = torch.randn((1, 1, 256), device="cuda", dtype=dtype)

# Simulate non-contiguous chunked tensors as in production

shift = torch.randn((1, 2, 256), device="cuda", dtype=dtype)[:, :1, :]

scale = torch.randn((1, 2, 256), device="cuda", dtype=dtype)[:, :1, :]

gemini-code-assist · 2026-07-05T15:48:05Z

+    gate = torch.randn((1, 1, 256), device="cuda", dtype=dtype)
+    shift = torch.randn((1, 1, 256), device="cuda", dtype=dtype)
+    scale = torch.randn((1, 1, 256), device="cuda", dtype=dtype)


Simulate non-contiguous chunked tensors in the test to match production conditions and ensure the fast path handles them correctly.

Suggested change

gate = torch.randn((1, 1, 256), device="cuda", dtype=dtype)

shift = torch.randn((1, 1, 256), device="cuda", dtype=dtype)

scale = torch.randn((1, 1, 256), device="cuda", dtype=dtype)

# Simulate non-contiguous chunked tensors as in production

gate = torch.randn((1, 2, 256), device="cuda", dtype=dtype)[:, :1, :]

shift = torch.randn((1, 2, 256), device="cuda", dtype=dtype)[:, :1, :]

scale = torch.randn((1, 2, 256), device="cuda", dtype=dtype)[:, :1, :]

BBuf · 2026-07-06T10:08:55Z

Closing this direction for now. The current fused AdaLN/residual-gate path is not safe for ERNIE-Image Turbo full-image consistency: even after reverting the two AdaLN RMSNorm fused helpers, keeping only the residual-gate fused path gave SSIM=0.736905 against the production baseline; a one-rounding CUDA variant improved only to SSIM=0.802430. Full eager / compile-fallback restored bitwise identical images (SSIM=1.000000, max_abs=0), so this should be reworked as a bitwise-first kernel task rather than merged as an optimization PR. I captured a follow-up KDA-Pilot task for a B200 bitwise ERNIE AdaLN/residual-gate fused kernel.

[diffusion] Fuse ERNIE AdaLN residual path

93792da

BBuf requested review from mickqian, ping1jing2, yhyang201 and yingluosanqian as code owners July 5, 2026 15:46

github-actions Bot added the diffusion SGLang Diffusion label Jul 5, 2026

BBuf added the run-ci CI: run the baseline test suite on this PR label Jul 5, 2026

gemini-code-assist Bot reviewed Jul 5, 2026

View reviewed changes

BBuf added run-ci CI: run the baseline test suite on this PR and removed run-ci CI: run the baseline test suite on this PR labels Jul 5, 2026

Simplify ERNIE residual gate fast path

8d2b755

BBuf closed this Jul 6, 2026

This was referenced Aug 5, 2026

[diffusion] ERNIE-Image bit-exact residual-gate fast path (H200 1024^2 e2e 16.17 -> 15.75 s) #33734

Merged

[diffusion] ERNIE-Image bit-exact fused RMSNorm+scale/shift (H200 1024^2 e2e 15.63 -> 15.00 s, denoise -3.3%) #33854

Merged

Navigation Menu

Repository navigation

[diffusion] Fuse ERNIE AdaLN residual path - #30170

[diffusion] Fuse ERNIE AdaLN residual path#30170
BBuf wants to merge 2 commits into
sgl-project:mainfrom
BBuf:bbuf/ernie-fused-adaln

BBuf commented Jul 5, 2026 •

edited by github-actions Bot

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jul 5, 2026

Uh oh!

gemini-code-assist Bot Jul 5, 2026

Uh oh!

gemini-code-assist Bot Jul 5, 2026

Uh oh!

gemini-code-assist Bot Jul 5, 2026

Uh oh!

gemini-code-assist Bot Jul 5, 2026

Uh oh!

BBuf commented Jul 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

		if can_use_residual_gate_add_cuda(residual, update, gate):
		return residual_gate_add_cuda(residual, update, gate)

-            if can_use_residual_gate_add_cuda(residual, update, gate):
-                return residual_gate_add_cuda(residual, update, gate)
+            gate_contig = gate.contiguous()
+            if can_use_residual_gate_add_cuda(residual, update, gate_contig):
+                return residual_gate_add_cuda(residual, update, gate_contig)

		shift = torch.randn((1, 1, 256), device="cuda", dtype=dtype)
		scale = torch.randn((1, 1, 256), device="cuda", dtype=dtype)

-    shift = torch.randn((1, 1, 256), device="cuda", dtype=dtype)
-    scale = torch.randn((1, 1, 256), device="cuda", dtype=dtype)
+    # Simulate non-contiguous chunked tensors as in production
+    shift = torch.randn((1, 2, 256), device="cuda", dtype=dtype)[:, :1, :]
+    scale = torch.randn((1, 2, 256), device="cuda", dtype=dtype)[:, :1, :]

Conversation

BBuf commented Jul 5, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

H200 Fused-vs-Unfused Accuracy

Result Image Comparison

Checklist

CI States

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jul 5, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 5, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 5, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 5, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 5, 2026

Choose a reason for hiding this comment

Uh oh!

BBuf commented Jul 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

BBuf commented Jul 5, 2026 •

edited by github-actions Bot

Loading