Request For Help: on exporting the SAM3 Vision Encoder to ONNX? (view_as_complex, SDPA, dynamic ops issues)

[Aravind-Sreenivas]

Hi — thanks for your work on EfficientSAM3 (great project).
I’m trying to export the SAM3 vision encoder from the official Meta implementation into ONNX, mainly for TensorRT deployment. However, the vanilla SAM3 code fails ONNX export with errors like:

UnsupportedOperatorError: aten::view_as_complex

and several dynamic-shape / SDPA-related issues.

Where the vanilla SAM3 export currently fails

• view_as_complex in rotary embeddings (not supported in ONNX)
• dynamic asserts on shape (assert size * size == xy_num)
• conditional padding logic (if pad_h > 0 or pad_w > 0)
• SDPA → requires replacing with matmul attention
• dynamic cumsum / dynamic position encoding

---

My Question to the EfficientSAM3 team

Since your repo deals deeply with SAM models and modifications, could you guide me on how to make SAM3’s vision encoder ONNX-exportable?

Specifically:

What is the correct approach to handle these?

1. How to remove/replace view_as_complex?
   Should rotary embedding be rewritten using real-valued sin/cos?
2. Should SDPA be replaced with a manual matmul attention block?
3. How to replace dynamic cumsum → static arange?
4. How to precompute/register positional embeddings as buffers?
5. Do FPN layers or padding need to be rewritten?
6. Any tips from your experience exporting SAM/SAM2/EfficientSAM variants?

Even a high-level guide or a minimal patch would be extremely helpful.

Thanks!

[SimonZeng7108]

Thanks for your request. We have not been able to make some time to look into the onnx conversion ATM as we are focusing on distilling the weights first. Your problems are valid and we will help to understand them all hopfully soon. At the mean time, if you were to find a solution, we encourage you to submit a pull request.

[rodneyML]

I had the same issue with EfficientTAM, the workaround is to use other ops that exclude the use of complex operations.
I believe following the same principle will work in your case @Aravind-Sreenivas . I'll post my workaround for EfficientTAM, you should just ensure that the outputs match when applying the workaround.

Original version:

def apply_rotary_enc(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
    repeat_freqs_k: bool = False,
):
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = (
        torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
        if xk.shape[-2] != 0
        else None
    )
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    if xk_ is None:
        # no keys to rotate, due to dropout
        return xq_out.type_as(xq).to(xq.device), xk
    # repeat freqs along seq_len dim to match k seq_len
    if repeat_freqs_k:
        r = xk_.shape[-2] // xq_.shape[-2]
        if freqs_cis.is_cuda:
            freqs_cis = freqs_cis.repeat(*([1] * (freqs_cis.ndim - 2)), r, 1)
        else:
            # torch.repeat on complex numbers may not be supported on non-CUDA devices
            # (freqs_cis has 4 dims and we repeat on dim 2) so we use expand + flatten
            freqs_cis = freqs_cis.unsqueeze(2).expand(-1, -1, r, -1, -1).flatten(2, 3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq).to(xq.device), xk_out.type_as(xk).to(xk.device)

Workaround version:

# for cross-attention
def apply_rotary_matenc_ca(xq, xk, rotmats, repeat_freqs_k=False): # with repeats
    
    bq, hq, nq, cq = xq.shape 
    bk, hk, nk, ck = xk.shape
    
    xq_reshape = xq.reshape(bq, hq, nq, cq // 2, 2, 1)
    q_out = torch.matmul(rotmats, xq_reshape)
    q_out_flattened = q_out.flatten(3)
    
    k_rotmat = rotmats.repeat(1, 1, nk // nq, 1, 1, 1)
    xk_reshape = xk.reshape(bk, hk, nk, ck // 2, 2, 1)
    k_out = torch.matmul(k_rotmat, xk_reshape)
    k_out_flattened = k_out.flatten(3)
    
    return q_out_flattened, k_out_flattened

# for self attention
def apply_rotary_matenc_sa(xq, xk, rotmats): # without repeats
    
    bq, hq, nq, cq = xq.shape 
    bk, hk, nk, ck = xk.shape
    
    xq_reshape = xq.reshape(bq, hq, nq, cq // 2, 2, 1)
    q_out = torch.matmul(rotmats, xq_reshape)
    q_out_flattened = q_out.flatten(3)
    
    k_rotmat = rotmats
    xk_reshape = xk.reshape(bk, hk, nk, ck // 2, 2, 1)
    k_out = torch.matmul(k_rotmat, xk_reshape)
    k_out_flattened = k_out.flatten(3)
    
    return q_out_flattened, k_out_flattened

Note two different functions, as one uses repeat_freqs_k, the other one doesn't.
You could probably merge them into one function with a conditional branch.

[Aravind-Sreenivas]

Thanks @rodneyML will try this workaround !

[rodneyML]

@Aravind-Sreenivas , did you manage to convert the full EfficientSAM3 to ONNX, including memory attention and memory encoding? I think a lot of poeple would benefit from that.

[SimonZeng7108]

@Aravind-Sreenivas @rodneyML , not sure if you guys saw this, https://github.com/wkentaro/sam3-onnx, by brilliant @wkentaro

[rodneyML]

@Aravind-Sreenivas @rodneyML , not sure if you guys saw this, https://github.com/wkentaro/sam3-onnx, by brilliant @wkentaro

Thanks @SimonZeng7108 , will have a look. I am however more interested in exporting the full model, i.e. including memory attention so that we could perform video inference.

[darkSuperman]

For exporting to ONNX models, you can refer to the scripts in this project. I've used their project to run a SAM3 model on the ONNX runtime. https://github.com/jamjamjon/usls/tree/main/scripts/sam3-image