CliMA · tapios · Jul 14, 2026 · Jul 12, 2026 · Jul 12, 2026 · Jul 12, 2026
diff --git a/NEWS.md b/NEWS.md
@@ -3,6 +3,36 @@ RRTMGP.jl Release Notes
 
 main
 ------
+- Internal compute buffers now use device-dependent physical layouts behind a
+  uniform column-first `(ncol, nlay/nlev)` indexing convention: optical
+  properties (`OneScalar`, `TwoStream`), source functions (`SourceLWNoScat`,
+  `SourceLW2Str`, `SourceSW2Str`), and the broadband flux buffers (`FluxLW`,
+  `FluxSW`). On the GPU the storage is column-first, giving coalesced access
+  (on an A100 at DYAMOND scale this cuts longwave kernel times by 25-36% and
+  shortwave all-sky by 14-20%); on the CPU the storage stays vertical-first
+  under a lazy `PermutedDimsArray` wrapper, preserving the stride-1 vertical
+  sweeps (and CPU performance) while running the same kernel code.
+  The Layer-2 flux getters (`lw_flux_up`, `net_flux`, ...) still present
+  plain `(nlev, ncol)` views: they read `Fluxes.FluxPresentation` arrays that
+  `update_fluxes!` fills from the compute buffers with one fused transposing
+  copy per call (so getters remain `Array`-materializable, broadcastable, and
+  reducible on the GPU). The opt-in per-band buffers (`FluxBand`) keep the
+  host-facing `(nlev, ncol, n_bnd)` layout. **Breaking for Layer-1
+  (functional-core) users**: code reading the workspace fields directly
+  (`slv.flux.flux_up`, `op.τ`, ...) must swap index order to `[gcol, ilev]`
+  or wrap the buffer in `Fluxes.lazy_transpose`. State inputs
+  (`AtmosphericState` arrays) and all boundary-condition arrays are
+  unchanged.
+- New `TransposedStateCache`: column-first copies of the hot
+  `AtmosphericState` arrays (layer pressure/temperature/dry-air column
+  amount/relative humidity, level temperatures), refreshed by one
+  `permutedims!` per spectral solve, so the gas-optics g-point loop reads
+  coalesced memory on GPUs. On by default on the GPU (the RTE workspace
+  constructors build one; the `RRTMGPSolver` shares a single cache between
+  the longwave and shortwave workspaces) and off (`nothing`) on the CPU,
+  where the state's vertical-first layout is already cache-friendly; pass
+  `state_cache = nothing` to opt out, or pass the same cache to several
+  workspaces to share the storage.
 - Documentation reorganized along tutorial / how-to / explanation / reference
   lines: two new executable (Literate.jl) tutorials — "A first radiation
   calculation" and "Radiative-convective equilibrium", which reproduces

diff --git a/docs/src/functional_core.md b/docs/src/functional_core.md
@@ -32,9 +32,18 @@ workspace's flux buffers in place, and you can read the results:
 
 ```julia
 solve_lw!(slv_lw, state, lookup_lw)   # or solve_lw!(slv_lw, state) for gray
-F = slv_lw.flux.flux_net              # (nlev, ncol) [W/m²]
+F = slv_lw.flux.flux_net              # (ncol, nlev) [W/m²]
 ```
 
+Note the axis order: the raw flux buffers are indexed column-first, `(ncol,
+nlev)`, on every device. On the GPU they are also stored that way, so that
+neighboring threads (one per column) access consecutive memory; on the CPU
+the storage is vertical-first under a lazy wrapper, which keeps the
+per-column sweeps stride-1 without changing the indexing. Only the Layer-2
+getters (`net_flux`, `lw_flux_up`, ...) present the vertical-first
+`(nlev, ncol)` orientation, from presentation copies that `update_fluxes!`
+refreshes at the end of each call.
+
 This is what RRTMGP's own tests drive, and what [`RRTMGPSolver`](@ref
 RRTMGP.RRTMGPSolver) wraps behind [`update_fluxes!`](@ref RRTMGP.update_fluxes!)
 and the named getters.
@@ -81,10 +90,10 @@ slv_sw = TwoStreamSWRTE(grid_params; cos_zenith, toa_flux,
 solve_lw!(slv_lw, as)
 solve_sw!(slv_sw, as)
 
-lw_net = slv_lw.flux.flux_net   # (nlev, ncol) [W/m²]
+lw_net = slv_lw.flux.flux_net   # (ncol, nlev) [W/m²]
 sw_net = slv_sw.flux.flux_net
 net = lw_net .+ sw_net
-net[end, 1]                     # net flux at the top of the atmosphere [W/m²]
+net[1, end]                     # net flux at the top of the atmosphere [W/m²]
 ```
 
 The one-liner [`solve_gray`](@ref RRTMGP.solve_gray) wraps this.

diff --git a/docs/src/getters.md b/docs/src/getters.md
@@ -53,9 +53,13 @@ F = RRTMGP.net_flux(solver)                      # (nlev, ncol) view of the resu
 With ClimaCore, this forms the bridge: `array2field(getter(solver), space)`
 wraps an input view as a `Field` (writes flow into the solver's buffer), and
 `field2array(dst) .= getter(solver)` copies an output into a destination field.
-On the GPU, the returned views are strided `SubArray`s of device arrays; to
+On the GPU, the returned views are plain `SubArray`s of device arrays; to
 bring one to the CPU, use `Array(getter(solver))` (or the element-wise `.=`
-copy above).
+copy above). One subtlety: the solver's compute buffers are indexed
+column-first, `(ncol, nlev)` (with the physical layout chosen per device for
+performance), and `update_fluxes!` ends by copying them into the
+`(nlev, ncol)` presentation arrays the flux getters expose — so read fluxes
+after `update_fluxes!`, not between manual Layer-1 solves.
 
 ## Responsibilities: what the host provides vs what RRTMGP derives
 

diff --git a/docs/src/howto/gpu.md b/docs/src/howto/gpu.md
@@ -50,7 +50,11 @@ outside hot loops.
 `Adapt.adapt(Array, solver)` produces a host copy (e.g., for serialization) and
 `Adapt.adapt(CuArray, solver)` moves it back. RRTMGP's custom `adapt_structure`
 methods preserve the internal view topology (scratch views into packed
-buffers), so an adapted solver keeps the zero-allocation property.
+buffers), so an adapted solver keeps the zero-allocation property. Adapt
+preserves each buffer's physical layout, so round-tripping a GPU solver
+through the host is lossless; adapting a CPU-*constructed* solver to the GPU
+runs correctly but keeps the CPU layout (uncoalesced) — for GPU work,
+construct the solver on the GPU.
 
 ## Reproducibility caveat: McICA cloud sampling
 
@@ -78,6 +82,18 @@ both backends:
   g-point's lookup-table slices hot in cache while sweeping the columns. With
   a single-threaded context, the same code runs serially.
 
+The memory layout matches the thread mapping on each device. On the GPU, the
+internal scratch and output buffers are stored column-first, `(ncol, nlev)`,
+so consecutive threads write consecutive addresses (coalesced access), and
+the frequently re-read state arrays (layer pressures and temperatures, level
+temperatures) are copied into a column-first `TransposedStateCache` once per
+solve so the g-point loop also reads coalesced memory. On the CPU, the same
+buffers are stored vertical-first under a lazy wrapper — each thread's
+vertical sweep is then stride-1 — and the kernels read the
+`AtmosphericState` directly, so both devices run identical kernel code on
+their preferred layout. The caller-owned `AtmosphericState` and the getter
+views keep their vertical-first `(nlev, ncol)` presentation throughout.
+
 ## Performance expectations
 
 CI benchmarks the kernels on GPUs (see `perf/benchmark_ratchet.jl`
@@ -105,20 +121,20 @@ Every configuration evaluated by the benchmark (see [`perf/benchmark_ratchet.jl`
 
 ### Single precision (`Float32`) — Primary target
 
-At `Float32` precision, memory traffic across lookup-table interpolations is halved compared to `Float64`, enabling large GPU speedups:
+At `Float32` precision, memory traffic across lookup-table interpolations and solver evaluations is halved compared to `Float64`, enabling large GPU speedups:
 
 | Configuration | CPU (Intel Xeon, 12 threads) | GPU (NVIDIA A100 40GB) | Speedup |
 |---|---|---|---|
-| Clear-sky (two-stream) | 77.95 s | 1.24 s | 63.1× |
-| All-sky, McICA clouds | 98.42 s | 1.82 s | 54.1× |
-| All-sky with aerosols | 99.39 s | 1.95 s | 51.0× |
+| Clear-sky (two-stream) | 80.05 s | 0.95 s | 84.3× |
+| All-sky, McICA clouds | 99.44 s | 1.40 s | 71.0× |
+| All-sky with aerosols | 106.77 s | 1.56 s | 68.4× |
 
 ### Double precision (`Float64`)
 
 For reference, running the identical uniform benchmark sweep in double precision (`Float64`) yields:
 
 | Configuration | CPU (Intel Xeon, 12 threads) | GPU (NVIDIA A100 40GB) | Speedup |
 |---|---|---|---|
-| Clear-sky (two-stream) | 79.64 s | 1.70 s | 46.8× |
-| All-sky, McICA clouds | 101.60 s | 2.43 s | 41.8× |
-| All-sky with aerosols | 106.33 s | 2.59 s | 41.1× |
+| Clear-sky (two-stream) | 81.88 s | 1.30 s | 63.0× |
+| All-sky, McICA clouds | 103.97 s | 1.92 s | 54.2× |
+| All-sky with aerosols | 108.38 s | 2.07 s | 52.4× |
diff --git a/docs/src/optics/AtmosphericStates.md b/docs/src/optics/AtmosphericStates.md
@@ -10,6 +10,7 @@ AtmosphericState
 GrayAtmosphericState
 CloudState
 AerosolState
+TransposedStateCache
 AbstractCloudMask
 MaxRandomOverlap
 AbstractGrayOpticalThickness

diff --git a/docs/src/optics/Fluxes.md b/docs/src/optics/Fluxes.md
@@ -8,6 +8,10 @@ CurrentModule = RRTMGP.Fluxes
 AbstractFlux
 FluxLW
 FluxSW
+FluxPresentation
+update_presentation!
+transpose_into!
+transpose_sum_into!
 set_flux_to_zero!
 compute_net_flux!
 ```
diff --git a/ext/RRTMGPCUDAExt.jl b/ext/RRTMGPCUDAExt.jl
@@ -23,6 +23,7 @@ import RRTMGP.AtmosphericStates: GrayAtmosphericState
 import RRTMGP.AtmosphericStates: AtmosphericState
 import RRTMGP.AtmosphericStates: CloudState
 import RRTMGP.AtmosphericStates: AerosolState
+import RRTMGP.AtmosphericStates: TransposedStateCache
 import RRTMGP.AtmosphericStates
 import RRTMGP.GrayAtmosphere: update_profile_lw!
 import RRTMGP.GrayAtmosphere: compute_gray_heating_rate!

diff --git a/ext/cuda/rte_longwave_2stream.jl b/ext/cuda/rte_longwave_2stream.jl
@@ -54,6 +54,7 @@ function rte_lw_2stream_solve!(
     bcs_lw::LwBCs,
     op::TwoStream,
     as::AtmosphericState,
+    state_cache::Union{TransposedStateCache, Nothing},
     lookup_lw::LookUpLW,
     lookup_lw_cld::Union{LookUpCld, Nothing} = nothing,
     lookup_lw_aero::Union{LookUpAerosolMerra, Nothing} = nothing,
@@ -71,6 +72,7 @@ function rte_lw_2stream_solve!(
         nlay,
         ncol,
         as,
+        state_cache,
         lookup_lw,
         lookup_lw_cld,
         lookup_lw_aero,
@@ -91,6 +93,7 @@ function rte_lw_2stream_solve_CUDA!(
     nlay,
     ncol,
     as::AtmosphericState,
+    state_cache::Union{TransposedStateCache, Nothing},
     lookup_lw::LookUpLW,
     lookup_lw_cld::Union{LookUpCld, Nothing},
     lookup_lw_aero::Union{LookUpAerosolMerra, Nothing},
@@ -115,6 +118,7 @@ function rte_lw_2stream_solve_CUDA!(
                 bcs_lw,
                 op,
                 as,
+                state_cache,
                 lookup_lw,
                 lookup_lw_cld,
                 lookup_lw_aero,

diff --git a/ext/cuda/rte_longwave_noscat.jl b/ext/cuda/rte_longwave_noscat.jl
@@ -60,6 +60,7 @@ function rte_lw_noscat_solve!(
     op::OneScalar,
     angle_disc::AngularDiscretization,
     as::AtmosphericState,
+    state_cache::Union{TransposedStateCache, Nothing},
     lookup_lw::LookUpLW,
     lookup_lw_cld::Union{LookUpCld, Nothing} = nothing,
     lookup_lw_aero::Union{LookUpAerosolMerra, Nothing} = nothing,
@@ -76,6 +77,7 @@ function rte_lw_noscat_solve!(
         nlay,
         ncol,
         as,
+        state_cache,
         lookup_lw,
         lookup_lw_cld,
         lookup_lw_aero,
@@ -96,6 +98,7 @@ function rte_lw_noscat_solve_CUDA!(
     nlay,
     ncol,
     as::AtmosphericState,
+    state_cache::Union{TransposedStateCache, Nothing},
     lookup_lw::LookUpLW,
     lookup_lw_cld::Union{LookUpCld, Nothing},
     lookup_lw_aero::Union{LookUpAerosolMerra, Nothing},
@@ -123,6 +126,7 @@ function rte_lw_noscat_solve_CUDA!(
                 Ds,
                 w_μ,
                 as,
+                state_cache,
                 lookup_lw,
                 lookup_lw_cld,
                 lookup_lw_aero,

diff --git a/ext/cuda/rte_shortwave_2stream.jl b/ext/cuda/rte_shortwave_2stream.jl
@@ -64,6 +64,7 @@ function rte_sw_2stream_solve!(
     bcs_sw::SwBCs,
     src_sw::SourceSW2Str,
     as::AtmosphericState,
+    state_cache::Union{TransposedStateCache, Nothing},
     lookup_sw::LookUpSW,
     lookup_sw_cld::Union{LookUpCld, Nothing} = nothing,
     lookup_sw_aero::Union{LookUpAerosolMerra, Nothing} = nothing,
@@ -81,6 +82,7 @@ function rte_sw_2stream_solve!(
         nlay,
         ncol,
         as,
+        state_cache,
         lookup_sw,
         lookup_sw_cld,
         lookup_sw_aero,
@@ -101,6 +103,7 @@ function rte_sw_2stream_solve_CUDA!(
     nlay,
     ncol,
     as::AtmosphericState,
+    state_cache::Union{TransposedStateCache, Nothing},
     lookup_sw::LookUpSW,
     lookup_sw_cld::Union{LookUpCld, Nothing} = nothing,
     lookup_sw_aero::Union{LookUpAerosolMerra, Nothing} = nothing,
@@ -126,6 +129,7 @@ function rte_sw_2stream_solve_CUDA!(
                     bcs_sw,
                     src_sw,
                     as,
+                    state_cache,
                     lookup_sw,
                     lookup_sw_cld,
                     lookup_sw_aero,

diff --git a/ext/cuda/rte_shortwave_noscat.jl b/ext/cuda/rte_shortwave_noscat.jl
@@ -58,11 +58,12 @@ function rte_sw_noscat_solve!(
     op::OneScalar,
     bcs_sw::SwBCs,
     as::AtmosphericState,
+    state_cache::Union{TransposedStateCache, Nothing},
     lookup_sw::LookUpSW,
 )
     nlay, ncol = AtmosphericStates.get_dims(as)
     tx, bx = _configure_threadblock(ncol)
-    args = (flux, flux_sw, op, bcs_sw, nlay, ncol, as, lookup_sw)
+    args = (flux, flux_sw, op, bcs_sw, nlay, ncol, as, state_cache, lookup_sw)
     @cuda always_inline = true threads = (tx) blocks = (bx) rte_sw_noscat_solve_CUDA!(
         args...,
     )
@@ -77,6 +78,7 @@ function rte_sw_noscat_solve_CUDA!(
     nlay,
     ncol,
     as::AtmosphericState,
+    state_cache::Union{TransposedStateCache, Nothing},
     lookup_sw::LookUpSW,
 )
     gcol = threadIdx().x + (blockIdx().x - 1) * blockDim().x # global id
@@ -95,6 +97,7 @@ function rte_sw_noscat_solve_CUDA!(
                         op,
                         bcs_sw,
                         as,
+                        state_cache,
                         lookup_sw,
                         n_gpt,
                         nlev,

diff --git a/perf/benchmark_baselines/nvidia_a100_sxm4_40gb.txt b/perf/benchmark_baselines/nvidia_a100_sxm4_40gb.txt
@@ -1,17 +1,17 @@
-# RRTMGP DYAMOND GPU benchmark measurement (minimum ns)
+# RRTMGP DYAMOND GPU benchmark baseline (minimum ns)
 # device: NVIDIA A100-SXM4-40GB
 # julia: 1.12.6
 # recorded: local run
 # regenerate with RRTMGP_BENCH_WRITE_BASELINE=1
-all_sky_lw_Float32 9.76149447e8
-all_sky_lw_Float64 1.331256434e9
-all_sky_sw_Float32 8.43344196e8
-all_sky_sw_Float64 1.097443624e9
-all_sky_with_aerosols_lw_Float32 1.05467407e9
-all_sky_with_aerosols_lw_Float64 1.412722056e9
-all_sky_with_aerosols_sw_Float32 8.94524846e8
-all_sky_with_aerosols_sw_Float64 1.173976877e9
-clear_sky_lw_Float32 7.7122294e8
-clear_sky_lw_Float64 1.049297832e9
-clear_sky_sw_Float32 4.64954059e8
-clear_sky_sw_Float64 6.53718113e8
+all_sky_lw_Float32 7.26150458e8
+all_sky_lw_Float64 9.92342112e8
+all_sky_sw_Float32 6.71148384e8
+all_sky_sw_Float64 9.40083761e8
+all_sky_with_aerosols_lw_Float32 7.97787842e8
+all_sky_with_aerosols_lw_Float64 1.09698575e9
+all_sky_with_aerosols_sw_Float32 7.28343007e8
+all_sky_with_aerosols_sw_Float64 9.66275699e8
+clear_sky_lw_Float32 4.89919056e8
+clear_sky_lw_Float64 6.82068605e8
+clear_sky_sw_Float32 4.55409919e8
+clear_sky_sw_Float64 6.39967283e8
diff --git a/src/RRTMGP.jl b/src/RRTMGP.jl
@@ -26,12 +26,12 @@ include(joinpath("api", "aerosols.jl"))
 # ── Layer 0: core containers — states, sources, optics, fluxes, BCs, lookups
 include(joinpath("optics", "VolumeMixingRatios.jl"))
 const Vmrs = VolumeMixingRatios # compat alias (deprecated name; remove in 0.24)
+include(joinpath("optics", "Fluxes.jl"))
 include(joinpath("optics", "LookUpTables.jl"))
 include(joinpath("optics", "AngularDiscretizations.jl"))
 include(joinpath("optics", "AtmosphericStates.jl"))
 include(joinpath("optics", "Sources.jl"))
 include(joinpath("optics", "Optics.jl"))
-include(joinpath("optics", "Fluxes.jl"))
 include(joinpath("optics", "GrayAtmosphere.jl"))
 const GrayUtils = GrayAtmosphere # compat alias (deprecated name; remove in 0.24)
 include(joinpath("optics", "BCs.jl"))

diff --git a/src/api/getters.jl b/src/api/getters.jl
@@ -57,17 +57,17 @@ _require_clear_sky(x) = x
 # Potentially views
 top_of_atmosphere_lw_flux_dn(s::RRTMGPSolver)         = _maybe_transpose(s.lws.bcs.inc_flux)
 top_of_atmosphere_diffuse_sw_flux_dn(s::RRTMGPSolver) = _maybe_transpose(s.sws.bcs.inc_flux_diffuse)
-lw_flux_up(s::RRTMGPSolver)                           = _domain_view(s, s.lws.flux.flux_up)
-lw_flux_dn(s::RRTMGPSolver)                           = _domain_view(s, s.lws.flux.flux_dn)
-lw_flux_net(s::RRTMGPSolver)                          = _domain_view(s, s.lws.flux.flux_net)
+lw_flux_up(s::RRTMGPSolver)                           = _domain_view(s, s.presented_flux_lw.flux_up)
+lw_flux_dn(s::RRTMGPSolver)                           = _domain_view(s, s.presented_flux_lw.flux_dn)
+lw_flux_net(s::RRTMGPSolver)                          = _domain_view(s, s.presented_flux_lw.flux_net)
 clear_lw_flux_up(s::RRTMGPSolver)                     = _domain_view(s, _require_clear_sky(s.clear_flux_lw).flux_up)
 clear_lw_flux_dn(s::RRTMGPSolver)                     = _domain_view(s, _require_clear_sky(s.clear_flux_lw).flux_dn)
 clear_lw_flux(s::RRTMGPSolver)                        = _domain_view(s, _require_clear_sky(s.clear_flux_lw).flux_net)
 surface_emissivity(s::RRTMGPSolver)                   = s.lws.bcs.sfc_emis
-sw_flux_up(s::RRTMGPSolver)                           = _domain_view(s, s.sws.flux.flux_up)
-sw_flux_dn(s::RRTMGPSolver)                           = _domain_view(s, s.sws.flux.flux_dn)
-sw_flux_net(s::RRTMGPSolver)                          = _domain_view(s, s.sws.flux.flux_net)
-sw_direct_flux_dn(s::RRTMGPSolver)                    = _domain_view(s, s.sws.flux.flux_dn_dir)
+sw_flux_up(s::RRTMGPSolver)                           = _domain_view(s, s.presented_flux_sw.flux_up)
+sw_flux_dn(s::RRTMGPSolver)                           = _domain_view(s, s.presented_flux_sw.flux_dn)
+sw_flux_net(s::RRTMGPSolver)                          = _domain_view(s, s.presented_flux_sw.flux_net)
+sw_direct_flux_dn(s::RRTMGPSolver)                    = _domain_view(s, s.presented_flux_sw.flux_dn_dir)
 clear_sw_flux_up(s::RRTMGPSolver)                     = _domain_view(s, _require_clear_sky(s.clear_flux_sw).flux_up)
 clear_sw_flux_dn(s::RRTMGPSolver)                     = _domain_view(s, _require_clear_sky(s.clear_flux_sw).flux_dn)
 clear_sw_direct_flux_dn(s::RRTMGPSolver)              = _domain_view(s, _require_clear_sky(s.clear_flux_sw).flux_dn_dir)