LesionSegmenter: GPU-resident export (no CPU round trip) - #121
Merged
Conversation
predict_logits_from_preprocessed_data returns logits on the CPU, so the export step moves the ~1.5 GB, 43-channel logit volume CPU->GPU to run the (already GPU-patched) resample, then argmaxes and comes back. The transfer and CPU-side handling dominate the stage. Keeping the export resident on the GPU -- same resample function, same spacings, same argmax, only the small uint8 label map returned -- removes that round trip. Measured on a size-stratified sample against the current convert_logits output: bit-identical (0 differing voxels across every case, including anisotropic ones where a naive trilinear shortcut would differ), and 3.6-25x faster on the export stage (median export ~5s -> ~1s). Guarded to plain softmax-argmax label spaces; region-based models fall back to convert_logits. Verified end to end by running the cold path on a real case (valid output, no errors) in addition to the per-stage bit-exactness check.
aperson30
added a commit
that referenced
this pull request
Aug 4, 2026
Two speed additions to the persistent LesionSegmenter service, both of which only make sense for a warm/persistent process: torch.compile (gated by LESIONSEG_COMPILE, default off): compiles the patch-level network once at startup with max-autotune-no-cudagraphs. The sliding-window patch is a constant [128,224,224], so one compile serves every patch of every case. Measured 1.18x on inference, accuracy-safe (voxel agreement 0.99996, lesion Dice 0.999 vs eager). The ~20s-3min compile is paid once at startup, behind the health gate -- while it runs, lesionseg_predict.py falls back to the cold path so the site keeps working. Env-gated so it can be dropped without a redeploy if a node shows the compile-latency variance seen on some shared GB10 nodes; the warm predictor's main win (cold-start elimination) is independent of it. GPU-resident export (default for softmax-argmax models): predict_logits returns logits on the CPU, so convert_logits moves the ~1.5 GB, 43-channel volume CPU->GPU for the (GPU-patched) resample and back. This runs the same resample, spacings and argmax GPU-resident and returns only the uint8 label map. Verified 0 differing voxels vs convert_logits across a size-stratified sample incl. anisotropic cases; 3.6-25x on the export stage. Region-based models fall back to convert_logits. Verified the service starts, compiles, and listens; export logic matches PR #121 (shipped) which was validated bit-exact.
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.
The LesionSegmenter export stage moves the ~1.5 GB, 43-channel logit volume off and back onto the GPU on every request, because
predict_logits_from_preprocessed_datareturns logits on the CPU while the export resample (already GPU-patched, PR #108) runs on the GPU. This keeps the whole export GPU-resident.What changes
The cold path now does the export in place on the GPU — same resample function, same spacings, same argmax — and returns only the small
uint8label map. No 1.5 GB round trip.Measured (size-stratified sample vs current
convert_logitsoutput)Bit-identical to the current output on every case (0 differing voxels), including anisotropic cases where a naive trilinear shortcut diverges. Median export ~5s → ~1s.
Safety
argmax(logits) == argmax(softmax(logits))); region-based models fall back toconvert_logits, unchanged.