[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 pr537/ (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 adaln_mod_projection: the per-stream modulation projection SiLU(temb) @ Wᵀ + b, 3072 → 18432, run separately for img_mod.1 and txt_mod.1, with forward and backward.

• Forward: native SiLU stored in the execution dtype, then a fixed-order projection with FP32 accumulation. FP32 uses one warp's fixed vector/FMA order; BF16 uses zero-initialized mma.sync K=16 partials summed in a fixed order. Neither depends on B or on a row's batch position.
• Backward: FP32 intermediates. The input gradient uses compensated dot partials with a fixed fold. d_weight uses an ascending FP32 Neumaier sum of products and their FMA residuals; d_bias uses an ascending Neumaier sum. Only the final result is cast.
• An independent CPU FP32 reference with explicit storage rounding (rl_engine/reference/model_specific/qwenimage/).
• Registry key qwenimage_adaln_mod_projection: CUDA slot only, no fallback; CPU maps to the reference; ROCm/NPU/MUSA fail closed. WS1 operator spec and inputs, binding, build source, and docs/operators/qwenimage-adaln-mod-projection.md.

Placement. The sources are under csrc/cuda/model_specific/qwenimage/, rl_engine/backends/cuda/model_specific/qwenimage/ and rl_engine/reference/model_specific/qwenimage/, following the repository layout guide's backends/<platform>/model_specific/<model>/ convention for operators that depend on one model. The contract is Qwen-Image-specific: fixed 3072 → 18432 shape, (shift, scale, gate) × 2 chunk order, model-qualified registry key. If you prefer the operator-family layout used by #499 (gemm/), I can move it.

The final_adaln_out row builds on this projection and follows as a stacked PR.

Prior art & reuse decision

Every candidate was measured on the full forward and backward path with the checkpoint's block-0 weights, on the same harness as this kernel (benchmarks/models/qwenimage/adaln_reuse_bi_audit.py). BI is checked bytewise: a singleton vs every row of B = 64/257/2048, three seeds, dense and power-of-two-boundary batch sweeps, row position and reversal, and the full 4096-sample workload in 5 chunkings.

Library	Implementation (file@commit)	BI measured	Verdict
PyTorch	F.silu + F.linear (cuBLAS), torch 2.13.0+cu130	--mode native --phase bi: 1026/1608 records fail, 26502/28428 singleton d_temb rows differ	Not BI
vLLM	vllm/model_executor/layers/batch_invariant.py@7797b60 (vendored as benchmarks/layers/vllm_batch_invariant_matmul.py)	--mode vllm_local --phase bi: 0/1608	BI; FP32 B=2048 d_weight 2 elements/stream over tolerance
vLLM (cuBLAS workspace settings only)	CUBLAS_WORKSPACE_CONFIG=:16:8, CUBLASLT_WORKSPACE_SIZE=1	--mode cublas_policy --phase bi: 1506/1608 fail	Not BI
SGLang	sglang/srt/batch_invariant_ops/batch_invariant_ops.py, sglang 0.5.21	--mode sglang_local --phase bi: 0/1608	BI; same FP32 d_weight failure
Megatron-LM (training path of Miles and Vime)	megatron/core/transformer/custom_layers/batch_invariant_kernels.py@52fbcbc	--mode megatron_local --phase bi: 0/1608	BI; same FP32 d_weight failure
CUTLASS	0b55a2f: GEMM templates/examples only	Not applicable	No SiLU + Linear operator with a BI contract or autograd
FlashInfer	34bb66e: no batch-invariant linear	Not applicable	No implementation
FlashAttention	flash_attn/ops/fused_dense.py@94e22c9: cuBLASLt fused_dense_lib, GELU/ReLU epilogues	BF16 linear_act_forward (GELU/ReLU epilogue, the only variants it has), singleton vs batch at in 3072 / out 6*3072, 3 seeds, B = 64/257/2048, torch 2.12.0+cu130 (the other rows use 2.13): GELU fails 9/9 records (e.g. 277 of 2048 rows differ at B = 2048), ReLU 0/9; probe. FP16/BF16 only, so the FP32 path cannot run	No SiLU epilogue, no FP32, no backward of the modulation sums; GELU variant is not BI
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	transformer_qwenimage.py@d961a38: nn.SiLU + nn.Linear	Same as PyTorch row	Not BI

Accuracy at unchanged tolerances (--phase accuracy/accuracy_large, FP64 reference): the three persistent GEMMs pass all BF16 cases and FP32 up to B = 256; at FP32 B = 2048, 2 d_weight elements per stream exceed atol = rtol = 1e-4. This kernel has no failures.

Image stream, forward + backward median ms:

	BF16 B=1	BF16 B=2048	FP32 B=1	FP32 B=2048
PyTorch (not BI)	0.177	0.468	0.615	11.640
vLLM / SGLang / Megatron GEMM	2.39 / 2.37 / 2.38	9.94 / 10.05 / 9.94	2.65 / 2.63 / 2.65	21.04 / 21.01 / 21.04
This kernel	0.787	175.3	0.772	197.9

Decision: I kept the kernel, but this is the maintainers' call (see Status below).

• For the kernel: it is the only candidate that passes every accuracy case, including FP32 d_weight at B = 2048. The persistent GEMMs are forward-only, so reusing them still needs a training adapter like the one in this PR's audit. Qwen-Image conditioning batches (prompts × CFG branches) are small, where the kernel is faster.
• For reuse: at B ≥ 256 the GEMMs are 6–18× faster, and they pass every BF16 case, which matches the BF16 checkpoint.
• Middle option: register the vLLM GEMM with the FP32 VJP adapter as the BF16 path and keep this kernel for FP32. Each dtype then has one implementation for every B, so batch invariance holds. The cost is two arithmetic contracts to maintain. Choosing between implementations by B is not an option, because it would change bytes across batch sizes.

Rejected along the way (historical, code removed): a K=32 partial-dot reducer failed cancellation accuracy and changed bytes at the B = 256 dispatch boundary. An ordered-tile Neumaier reducer passed every numerical gate but was 17–26% slower than the fixed reducer.

Results

Measured at c88f7c0fe35d6544e092a16097a094d323b35581. The evidence commit adds only docs/usage/evidence/qwenimage-adaln-mod-projection-b200/.

fixed/acceptance.json is qualified with zero failures:

Gate	Result
Real-weight accuracy vs CPU FP32 and FP64, seeds 0/127/2026, BF16/FP32, B=1/3/256/2048	48/48
Batch-invariance records (raw bytes)	1608/1608
Full 4096-sample workload, 5 chunkings	60/60
Shared-parameter repeat and ordered zero-padding, B=1/255/256/257/2048	84/84
Parameter-gradient cancellation screen	8/8

Accuracy (strict FP64 acceptance; this kernel vs the batch-invariant GEMMs):

Batch invariance (raw-byte mismatches per candidate):

Performance (forward + backward, all candidates; and this kernel vs native PyTorch at B=1):

Peak memory at B=1 is 111.6 MiB vs 108.1 MiB for native PyTorch. 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 a73c66c5)

Measured at a73c66c5e2db7ec844e30ac763ce85d410179259 on H100-94GB (SM90). Raw logs and JSON are on pr_assets under pr537/hopper/ (commit 42eb4190).

fixed/acceptance.json is qualified with zero failures (1808 checks):

Gate	Result
Real-weight accuracy vs CPU FP32 and FP64, seeds 0/127/2026, BF16/FP32, B=1/3/256/2048	48/48
Batch-invariance records (raw bytes)	1608/1608
Full 4096-sample workload, 5 chunkings	60/60
Shared-parameter repeat and ordered zero-padding, B=1/255/256/257/2048	84/84
Parameter-gradient cancellation screen	8/8

Within-H100 repeat/batch invariance only (BF16 mma.sync); raw bytes are not claimed equal to B200 (RFC #428 §2.2). Model: Qwen/Qwen-Image @ 75e0b4be04f60ec59a75f475837eced720f823b6. Environment: H100 (SM90), PyTorch 2.13.0+cu130, CUDA 13.0, driver 580.173.02, TF32 off, TORCH_CUDA_ARCH_LIST=9.0, no fast-math. Summary: SUMMARY.md, report.json, env_header.txt.

Status

No known failures at d0708cddd0587755b8fb77a211bce5d0f4df84a6. Evidence was produced from a clean checkout of the code commit c88f7c0fe35d6544e092a16097a094d323b35581, which is this PR's parent of the evidence commit.

Open question for the maintainers: keep this kernel, or reuse a batch-invariant GEMM? vLLM's, SGLang's and Megatron's persistent GEMMs are batch-invariant for this projection, pass every BF16 case, and are 6–18× faster at B ≥ 256. They need an FP32 VJP adapter (the GEMMs are forward-only), and they exceed the FP32 d_weight tolerance at B = 2048 (2 elements per stream). This kernel passes every case and is 2.6–3.4× faster at B = 1–3. Details are in Prior art & reuse decision above.

Re-run at head a73c66c5 (B200, torch 2.13). Tests: 101 passed, 0 failed, 0 skipped on two GPUs. Fixed-mode acceptance from a clean checkout of this commit: qualified, 1808 checks, 0 failures; tensor hashes of the real weights are identical to the earlier evidence. The reuse-survey modes that need vLLM / SGLang / Megatron sources were not rerun. Raw data: pr_assets pr537/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

101 passed, 0 failed, 0 skipped on 2× B200 (log); the same command with TORCH_CUDA_ARCH_LIST=9.0 on 2× H100: 101 passed (log). Two GPUs are needed for the device-isolation test. The tests cover:

• CPU-reference accuracy of all gradients.
• Batch position and repeat bytes, and same-order parameter contributions.
• Noncontiguous and invalid inputs, readback, registry/gtest, and the fail-closed missing-extension and fast-math paths.

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: isolated Phase A operator with first-order autograd. No LayerNorm/modulation, gated residual, full-model, LoRA or distributed claim. Block-0 weights of both streams were used, not all 60 blocks.
• Performance: the unfused, correctness-first kernel is slower than native PyTorch at B=1 and much slower at large B (see the reuse question above).
• Extension: requires a build without fast-math; a missing or stale extension raises instead of falling back.

[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: e51b09e5-83a3-4009-b7dd-7cb269b24384

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

[Flink-ddd]

FAST_MATH doesn't reliably detect NVCC's --use_fast_math, so this can incorrectly report a strict build. Please enforce this in the build configuration and test an actual fast-math build; the current test only mocks the result.

[fusheng-ji]

Thanks, you are right: nvcc defines neither __FAST_MATH__ nor __USE_FAST_MATH__ for device code (I checked with CUDA 13 and --use_fast_math), so the old check always reported a strict build. Fixed in 274b6c5: qwenimage_strict_math_compiled() now runs a small probe kernel in the same translation unit, so it sees the real flags. It checks that a subnormal survives a multiply and that one division matches IEEE, and both change under --use_fast_math. The new tests/models/qwenimage/test_strict_math_build.py compiles the real det_linear.cu with and without --use_fast_math and asserts False/True.

[github-actions]

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

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

⚠️ 5 warning(s) occurred during review.

---

📄 rl_engine/runtime/registry.py (L800-L802)

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

On cuda the strict CUDA op is the only candidate, but its constructor (QwenImageAdaLNModProjectionCudaOp.__init__ -> require_extension()) raises on builds without the new _C symbols or with fast math. _get_or_create_backend swallows that exception and records the backend in _failed_backends, so get_op ends with RuntimeError("No functional backend found for ... on cuda") — the diagnostic no longer mentions the actual cause (rebuild _C / fast-math build), unlike the direct constructor path which says "rebuild rl_engine._C ...". Consider either letting these strict load/instantiation failures propagate, or including _failed_backends/the captured reason in the get_op error for this op type so a strict-operator misconfiguration is diagnosable.

---

📄 rl_engine/backends/cuda/model_specific/qwenimage/adaln_mod_projection.py (L39-L46)

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

aligned() assumes a 16-byte alignment unit, but check_matrix in det_linear.cu only requires the start pointer to be 16-byte aligned and the row length K to be a multiple of 16 / element_size(). For a float32 tensor, each row test uses a fixed 16-byte vector stride (kStride = 32 * kVec elements), so a view that is contiguous but starts at a storage offset that is 16-byte aligned while its row-to-row distance would need to stay consistent is fine — however, for a BF16 tensor whose storage dtype is BF16 but which is passed to the kernel via float() casts in projection_backward this is easy to get wrong. Concretely: grad = aligned(grad.float()) in projection_backward produces a float32 tensor whose element count is unchanged, so if grad was a misaligned BF16 view, alignment is re-checked correctly there. The fragile case is temb: projection_backward uses the saved (original, possibly misaligned/offset) temb directly with torch.sigmoid, which is fine numerically. The alignment helper itself is sound; the only remaining concern is that the copy branch drops autograd history (.clone() on a leaf-with-grad tensor inside forward under no_grad semantics of Function.forward is fine) — so no action needed. Please double-check that aligned(weight) inside projection_backward is not used for the returned dw (it is not).

---

📄 rl_engine/_C.pyi (L342-L347)

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

The stub advertises out_dtype, but the bound C++ signature (csrc/bindings/ops.cpp / det_linear.cu) names the third parameter w_dtype, and pybind registers it without explicit py::arg names. Callers that rely on the stub and pass out_dtype= as a keyword will fail at runtime with an unexpected-keyword error. Rename the stub parameter to w_dtype (or add py::arg names to the pybind registration).

💡 Suggested Change

Before:

def qwenimage_det_linear_backward_weight(
    grad: torch.Tensor,
    x: torch.Tensor,
    out_dtype: torch.dtype,
    with_bias: bool,
) -> list[torch.Tensor]: ...

After:

def qwenimage_det_linear_backward_weight(
    grad: torch.Tensor,
    x: torch.Tensor,
    w_dtype: torch.dtype,
    with_bias: bool,
) -> list[torch.Tensor]: ...

---

📄 csrc/cuda/model_specific/qwenimage/det_linear.cu (L658-L664)

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

The db launch path differs from the dW dispatch above it: det_linear_dweight_kernel is launched for all dtype combinations, but the db kernel here is gated on w_dtype only. This part is fine, but note bias_blocks is computed from n_out in the host code — when with_bias is true and n_out is large that is correct; however the zero-k_in case is not short-circuited, so with k_in == 0 and n_out > 0 total is 0 and det_linear_dweight_kernel is still launched with blocks from grid_for(0, threads) (=1). It writes nothing and is harmless, but dw of shape [n_out, 0] has no elements — acceptable. No defect confirmed here.

---

📄 csrc/cuda/model_specific/qwenimage/det_linear.cu (L541-L544)

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

pre_act is allocated with x.options().dtype(at::kFloat) and, in the BF16 branch, written by det_linear_forward_bf16_mma_kernel only for n < n_out && t < rows. Indices outside that condition are never written, and since the buffer is torch::empty, any element in the covered range is fine — but the kernel's n < n_out / t < rows guards mean every element of the [rows, n_out] buffer is written exactly once, so no uninitialized reads. Verified, no issue.

---

📄 csrc/cuda/model_specific/qwenimage/det_linear.cu (L71-L78)

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

The probe hardcodes torch::kCUDA (current default device) and caches the verdict in a process-wide static, while every entrypoint calls check_matrix after creating a CUDAGuard on the input device. On multi-GPU runs where the current device is not the tensor's device, the probe launches on whatever device is current at first call and the cached verdict is then reused for all devices. Since the verdict is a property of the compiled binary (nvcc flags), the semantics happen to hold, but the launch can fail or contend on an unrelated device (e.g. when the default device is not visible/initialized). Use at::cuda::getCurrentCUDAStream()'s device consistently or probe on the device already selected by the guard.

---

📄 rl_engine/validation/operators/operator_inputs.py (L57-L57)

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

The shape name is hardcoded as f"{batch}x3072x18432" and does not reflect multiplier, whereas the actual input builder produces weight of shape (multiplier * 3072, 3072). If multiplier changes (e.g. a parameter-screening variant), the reported shape name will silently disagree with the tensors actually validated, mislabeling evidence. Derive the dimension from multiplier (e.g. f"{batch}x3072x{6*3072}" built from the same constant used in the builder) so the name cannot drift from the inputs.

---

📄 rl_engine/validation/operators/operator_inputs.py (L402-L402)

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

_qwenimage_inputs and _make_qwenimage_projection_inputs lack type annotations, unlike the surrounding builder helpers in this module (_make_rms_norm_inputs etc.). More importantly, the magic numbers 6 / 3072 are duplicated here and in operator_shape_name; extract them into module-level constants so the shape name and tensor construction stay in sync.

---

⚠️ 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/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/runtime/registry.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/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/runtime/registry.py" mid-review: used 607530 tokens exceeds budget 500000
• benchmarks/models/qwenimage/__init__.py (token_budget_reached): stopped dispatch: used 607530 tokens + group estimate 325640 = projected 933170 exceeds budget 500000
• build_tools/extensions.py,csrc/bindings/ops.cpp,csrc/cuda/model_specific/qwenimage/det_linear.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/validation/operators/operator_inputs.py,rl_engine/validation/operators/operator_specs.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.