Reduce recognizer VRAM by discarding logits per chunk

Computes argmax + softmax probability of the winning token inside
the recognizer chunk loop and discards the full [N, T, num_tokens]
logits tensor immediately, reducing peak VRAM.

Also converts region vertices to numpy during graph_build to avoid
per-region .cpu() calls in post-processing.

nemotron-ocr/src/nemotron_ocr/inference/pipeline_v2.py

@@ -409,23 +409,28 @@ class NemotronOCRV2(NemotronOCR):
         """Run recognizer in fixed-size region chunks.
 
         Every chunk is padded to exactly ``recognizer_chunk_size`` regions
-        so CUDA always sees the same tensor shape — eliminates kernel
-        recompilation and cuBLAS workspace re-allocation.
+        so CUDA always sees the same tensor shape.
+
+        Computes argmax + softmax probability of the winning token inside
+        each chunk loop and discards the full logits immediately, reducing
+        peak VRAM from ``[N, T, num_tokens]`` to ``[chunk_size, T, num_tokens]``.
 
         Returns:
-            rec_output: (N_regions, T, C)
-            rec_features: (N_regions, T, D)
+            rec_ids: (N_regions, T) int64 — argmax token indices
+            rec_probs: (N_regions, T) float32 — softmax probability of winning token
+            rec_features: (N_regions, T, D) — transformer features for relational model
         """
         n = rec_quads.shape[0]
         if n == 0:
             device = rec_quads.device
             return (
-                torch.empty(0, self.max_width, self.num_tokens, dtype=torch.float16, device=device),
+                torch.empty(0, self.max_width, dtype=torch.int64, device=device),
+                torch.empty(0, self.max_width, dtype=torch.float32, device=device),
                 torch.empty(0, self.max_width, self.recognizer.feature_depth, dtype=torch.float16, device=device),
             )
 
         cs = self.recognizer_chunk_size
-        out_parts, feat_parts = [], []
+        ids_parts, probs_parts, feat_parts = [], [], []
         with amp.autocast("cuda", enabled=True), torch.inference_mode():
             for start in range(0, n, cs):
                 end = min(start + cs, n)
@@ -434,13 +439,21 @@ class NemotronOCRV2(NemotronOCR):
                 if real_n < cs:
                     pad = torch.zeros(cs - real_n, *chunk.shape[1:], dtype=chunk.dtype, device=chunk.device)
                     chunk = torch.cat([chunk, pad], dim=0)
-                o, f = self.recognizer(chunk)
-                out_parts.append(o[:real_n])
-                feat_parts.append(f[:real_n])
+                logits, feats = self.recognizer(chunk)
+
+                ids = logits.argmax(dim=2)
+                probs = torch.softmax(logits, dim=2).gather(
+                    2, ids.unsqueeze(2)
+                ).squeeze(2).float()
+
+                ids_parts.append(ids[:real_n])
+                probs_parts.append(probs[:real_n])
+                feat_parts.append(feats[:real_n])
 
-        if len(out_parts) == 1:
-            return out_parts[0], feat_parts[0]
-        return torch.cat(out_parts, dim=0), torch.cat(feat_parts, dim=0)
+        def _cat_or_single(parts):
+            return parts[0] if len(parts) == 1 else torch.cat(parts, dim=0)
+
+        return _cat_or_single(ids_parts), _cat_or_single(probs_parts), _cat_or_single(feat_parts)
 
     # ------------------------------------------------------------------
     # Phase 6 — Relational model + build output dict
@@ -453,12 +466,17 @@ class NemotronOCRV2(NemotronOCR):
         region_counts,
         region_counts_cpu,
         rec_features,
-        rec_output,
+        rec_ids,
+        rec_probs,
         e2e_det_conf,
         confidence,
         padded_lengths,
     ):
-        """Run relational model and assemble the output dict used by post-processing."""
+        """Run relational model and assemble the output dict.
+
+        Accepts pre-computed argmax ids and per-token probabilities from the
+        recognizer (avoids materializing the full logits tensor).
+        """
         with amp.autocast("cuda", enabled=True), torch.inference_mode():
             rel_output = self.relational(
                 rel_quads,
@@ -481,12 +499,9 @@ class NemotronOCRV2(NemotronOCR):
             scale_per_region = torch.repeat_interleave(scale_factors, counts_long, dim=0)
             quads_scaled = quads_cuda * scale_per_region.view(-1, 1, 1)
 
-            rec_logits = rec_output.float()
-            log_norm = rec_logits.logsumexp(dim=2, keepdim=True)
-            seq_ids = rec_logits.argmax(dim=2)
-            seq_probs = (rec_logits.gather(2, seq_ids.unsqueeze(2)) - log_norm).squeeze(2).exp()
+            seq_ids = rec_ids
+            seq_probs = rec_probs
 
-            # Pre-compute text confidence on GPU (geometric mean of per-token probs)
             _before_eos = (seq_ids == 1).cumsum(dim=1) == 0
             _real = (seq_ids != 0) & _before_eos
             _counts = _real.sum(dim=1).clamp(min=1).float()
@@ -510,48 +525,19 @@ class NemotronOCRV2(NemotronOCR):
                 "e2e_recog_features": rec_features,
             }
 
-        return output, rec_output
+        return output
 
     # ------------------------------------------------------------------
     # Phase 7 — Post-process
     # ------------------------------------------------------------------
 
-    def _decode_with_fallback(self, output, rec_output):
-        """Call recog_encoder, falling back to full softmax if needed.
-
-        The local encoder accepts pre-computed argmax indices (dim=2) for
-        efficiency.  The HuggingFace encoder only accepts full softmax
-        probabilities (dim=3).  We try the fast path first and cache the
-        result so subsequent calls skip the try/except.
-        """
-        if getattr(self, "_use_full_softmax", False):
-            output = dict(output)
-            output["sequences"] = F.softmax(rec_output.float(), dim=2)
-            output.pop("sequence_probs", None)
-            return self.recog_encoder.convert_targets_to_labels(
-                output, image_size=None, is_gt=False,
-            )
-
-        try:
-            batch = self.recog_encoder.convert_targets_to_labels(
-                output, image_size=None, is_gt=False,
-            )
-            self._use_full_softmax = False
-            return batch
-        except (ValueError, RuntimeError):
-            logger.info(
-                "Encoder does not support argmax sequences; "
-                "falling back to full softmax (one-time detection)."
-            )
-            self._use_full_softmax = True
-            output = dict(output)
-            output["sequences"] = F.softmax(rec_output.float(), dim=2)
-            output.pop("sequence_probs", None)
-            return self.recog_encoder.convert_targets_to_labels(
-                output, image_size=None, is_gt=False,
-            )
+    def _decode_with_fallback(self, output):
+        """Decode recognized sequences using pre-computed argmax indices."""
+        return self.recog_encoder.convert_targets_to_labels(
+            output, image_size=None, is_gt=False,
+        )
 
-    def _postprocess_batch(self, output, rec_output, original_shapes, merge_level, include_invalid, timings=None):
+    def _postprocess_batch(self, output, original_shapes, merge_level, include_invalid, timings=None):
         """Decode sequences, build relation graphs, format per-image results."""
         import gc
         gc_was_enabled = gc.isenabled()
@@ -559,17 +545,17 @@ class NemotronOCRV2(NemotronOCR):
 
         try:
             return self._postprocess_batch_inner(
-                output, rec_output, original_shapes, merge_level, include_invalid, timings,
+                output, original_shapes, merge_level, include_invalid, timings,
             )
         finally:
             if gc_was_enabled:
                 gc.enable()
 
-    def _postprocess_batch_inner(self, output, rec_output, original_shapes, merge_level, include_invalid, timings=None):
+    def _postprocess_batch_inner(self, output, original_shapes, merge_level, include_invalid, timings=None):
         _t = time.perf_counter
 
         t0 = _t()
-        batch = self._decode_with_fallback(output, rec_output)
+        batch = self._decode_with_fallback(output)
         recog_decode_ms = (_t() - t0) * 1000
 
         t0 = _t()
@@ -595,7 +581,8 @@ class NemotronOCRV2(NemotronOCR):
 
         for example in batch:
             for text_region in example:
-                text_region.region = text_region.region.vertices
+                v = text_region.region.vertices
+                text_region.region = v.cpu().numpy() if hasattr(v, 'cpu') else np.asarray(v)
         graph_build_ms = (_t() - t0) * 1000
 
         t0 = _t()
@@ -679,7 +666,12 @@ class NemotronOCRV2(NemotronOCR):
 
         # Phase 2.5: centerness + peak prefilter
         if self._use_prefilter:
+            t0 = time.perf_counter()
             det_conf = self._prefilter_detections(det_conf, det_rboxes)
+            torch.cuda.synchronize()
+            T["prefilter"] = (time.perf_counter() - t0) * 1000
+        else:
+            T["prefilter"] = 0.0
 
         # Phase 3: NMS
         t0 = time.perf_counter()
@@ -709,22 +701,23 @@ class NemotronOCRV2(NemotronOCR):
         torch.cuda.synchronize()
         T["rectify"] = (time.perf_counter() - t0) * 1000
 
-        # Phase 5: recognizer
+        # Phase 5: recognizer (argmax + softmax prob per chunk, discards full logits)
         t0 = time.perf_counter()
-        rec_output, rec_features = self._run_recognizer_chunked(rec_quads)
+        rec_ids, rec_probs, rec_features = self._run_recognizer_chunked(rec_quads)
         del rec_quads
         torch.cuda.synchronize()
         T["recognizer"] = (time.perf_counter() - t0) * 1000
 
         # Phase 6: relational model
         t0 = time.perf_counter()
-        output, rec_output_raw = self._run_relational_and_build_output(
+        output = self._run_relational_and_build_output(
             rel_quads,
             quads_cuda,
             region_counts,
             region_counts_cpu,
             rec_features,
-            rec_output,
+            rec_ids,
+            rec_probs,
             e2e_det_conf,
             confidence,
             padded_lengths,
@@ -736,7 +729,7 @@ class NemotronOCRV2(NemotronOCR):
         # Phase 7: postprocess (broken down into sub-phases)
         t0 = time.perf_counter()
         results = self._postprocess_batch(
-            output, rec_output_raw, original_shapes, merge_level, include_invalid, timings=T,
+            output, original_shapes, merge_level, include_invalid, timings=T,
         )
         T["post_total"] = (time.perf_counter() - t0) * 1000
 
@@ -747,13 +740,13 @@ class NemotronOCRV2(NemotronOCR):
 
         logger.info(
             "batch=%d  regions=%d(%s)  "
-            "img_load=%.1f  gpu_pre=%.1f  det=%.1f  nms=%.1f  "
+            "img_load=%.1f  gpu_pre=%.1f  det=%.1f  prefilt=%.1f  nms=%.1f  "
             "rectify=%.1f  recog=%.1f  rel=%.1f  "
             "post[decode=%.1f rel=%.1f graph=%.1f fmt=%.1f]=%.1f  "
             "total=%.1f ms",
             num_images, total_regions, counts_str,
             T.get("img_load", 0), T.get("gpu_preproc", 0),
-            T["detector"], T["nms"],
+            T["detector"], T["prefilter"], T["nms"],
             T["rectify"], T["recognizer"], T["relational"],
             T.get("post_recog_decode", 0), T.get("post_rel_decode", 0),
             T.get("post_graph_build", 0), T.get("post_format", 0),