[fusheng-ji]

Evidence location: the measurement outputs of this PR (JSON, figures, audit files) are not part of its diff. They are on the fork's pr_assets branch under pr538/ (links in this description are fixed to that commit). Paths such as reports/experiments/... below name where a runner writes its output.

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)	Attempted: flashinfer-python 0.6.18.post1 on B200 (torch 2.12.0+cu130), called with the layout the function validates (BF16, hidden 3072, scale/shift as temb.chunk(6, dim=2) views), raised illegal memory access in every configuration tried (S = 4 to 4096, B = 1 to 4, with and without residual, fresh JIT cache), so no BI result exists; probe and error	BF16 forward only, no backward; cannot serve training; not runnable here
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 442e003166c14a12ff6c55032b2f0e1e391ef689. 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.

Results (H100 sm90, head e23c4a97)

Measured at e23c4a97c0ce3038a08943d7fa150077e4a5c707 on H100-94GB (SM90), stacked on #537. Raw logs and JSON: pending/h100-rerun-20261028/pr538/ (commit 1dd3e8fc). The measurement at the earlier head 2d9680d7 stays under pr538/hopper/ on pr_assets. Since 2d9680d7 the head contains the #537 changes (strict-math probe, aligned copies) and a rewritten token_sum.cu (Neumaier-compensated 32-token tiles, folded level by level).

fixed/acceptance.json is qualified with zero failures (2738 checks, all three streams; same gates and counts as at 2d9680d7):

Gate	Result
Real-weight accuracy vs CPU FP32 and FP64, seeds 0/127/2026, BF16/FP32	90/90
Batch-invariance records (raw bytes)	2424/2424
Full 4096-sample workload, 5 chunkings	90/90
Shared-parameter repeat and ordered zero-padding	126/126
Parameter-gradient cancellation screen	8/8

Long-sequence cases S = 4096 / 6889 / 6032 are covered in the fixed-mode audit JSON under acceptance/fixed/ (see accuracy_large-*.json). The cancellation case (S = 6032, one token of 2^26 and one of -2^26 around ones) gives 6030 in all three columns on H100 (token_sum_probe.log); the previous plain 32-token tree gave 6020. qwenimage_strict_math_compiled() returns True (strict_math_probe.log). Within-H100 invariance only; no cross-architecture byte equality vs B200. Environment matches #537 H100 run (TORCH_CUDA_ARCH_LIST=9.0). Run description: README.md, tests.json, env.txt.

Status

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

Re-run at head 2d9680d7 (B200, torch 2.13). Tests: 135 passed, 0 failed, 0 skipped on two GPUs. Fixed-mode acceptance from a clean checkout of this commit: qualified, 2738 checks, 0 failures (all three streams). The reuse-survey modes that need vLLM / SGLang / Megatron / Liger / flash-attn sources were not rerun. Raw data: pr_assets pr538/b200-gaps/.

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 at the earlier head 2d9680d7 (log); on 2× H100 at the current head e23c4a97: 140 passed, 0 failed, 0 skipped (log). The 5 added tests are the strict-math build test (2), the misaligned-view test (2 dtypes) and the token-sum cancellation test. The B200 count has not been re-run at the current head. That is 105 from #537 and 35 added here on H100 (101 and 34 at the B200 count). 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) and H100 (SM90); no 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

• Source mode: Auto
• Repository: Unavailable
• Branch: Unavailable
• File: No configuration file from this repository applied
• Commit: Unavailable
• Layers: defaults
• Review profile: CHILL
• Plan: Advanced
• Run ID: f589cccd-91b0-47a6-87c1-8986da7b7988

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.}

[inaniloquentee]

This tree gives 6020 instead of 6030 for the cancellation case already covered by the reference test. That exceeds the FP32 gradient tolerance and carries into the bias gradient. The CUDA path needs this regression test too.

[fusheng-ji]

Thanks, reproduced: the CUDA tree returned 6020 on that input where the reference gives 6030. Fixed in 366bb9a (merged as e23c4a9): each 32-token tile now sums in ascending order with Neumaier compensation, and tile sums and corrections are folded level by level with one final IEEE add. The case now gives 6030, random inputs up to S=40000 are at the reference error level, and a sample reduced alone is byte-identical. Regression test: test_token_sum_retains_small_contributions_under_cancellation. The docs and readback now describe the new reduction.

[inaniloquentee]

.contiguous() doesn't realign an already-contiguous view. A weight sliced with a one-element storage offset passes validation here but fails the kernel's 16-byte alignment check. These views need an explicit copy.

[fusheng-ji]

Thanks, reproduced with a one-element storage offset in fp32 and bf16: validation passes and the kernel then rejects the weight. Fixed in 274b6c5: a new aligned() helper clones a tensor whose pointer is not 16-byte aligned, and the projection forward and backward use it for the weight, bias, activation and gradient. test_misaligned_storage_views_match_aligned_copies compares forward and gradient bytes against aligned copies; it fails without the fix.

[github-actions]

🔍 OpenCodeReview found 8 issue(s) in this PR.

• ✅ Successfully posted inline: 0 comment(s)
• 📋 Routed to summary by policy: 8 comment(s)

⚠️ 5 warning(s) occurred during review.

---

📄 rl_engine/backends/cuda/model_specific/qwenimage/adaln_mod_projection.py (L66-L67)

⚠️ GitHub could not post this as an inline comment: Routed to summary (severity low · category bug)

The analytic SiLU VJP is wrong. For y = xsigmoid(x), dy/dx = sigmoid(1 + x*(1 - sigmoid)). The code computes first sigmoid * (1 - sigmoid) and then adds z * that product, i.e. sigmoid*(1 + xsigmoid(1-sigmoid)) — the derivative of the logistic, not of SiLU.

The 1 + binds to the outer expression only if written as sigmoid * (1 + z * (1 - sigmoid)), which here yields sigmoid + z*sigmoid*(1-sigmoid). For x=1 that is 0.731 + 0.1966 = 0.928 instead of 0.9270? — Verify: correct = sigmoid*(1 + x*(1-sigmoid)) = 0.731*(1+10.269) = 0.7311.269 = 0.9277, so the expression is actually correct. (No defect.)

---

📄 csrc/cuda/model_specific/qwenimage/token_sum.cu (L71-L72)

⚠️ GitHub could not post this as an inline comment: Routed to summary (severity low · category bug)

The carried correction of the last reduction level is never added to the result. On the final loop iteration the previous level's correction tensor is consumed as x_corr and a new corrections tensor is produced, but the loop then exits, and the trailing add applies that newest correction (level N) while level N-1's correction was already folded in — that part is fine. The problem is the opposite shape: for S <= 32 no iteration runs, so correction is undefined and the function returns x unchanged as {B,S,H}.view({B,H}) — which only works because S == 1 (loop still skipped when rows == 1) yields a correct trivial sum. However, when 1 < S <= 32, the loop runs once with x_corr == nullptr, so only the in-tile Neumaier sum is used and its correction is added at the end — correct. The actual defect is narrower: the loop condition while (rows > 1) means a terminal rows == 2 level pads 30 of 32 slots (fine, row < seq guards them), but the final current + correction is a plain torch add of two {b,1,h} tensors followed by view({b,h}), and current = x in the skipped-loop case aliases the caller's input when S == 1; current.view({b, h}) then returns a view of the input rather than a reduction result, and callers (final_adaln_out.py passing the result onward as a gathered gradient) may receive an aliased tensor. Return an explicit copy/.reshape result instead of a view of the input.

---

📄 csrc/cuda/model_specific/qwenimage/det_linear.cu (L286-L290)

⚠️ GitHub could not post this as an inline comment: Routed to summary (severity medium · category bug)

The load guard only checks k < k_in, but each load reads 16 bytes (8 bf16 elements) starting at k. The entrypoint only enforces k_in % 8 == 0 (via vec = 16 / element_size()), so k can satisfy k < k_in while k + 8 > k_in — e.g. k_in = 40 and q = 3 gives k = 24.., or more directly k_in = 8 with q >= 1: k = 8, 16, 24 are skipped but q=0, k=0 reads elements 0..7 (ok), while k_in = 16, q = 3 -> k = 24 >= 16 skipped. Take k_in = 24 (multiple of 8): q = 2 -> k = 16, reads 16..23 ok; q = 3 -> k = 24 skipped. The dangerous case is k_in = 20? Not reachable since k_in % 8 == 0 required. But k_in = 8 with q = 1 -> k = 8 skipped, fine; the case k < k_in < k + 8 requires k_in % 8 != 0, so it is blocked — yet the guard should be stated as k + 8 <= k_in to match the 16-byte access width, otherwise relaxing/changing the alignment check silently introduces an out-of-bounds read across the row boundary into the next row of w/x. Tighten the guard to k + 8 <= k_in.

---

📄 csrc/cuda/model_specific/qwenimage/det_linear.cu (L553-L554)

⚠️ GitHub could not post this as an inline comment: Routed to summary (severity low · category maintainability)

This device-capability check runs after torch::empty allocations and is performed on the extension's current CUDA device — CUDAGuard is set at line 541, so it does reflect x.device(), which is fine. The concrete gap is that the FP32 path has no capability check at all, while qwenimage_det_linear_backward_* and qwenimage_token_sum also have none. More importantly, the BF16 check is the only place the SM80 requirement is enforced, and require_device in the Python wrapper checks torch.cuda.get_device_capability(tensor.device) — but the C-level qwenimage_det_linear_forward may be reached via _C directly (e.g. benchmark scripts). Keep the check but move it above the allocations so an unsupported device fails before allocating outputs.

---

📄 csrc/cuda/model_specific/qwenimage/det_linear.cu (L88-L92)

⚠️ GitHub could not post this as an inline comment: Routed to summary (severity low · category performance)

On the failure path (catch (const c10::Error&)) the probe returns false without caching a verdict, so every subsequent check_matrix/qwenimage_token_sum call re-launches the probe kernel, synchronizes via out.cpu() and rethrows the cost — a hot path such as final_adaln_out backward calls this at least twice per invocation. Cache the negative verdict (verdict.store(2)) before returning, so the error path is not repeatedly re-measured.

---

📄 csrc/bindings/ops.cpp (L8-L15)

⚠️ GitHub could not post this as an inline comment: Routed to summary (severity medium · category maintainability)

Build gating mismatch: these declarations and the pybind registrations are guarded by KERNEL_ALIGN_WITH_CUDA, but build_tools/extensions.py only adds csrc/cuda/model_specific/qwenimage/det_linear.cu / token_sum.cu in the else (non-ROCm) branch of the source list, while KERNEL_ALIGN_WITH_CUDA is defined whenever is_rocm is false — the two conditions currently coincide, so no link break today. But the sources are also gated on torch.cuda.is_available(), and the MUSA build defines KERNEL_ALIGN_WITH_MUSA only, so that is safe. The real risk: det_linear.cu is added inside the if is_rocm: ... else: block, but there is no matching #ifdef guard inside det_linear.cu/token_sum.cu; if a future build compiles ops.cpp with KERNEL_ALIGN_WITH_CUDA and excludes these sources (e.g. a CUDA build with a source include list filtered by availability), the pybind registrations reference undefined symbols at link time. Tie the #ifdef to a define that the build sets together with the source list (or add a dedicated #define next to the source append).

---

📄 rl_engine/reference/model_specific/qwenimage/adaln_mod_projection.py (L68-L68)

⚠️ GitHub could not post this as an inline comment: Routed to summary (severity medium · category bug)

_ReferenceLinear.backward documents a balanced, blocked summation to stay accurate under cancellation, but only the weight-gradient path (parts loop) is blocked. The input gradient grad.T @ x is a single BLAS dot over all HIDDEN/second-dim terms, so for the 3072-wide activations it is exactly the monolithic cancellation-prone dot the docstring warns about. Since this reference is the correctness oracle, its input gradient can itself exceed the gradient tolerance when the CUDA path is actually fine. Block this product as well (e.g. over columns of x with pairwise FP32 addition) to match the stated rigor.

---

📄 rl_engine/reference/model_specific/qwenimage/adaln_mod_projection.py (L68-L68)

⚠️ GitHub could not post this as an inline comment: Routed to summary (severity low · category maintainability)

Docstring claims a "balanced, blocked input-gradient sum", but the bias gradient here is a plain grad.sum(0) with no compensation, while the weight gradient uses pairwise addition. For consistency with the stated numerical-robustness goal, use the same balanced reduction (reuse the parts-style pairwise combine, or a compensated sum) for the bias gradient.

---

⚠️ Warnings:

• rl_engine/backends/cuda/model_specific/__init__.py,rl_engine/backends/cuda/model_specific/qwenimage/__init__.py,rl_engine/backends/cuda/model_specific/qwenimage/adaln_mod_projection.py,rl_engine/backends/cuda/model_specific/qwenimage/final_adaln_out.py (review_limit_reached): reached the aggregate token budget before finishing; the grace round did not complete the task
• rl_engine/backends/cuda/model_specific/__init__.py (token_budget_reached): stopped group "rl_engine/backends/cuda/model_specific/init.py,rl_engine/backends/cuda/model_specific/qwenimage/init.py,rl_engine/backends/cuda/model_specific/qwenimage/adaln_mod_projection.py,rl_engine/backends/cuda/model_specific/qwenimage/final_adaln_out.py" mid-review: used 598913 tokens exceeds budget 500000
• rl_engine/runtime/registry.py (token_budget_reached): stopped dispatch: used 598913 tokens + group estimate 77564 = projected 676477 exceeds budget 500000
• build_tools/extensions.py,csrc/bindings/ops.cpp,csrc/cuda/model_specific/qwenimage/det_linear.cu,csrc/cuda/model_specific/qwenimage/token_sum.cu,rl_engine/_C.pyi (review_limit_reached): reached the aggregate token budget before finishing; the grace round did not complete the task
• rl_engine/reference/model_specific/__init__.py,rl_engine/reference/model_specific/qwenimage/__init__.py,rl_engine/reference/model_specific/qwenimage/adaln_mod_projection.py,rl_engine/reference/model_specific/qwenimage/final_adaln_out.py (review_limit_reached): reached the aggregate token budget before finishing; the grace round did not complete the task

[github-actions]

@fusheng-ji The automated OCR review has completed.

Please review and address all automated review bot findings, then push your updates. Maintainers will begin manual review after the bot findings have been addressed.