GPU rules for `vcat`, `cat`, `hcat`, `permutedims`. by AstitvaAggarwal · Pull Request #1225 · chalk-lab/Mooncake.jl

AstitvaAggarwal · 2026-06-23T18:38:10Z

CI Summary — GitHub Actions

Documentation Preview

Mooncake.jl documentation for PR #1225 is available at:
https://chalk-lab.github.io/Mooncake.jl/previews/PR1225/

Performance

Performance Ratio:
Ratio of time to compute gradient and time to compute function.
Warning: results are very approximate! See here for more context.

┌────────────────────────────┬──────────┬──────────┬─────────────┬─────────┬─────────────┬────────┐
│                      Label │   Primal │ Mooncake │ MooncakeFwd │  Zygote │ ReverseDiff │ Enzyme │
│                     String │   String │   String │      String │  String │      String │ String │
├────────────────────────────┼──────────┼──────────┼─────────────┼─────────┼─────────────┼────────┤
│                   sum_1000 │ 180.0 ns │     1.62 │        1.72 │   0.667 │        3.28 │   6.23 │
│                  _sum_1000 │ 971.0 ns │      6.8 │        1.01 │  1630.0 │        41.6 │   1.06 │
│               sum_sin_1000 │  7.23 μs │     3.51 │        1.33 │    1.46 │        10.9 │   1.72 │
│              _sum_sin_1000 │  6.58 μs │     3.46 │        1.72 │   211.0 │        12.2 │   1.98 │
│                   kron_sum │ 219.0 μs │     12.8 │        3.26 │    8.21 │       363.0 │   18.6 │
│              kron_view_sum │ 301.0 μs │     10.7 │        4.92 │    23.7 │       308.0 │   11.7 │
│      naive_map_sin_cos_exp │  2.39 μs │     2.99 │        1.49 │ missing │         7.5 │   2.08 │
│            map_sin_cos_exp │  2.28 μs │      3.7 │        1.68 │    1.56 │         6.7 │   2.69 │
│      broadcast_sin_cos_exp │  2.34 μs │     3.15 │        1.67 │    4.19 │        1.46 │   2.13 │
│                 simple_mlp │ 375.0 μs │     4.51 │        2.52 │     1.4 │        8.24 │   2.77 │
│                     gp_lml │ 389.0 μs │     5.61 │        1.66 │    4.69 │     missing │   3.25 │
│ turing_broadcast_benchmark │  1.66 ms │     6.94 │        3.58 │ missing │        43.5 │   2.36 │
│         large_single_block │ 501.0 ns │     7.24 │        1.72 │  3930.0 │        26.7 │   1.72 │
└────────────────────────────┴──────────┴──────────┴─────────────┴─────────┴─────────────┴────────┘

…x,NNlib,Lux

yebai · 2026-06-23T19:00:18Z

        SkipConnection(Dense(2 => 2), vcat),
        randn(Float32, 2, 3),
        "SkipConnection(Dense(2 => 2), vcat)",
-    ),  # vcat tangent mixes CPU/GPU: Illegal conversion of CUDA.DeviceMemory to Ptr{Float32}


The rules for vcat/hcat/cat on CuArrays are left unimplemented on purpose — if the underlying issue has been fixed, let's investigate what changed.

So I tested further by disabling the rules, the three functions being diffed fail as the following places:

vcat -> Base._typed_vcat -> CUDAKernels.##_#6 -> kernel_compile -> cufunction (contains Base.@lock (try/finally)

hcat -> Base._typed_hcat -> CUDAKernels.##_#6 -> kernel_compile -> cufunction

cat -> Base._copy_or_fill! -> CUDAKernels.##_#6 -> kernel_compile -> cufunction

The root cause is Base.@lock inside cufunction (used for thread-safe kernel caching), which Mooncake can't handle in reverse mode due to the try/finally block. I tried adding a @zero_derivative rule for cufunction (to skip past the @lock), and then prepare_cuda_state (skip cuCtxGetCurrent ccall inside). After both, CUDAKernels.##_#6 itself then hits cuOccupancyMaxPotentialBlockSize (raw ccall). (this continues)

The explicit rules at the vcat/hcat/cat level might be the right intervention point. But in case Im seeing something incorrectly, happy to be corrected to the right approach ahead.

for clarity, I'm of the understanding that the underlying issue was handling the relevant CUDA ccalls along these paths. Is it preferred to mark all these calls as non differentiable/with their respective rules?

codecov · 2026-06-23T19:14:21Z

Codecov Report

❌ Patch coverage is 89.51613% with 13 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
ext/MooncakeCUDAExt/MooncakeCUDAExt.jl	89.51%	13 Missing ⚠️

📢 Thoughts on this report? Let us know!

sunxd3 · 2026-06-29T12:25:35Z

+        _throw_mixed_cat_error($_fn)
+    end
+    @eval function rrule!!(
+        ::CoDual{typeof($_fn)}, ::CoDual{<:CuMaybeComplexArray}, ::CoDual{<:AbstractArray}


this only cover the two argument case

sunxd3 · 2026-06-29T12:29:52Z

+function rrule!!(
+    ::CoDual{typeof(Core.kwcall)},
+    kw::CoDual{<:NamedTuple},
+    ::CoDual{typeof(cat)},
+    args::CoDual{<:CuMaybeComplexArray}...,
+)
+    pkw = primal(kw)
+    dim = pkw.dims
+    primals = map(primal, args)
+    fdatas = map(tangent, args)
+    out = cat(primals...; pkw...)
+    dy_out = zero(out)
+    pb!!(::NoRData) = (
+        _cu_concat_pb!(fdatas, dy_out, dim);
+        (NoRData(), NoRData(), NoRData(), map(_ -> NoRData(), args)...)
+    )
+    return CoDual(out, dy_out), pb!!
+end


cat's dim argument could take Tuple and Val, this assumes dims to be a single integer?

sunxd3 · 2026-06-29T12:36:00Z

    end
 end

+@is_primitive(MinimalCtx, Tuple{typeof(vcat),Vararg{CuMaybeComplexArray}})


old error message functions would match Numbers, consider add them back for better error message?

sunxd3 · 2026-06-30T06:54:35Z

+# Vararg{CuMaybeComplexArray} rules above, so pure-GPU calls still dispatch to those rules.
+# For mixed GPU/CPU calls _throw_if_any_cu detects the CuArray at any position and throws.
+# For pure-CPU calls the rules compute correct pullbacks using the same _cu_concat_pb! helper.
+@is_primitive(MinimalCtx, Tuple{typeof(vcat),Vararg{AbstractArray}})


I am not sure AbstractArray is a good idea: could you verify view, adjoints etc. not breaking on CPU with the extension loaded?

The PR is still WIP, will request a review when ready

changed the PR to draft as well for this, I think this needs more thought. sorry for the delay

…max JET failures

- Restore the mixed CPU/GPU device guard (deleted in b26d1a0), this time only intercepting genuinely mixed calls, so pure-CPU/pure-GPU are unaffected. - Support cat(...; dims=(1,2)) (block-diagonal), not just Integer/Val. - Support Adjoint/Transpose/SubArray wrapping a CuArray. - Support Integer/Bool-eltype CuArrays. - Add matching tests.

@generated

…A is loaded - Replace @generated _is_primitive with a plain, defensive method (handles Vararg/UnionAll call-site types safely); widen it to also cover GPU+scalar mixing at any arity, fixing the N-arg regression and dropping ~18 now-dead methods. - Exclude Float16 from CuMaybeWrappedArray's SubArray branch (arrayify doesn't support it); reword the error message; fix a zero-arg edge case. - Restore deleted flux.jl comments; add a regression test; bump version.

sunxd3

Great work, we are really close.

Two smaller items left worth fixing:

1. `CuMaybeWrappedArray` declares `ComplexF16` wrappers the rule can't handle

ext/MooncakeCUDAExt/MooncakeCUDAExt.jl:127-131

The Adjoint/Transpose arms use CuFloatOrComplex, which includes Complex{Float16}:

Adjoint{<:CuFloatOrComplex,<:CuMaybeComplexArray},
Transpose{<:CuFloatOrComplex,<:CuMaybeComplexArray},

But these rules canonicalise through arrayify, whose Adjoint/Transpose methods
(src/rules/blas.jl) are bounded Union{IEEEFloat,BlasFloat}, and ComplexF16 is not in
that set. So Adjoint{ComplexF16}/Transpose{ComplexF16} — and bare CuArray{ComplexF16}
(CuMaybeComplexArray on L128) — are declared primitives but then fail inside the rule
(MethodError from arrayify, or a kernel exception), instead of falling through to the
interpreter like other unsupported cases.

Fix: bound the complex arms to ComplexF32/ComplexF64 so the primitive declaration matches
what the rule can actually canonicalise, mirroring the BlasFloat bound already used on the
SubArray arm (L131).

Also fix the comment at L119-126, which is inaccurate on two points:

The example vcat(view(cu_f16_vec, 1:4), y) — a 1-D contiguous view of a CuArray
returns a CuArray, not a SubArray, so it is unaffected by the SubArray bound and
works. Use a strided/2-D view (e.g. view(cu_f16_mat, 1:2, :)) to illustrate the point.
"excludes Float16/Complex{Float16}, unlike its Adjoint/Transpose/Symmetric siblings" —
the siblings do handle Float16 but not Complex{Float16} (same Union{IEEEFloat, BlasFloat} bound). Tightening the complex arms above makes this consistent.

2. The `f_view` test does not exercise what it claims

test/ext/cuda/cuda.jl:1325-1327

# Float16 SubArrays are excluded from CuMaybeWrappedArray (Finding 3);
# confirm this now falls through cleanly instead of erroring.
f_view(x, y) = sum(vcat(view(x, 1:2), y))

view(x16, 1:2) on a 1-D contiguous CuArray returns a CuArray, not a SubArray, so this
hits the normal pure-GPU vcat rule and duplicates the "Bare Float16" test just above it
(L1314). It never touches the SubArray-exclusion path the comment describes.

Fix: to test the exclusion, use a strided view that stays a SubArray
(e.g. view(cu_f16_mat, 1:2, :)) and assert it errors (@test_throws); otherwise remove
the test as redundant.

sunxd3

LG, thanks!

sunxd3 · 2026-07-09T10:11:53Z

Don't release yet until we look into the CI failures
c.f. TuringLang/Bijectors.jl#482

AstitvaAggarwal and others added 3 commits June 23, 2026 19:33

multihead attention, batchnorm, and skipconnection models for GPU flu…

c41d201

…x,NNlib,Lux

.

e5679ba

Merge branch 'main' into cuda-gpu-rules

95c8b4c

github-actions Bot assigned AstitvaAggarwal Jun 23, 2026

fix comment

1272219

yebai reviewed Jun 23, 2026

View reviewed changes

AstitvaAggarwal and others added 2 commits June 24, 2026 16:02

remove nnlib ext

42966cd

Merge branch 'main' into cuda-gpu-rules

b19ea44

AstitvaAggarwal changed the title ~~GPU rules for vcat, cat, hcat, permutedims and NNLib.batchnorm~~ GPU rules for vcat, cat, hcat, permutedims. Jun 24, 2026

AstitvaAggarwal added 3 commits June 26, 2026 13:33

mixed case guards

d5d6a41

format, frule guards

d25f8b8

more tests

54b1fcf

sunxd3 reviewed Jun 29, 2026

View reviewed changes

changes from reviews, mixed device checks for N args

f3df3ce

sunxd3 reviewed Jun 30, 2026

View reviewed changes

AstitvaAggarwal marked this pull request as draft June 30, 2026 11:31

Remove overly-broad vcat/hcat/cat AbstractArray rules; fix NNlib soft…

b26d1a0

…max JET failures

AstitvaAggarwal force-pushed the cuda-gpu-rules branch from 94960b3 to b26d1a0 Compare June 30, 2026 12:03

AstitvaAggarwal and others added 9 commits June 30, 2026 14:26

Merge branch 'main' into cuda-gpu-rules

6e191c7

rebase

03c1c02

Merge branch 'main' into cuda-gpu-rules

7ec09bf

trim int handling

4c546d7

Merge branch 'main' into cuda-gpu-rules

ff5907f

Merge branch 'main' into cuda-gpu-rules

f945b13

more tests fro codecov

27e52bd

Merge branch 'main' into cuda-gpu-rules

798a6cb

AstitvaAggarwal marked this pull request as ready for review July 7, 2026 09:35

AstitvaAggarwal requested a review from sunxd3 July 7, 2026 09:35

AstitvaAggarwal added 4 commits July 7, 2026 14:23

Add review-requested nice-to-have test coverage

ad85d19

format

29163c0

trim comments

61d5fc5

sunxd3 approved these changes Jul 8, 2026

View reviewed changes

AstitvaAggarwal and others added 3 commits July 8, 2026 22:18

changes from reviews

eaff764

Merge branch 'main' into cuda-gpu-rules

5cb3117

Merge branch 'main' into cuda-gpu-rules

2cd64f1

AstitvaAggarwal requested a review from sunxd3 July 9, 2026 09:07

sunxd3 approved these changes Jul 9, 2026

View reviewed changes

sunxd3 merged commit 71b34d3 into main Jul 9, 2026
94 of 97 checks passed

sunxd3 deleted the cuda-gpu-rules branch July 9, 2026 10:12

Uh oh!

Conversation

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

Uh oh!

CI Summary — GitHub Actions

Documentation Preview

Performance

Uh oh!

yebai Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AstitvaAggarwal Jun 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sunxd3 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

1. CuMaybeWrappedArray declares ComplexF16 wrappers the rule can't handle

2. The f_view test does not exercise what it claims

Uh oh!

sunxd3 left a comment

Choose a reason for hiding this comment

Uh oh!

sunxd3 commented Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

AstitvaAggarwal commented Jun 23, 2026 •

edited by github-actions Bot

Loading

yebai Jun 23, 2026 •

edited

Loading

AstitvaAggarwal Jun 24, 2026 •

edited

Loading

codecov Bot commented Jun 23, 2026 •

edited

Loading

sunxd3 left a comment •

edited

Loading

1. `CuMaybeWrappedArray` declares `ComplexF16` wrappers the rule can't handle

2. The `f_view` test does not exercise what it claims

sunxd3 commented Jul 9, 2026 •

edited

Loading