[dandan20031031]

Qwen-Image rotates adjacent channel pairs using centered three-axis positions; the existing GPT-NeoX half-split RoPE cannot implement that contract. This adds the unfused CUDA multi_axis_rope row of #386, with FP32/FP16/BF16 BSHD inputs, fixed FP32 multiply/add order, input gradients, per-sample tables and actual-backend traces.

The Triton backend adapts FlashAttention's rotary kernel, preserving its BSD license and attribution. The adaptation reads independent table/input strides and disables contraction. Native CUDA supplies the same arithmetic without a Triton dependency. Position semantics are pinned to Diffusers@40f63690.

Both GPU paths support packed-QKV views without materialization, fake/autograd registration, higher input VJPs, CUDA Graph replay and fullgraph forward/backward. OP_SPECS, the standard input builder and a gradient adapter expose the operator to existing accuracy/invariance APIs.

The ordinary CLI and WS1 forward/gradient adapters share a centered image-grid prefix derived from qwen_image_positions. Full logical sample length and (sample_id, token_position) select coordinates consistently across packing, permutation, chunks and padding. Regressions cover those layouts, non-CUDA reference routing, load-error propagation and arbitrary head-dimension strides. Backward telemetry documents its CUDA Graph replay limit.

Validation

Two separately recorded environments use the same RTX 5080/SM120 and PyTorch 2.9.1+cu128:

• Ubuntu 24.04.5 on WSL2, Python 3.12.3, CUDA 12.8.93, GCC 13.3, Triton 3.5.1, driver 617.42.
• Windows, Python 3.13.9, CUDA 12.8/MSVC, triton-windows 3.5.1.post24; retained as historical evidence.

Latest WSL2 results:

Check	Result
Ordinary complete rl_engine._C wheel, all 10 configured translation units, and import	Passed
Operator suite, including independent accuracy, byte invariance, stride/no-copy, opcheck and fullgraph backward	105 passed, one two-GPU guard skip
CUDA/Triton WS1 forward/gradient accuracy and invariance	Both passed with explicit provenance
Related regressions	206 passed, seven Ascend skips
Repository layout with extension enforcement	37 passed, zero skips/failures
Fresh installed full wheel, outside checkout	Both backends' eager/compiled forward and input-gradient bytes exact
Actual CPU-only PyTorch 2.9.1+cpu package/reference check, Python 3.10.22 and 3.12.3	Each: 80 passed, 64 expected GPU skips
Pinned actual Diffusers symbols	30 cases passed; image/text tables exact
Compute Sanitizer	52 selected cases passed; zero errors
Changed-file formatting/static checks and strict docs	Passed locally

The complete WSL2 build used the normal wheel path with RL_KERNEL_ROPE_ONLY unset. The installed-package probe records the actual loaded full _C path and its fingerprint. The existing optional operator-only wheel profile remains available. Both wheel paths retain the third-party license.

Acceptance entry points and CI

ci/scripts/run_qwenimage_rope_ws1.sh requires a real NVIDIA GPU, Triton and the complete _C RoPE registration. It records binary/backend identity and allows only the two-device guard to skip.

The new Qwen-Image RoPE workflow covers PRs targeting test-qwenimage, with CPU contracts on Python 3.10/3.12, changed-file checks and strict documentation build. Hosted CI is awaiting upstream authorization (action_required). Both Python versions were checked locally with actual CPU-only PyTorch wheels, including installed-package reference forward/backward. These are not recorded as passing hosted CI results. GPU qualification uses the explicit gate.

Standalone benchmark

BF16 [1,4096,24,128], packed QKV, CUDA Graph forward/forward+backward medians in ms/call:

• Native CUDA: 0.05868/0.13079.
• Adapted Triton: 0.04731/0.11841.
• Original FlashAttention with only contraction disabled: 0.04722/0.11847.
• PyTorch complex baseline: 0.24754/0.51583.

The benchmark includes default FlashAttention, all three reference image grids plus a short case, both layouts, eager wall/graph timings and extra peak allocation. All 40 pre-timing comparisons passed; strict CUDA/Triton/no-FMA output and input-gradient bytes agree exactly. Forward extra peak at 4096 tokens is 24 MiB, the returned output size. These one-process measurements establish no advantage over FlashAttention or model-throughput speedup.

Scope

This qualifies Linux userspace/compiler/package behavior on WSL2/SM120. Bare-metal Linux, A100/H100, multi-GPU guards, QK-norm fusion, model/block assembly and trajectories remain unqualified. Building the complete library does not certify every native operator or the Dense Qwen3 chain.

Related Ascend work is #411. Its audited head 285c4dccab7e9048956584caf431946d82405c09 uses a different half-split/uncentered reference, so a shared interface requires semantic alignment.

Validation summary and reproduction: SM120 validation report. Detailed per-case logs, benchmark data and source/binary SHA256 manifests remain at immutable commit 84751473006401c19c1c8c17844e8ef8b1c44375: WSL2 full-extension archive, historical Windows archive. The final diff retains one summary instead of the raw run files. This addresses one platform row, not all of #386.

[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: 3ddf3b18-2e2f-43bf-b50d-0af245fa61a9

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

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

---

📄 csrc/cuda/rope/multi_axis_rope.cu (L40-L41)

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

The output index is derived from the linear thread counter i (y[2*i]), which is only correct when y is contiguous in [B,S,H,128] order — that holds. The real problem is that n = x.numel() / 2 and the decomposition pair = i % 64 / row = i / 64 mixed with head = row % heads, token = row / heads, batch = token / seq is only meaningful if the logical row-major order of x matches the flat counter, which it does — but the read side also assumes x.stride(3) separates even/odd elements (2*pair*strides.xd and xi + strides.xd). For a negative or zero strides.xd (e.g. a flipped/expanded head dim, which validate_inputs does not reject) the pair read collapses or aliases, silently producing wrong values instead of erroring. Consider rejecting strides.xd <= 0 in multi_axis_rope_apply.

---

📄 csrc/cuda/rope/multi_axis_rope.cu (L67-L67)

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

For the strongly-preferred case where B*S is large, 65535 cap is fine because of the grid-stride loop. However the loop step uses static_cast<int64_t>(blockDim.x) * gridDim.x — correct. The remaining issue is that 65535 is a pre-Volta limit; on modern devices gridDim.x supports 2^31-1, but capping at 65535 is still safe. No change required; flagging only that the cap makes the effective grid-stride path the common path for Qwen-Image (S ~ several thousand, H ~ 24-40 → n > 1e8/2), so the %// integer divisions per element dominate. Not a correctness defect.

---

📄 csrc/cuda/rope/multi_axis_rope.cu (L63-L63)

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

n = x.numel() / 2 is computed after y = at::empty(...), but the early return skips the launch only when x.numel()==0. Because x.size(2) > 0 is checked, a zero element count can still come from B==0 or S==0, which is handled. Fine. However note that x.size(1) (seq) is passed to the kernel as the divisor for token / seq; if x.size(1) == 0 but x.size(0) > 0, n is 0 and we return early, so no division by zero. OK.

---

📄 rl_engine/backends/shared/triton/rope/multi_axis_rope.py (L15-L15)

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

_apply unpacks x.shape as BSHD and forwards *x.stride() positionally to the kernel's stride_x_batch, stride_x_seqlen, stride_x_nheads, stride_x_headdim. The kernel only respects that ordering for a 4-D BSHD tensor. validate_inputs accepts any 4-D [B, S, H, 128], so shape itself is enforced, but transpose is forwarded from the caller while the launch still assumes the non-transposed layout — the CUDA path is called with transpose=False from MultiAxisRoPECudaOp, so both backends are consistent today, yet nothing in this wrapper rejects or adapts a transpose=True invocation of _rotation (e.g. via the exported _rotation used in test_dispatcher_fake_autograd_and_fullgraph_compile). Consider asserting not transpose (or documenting/handling the layout) here so a future transpose caller cannot silently get wrong-indexed access.

---

📄 rl_engine/backends/shared/triton/rope/multi_axis_rope.py (L30-L36)

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

For 3-D tables, cos.stride(0) is passed as the batch stride, and the kernel then computes COS + pid_batch * stride_cos_batch + .... validate_inputs accepts any 3-D table of shape [B, S, 64], which includes a non-contiguous / batch-broadcast table such as t.expand(B, -1, -1) whose stride(0) is 0 — that case is fine — but it also accepts a table whose batch stride is nonzero while batch was produced by a broadcast over a smaller first dim. More importantly, unlike the native kernel (cos.dim()==3 ? cos.stride(0) : 0 after the same checks), there is no guard that a 3-D table's first dim equals x.size(0); the shape check in validate_inputs does enforce (x.shape[0], x.shape[1], 64), so this matches. No defect proven; flagging only that 3-D shared tables rely entirely on that check.

---

📄 rl_engine/backends/cuda/rope/multi_axis_rope.py (L21-L23)

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

native_available() is called during backend selection (registry._get_or_create_backend, and directly by tests/tools). find_spec can resolve but library.origin may be None (namespace packages) or load_library may raise on an ABI-mismatched wheel; either path propagates an exception out of native_available() instead of returning False, so the registry's "backend could not be loaded → fall back to Triton/reference" logic never runs and get_multi_axis_rope_op('auto') fails hard. Wrap the load in try/except (ImportError, OSError, RuntimeError) and return False (keeping the hasattr re-check) so availability probing is total, matching the documented "conditional backend registers itself only when the stack loads" policy.

---

📄 rl_engine/ops/rope/_registration.py (L21-L26)

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

record_backward is called from inside the custom_op body but is labelled/monitored as a backward event, while this executes during the forward (transpose=True forward). Combined with the backward calling rotation(..., not ctx.transpose), this counts a backward only on forward-transpose calls and also fires during autograd replay, so telemetry attribution is inverted. It should be emitted from the backward/VJP path (or keyed on an explicit direction flag), not from the forward body.

---

📄 rl_engine/ops/rope/_registration.py (L29-L31)

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

register_fake calls validate_inputs, which rejects tensors with requires_grad cos/sin and inspects device/dtype. Under fake/meta tracing during torch.compile or export, x is a meta tensor with device='meta'; the x.device != cos.device check will pass only if tables were also placed on meta, otherwise this raises during tracing. The fake kernel should mirror shape/dtype constraints only, not the same runtime device/grad checks as the eager body.

---

⚠️ Warnings:

• csrc/bindings/rope_module.cpp,csrc/cuda/rope/multi_axis_rope.cu (review_limit_reached): reached the aggregate token budget before finishing; the grace round did not complete the task
• csrc/bindings/rope_module.cpp (token_budget_reached): stopped group "csrc/bindings/rope_module.cpp,csrc/cuda/rope/multi_axis_rope.cu" mid-review: used 577375 tokens exceeds budget 500000
• rl_engine/validation/operators/gradient_adapters.py (token_budget_reached): stopped dispatch: used 577375 tokens + group estimate 130140 = projected 707515 exceeds budget 500000
• rl_engine/backends/cuda/rope/multi_axis_rope.py,rl_engine/backends/shared/triton/rope/FLASH_ATTENTION_LICENSE,rl_engine/backends/shared/triton/rope/_flash_attn_rotary.py,rl_engine/backends/shared/triton/rope/multi_axis_rope.py (review_limit_reached): reached the aggregate token budget before finishing; the grace round did not complete the task
• rl_engine/ops/rope/_registration.py,rl_engine/ops/rope/multi_axis_rope.py,rl_engine/reference/rope/multi_axis_rope.py,rl_engine/runtime/registry.py (review_limit_reached): reached the aggregate token budget before finishing; the grace round did not complete the task

[github-actions]

@dandan20031031 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.

[dandan20031031]

OCR review follow-up in 9e2e50f.

The CLI and WS1 adapters now share canonical centered image-grid prefixes keyed by complete logical sample/token identity. New regressions cover physical layouts, non-CUDA routing, load failures and head-dimension strides. The validation summary records the new runs separately from the initial benchmark measurements.

Local validation on WSL2 / RTX 5080 / PyTorch 2.9.1+cu128	Result
Rebuilt ordinary complete CUDA wheel, normal installation, fresh-process installed-package eager/fullgraph forward and input gradients	Passed
Operator suite	105 passed, one two-GPU guard skip
CUDA/Triton WS1 four judgments and ordinary CLI checks	Both passed with backend provenance
Related regressions	206 passed, seven Ascend skips
Repository layout with extension enforcement	37 passed
Compute Sanitizer, expanded to include the new stride cases	52 passed, zero errors
Actual CPU-only torch 2.9.1+cpu, Python 3.10.22 and 3.12.3	Each: 80 passed, 64 expected GPU skips; installed reference forward/backward passed
Changed-file checks and strict documentation build	Passed

Disposition of the nine findings routed to the summary:

1. Gradient-adapter coordinates: addressed by the shared grid-prefix helper and the packed/permuted/chunked/padded regressions. Full logical sequence lengths, rather than current physical span lengths, determine the grid.
2. Shared-table broadcasting: [S,1,64] right-aligns as [1,S,1,64] against [B,S,H,64]. The accuracy oracle now explicitly expands shared tables to a batch dimension, while the candidate still consumes the shared table; B=2/S=7 coverage passes.
3. Backward telemetry placement: retained inside the opaque custom-op implementation, outside fake/AOT tracing. Documentation and source comments clarify that Python execution_count does not count CUDA Graph replays. Fullgraph forward/backward and graph replay tests pass.
4. Recursive backward / gradient layouts: transpose toggling implements the appropriate adjoint at each input VJP; backward is only requested by autograd when needed. Higher-VJP tests pass, and the added stride tests also use noncontiguous upstream gradients.
5. Library-load errors: the cited function does not catch load_library errors. New failure-injection tests verify that both OSError and RuntimeError propagate unchanged.
6. CUDA strided indexing: each logical pair is read at 2*p*xd and +xd, then written in contiguous logical order. Added FP64-oracle forward/VJP tests cover last-dimension stride 2 and last-dimension outermost storage for three dtypes and all three implementations, including 12 CUDA/Triton cases.
7. 64/128 constants: accepted; native indexing and shape checks now share named kHeadDim / kPairsPerHead constants. The complete CUDA wheel was rebuilt and requalified after this change.
8. Explicit operator-only profile: fail-closed behavior is intentional and documented. Existing build-profile tests reject missing Torch and CPU/ROCm Torch for an explicitly requested CUDA operator wheel; those tests pass in both CPU matrix versions.
9. License packaging: reopened the newly built full wheel and confirmed rl_engine/backends/shared/triton/rope/FLASH_ATTENTION_LICENSE contains the BSD text. The wheel also contains the new validation helper. Wheel SHA256: 9feb8200b948021c268c933f93efcd424b6b68d6c55bc24c3004d6336deb3a13.

The two inline findings have separate replies with implementation/test references. Hosted CPU CI still awaits upstream workflow authorization; the local CPU runs above are not presented as passing hosted CI. Bare-metal Linux, A100/H100, two-GPU execution and model/trajectory integration remain unqualified.

[dandan20031031]

Follow-up on OCR run 40717454691-1 (the refreshed nine-finding summary).

The inline VJP finding has a separate reply with source references and 12 additional independent FP64-oracle output/VJP checks. The eight summary findings were checked as follows; no further implementation change was indicated by these reproductions.

Finding	Verification
Zero/negative last-dimension strides	An expanded tensor with stride(-1)==0 has well-defined repeated logical values. Six added local FP32/FP16/BF16 CUDA/Triton forward/VJP probes, also using batch-broadcast tables, agree with the FP64 oracle. On this PyTorch 2.9.1 installation, as_strided rejects negative strides before the operator is called. Requiring positive strides would unnecessarily exclude a valid read-only expanded input.
Grid cap / integer division	The grid-stride cap is safe, as the comment notes. At B=1/H=24, the three recorded image grids need 24,576, 41,334 and 36,192 blocks, all below 65,535. No new performance claim or launch-policy change is made.
Empty input / sequence divisor	The zero-element return precedes launch. Existing B=0/S=0 forward/backward cases pass in the 105-pass operator run.
Triton transpose=True	This flag conjugates the rotation, not the input memory layout; it reaches CONJUGATE and negates sine. Both flag values pass independent forward/VJP checks for every GPU backend and dtype.
Batched table shape/stride	Validation requires exactly [B,S,64] for 3-D tables. Batch-broadcast stride-zero tables pass the additional six probes; the review itself notes the shape check is present.
Automatic fallback on library-load errors	KernelRegistry._get_or_create_backend catches construction exceptions and lets selection proceed. I injected ImportError, OSError and RuntimeError from the native availability probe, then called the actual get_multi_axis_rope_op('auto', device='cuda') path with a fresh registry: all three selected the Triton backend. Explicit CUDA selection continues to preserve the original load error.
Backward provenance	The public forward always calls the primitive with transpose=False; its first input VJP calls transpose=True, which records the declared transpose-kernel provenance. The primitive is an internal shared forward/adjoint implementation. The documentation states that its Python count is not a CUDA Graph replay counter. Both backend WS1 gates, recursive input-VJP and fullgraph/graph tests pass.
Fake/meta device validation	FakeTensor retains a logical device; in these CUDA traces x/cos/sin remain logically cuda:0, although storage is meta-backed. Explicit FakeTensorMode calls and torch.export.export followed by exported-module execution pass for both backends. Existing opcheck/fullgraph cases also pass. Retaining the same device/fixed-table constraints keeps fake and eager contracts consistent.

These follow-up probes used the same validated source snapshot published as 9e2e50f. The committed validation report distinguishes the initial measurements from the new local qualification results; hosted CPU CI is still awaiting upstream authorization.