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

[cptnm3]

Pull Request

Issue

Fixes #1331
Fixes #1549

Description

Each ModelOutputThunk returned by _generate_from_raw previously left mot.raw.response as None, making the HF raw path inconsistent with all other backends and with the HF chat path.

Changes:

• Inside the per-MOT loop, construct a GenerateDecoderOnlyOutput slice for row i using tensor views (no .clone()) so no additional GPU memory is allocated. sequences, scores, and logits are sliced; past_key_values, attentions, and hidden_states are set to None with a one-time debug log.
• After the loop, drop the shared outputs object — null out sequences, scores, logits, past_key_values, attentions, and hidden_states with hasattr guards, then call gc.collect() and torch.cuda.empty_cache(). Debug-log GPU memory before and after. Per-MOT views keep the underlying tensor storage alive via refcounting.

Tests added:

• test_generate_from_raw_raw_response_set_per_mot — asserts raw.response is set, sequences shape is (1, seq_len), storage is shared (view not clone), and omitted fields are None.
• test_generate_from_raw_raw_response_scores_are_views_when_logits_requested — asserts raw.response.scores is a tuple of views sharing storage with the original batch scores when LOGITS=True.
• test_generate_from_raw_raw_response_scores_none_when_logits_not_requested — asserts raw.response.scores is None when model.generate() returns no scores.
• New file test_huggingface_raw_response_copy.py with five tests covering shallow copy (shared raw.response identity and storage) and deepcopy (distinct object, broken storage sharing, preserved values).

Testing

• Tests added to the respective file if code was changed
• New code has 100% coverage if code was added
• Ensure existing tests and github automation passes (a maintainer will kick off the github automation when the rest of the PR is populated)

Attribution

• AI coding assistants used

---

Adding a new component, requirement, sampling strategy, or tool?

If your PR adds or modifies one of the types below, check the matching box. A checklist of type-specific review items will be posted as a comment.

• Component
• Requirement
• Sampling Strategy
• Tool

NOTE: Please ensure you have an issue that has been acknowledged by a core contributor and routed you to open a pull request against this repository. Otherwise, please open an issue before continuing with this pull request.

[jakelorocco]

Can you please double check the claims that we aren't storing unnecessary tensors / rows? I did a quick investigation (and then had Claude write some tests in https://github.com/jakelorocco/mellea/tree/test/backend-memory-regressions); I believe the whole tensor is saved when we have a single view into it.

Here's the corresponding analysis of issue/test:

Review finding	Test	On PR code
Views pin the whole batch (sequences)	test_raw_response_sequences_retain_only_their_own_row	FAIL — retains 256 B for a 32 B row
Views pin the whole batch (scores)	test_raw_response_scores_retain_only_their_own_row	FAIL — retains 2048 B for a 256 B row
Same, logits/RAW_LOGITS branch (untested by the PR)	test_raw_response_raw_logits_retain_only_their_own_row	FAIL — 1024 B for a 256 B row
Holding one MOT keeps the entire batch alive	test_batch_tensors_are_freed_once_only_mots_are_held	FAIL — batch tensor still alive via weakref
deepcopy duplicates the batch	test_deepcopy_of_result_does_not_duplicate_the_batch	FAIL — 256 B allocated for a 32 B row
Per-call gc.collect() + empty_cache()	test_generate_from_raw_does_not_force_gc_or_cuda_flush_without_cuda	FAIL — 1 full GC pass with no CUDA
Dead isinstance branch emits sequences=None	test_raw_response_is_never_emitted_with_null_sequences	FAIL
del out.f / out.f = None never frees (chat path)	test_post_processing_clearing_raw_logits_actually_releases_them	FAIL — tensors alive after clear
Mislabeled invariant in the PR's test	test_raw_response_scores_follow_generate_output_not_the_logits_option	PASS (documents the real rule)
One-time notice shouldn't spam per item/call	test_omitted_fields_notice_is_logged_once_per_backend	PASS (regression guard)

Maybe we are fine with one mot causing the full tensor to be saved, but that seems excessive to me (unless the fix is complicated / messy).

[planetf1]

Claude review.

Confirming @jakelorocco's finding, and correcting one row of the table.

The retention is real. A row view keeps the whole batch storage alive, not just its row, and nothing releases it: generate_from_raw (mellea/core/backend.py:232) never calls post_processing, and unlike the chat path at :1617 the raw path never sets raw.response = None. This PR's own test shows the size: at batch 2, vocab 32000, fp32, the view's storage is 256000 bytes where a per-row clone is 128000. The ratio is the batch size, so "generate N, keep the best one" retains N times what it needs.

On "maybe we are fine with one mot causing the full tensor to be saved": the file already decided it twice. :1798 reads # Clone each slice so this MOT does not hold a view into the shared batch allocation., and test_generate_from_raw_logits_sliced_per_item at :561 requires logits must be a clone, not a view for this same function. Nor is the fix messy: .detach().clone() on the three slices passes all 64 tests across both HF unit files, ruff format clean.

Correction: the isinstance branch is not dead. Removing it fails four test_multimodal_blocks_in_raw_ctx_not_checked cases, because test_huggingface_unit.py:904 mocks sequences as a plain list, so outputs.sequences[i : i + 1, :] raises TypeError: list indices must be integers or slices, not tuple. Keep the guard, clone inside it. It does read as dead code from the diff; only running it showed otherwise.

Two smaller additions on the cleanup block:

• memory_allocated() can't move here. It reports memory occupied by tensors, while empty_cache() releases unoccupied cached memory. memory_reserved() is the one that would move.
• It's CUDA-only twice over: torch.cuda.empty_cache() is a no-op off CUDA and both debug logs sit behind torch.cuda.is_available(), yet the backend selects cuda, then mps, then cpu at :399-405 and this function already special-cases mps at :1681. On mps or cpu it reduces to a gc.collect().

The two inline comments are on the new tests, which pin the view-sharing as the contract and would invert alongside the fix.

[planetf1]

Correcting my own comment above: "keep the guard, clone inside it" is not sufficient on its own.

[cptnm3]

Thanks @jakelorocco, @planetf1 for the comments. I'm working on incorporating requested changes.

[cptnm3]

Hi @jakelorocco @planetf1,
I have made the requested changes. @jakelorocco All the tests provided by you are green now, except test_post_processing_clearing_raw_logits_actually_releases_them. I have left out the changes needed for that test as my changes were purely on the generate_from_raw path and the required changes would affect the chat path. Please let me know if you would like me to add changes for the test test_post_processing_clearing_raw_logits_actually_releases_them to pass in this PR
Please take a look and let me know if there are any more changes required.
Thanks!

[jakelorocco]

This output deletion code apparently doesn't quite work due to transformer's implementation:

GenerateDecoderOnlyOutput subclasses transformers.utils.generic.ModelOutput, which subclasses OrderedDict. It overrides __setattr__/__setitem__ to keep the attribute and the dict entry in sync, but does not override __delattr__. So del outputs.sequences removes only the __dict__ slot; the OrderedDict entry keeps a strong reference. Verified locally:

del o.sequences
after del: has attr? False
after del: dict entry? True True
o["sequences"] is t: True        # tensor still alive
refcount: 4 -> 3                 # one slot dropped, dict ref retained

[jakelorocco]

I believe you need to do something like outputs["sequences"] = None.

[jakelorocco]

I think you also have to delete previous references like sequences_to_decode from above before the gc / cache clear will work.

So the final version would be something like:

del sequences_to_decode  # views into outputs.sequences
outputs = None           # drops the ModelOutput and its dict entries

[jakelorocco]

I think a test like this would work:

async def test_post_processing_clearing_raw_logits_actually_releases_them():
    """Clearing `hf_output.logits` must drop the tensors, not just the attribute.

    `GenerateDecoderOnlyOutput` is a `ModelOutput`, i.e. an `OrderedDict` subclass
    that mirrors every field into the mapping. `ModelOutput.__setattr__` skips the
    mapping write when the value is `None`, and `ModelOutput` defines no
    `__delattr__`, so `out.logits = None` and `del out.logits` both leave the
    mapping entry — and therefore the tensors — in place. Any code that nulls a
    field to free memory while keeping the container has to clear the mapping too.
    """
    backend = _make_backend(1)
    backend._use_caches = True  # keeps raw.response, so the container survives

    mot, refs = await _post_process_holding_only_weakrefs(backend, n_steps=2)
    gc.collect()
    gc.collect()

    assert mot.raw.response is not None, "test setup: raw.response should be retained"
    assert mot.raw.response.logits is None, "test setup: logits attribute was cleared"
    for step, ref in enumerate(refs["logits"]):
        assert ref() is None, (
            f"raw logits tensor for step {step} is still alive after hf_output.logits "
            "was set to None — the ModelOutput mapping entry still references it"
        )

[cptnm3]

This issue existed in existing chat path cleanup as well. I have made changes to cover both chat and batch path to do the cleanup properly.
tested test_post_processing_clearing_raw_logits_actually_releases_them and it was green.

[jakelorocco]

I didn't realize that we are already duplicating logits here. Can we just offer a view into this or do some common extraction so that we don't reduplicate them above?

[cptnm3]

I did originally want to merge these but did not, as they needed different structures. But now I have tried to implement the slicing and share the same clones across both raw.response and result.generation. Please take a look and share your thoughts

[jakelorocco]

We should probably also gate this saving / copying logic on _use_caches=True. Otherwise, everyone pays unconditionally for this. I believe this is how the chat completions api is treated.

[cptnm3]

Added

[planetf1]

All my threads are now resolved — verified at eb0ea8cd (ran test_huggingface_unit.py there: 64/64 pass).

Three nits, all non-blocking, for this PR or a follow-up:

1. The docstring on test_generate_from_raw_raw_response_set_per_mot (test_huggingface_unit.py:1342) still says the sequences slice is a "view, not clone" — the body now asserts the opposite. One line.
2. The two new "raw.response stays None" paths — non-tensor sequences and beam search — have no tests in this PR. Both invariants are pinned by the code alone.
3. The post_processing release fix: the regression suite referenced in the thread (test_huggingface_raw_response_memory.py) lives on a separate branch (test/backend-memory-regressions), not in this PR. Worth stating whether it lands here or ships separately.

Nothing above is a blocker, but I am holding approval until the open items in @jakelorocco's threads are addressed.

[AngeloDanducci]

Edited the top level comment to include Fixes #1549 since that is also covered by this PR.

[planetf1]

Two questions:

• Will @jakelorocco's regression tests (the test/backend-memory-regressions branch) go into this PR? As it stands, nothing in the diff proves the tensors are actually released.
• With the _use_caches gate, raw.response is still None when caching is off (same as the chat path, as requested in-thread). Is that the intended scope for #1331, or a gap we should close?

A few small things inline. The one on the non-cached clear list I'd like resolved before merge: "Fixes #1549" will auto-close the issue, and that corner would still be open. No blockers.

[planetf1]

One thing this list lets through: on the intrinsic path with caching off, the captured output (the raw_hf_output_cell at :752) is still reachable through mot._gen.process (the partial at :828) even after mot.raw.response = None below — and this list doesn't touch past_key_values. So with logits requested, the KV cache can still linger on the held MOT. Not a regression — before this PR the whole output was pinned there, since the old del never reached the mapping — but "Fixes #1549" closes the issue on merge with this corner still open and untracked. Either fix it here (add past_key_values to the list and drop the cell reference, with a weakref test that holds the MOT) or open a tracking issue.

[planetf1]

One that predates the PR: if batch_decode, the clone loop or action.parse raises, this cleanup never runs and the traceback keeps outputs — the whole batch — alive for as long as the exception lives. The PR improves this path overall, so just noting it; a try/finally would close the gap.

[planetf1]

Still says "view, not clone" here while the assertion below it requires a clone — my open nit from the 2026-08-19 comment. One word.

[planetf1]

This deliberately builds a view to test copy semantics (fine) — but the comment reads as though the raw batch path produces a view, and it clones now. A word or two so nobody chases a view that no longer exists.

[planetf1]

The docstring promises a 4-tuple; the helper returns a 3.

[planetf1]

This file skips on llguidance, but nothing it imports needs it — an install without llguidance silently skips these five tests for no reason.