[WS1][Qwen-Image][CUDA] Final AdaLN output (final_adaln_out)

[fusheng-ji]

What

RFC #386 WS1 row final_adaln_out: the image-stream final modulation norm_out:

scale, shift = (SiLU(temb) @ Wᵀ + b).chunk(2)        # 3072 → 6144, (scale, shift) order
out = LayerNorm(x, eps=1e-6, no affine) * (1 + scale) + shift

Forward and backward are implemented. The later 3072 → 64 proj_out belongs to latent_proj_gemm and is not included.

• Projection: reuses adaln_mod_projection's fixed-order CUDA projection and backward ([WS1][Qwen-Image][CUDA] AdaLN modulation projection (adaln_mod_projection) #537) unchanged.
• LayerNorm: PyTorch's native CUDA LayerNorm forward and backward (aten.native_layer_norm and native_layer_norm_backward), run on FP32 input with a declared execution-dtype store. Both are per-row, so batching does not change them; the BI checks above confirm this.
• Modulation parameter gradients: d_scale and d_shift are token sums, reduced by a new FP32 kernel (csrc/cuda/model_specific/qwenimage/token_sum.cu) that applies fixed 32-way balanced trees recursively until one value remains. The tree depends only on the sequence length, not on B or batch position.
• Reference: independent CPU FP32 reference with sequential Neumaier token sums.
• Registration: registry key qwenimage_final_adaln_out (CUDA only, no fallback; CPU maps to the reference; ROCm/NPU/MUSA fail closed), plus the WS1 operator spec/inputs, binding, build source and docs/operators/qwenimage-final-adaln-out.md.
• Runners: the evidence runners, plots and tests gain the final stream (benchmarks/models/qwenimage/streams.py), plus two reuse-survey modes for the LayerNorm part: Liger and flash-attention.

Placement follows #537: model_specific/qwenimage/.

Prior art & reuse decision

Each candidate composes an existing batch-invariant projection with an existing LayerNorm, and is measured on the full forward and backward path with the checkpoint's norm_out.linear weights. BI is checked bytewise: a singleton vs every row of B = 64/257/2048, three seeds, dense and boundary batch sweeps, the 4096-sample workload in 5 chunkings, and B = 2 vs B = 1 at S = 31/32/33/4096/6889/6032.

Library	Implementation (file@commit)	BI measured (final stream)	Verdict
PyTorch	F.layer_norm + F.linear (cuBLAS), torch 2.13.0+cu130	--mode native --phase bi: 581/816 records fail	Not BI
vLLM	batch_invariant.py@7797b60 GEMM + F.layer_norm	--mode vllm_local: 4/816 fail (long sequences)	BF16 d_weight 40,442 elements over tolerance at S=6889
vLLM (cuBLAS workspace settings only)	CUBLAS_WORKSPACE_CONFIG=:16:8, CUBLASLT_WORKSPACE_SIZE=1	761/816 fail	Not BI
SGLang	batch_invariant_ops.py, sglang 0.5.21	4/816 fail (long sequences)	Same accuracy failures as vLLM
Megatron-LM (training path of Miles and Vime)	batch_invariant_kernels.py@52fbcbc	4/816 fail (long sequences)	Same accuracy failures as vLLM
Liger	liger_kernel/ops/layer_norm.py, liger-kernel 0.8.4, with the vLLM GEMM	5/816 fail (long sequences)	BF16 d_weight 40,459 elements over tolerance at S=6889
FlashAttention	flash_attn/ops/triton/layer_norm.py@94e22c9 (Triton, with backward), with the vLLM GEMM	5/816 fail (long sequences)	BF16 d_weight 40,442 elements over tolerance at S=6889
FlashInfer	flashinfer/norm/__init__.py@34bb66e fused_dit_residual_layernorm_scale_shift (+ gated variants)	Not measured	BF16 forward only, no backward; cannot serve training
CUTLASS	0b55a2f: examples only (e.g. 37_gemm_layernorm_gemm_fusion)	Not applicable	No LayerNorm + modulation operator with autograd
verl / Miles / Vime	1d367a4 / 05bd9cb / b9ecdf5: git grep -iE 'img_mod|txt_mod|AdaLayerNormContinuous|adaln' → no files	Not applicable	No operator of their own; rollout uses vLLM/SGLang and training uses Megatron or FSDP (PyTorch), all measured above
Diffusers	AdaLayerNormContinuous in normalization.py@d961a38	Same as PyTorch row	Not BI

All five batch-invariant compositions also exceed the BF16 tolerance on d_weight/d_bias:

• at S = 4096/6889/6032 (about 24,000–40,000 elements);
• at B = 256 and B = 2048 (up to 239,278 elements).

At FP32 B = 2048 they fail by 70–75 elements. In these compositions, the modulation backward and its token sums run through native autograd; this operator accumulates those sums in a fixed FP32 tree and passes every case.

Final stream, forward + backward median ms:

	BF16 B=1, S=6889	BF16 B=2048, S=3	FP32 B=2048, S=3
PyTorch (not BI)	0.814	0.780	4.485
vLLM / SGLang / Megatron GEMM + F.layer_norm	1.48 / 1.48 / 1.49	4.00 / 4.03 / 4.00	8.01 / 7.98 / 8.01
vLLM GEMM + Liger / flash-attention LayerNorm	1.60 / 1.91	4.17 / 4.45	8.06 / 8.02
This operator	0.968	61.9	68.6

Decision: write our own. No existing implementation is both batch-invariant and within the declared tolerances:

• PyTorch and the cuBLAS policy are not BI.
• Every BI composition fails BF16 parameter-gradient accuracy at the RFC's long token counts.
• FlashInfer's fused kernel has no backward.

The operator is slower than these candidates at large conditioning batches; the projection reuse question in #537 applies to its projection part as well.

Results

Measured at 55f83d2b9043f41dca84440d9002ca30f5d3fac4. The evidence commit adds only docs/usage/evidence/qwenimage-final-adaln-out-b200/.

fixed/acceptance.json is qualified with zero failures (all three streams are revalidated at this commit):

Gate	Final stream	All streams
Real-weight accuracy vs CPU FP32 and FP64, seeds 0/127/2026, BF16/FP32	42/42	90/90
Batch-invariance records (raw bytes)	816/816	2424/2424
Full 4096-sample workload, 5 chunkings	30/30	90/90
Shared-parameter repeat and ordered zero-padding	42/42	126/126
FP64 diagnostic of CUDA and CPU FP32 gradients (FP32 final)	4/4	—

Accuracy (strict FP64 acceptance; this operator vs the BI compositions):

Batch invariance:

Performance (all candidates; and this operator vs native PyTorch at B=1):

At B=1, S=6889 (BF16) the operator takes 0.979 ms forward + backward vs 0.802 ms for native PyTorch. Peak memory is 605.5 MiB vs 350.4 MiB, because the operator keeps FP32 normalization state. Environment: B200 (SM100), PyTorch 2.13.0+cu130, CUDA 13.0, driver 580.126.20, TF32 off, no fast-math. Full evidence, pins and reproduction: README.

Status

No known failures at efeb3b9d57668371a2bb394ae3d49a5c3e794af3. Evidence was produced from a clean checkout of the code commit 55f83d2b9043f41dca84440d9002ca30f5d3fac4, which is this PR's parent of the evidence commit. Stacked on #537; review that first.

Tests

KERNEL_ALIGN_USE_FAST_MATH=0 TORCH_CUDA_ARCH_LIST=10.0 python setup.py build_ext --inplace
CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=. python -m pytest tests/models/qwenimage \
  tests/runtime/test_kernel_registry.py tests/contracts/test_tolerance_contract.py \
  tests/build/test_build_platform_collectives.py tests/validation/operators/test_ws1_gtest_gpu.py \
  tests/runtime/test_layout_compatibility.py -q -rs

135 passed, 0 failed, 0 skipped on 2× B200 (log). That is 101 from #537 and 34 added here. The added tests cover:

• Long-sequence gradients at S = 4096/6889/6032 over three seeds.
• (scale, shift) order and LayerNorm.
• An FP64 gradcheck of the reference modulation backward.
• Retention of small token contributions in the reference.
• The final stream in every shared test.

Scope / Known limitations

• Hardware: NVIDIA SM80+ code, validated on B200 (SM100) only; no H100/SM90 or ROCm run. Other platforms' registry slots are empty and fail closed.
• Operator scope: norm_out only; proj_out belongs to latent_proj_gemm. First-order autograd only. Long token counts and large batches are covered separately, not every B×S combination.
• Cost: unfused; uses more memory than native PyTorch and is much slower at large conditioning batches.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

• Configuration used: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: 078f54a4-bfb0-4331-a213-ef7a8e8e69ee

---

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

_{Comment @coderabbitai help to get the list of available commands.}