refactor: chunk the kernel-CDF forward transform (512-query blocks)#378
Merged
Conversation
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>
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
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.
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 fixedKERNEL_FORWARD_BLOCK = 512query blocks:lax.mapover 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
lax.mapValidation checklist (--auto run — plan was not pre-approved)
🤖 Generated with Claude Code