Skip to content

perf(vit): fuse ViT axial RoPE via ggml GGML_ROPE_TYPE_VISION kernel - #14

Open
gabilan wants to merge 1 commit into
PABannier:mainfrom
gabilan:perf/rope-fusion
Open

perf(vit): fuse ViT axial RoPE via ggml GGML_ROPE_TYPE_VISION kernel#14
gabilan wants to merge 1 commit into
PABannier:mainfrom
gabilan:perf/rope-fusion

Conversation

@gabilan

@gabilan gabilan commented Aug 1, 2026

Copy link
Copy Markdown

The Hiera ViT encoder applied 2D axial RoPE with a hand-rolled complex multiply (sam3_apply_rope): reshape + strided (real,imag) views + 4 muls

  • sub + add + concat + cont, ~8 primitive ggml ops per Q/K per block, 64 calls per encode. On Metal this dominated the encode: ~828ms / ~38% of a ~2170ms 1008x1008 encode.

ggml already ships a single fused vision-RoPE kernel (ggml_rope_multi with GGML_ROPE_TYPE_VISION), but it computes rotary frequencies internally from theta+positions and pairs head_dim elements split-half (i, i+n_dims), whereas the model stores a precomputed freqs_cis table paired interleaved ((2i, 2i+1)). This bridges the two conventions:

  • Reorder head_dim once from interleaved to split-half before the kernel (new[c]=old[2c], new[c+n_dims]=old[2c+1]) so the kernel rotates the same (real,imag) pair the stored table does.
  • n_dims = head_dim/2, sections {n_dims/2, n_dims/2, 0, 0}: X-axis frequencies land on head_dim pairs 0..15 (position block 0), Y-axis on 16..31 (block 1), freq = base^(-c/16) -- identical to the stored compute_axial_cis layout.
  • freq_scale interpolates global-block positions back to the WS x WS pretrain grid (WS/n_img_embd = 24/72 = pretrain_image_size/image_size); windowed blocks are native (1.0). This reproduces the stored table's scale_pos.

The head_dim reorder is deliberately NOT undone: RoPE touches only Q and K, and the attention score Q.K^T is invariant under a head_dim permutation applied consistently to both, while V (and the output basis) is untouched -- so the attention result is unchanged while saving the reorder-back copy.

Validation (real q8_0 model, Metal, cat.jpg):

  • The live-computed rotary frequencies were checked element-wise against every block's stored freqs_cis table: max |stored - computed| = 1.19e-07 (f32 epsilon) across all 32 blocks, both windowed and global.
  • Golden mask regression vs the pre-change build: point/box/text-detection masks agree at IoU 0.999986 / 0.999996 / 0.999992 (3-12 disagreeing pixels out of 1.44M, all on mask boundaries); confidence deltas ~6e-6.
  • Full downstream test suite (395 assertions, 17 cases) passes unchanged.

Isolated encode wall-clock: 2173.5ms -> 1378.8ms on Metal (1.58x).

The Hiera ViT encoder applied 2D axial RoPE with a hand-rolled complex
multiply (sam3_apply_rope): reshape + strided (real,imag) views + 4 muls
+ sub + add + concat + cont, ~8 primitive ggml ops per Q/K per block, 64
calls per encode. On Metal this dominated the encode: ~828ms / ~38% of a
~2170ms 1008x1008 encode.

ggml already ships a single fused vision-RoPE kernel (ggml_rope_multi with
GGML_ROPE_TYPE_VISION), but it computes rotary frequencies internally from
theta+positions and pairs head_dim elements split-half (i, i+n_dims),
whereas the model stores a precomputed freqs_cis table paired interleaved
((2i, 2i+1)). This bridges the two conventions:

  * Reorder head_dim once from interleaved to split-half before the kernel
    (new[c]=old[2c], new[c+n_dims]=old[2c+1]) so the kernel rotates the
    same (real,imag) pair the stored table does.
  * n_dims = head_dim/2, sections {n_dims/2, n_dims/2, 0, 0}: X-axis
    frequencies land on head_dim pairs 0..15 (position block 0), Y-axis on
    16..31 (block 1), freq = base^(-c/16) -- identical to the stored
    compute_axial_cis layout.
  * freq_scale interpolates global-block positions back to the WS x WS
    pretrain grid (WS/n_img_embd = 24/72 = pretrain_image_size/image_size);
    windowed blocks are native (1.0). This reproduces the stored table's
    scale_pos.

The head_dim reorder is deliberately NOT undone: RoPE touches only Q and K,
and the attention score Q.K^T is invariant under a head_dim permutation
applied consistently to both, while V (and the output basis) is untouched --
so the attention result is unchanged while saving the reorder-back copy.

Validation (real q8_0 model, Metal, cat.jpg):
  * The live-computed rotary frequencies were checked element-wise against
    every block's stored freqs_cis table: max |stored - computed| = 1.19e-07
    (f32 epsilon) across all 32 blocks, both windowed and global.
  * Golden mask regression vs the pre-change build: point/box/text-detection
    masks agree at IoU 0.999986 / 0.999996 / 0.999992 (3-12 disagreeing
    pixels out of 1.44M, all on mask boundaries); confidence deltas ~6e-6.
  * Full downstream test suite (395 assertions, 17 cases) passes unchanged.

Isolated encode wall-clock: 2173.5ms -> 1378.8ms on Metal (1.58x).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant