Fix template offload helpers for long-sequence inference

[taivu1998]

Summary

Fixes #489.

This PR repairs the monomer template memory-saving paths used by long-sequence inference. embed_templates_offload and embed_templates_average now call template submodules through the actual owner, model.template_embedder, instead of looking for stale direct attributes on AlphaFold.

Root Cause

AlphaFold constructs monomer template machinery as a TemplateEmbedder instance at self.template_embedder. The normal template path already uses that object, but the offload and average helper functions still called model.template_pair_embedder, model.template_pair_stack, model.template_pointwise_att, and model.template_single_embedder directly. Enabling long_sequence_inference=True turns on template offload, so the stale direct access raises:

AttributeError: 'AlphaFold' object has no attribute 'template_pair_embedder'

Changes

• Route embed_templates_offload and embed_templates_average through model.template_embedder.*.
• Keep runtime flags on model.globals and template-specific config on the template embedder config.
• Fix final template-mask broadcasting for batched template inputs.
• Correct the offload helper's inplace_safe final mask branch.
• Remove the brittle sys.getrefcount assertion from the offload helper.
• Update trace_utils.py to use model.template_embedder.template_pair_stack.
• Add CPU-friendly regression tests for offload and average template paths, including no-batch and batched template dimensions.

Validation

Validated in an isolated worktree. The local macOS environment does not have the compiled OpenFold attention extension installed, so I used a disposable /private/tmp test venv and a test-only stub for the optional attn_core_inplace_cuda import; the tested code path does not call that extension.

Commands run:

PYTHONPATH=/private/tmp/openfold-issue-489-test-stubs:/private/tmp/openfold-issue-489-template-offload /private/tmp/openfold-issue-489-test-venv/bin/python -m unittest tests.test_template.TemplateOffloadRegression
PYTHONPATH=/private/tmp/openfold-issue-489-test-stubs:/private/tmp/openfold-issue-489-template-offload /private/tmp/openfold-issue-489-test-venv/bin/python -m unittest tests.test_template
python -m py_compile openfold/model/template.py openfold/utils/trace_utils.py tests/test_template.py
git diff --check
rg -n "model\\.template_(pair_embedder|pair_stack|pointwise_att|single_embedder)|getrefcount|import sys" openfold/model/template.py openfold/utils/trace_utils.py

Results:

• TemplateOffloadRegression: 3 tests passed
• tests.test_template: 7 tests passed, 2 skipped
• Syntax compile passed
• Diff whitespace check passed
• Static stale-reference search returned no matches