feat(backends): populate mot.raw.response on the HD generate_from_raw

mellea/backends/huggingface.py

@@ -1525,12 +1525,25 @@ class used during generation, if any.
             self.cache_put(cache_key, cache_info)
 
             # Clear KV cache and scores from HF output; retained via LRU cache above.
-            hf_output.past_key_values = None
-            hf_output.scores = None
+            # `ModelOutput` (`OrderedDict` subclass) does not sync `None` writes back
+            # to the mapping, so plain attribute assignment leaves the dict entry — and
+            # its tensor — alive.
+            for _field in ("past_key_values", "scores"):
+                if _field in hf_output:
+                    dict.__delitem__(hf_output, _field)
+                try:
+                    object.__delattr__(hf_output, _field)
+                except AttributeError:
+                    pass
 
         # Clear the raw logits tensor (scores already cleared above if cached).
         if isinstance(hf_output, GenerateDecoderOnlyOutput):
-            hf_output.logits = None
+            if "logits" in hf_output:
+                dict.__delitem__(hf_output, "logits")
+            try:
+                object.__delattr__(hf_output, "logits")
+            except AttributeError:
+                pass
 
         # Only scan for tools if we are not doing structured output and tool calls were provided to the model.
         if _format is None and tool_calls:
@@ -1608,20 +1621,26 @@ class used during generation, if any.
         if not self._use_caches and isinstance(
             mot.raw.response, GenerateDecoderOnlyOutput
         ):
-            import gc
-
             hf_out = mot.raw.response
-            if hasattr(hf_out, "sequences") and hf_out.sequences is not None:
-                del hf_out.sequences
-            if hasattr(hf_out, "scores") and hf_out.scores is not None:
-                del hf_out.scores
-            if hasattr(hf_out, "logits") and hf_out.logits is not None:
-                del hf_out.logits
+            # GenerateDecoderOnlyOutput is a ModelOutput (OrderedDict subclass).
+            # ModelOutput.__setattr__ skips the dict write when value is None, and
+            # ModelOutput defines no __delattr__, so both `out.f = None` and
+            # `del out.f` leave the mapping entry — and its tensor — alive.
+            # Clear both the dict entry and the instance attribute to release them.
+            for field in ("sequences", "scores", "logits"):
+                if field in hf_out:
+                    dict.__delitem__(hf_out, field)
+                try:
+                    object.__delattr__(hf_out, field)
+                except AttributeError:
+                    pass
             mot.raw.response = None
 
-            # Force Python GC and return CUDA memory to device
-            gc.collect()
-            torch.cuda.empty_cache()
+            if torch.cuda.is_available():
+                import gc
+
+                gc.collect()
+                torch.cuda.empty_cache()
 
         # Generate the log for this ModelOutputThunk.
         generate_log = GenerateLog()
@@ -1766,16 +1785,75 @@ async def _generate_from_raw(
             result.generation.provider = self._provider
             result.raw.provider = self._provider
 
-            if want_logits and outputs.scores is not None:
-                # Clone each slice so this MOT does not hold a view into the shared batch allocation.
-                result.generation.logits = tuple(
-                    step_scores[i].detach().clone() for step_scores in outputs.scores
+            # Extract per-MOT tensors once. These are shared by generation
+            # and, where supported, raw.response.
+            mot_scores: tuple[torch.Tensor, ...] | None = None
+            if outputs.scores is not None:
+                mot_scores = tuple(
+                    score[i].detach().clone() for score in outputs.scores
+                )
+
+            mot_logits_raw: tuple[torch.Tensor, ...] | None = None
+            if outputs.logits is not None:
+                mot_logits_raw = tuple(
+                    logits[i].detach().clone() for logits in outputs.logits
+                )
+
+            # Construct a per-MOT GenerateDecoderOnlyOutput slice holding clones of row i.
+            # past_key_values, attentions, and hidden_states are shared across the batch
+            # and cannot be sliced per MOT, so they are omitted.
+            # Note: beam-search outputs (GenerateBeamDecoderOnlyOutput) are not supported here;
+            # mot.raw.response will be None.
+            if (
+                self._use_caches
+                and isinstance(outputs, GenerateDecoderOnlyOutput)
+                and isinstance(outputs.sequences, torch.Tensor)
+            ):
+                response_scores = (
+                    tuple(score.unsqueeze(0) for score in mot_scores)
+                    if mot_scores is not None
+                    else None
                 )
+                response_logits = (
+                    tuple(logits.unsqueeze(0) for logits in mot_logits_raw)
+                    if mot_logits_raw is not None
+                    else None
+                )
+
+                if "raw_batch_response_fields_omitted" not in self._warned_about:
+                    self._warned_about.add("raw_batch_response_fields_omitted")
+                    MelleaLogger.get_logger().debug(
+                        "mot.raw.response.past_key_values, .attentions, and "
+                        ".hidden_states are not available on the raw batch path "
+                        "and will always be None."
+                    )
 
-            if want_raw_logits and outputs.logits is not None:
-                result.generation.raw_logits = tuple(
-                    step_logits[i].detach().clone() for step_logits in outputs.logits
+                result.raw.response = GenerateDecoderOnlyOutput(
+                    sequences=cast(
+                        "torch.LongTensor",
+                        outputs.sequences[i : i + 1, :].detach().clone(),
+                    ),
+                    scores=cast("tuple[torch.FloatTensor] | None", response_scores),
+                    logits=cast("tuple[torch.FloatTensor] | None", response_logits),
+                    attentions=None,
+                    hidden_states=None,
+                    past_key_values=None,
                 )
+            elif self._use_caches:
+                # Beam search slicing is not feasible — raw.response is not populated.
+                warn_key = "raw_batch_beam_search_unsupported"
+                if warn_key not in self._warned_about:
+                    self._warned_about.add(warn_key)
+                    MelleaLogger.get_logger().debug(
+                        "mot.raw.response is not available on the raw batch and will be None."
+                    )
+
+            # Reuse the cloned tensors.
+            if want_logits and mot_scores is not None:
+                result.generation.logits = mot_scores
+
+            if want_raw_logits and mot_logits_raw is not None:
+                result.generation.raw_logits = mot_logits_raw
 
             action = actions[i]
             result.parsed_repr = (
@@ -1794,6 +1872,28 @@ async def _generate_from_raw(
             result._generate_log = generate_log
             results.append(result)
 
+        # Drop all references that might pin the shared batch tensors.
+        # `sequences_to_decode` holds slice views of `outputs.sequences`, and
+        # `outputs` is a `GenerateDecoderOnlyOutput` — a `ModelOutput`, which
+        # subclasses `OrderedDict`. `del obj.attr` only removes the `__dict__`
+        # slot; the `OrderedDict` entry keeps a strong reference to the tensor.
+        # Setting `outputs = None` drops the whole container at once
+        del sequences_to_decode
+        outputs = None
+        if torch.cuda.is_available():
+            import gc
+
+            MelleaLogger.get_logger().debug(
+                "GPU memory before raw batch cleanup: %d bytes reserved",
+                torch.cuda.memory_reserved(),
+            )
+            gc.collect()
+            torch.cuda.empty_cache()
+            MelleaLogger.get_logger().debug(
+                "GPU memory after raw batch cleanup: %d bytes reserved",
+                torch.cuda.memory_reserved(),
+            )
+
         usage: dict[str, Any] | None = (
             {
                 "prompt_tokens": agg_prompt,

test/backends/test_huggingface_raw_response_copy.py

@@ -0,0 +1,90 @@
+# Copyright IBM Corp. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Unit tests for copy/deepcopy semantics of ModelOutputThunks that hold an HF
+GenerateDecoderOnlyOutput in raw.response — the raw batch path specific case."""
+
+import copy
+from typing import Any
+
+import pytest
+
+torch = pytest.importorskip("torch", reason="torch not installed — install mellea[hf]")
+pytest.importorskip(
+    "transformers", reason="transformers not installed — install mellea[hf]"
+)
+pytest.importorskip(
+    "llguidance", reason="llguidance not installed — install mellea[hf]"
+)
+
+from transformers.generation.utils import GenerateDecoderOnlyOutput
+
+from mellea.core import ModelOutputThunk
+
+
+def _make_mot_with_hf_raw_response() -> tuple[ModelOutputThunk, Any]:
+    """Build a MOT whose raw.response is a CPU-only GenerateDecoderOnlyOutput with a view."""
+    full_batch = torch.arange(6, dtype=torch.long).reshape(2, 3)
+    # Simulate the view produced by the raw batch path for batch item 0.
+    row_view = full_batch[0:1, :]
+    hf_out = GenerateDecoderOnlyOutput(
+        sequences=row_view,
+        scores=None,
+        logits=None,
+        attentions=None,
+        hidden_states=None,
+        past_key_values=None,
+    )
+    mot = ModelOutputThunk(value="hello")
+    mot.raw.response = hf_out
+    return mot, full_batch
+
+
+def test_shallow_copy_raw_response_is_same_object():
+    """copy.copy(mot).raw.response is the same object as mot.raw.response."""
+    mot, _ = _make_mot_with_hf_raw_response()
+    copied = copy.copy(mot)
+    assert copied.raw.response is mot.raw.response, (
+        "shallow copy must keep raw.response as the same object"
+    )
+
+
+def test_shallow_copy_raw_response_sequences_shares_storage():
+    """Shallow-copied MOT preserves shared tensor storage for raw.response.sequences."""
+    mot, full_batch = _make_mot_with_hf_raw_response()
+    copied = copy.copy(mot)
+    assert (
+        copied.raw.response.sequences.untyped_storage().data_ptr()
+        == full_batch.untyped_storage().data_ptr()
+    ), (
+        "shallow copy: raw.response.sequences must still share storage with the original batch"
+    )
+
+
+def test_deepcopy_raw_response_is_distinct_object():
+    """copy.deepcopy(mot).raw.response is a distinct object from mot.raw.response."""
+    mot, _ = _make_mot_with_hf_raw_response()
+    deep = copy.deepcopy(mot)
+    assert deep.raw.response is not mot.raw.response, (
+        "deepcopy must produce a new raw.response object"
+    )
+
+
+def test_deepcopy_raw_response_sequences_does_not_share_storage():
+    """Deepcopy breaks tensor storage sharing for raw.response.sequences."""
+    mot, full_batch = _make_mot_with_hf_raw_response()
+    deep = copy.deepcopy(mot)
+    assert (
+        deep.raw.response.sequences.untyped_storage().data_ptr()
+        != full_batch.untyped_storage().data_ptr()
+    ), "deepcopy: raw.response.sequences must NOT share storage with the original"
+
+
+def test_deepcopy_raw_response_sequences_preserves_values():
+    """Deepcopy preserves tensor values in raw.response.sequences despite storage isolation."""
+    mot, _ = _make_mot_with_hf_raw_response()
+    original_values = mot.raw.response.sequences.clone()
+    deep = copy.deepcopy(mot)
+    assert torch.equal(deep.raw.response.sequences, original_values), (
+        "deepcopy must preserve tensor values in raw.response.sequences"
+    )