aymous/sam3-rotary-embedding-2d

SAM3 interleaved 2D RoPE

SAM3's Perception-Encoder rotary apply (apply_rotary_pos_emb_2d): one launch rotates the interleaved pairs of both q and k against the per-layer axial cos/sin tables, fp32 compute with cast back to the input dtype (the reference's upcast contract). Everything moves as aligned 8-byte quad transactions; q/k are consumed in the model's native view(B,S,H,D).transpose(1,2) strided layout, so no .contiguous() copy is reintroduced. Eager instead rotates q and k separately through rotate_pairwise (view/unbind/stack/flatten), materializing ~10 fp32 intermediates per tensor per layer. The cos/sin tables may be fp32 (as the model registers them — read unrounded) or match the q/k dtype.

Built with kernel-builder as a native torch.library op (AOT, loadable via the kernels library). Reference op: transformers apply_rotary_pos_emb_2d (modeling_sam3).

Benchmark (NVIDIA GB10 / sm_121, aarch64)

Mean latency per workload — op_eager / op_compile are the torch reference (eager and torch.compile); custom is this kernel. Both in-model shapes of facebook/sam3 (image 1008, patch 14 → 72×72 tokens, ViT-L 16 heads × head_dim 64, bf16): the windowed layers (28 of 32 — 3×3 windows of 24×24) and the global layers (4 of 32 — the full grid).

shape	op_eager (ms)	op_compile (ms)	custom (ms)	custom vs op_compile
windowed: B=9, seq 576	3.261	0.235	0.187	1.25×
global: B=1, seq 5184	3.200	0.239	0.186	1.29×

17.4× / 17.2× vs op_eager. In-model (facebook/sam3, single forward at image 1008): swapping the kernel into all 32 Sam3ViTRoPEAttention layers takes prefill from 763 ms (stock eager) to 678 ms — 1.13× end-to-end, with 99.9995% binarized-mask agreement on the stock top-10 queries (accumulated bf16 drift on boundary pixels only).

Usage

import torch
from kernels import get_kernel

kernel = get_kernel("aymous/sam3-rotary-embedding-2d", version=1, trust_remote_code=True)
# q, k: [B, heads, seq, head_dim] (strided view+transpose layouts accepted, head_dim innermost)
# cos, sin: contiguous [seq, head_dim], fp32 or q's dtype
q_rot, k_rot = kernel.sam3_rotary_embedding_2d(q, k, cos, sin)

Note: the published build targets aarch64 / CUDA / sm_121 (NVIDIA GB10). Other platforms need the full build matrix.