Add GPU capability probe (standalone project + kkt workflow)#326
Merged
Conversation
A diagnostic harness that measures, on real NVIDIA/CUDA hardware, which CTFlows GPU combinations already work against the current, unmodified packages — and, for the failures, with which error. It exists to turn the GPU design hypotheses (AD is the gate; AD-free flows are GPU-ready first; the pv0 init-condition bug) into measured facts before any code is written. - probe/gpu/probe_gpu.jl: verbose, try/catch-per-experiment, never asserts (so it cannot fail CI); prints a capability matrix + interpretation guide. Activates its own project and Pkg.develops the checked-out CTFlows (measures the PR's source). - probe/gpu/Project.toml: self-contained env (CUDA, Zygote, ForwardDiff, OrdinaryDiffEqTsit5, SciMLBase, ADTypes) — CTFlows' own deps are untouched. - .github/workflows/GPUProbe.yml: dedicated workflow replicating the CTActions setup steps (checkout, setup-julia, add ct-registry via SSH_KEY) and running the probe on runs-on: [kkt]. Triggered by workflow_dispatch or the 'run ci gpu' label. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
First H200 run (PR #326) rewrote two assumptions and exposed one harness bug: - Flow(ODEProblem) point-call returns the final state directly (a CuArray); the final_state() wrapper was wrong and caused a MethodError. Drop the wrapper — the flow itself works on device. - host/device vcat(x0, p0, zeros(1)) does NOT fail (vcat promotes), so the pv0 `zeros` is hygiene, not a hard bug. Relabel B0 accordingly. The real augmented blocker is B4 (variable_costate) failing with a GPUCompiler.KernelError — its own investigation item. - Add expect_fail marking (:xfail) so the AutoForwardDiff-on-GPU cases are recorded as expected failures, separating them from real limitations in the summary; an unexpected pass is now flagged loudly. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The capability matrix is now sharp (13 ok / 3 xfail / 1 real fail); the only genuine limitation is B4 (variable_costate augmented flow, a GPUCompiler.KernelError whose reason is truncated in the log). Keep the full showerror+backtrace for unexpected :fail results and print it in a dedicated "UNEXPECTED-FAILURE DETAIL" section, so the next run reveals B4's root cause (which augmented-RHS operation is not GPU-compilable). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… cause Run 3's full trace pinned B4 (variable_costate on GPU): _aug_assign! (hamiltonian_vector_field_system.jl:221) broadcasts a HOST ∂pv into the DEVICE du, because the variable v transits host-side so variable_gradient returns a host array. The state derivatives ∂x/∂p are already device — only the variable-costate block is host. Add B4b: the same augmented flow with a DEVICE variable (variable=_dev([0.5])). If it compiles and runs, it confirms the root cause and validates the fix direction (device-adapt the variable-costate block, or carry v on-device). B4a (host variable) is kept as the reproducer. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Run 4's device-variable check was inconclusive: it failed at H evaluation, not at _aug_assign!, because the test Hamiltonian used `+ v[1]` and `v[1]` scalar-indexes a device v inside hamiltonian_gradient (Zygote pullback at probe_gpu.jl:270), never reaching the augmented assign. Switch to `+ sum(v)` (GPU-friendly, same value for a 1-vector) so B4b actually exercises the device→device _aug_assign! path. This also surfaced a real user-facing constraint worth documenting: Hamiltonians must avoid scalar indexing (v[i], x[i]) on GPU — use sum/dot/broadcasts. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Run 5's full trace (device v, sum(v) Hamiltonian) shows the augmented RHS now runs clean device->device (_aug_assign!/hamiltonian_gradient no longer fail) -- the failure moved downstream to Trajectories._aug_split_solution (building.jl:92-96): Systems._coerce_state(pv0) returns `only` when length(pv0)==1 (the "1-D=scalar" convention, hamiltonian_vector_field_system.jl:97), and only(::CuArray) scalar- indexes via iterate. This is the same general GPU gap already flagged for OCP/Model flows in the design report (S5), not an augmented-flow-specific defect. Add B4c: variable of length 2, so _coerce_state is `identity` on every dimension (x0, p0, pv0 all length >= 2), avoiding `only` entirely. A pass confirms the augmented RHS + solution-splitting pipeline is fully GPU-clean once no dimension collapses through the 1-D=scalar coercion -- isolating that coercion as the single remaining gap across the whole probed surface. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Two experiments on the probe branch:
1. Implement the chosen "1-D=scalar on GPU" fix (_safe_only): Systems._coerce_state now
returns `_safe_only` instead of `only`, where
_safe_only(x::GPUArraysCore.AbstractGPUArray) = GPUArraysCore.@allowscalar only(x)
_safe_only(x) = only(x)
Dispatch is on the live array at the call site, so the "1-D=scalar" contract returns an
identical host scalar on CPU and GPU. Adds GPUArraysCore (tiny, backend-agnostic) to
CTFlows deps. Probe B4b (device variable, length 1) should flip FAIL→OK, validating the
fix end-to-end; B4a (host variable) still fails (that's the separate _aug_assign! fix).
2. Add DiffEqGPU + StaticArrays to the probe project and a Block 6 ensemble tooling PoC:
EnsembleGPUArray (in-place, permissive) and EnsembleGPUKernel (OOP + SVector,
kernel-compiled), N=8 trajectories of ẋ=-x, compared to the analytic solution. Pins the
ensemble tooling on kkt (§3 model 2 / §7 Family C), not yet wired to CTFlows flows.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Run-7 change was incomplete: _coerce_state now returns _safe_only, but the RHS
functor structs bound their coercion type params to Union{typeof(only),typeof(identity)},
which typeof(_safe_only) does not satisfy — a latent regression for length-1 states
(the probe uses length-2, so it wouldn't catch it).
- Update all Systems RHS-functor bounds Union{typeof(only),...} → Union{typeof(_safe_only),...}
(hamiltonian/hvf/pseudo/constrained_pseudo functors). The Flows OCP bounds are left
untouched: they are fed by the Flows-local _dim_coerce/_finalize_vf, not Systems._coerce_state
(the OCP-on-GPU fix is separate, report §5).
- Move _safe_only into src/Systems/coercion.jl included before the functor files, because
typeof(_safe_only) in a struct bound is evaluated at struct-definition (include) time.
Validated on CPU (temp env): CTFlows precompiles; a length-1 Hamiltonian flow constructs
and runs, returning scalars cos(1)/-sin(1) — the "1-D=scalar" contract is preserved,
no regression. GPU path awaits kkt (run 8).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Run 8 confirmed the _safe_only fix (B4 device-var length-1 now OK on H200) but B6's ensemble PoC failed: this SciMLBase version replaced the classic prob_func(prob, i, repeat) with prob_func(prob, ctx), where ctx.sim_id is the trajectory index. Both EnsembleGPUArray and EnsembleGPUKernel closures updated accordingly.
Run 9 confirmed the ctx.sim_id prob_func fix (both ensemble solves now reach the GPU kernel), but the probe's own maxerr computation then failed: EnsembleSolution <: AbstractVectorOfArray, so sol[i] performs tensor-style component indexing, not trajectory selection — it silently returns scalar Float64s instead of the i-th ODESolution. sol.u[i] bypasses that overload and gets the right trajectory.
Run 10: ODESolution is itself an AbstractVectorOfArray (component x time) with IndexCartesian getindex, so sol.u[i][end] does linear Cartesian indexing over the whole solution and returns a single scalar component, not the final state vector. Go through the raw .u field at both nesting levels to get plain, unambiguous indexing.
…unds
The RHS-functor struct type bounds now require
Union{typeof(identity), typeof(_safe_only)} instead of
Union{typeof(identity), typeof(only)} (part of the GPU 1-D=scalar fix).
These tests constructed the structs directly with literal `only`,
bypassing _coerce_state, and no longer type-checked.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
A diagnostic probe that measures, on the self-hosted NVIDIA/CUDA runner (
kkt), which CTFlows GPU combinations already work against the current, unmodified packages — and, for the failures, with which error. It exists to turn the GPU design report's hypotheses into measured facts before any code is written.It is not a test suite: nothing asserts pass/fail. Every experiment is wrapped in
try/catch, so a single failure never stops the run or fails CI; each outcome is printed and collected into a final capability matrix (read theSUMMARYat the bottom of the job log).Files
probe/gpu/probe_gpu.jl— the diagnostic. Activates its own project andPkg.develops the checked-out CTFlows, so it measures this PR's source.probe/gpu/Project.toml— self-contained env (CUDA, Zygote, ForwardDiff, OrdinaryDiffEqTsit5, SciMLBase, ADTypes). CTFlows' own deps are untouched — the heavy GPU deps live only in the probe project.probe/gpu/README.md— what it measures and how to run it..github/workflows/GPUProbe.yml— dedicated workflow replicating the three load-bearing CTActions steps (checkout → setup-julia → add ct-registry viaSSH_KEY) then running the probe onruns-on: ["kkt"].What it probes (blocks B0–B5)
Array primitives incl. the
pv0init-condition bug-vs-fix;DifferentiationInterface.gradienton aCuArrayforAutoForwardDiffvsAutoZygote(the core "AD is the GPU gate" question); the realhamiltonian_gradientcall; flows on device (VectorField/HamiltonianVectorField/ODEProblemAD-free, vsHamiltoniandefault-vs-Zygote backend);variable_costate(thepv0path); andFloat32end-to-end.How to run it
run ci gpulabel.Once the works/doesn't-work boundary is known from the log, the probe is superseded by real assertions in the CTFlows test suite (design report §7). It is a throwaway measurement harness.
🤖 Generated with Claude Code