perf(vit): fuse ViT axial RoPE via ggml GGML_ROPE_TYPE_VISION kernel - #14
Open
gabilan wants to merge 1 commit into
Open
perf(vit): fuse ViT axial RoPE via ggml GGML_ROPE_TYPE_VISION kernel#14gabilan wants to merge 1 commit into
gabilan wants to merge 1 commit into
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The Hiera ViT encoder applied 2D axial RoPE with a hand-rolled complex multiply (sam3_apply_rope): reshape + strided (real,imag) views + 4 muls
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:
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):
Isolated encode wall-clock: 2173.5ms -> 1378.8ms on Metal (1.58x).