Skip to content

refactor: chunk the kernel-CDF forward transform (512-query blocks)#378

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/kernel-forward-chunking
Jul 10, 2026
Merged

refactor: chunk the kernel-CDF forward transform (512-query blocks)#378
Jammy2211 merged 1 commit into
mainfrom
feature/kernel-forward-chunking

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Summary

Removes the kernel-CDF meshes' O(M×N) memory wall (#376): the exact forward F(q) = Σᵢ wᵢ·Φ((q−xᵢ)/h) previously broadcast an (M, N, 2) array — ~60 GB at production imaging scale (M≈246k over-sampled queries × N≈15.4k traced points, observed OOM). It now evaluates in fixed KERNEL_FORWARD_BLOCK = 512 query blocks: lax.map over padded blocks under jax (static shapes, AD flows through the scan), a plain block loop under numpy. Peak block memory ≈ 126 MB at that scale; measured end-to-end: 1.06 GB peak RSS where the unblocked code failed to allocate. CPU wall-time at that scale is ~10 min/eval (2×10⁹ erf evaluations — the arithmetic, not the blocking); GPU remains the production target.

API Changes

None — internal changes only (one additive module constant, KERNEL_FORWARD_BLOCK).

Test Plan

  • Refactor invariant — float-identical values: block-invariance unit test (batched == row-by-row across a non-multiple-of-block query count); full suite 894 passed
  • Both jax_grad certification scripts pass on this branch with byte-identical FoM/parity values to the unblocked run (imaging 7 variants, interferometer 5)
  • Function-level JAX sanity: eager == JIT; AD == FD through lax.map
  • Previously-OOM scale probe (M=246,080 × N=15,380): completes at 1.06 GB peak RSS

Validation checklist (--auto run — plan was not pre-approved)

  • Effective level: safe (refactor cap; human-directed follow-up "do these two")
  • Plan: on issue refactor: chunk the kernel-CDF forward transform (remove the O(M×N) memory wall) #376
  • Gate: tests 894 + certification byte-identical · smoke n/a (jax_grad scripts are the downstream surface, both re-run) · review CLEAN · Heart YELLOW acked in-session (same 6 stale reasons)
  • Human: plan sound in hindsight?
  • Human: diff matches plan (no scope creep)?
  • Human: merge, amend, or reject — then log the outcome

🤖 Generated with Claude Code

The exact kernel forward broadcast queries x points into an (M, N, 2) array —
~60 GB at production imaging scale (M≈246k, N≈15.4k; observed OOM). It now
evaluates in fixed KERNEL_FORWARD_BLOCK=512 query blocks: lax.map over padded
blocks under jax (static shapes, AD flows through), a plain block loop under
numpy. Values are float-identical (the sum over points is unchanged; blocks
only tile the query axis): block-invariance unit test, full suite 894, both
jax_grad certification scripts pass with byte-identical FoM/parity values,
and the previously-OOM scale runs at ~1.1 GB peak RSS.

#376.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@Jammy2211 Jammy2211 added the pending-release PR queued for the next release build label Jul 10, 2026
@Jammy2211 Jammy2211 merged commit 8d9b1d4 into main Jul 10, 2026
4 checks passed
@Jammy2211 Jammy2211 deleted the feature/kernel-forward-chunking branch July 10, 2026 13:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pending-release PR queued for the next release build

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant