Choosing Between Custom CUDA, Libraries, and Framework Operators
A medical imaging startup has a preprocessing step that resamples 3D volumes. The current implementation uses a Python loop calling a deep learning framework's interpolation operator once per slice. Profiling shows the preprocessing takes 60% of total training step time, and the GPU is idle during much of it because each slice launch is small. The team has three options: (a) write a custom CUDA kernel, (b) restructure the code to use batched library calls on the framework's tensor operations, or (c) keep the Python loop but enable more aggressive framework compilation. The team has two GPU engineers with strong Python skills but limited CUDA experience, and a 5-week deadline. They ask which path to take and how to validate the choice.
1. Evaluate the three options against the constraints. Which is the best default, and under what conditions would you escalate to a custom kernel?
Read the worked answer
Option (b), restructuring to batched library calls, is the best default given the constraints.
The symptom is small per-slice launches with the GPU idle between them. That is a batching problem, not a compute-kernel problem. Reshaping the loop into a single batched operation over the volume, or over a chunk of slices, lets the framework issue one larger kernel launch that fills the GPU. This uses existing, well-tested operators and requires no CUDA expertise. It is also the lowest-risk path against a 5-week deadline.
Option (c), more aggressive compilation, is worth trying as a quick experiment but is not a strategy. Compilation can fuse and reduce overhead, but if the fundamental structure is a Python loop with per-slice dispatch, compilation gains are bounded by the launch pattern. Treat it as a cheap test, not the plan.
Option (a), a custom CUDA kernel, should be escalated only if batching does not close the gap. Conditions that justify it: the resampling math is not expressible with existing operators, the batched version still leaves the GPU underutilized because of memory access patterns, or the operation is on the critical path after batching and still dominates step time. A custom kernel also carries maintenance cost and requires the two GPU engineers to ramp up on CUDA, which is risky in 5 weeks.
Recommended sequence: implement batching first, measure step time and GPU utilization, then decide whether a kernel is needed. Keep the custom kernel as a contingency with a defined trigger threshold.
2. What validation would you run to prove the chosen option actually helps, and what metrics would you report?
Read the worked answer
Validation should be a controlled comparison, not an impression.
Baseline: record current preprocessing time per step, total step time, GPU utilization during preprocessing, and end-to-end training throughput. Use a fixed set of volumes and a fixed number of steps so runs are comparable.
Experiment: run the batched implementation on the same inputs. Vary only the batching factor and measure how step time and utilization change. Include a correctness check: compare resampled outputs against the baseline within a numerical tolerance, because batching can change interpolation boundaries and edge handling.
Metrics to report:
- Preprocessing wall-clock time per step and its share of total step time.
- GPU utilization during the preprocessing phase, not just overall.
- End-to-end throughput in samples per second.
- Numerical agreement with the baseline, expressed as maximum and mean absolute difference.
- Memory footprint, since larger batches increase peak memory.
Decision rule: adopt the batched version if it materially reduces preprocessing share and passes the correctness tolerance. If preprocessing still dominates after batching, that is the trigger to evaluate a custom kernel, and the same validation harness should be reused. Report the comparison to stakeholders with the baseline and the new numbers side by side so the choice is auditable.