[draft]Optimize Nemotron OCR GPU execution for stream-aware serving
Summary
This PR optimizes the NemotronOCRV2 GPU execution path for sustained, queued OCR serving.
Model and CUDA-extension changes live here. The native vLLM pooling model, JPEG-byte I/O path, queueing, sweep tooling, and benchmark drivers are in VibhuJawa/vllm#1.
Optimizations
- Launch custom CUDA kernels on PyTorch’s current CUDA stream instead of the implicit default stream.
- Remove avoidable CUDA scalar reads and synchronization in local IPS and NMS.
- Disable NMS timing synchronization unless verbose timing is requested.
- Batch relational GPU-to-CPU transfers through pinned memory, with one wait per CUDA stream instead of synchronizing each
.cpu()operation. - Skip inference-only finite-value assertions that force GPU synchronization.
- Fuse recognizer argmax and selected softmax probability extraction into one OpenAI Triton GPU kernel, with a PyTorch fallback.
- Add opt-in channels-last, cuDNN benchmarking, and
torch.compilecontrols. - Add CUDA regression tests for relational staging and graph construction.
“OpenAI Triton” above refers to the GPU kernel language, not NVIDIA Triton Inference Server.
Matched A100 benchmark
This model PR is one component of the complete optimized deployment; the speedup is not attributed to this patch alone.
All systems processed the same ordered 1,000 Digital Corpora Bo767 document pages, encoded as JPEG Q100 4:4:4 bytes and replayed 30 times:
- 30,000 timed requests
- Paragraph merge
infer_length=1024- One NVIDIA A100-SXM4-80GB
- Zero recorded request failures
| System | Images/s | Versus official HF |
|---|---|---|
| Official NVIDIA/HF in-process | 31.2457 | 1.00x |
| Isolated native vLLM baseline, one replica | 43.8927 | 1.40x |
| Optimized native vLLM, eight replicas | 70.1156 | 2.24x |
The optimized deployment is 1.60x faster than the isolated vLLM baseline and averaged 99.74% GPU utilization.
The conservative recognizer-chunk-128 control reached 69.5380 images/s. The slightly faster recognizer-chunk-64 configuration has not yet been evaluated against labeled OCR ground truth and should be treated as throughput-tuned rather than proven accuracy-preserving.
Validation
- A100 (
sm_80) CUDA extension rebuilt successfully. - Side-stream smoke testing confirmed custom kernels execute on the current non-default CUDA stream.
- Relational staging regression tests: 2 passed.
- Ruff checks passed for changed Python.
- A controlled 32-image optimized-kernel comparison recorded zero text or region-count mismatches.
Related code and results
- Model commit:
bb392d4 - vLLM integration: VibhuJawa/vllm#1
- Optimized vLLM implementation commit:
267b6f6d6 - Results and reproduction guide: VibhuJawa/nemotron-vllm-ocr
- Raw telemetry, exact 1,000-page manifest, provenance, and charts: A100 benchmark results