facebookresearch / facebookresearch/sam2
Per-frame streaming inference is launch-bound at 512px — CUDA-graph capture gives 5.9×. Is this expected, and is there a supported path?
- Dominant language
- Jupyter Notebook
- Stars
- 19.9k
- Forks
- 2.5k
- PR merge metrics
- No merged PRs in 30d
Description
## Summary
We built a real-time laparoscopic-surgery tracker on SAM 2 and found that per-frame streaming inference at 512px is almost entirely launch-bound, not compute-bound. Capturing the per-frame step into a CUDA graph gave a 5.9× speed-up with numerically identical masks (IoU 0.999). We'd like to understand whether this overhead is expected, and whether our workaround is sound or is relying on internals that may break.
## Environment
RTX A6000 48GB · Xeon E5-2630 v4 · torch 2.6.0+cu124 · CUDA 12.4 · SAM 2.1 checkpoints · bf16 · single object, 1280×720 source
## Measurements
| Model | Res | Eager fps | CUDA-graph fps | Speed-up |
|-----------|------|-----------|----------------|----------|
| tiny | 512 | 19.0 | 112.0 | 5.91× |
| small | 512 | 17.7 | 101.9 | 5.76× |
| base_plus | 512 | 14.2 | 78.8 | 5.56× |
| large | 512 | 10.9 | 50.9 | 4.66× |
| tiny | 1024 | 18.8 | 35.5 | 1.89× |
| large | 1024 | 10.9 | 14.9 | 1.37× |
For tiny@512: eager 52.6 ms/frame vs graph 8.9 ms/frame. The graph replays the identical kernels, so the 43.7 ms difference — 83% of wall time — was launch/dispatch/Python overhead. At 1024 the gain drops to 1.4–1.9× because pixel work dominates, which is consistent with the launch-bound explanation.
## What we're actually asking
**1. Is the 512px regime expected to be launch-bound?** Is SAM 2's per-frame step known to be dominated by launch overhead at lower resolutions, or does this indicate we're driving the model wrong? 83% seems high enough that we suspect we're missing something.
**2. Is there a supported low-latency / streaming path we've missed?** The documented API is `init_state()` → `propagate_in_video()`, which wants the whole video up front. We drive `track_step()` frame-by-frame with our own `output_dict` and prune entries older than `num_maskmem + max_obj_ptrs_in_encoder + 2` to keep per-frame cost flat. Is there an intended way to do this that we overlooked?
**3. Is our CUDA-graph approach sound, or are we depending on internals that will break?** Graphs need static shapes and addresses, but memory attention reads a Python dict keyed by frame index that changes every frame. Four things made it static:
- Freezing the frame index to a constant `T = 100000` during the graph phase so index arithmetic in `_prepare_memory_conditioned_features` constant-folds.
- Fixed-address ring buffers (`RING = 16`, never reallocated); each frame shifts only their contents via in-place `copy_()` outside the graph, with a frozen dict mapping `{T-1-j: ring[j]}` so the graph always reads the same addresses.
- Caching the one CPU→GPU op (object-pointer temporal encoding, `torch.tensor(pos_list).to(device)`), which is constant once the index is frozen. We gate this so the eager fallback still recomputes it — leaking it caused a shape-mismatch crash.
- `cudnn.benchmark=False`, 3 warm-up iterations in a side stream, then `torch.cuda.graph`.
We capture both `forward_image` (Hiera + FPN) and `track_step` (memory attention + mask decoder + memory encoder); host work stays eager. Multi-object reuses the single-object graph by swapping each object's memory into the graph buffers before replay. Every failure path falls back to eager.
Specifically: is freezing the frame index to a sentinel constant safe in general, or does it break any behavior that depends on true temporal distance (e.g. `max_obj_ptrs_in_encoder` windowing) in cases our clips don't exercise?
**4. One concrete finding worth recording:** multi-graph capture works only with PyTorch's default graph pool. Using a private `graph_pool_handle()` hit `use_count > 0` on re-capture. With the default pool we captured tiny+small+base_plus+large, kept all resident, and replayed each correctly.
## Verification
CUDA-graph vs eager: IoU 0.999 across all four models @512 (the 0.001 gap is bf16 non-determinism). Eager streaming vs reference `propagate_in_video`: IoU 1.000.
Happy to share the implementation or run any experiment that would help.
Contributor guide
Research direction
Start by reading the documented init_state() and propagate_in_video() flow alongside track_step(), output_dict, and _prepare_memory_conditioned_features; inspect the reported forward_image and torch.cuda.graph capture path. Validate whether freezing the frame index preserves temporal-window behavior and whether graph_pool_handle() is supported; done means the project confirms a supported streaming path or documents the constraints and safe behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- computer-vision, machine-learning, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100