Expand description
EmbeddingGemma 2 (google/embeddinggemma-2) image and text embeddings
on an AMD XDNA NPU.
The bundle iron/applications/embeddinggemma2/export_eg2.py writes
holds every compiled IRON kernel and the weights (the NPU ones
pre-packed); this crate replays the forward the Python app runs
(eg2_common.py / eg2_npu.py):
| NPU | host (here) | |
|---|---|---|
| vision tower (16 layers, 768 wide) | every projection (flm.GEMMs: patch embedding, qkv, o, GeGLU gate+up, down, embed_vision) and the attention (the MHA operator) | resize + patchify (preprocess.rs), position embeddings, RMSNorms, q / k / v norms, 2D RoPE, residual adds, 3 x 3 pooling |
| text encoder (24 layers, 512 wide) | every projection (per-layer-input projection, qkv, o, GeGLU, down, PLE gate + projection) | norms, RoPE, attention (~270 tokens), per-layer gating, mean pooling, the 512 -> 768 projection |
EmbeddingGemma2::embed_rgb and EmbeddingGemma2::embed_text
give the L2-normalized 768-d embeddings sentence-transformers computes
for an image and for a text (with its task prompt, e.g. SearchQuery
or Document); both live in the model’s one space, compared by cosine.
Text runs through the same text encoder as an image’s soft tokens:
Gemma’s tokenizer (tokenizer.rs), the token embeddings, then every
projection on the NPU and the attention (a 512-token sliding window on
20 of the 24 layers) on the host.
Modules§
- model
- The forward, as
eg2_common.pyruns it witheg2_npu.NpuBackend: every projection anflm.GEMMdispatch (all of an image’s rows in one), the vision attention an MHA dispatch, the rest here in f32. - npu
- The NPU side: the bundle’s kernels (kernels naming the same xclbin share its hardware context: the vision GEMMs one, the text GEMMs another, the MHA length buckets a third – 3 of NPU2’s 16, shared by every process), device row buffers, and the GEMM / MHA dispatches.
- preprocess
- Image preprocessing, as HF’s
Gemma4ImageProcessor(torchvision backend): - tokenizer
- EmbeddingGemma 2’s tokenizer (Gemma’s
tokenizer.json), as Hugging Facetokenizersruns it and sentence-transformers calls it:
Structs§
- Config
- The model’s constants (the manifest’s params).
- Embedding
Gemma2 - Timing
- Wall time per stage, in first-seen order; NPU dispatch time is kept
under
npu:<kernel>.