Enable torch.export of the full SAM3 grounding pipeline

[rbavery]

Summary

Bundle the SAM3 image encoder, text encoder, encoder fusion, and decoder
into a single ExportedProgram so consumers can ship one .pt2. The
graph accepts dynamic batch (B>=1) and prompt (P>=1) dims; image H/W and
CLIP context length are fixed at 1008 and 32.

This is a minimal redo of #3: only the patches strictly needed for
torch.export of the full pipeline + a single test, on top of latest main.
No artifact/sub-module export scripts, no benchmarks, no uv.lock.

The exported contract preserves what wherobots-rasterflow relies on
(/home/rave/wherobots-rasterflow/.../adapters/torch/geometry_actors.py
calls self.model(patch_batch, token_ids) and unpacks 4 outputs).

Verified on RTX 3090 (torch 2.10.0+cu128)

range_constraints: {s34: VR[1, int_oo], s69: VR[1, int_oo]}

$ uv run pytest tests/export/test_full_pipeline_export.py -m slow
test_full_pipeline_export_matches_eager                          PASSED  (atol=0)
test_full_pipeline_export_save_load_roundtrip                    PASSED
test_full_pipeline_export_supports_dynamic_shapes[1-1]           PASSED
test_full_pipeline_export_supports_dynamic_shapes[1-2]           PASSED
test_full_pipeline_export_supports_dynamic_shapes[3-2]           PASSED
test_full_pipeline_export_supports_dynamic_shapes[2-4]           PASSED
======== 6 passed in ~3 min ========

Eager-vs-exported max abs diff is 0 on all 4 outputs (pred_logits,
pred_boxes, pred_masks, presence_logit_dec).

Required model patches

File	Change
sam3/model/encoder.py	Fix x.dim → x.dim() (called the method object instead of invoking it)
sam3/model/geometry_encoders.py	Skip pin_memory + non_blocking copy under torch.compiler.is_compiling() and on CPU. Convert ROIAlign input from a list-of-tensors to the [N, 5] batched form (lists of tensors aren't traceable).
sam3/model/position_encoding.py	Skip the _encode_xy ndim/length assert when tracing — symbolic dims trip GuardOnDataDependentSymNode. Skip the shape-keyed forward cache when shape items are SymInts.
sam3/model/decoder.py	(1) Allow 4D (B, H, Q, K) cross_attn_mask through TransformerDecoderLayer and add _cross_attn_with_rpb, a manual SDPA path that preserves a per-head additive bias (nn.MultiheadAttention rejects 4D masks). Casts proj weights to activation dtype to coexist with bf16 autocast. (2) Forward the box-RPB matrix in 4D (drop the flatten(0, 1) that produced the MHA-friendly 3D shape). (3) _make_box_rpb_relative: use torch.compiler.is_compiling() instead of is_dynamo_compiling() and drop the SymInt tuple-equality cache check + eager-only shape asserts on deltas/B.
sam3/model_builder.py	Thread num_feature_levels through build_sam3_image_model → _create_sam3_transformer / _create_sam3_model so callers can pin it explicitly.

New files

• scripts/export_sam3_full_pipeline.py — FullSam3PipelineWrapper that unifies the four sub-modules into one forward. Wraps the model call in torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16) because upstream's ViT MLP uses sam3.perflib.fused.addmm_act (added in the SAM 3.1 release) which forces bf16 internally and Sam3TrackingPredictor enters the same autocast in __init__. Outputs are cast back to fp32 so consumers see the same input/output dtype contract production already expects. CLI saves a .pt2 with Dim("batch", min=1) and Dim("num_prompts", min=1).

• tests/export/test_full_pipeline_export.py — slow pytest marked with the slow marker (registered in pyproject.toml). Three tests:

  1. test_full_pipeline_export_matches_eager: exact-equality (atol=0) check between the eager wrapper and the exported program on the same input.
  2. test_full_pipeline_export_save_load_roundtrip: torch.export.save → torch.export.load → run on fresh inputs.
  3. test_full_pipeline_export_supports_dynamic_shapes: parametric run across (1,1), (1,2), (3,2), (2,4) (batch, num_prompts) pairs.

Notes / caveats

• The 2.11 release notes review (from Bump torch to 2.11 and torchvision to 0.26 #8) holds: none of these workarounds have a clean 2.11-specific revert. After merging Bump torch to 2.11 and torchvision to 0.26 #8 we can re-test whether the manual SDPA RPB path can fall back to plain nn.MultiheadAttention; for now it's load-bearing.
• AOTI (issue pytorch/pytorch#174608) is the next step; this PR sets up a clean base for that experiment.
• Loading the saved .pt2 requires import torchvision.ops first to register torch.ops.torchvision.roi_align.default — same as the production deploy.