Conversation
…pooled buffer Clearing the auxiliary-polynomial pool destroys every pooled RNSPoly; each pooled limb buffer freed through GPUfree records an event on the freeing stream and makes the pool stream wait on it. With N = 2^16 and a pool of 104 polys that is 4816 cudaEventRecord + 4816 cudaStreamWaitEvent calls per clear, 12.6 ms of host time on an H200 (nsys). After one device-wide synchronize no buffer is in use, so the waits are redundant: set a thread-local flag that lets GPUfree skip them while the pool is cleared. Co-Authored-By: Claude Fable 5.1 <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
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
ContextData::clearAuxilarPoly()now synchronizes the device once and then letsGPUfreeskip the per-buffer "pool stream waits for the freeing stream" event pair while the pool is being destroyed.Why
Every pooled limb buffer freed through
GPUfreerecords an event on the freeing stream and makes the memory-pool stream wait on it (s[id].wait(stream)inCudaUtils.cu). Clearing the auxiliary-polynomial pool destroys every pooledRNSPolyat once, so a clear costs thousands of those event operations even though nothing is in flight any more.Measured with Nsight Systems on an H200 (CUDA 12, N = 2^16, 25-limb chain, pool of 104 polynomials, one clear per MoE layer):
cudaEventRecord+ 4816cudaStreamWaitEventcudaDeviceSynchronize)The rest of the clear is the ~100
LimbPartitiondestructors and the pool bookkeeping itself.Why it is safe
After
cudaDeviceSynchronize()returns, no stream can still be using a buffer that is about to be returned to the pool, which is exactly what the skipped event pair was protecting against. The flag isthread_localand is set only for the duration ofprecom.auxPoly.clear(), so frees issued from other threads keep their waits.Validation: a 32-layer CKKS MoE chain (N = 2^16, depth 25, one
clearAuxilarPolyper layer, 16384 and 32768 slots) passes its regression gates with the change in four independent runs (expert routing 32/32 layers, output error class, noise band), and the per-stage timings of everything except the stage containing the clear are unchanged to 0.1 ms.Notes
clearAuxilarPolyis a coarse operation (callers use it between layers or at shutdown), so one device-wide synchronize inside it is a reasonable price. If you would rather keep the old behaviour available, the flag could be gated behind a parameter; happy to adjust.Fable 5.1 on behalf of Seyfal