Expose minimal Gemma3 surface for frozen-text-encoder use

[xocialize]

Motivation

Multimodal generators increasingly use a frozen LLM as a text encoder, conditioning on the hidden states of every layer rather than on generated tokens. A concrete example: Lightricks LTX-2 conditions its video DiT on all 49 hidden states of Gemma-3-12B-it. We've been carrying this as a small fork while building a Swift port of LTX-2 — upstreaming so the next consumer of Gemma-3-as-encoder doesn't have to.

Shape (reworked per review)

Originally this PR shipped an allHiddenStates(_:mask:) extension in MLXLLM. Per review, that tap is pipeline-specific (its uniform mask matches the LTX-2 reference encoder, not generic Gemma3 per-layer sliding/global semantics) — so it now lives in client code, and this PR only widens visibility, scoped to @_spi(GemmaEncoder) so it isn't advertised public API:

• Gemma3TextConfiguration.hiddenSize, .hiddenLayers
• Gemma3TransformerBlock + its callAsFunction(_:mask:cache:)
• Gemma3Model.embedTokens, .layers, .config

Clients opt in with @_spi(GemmaEncoder) import MLXLLM. Zero behavior change — attributes and doc comments only.

Testing

Gemma3EncoderAccessTests locks the contract: it implements the client-side all-hidden-states tap using @_spi(GemmaEncoder) import MLXLLM — no @testable — on a tiny randomly-initialized model (no downloads), asserting state count (hiddenLayers + 1), shapes, first-state == scaled embedding, layer transformation, and determinism. 3/3 pass via xcodebuild test. If the exposed surface regresses, this test stops compiling.

Branch is merged up to current main; scripts/verify-docs.sh builds MLXLLM docs clean under --warnings-as-errors.

🤖 Generated with Claude Code

[davidkoski]

This looks like a really cool idea, but I have some questions & concerns about it.

• should this live in MLXEmbedders? that is where models that do this type of encoding live. I realize this is a bit different than what they do, but would it make more sense to do this on an embedding model (less weights to load, etc)?
• the embedding model may not be the same as the full gemma3 model so it might not produce the same results, but I wonder if an ability to load the full weights into the encoding model would do the trick? it should be the same structure and the sanitize should discard the unused pieces
• is this something that would be widely used? It seems like it might be a bit specialized and maybe is better suited as an extension (like this) in the client code. perhaps we just need to make Gemma3Model open (plus expose the bits needed)?

and separately I think there is a correctness issue:

• the gemma3 mask is a bit complicated with global and local masks for different layers, e.g.

            fullMask = createAttentionMask(h: h, cache: cacheArray[firstFullIdx])

            let slidingWindow = config.slidingWindow > 0 ? config.slidingWindow : 4096
            slidingWindowMask = createAttentionMask(
                h: h, cache: cacheArray[firstSlidingIdx], windowSize: slidingWindow)

It could be that coincidentally, if your text is <= than the prefill window size (512) then it would produce the same result, but it is complicated and I think not quite correct. The mean cosine you report for 1024 might be caused by this.

[xocialize]

Thanks for the careful look @davidkoski — I think you've landed on the right answer, and your mask concern is actually the strongest argument for it.

On the mask: you're right about the semantics. The single uniform mask is deliberate — it matches the reference LTX-2 encoder implementation (which is what the video DiT was trained against), not generic Gemma3 generation semantics. The parity goldens in the PR description come from that same reference, so the 0.999866 at T=1024 is 4-bit quant/accumulation noise rather than mask divergence — but your underlying point stands: as a generic Gemma3Model feature, a mask that silently diverges from the model's own per-layer sliding/global selection for T > window is surprising API. That's pipeline-specific behavior, and it shouldn't live in MLXLLM.

Which is your option 3. I've reworked the PR accordingly: the allHiddenStates extension and its LTX-specific masking are gone; the PR is now just the minimal visibility changes needed so a client module can implement the tap itself, plus a test that proves sufficiency — it implements the same tap using only import MLXLLM (no @testable) against a tiny randomly-initialized model, so the exposed surface can't silently regress. We'll carry the actual encoder tap (uniform mask + the per-layer eval that keeps 12B-class encodes under the GPU watchdog) in our LTX package where it belongs.

On MLXEmbedders (your first two questions): we considered it, but the encoder must reproduce exactly the hidden states the pipeline was trained against — full Gemma-3-12B-it weights, all 49 states, the reference's masking. An embedders-side model would still need the same specialized tap and masking to be exact, so it moves the specialization rather than removing it. If a general "LLM-as-encoder" story lands in MLXEmbedders someday we'd happily migrate; for now the open/expose route keeps upstream clean.

The branch is also merged up to current main (which brings in the #404 CI fix, so mac_build_and_test should actually run now). Let me know if the exposure set looks right or if you'd prefer any of it shaped differently.

[davidkoski]

See what you think about my questions / concern. Thanks!

[davidkoski]

Awesome, i will take a look! What you said about exact repro makes sense.

For exposing the innards, I wonder if using @_spi would make sense:

• https://github.com/ml-explore/mlx-swift/pull/418/changes#diff-1fd21695c64d825d44de72085bac30d18662ee7815e766e1848414bf4ce593c5R129

    /// Usable by extensions to implement `callAsFunction()`
    @_spi(MaterializedModule)
    public let _base: LayerType  

Then in other packages you can explicitly expose access to the symbol like this:

• https://github.com/ml-explore/mlx-swift-lm/pull/335/changes#diff-0374a7a4063b7d76b00e59997224676afc69b82a4fd4a53650abf2efaf95f542R6

@_spi(MaterializedModule) import MLXNN

extension MaterializedModule: EmbeddingModel where LayerType: EmbeddingModel {

    public func callAsFunction() {
...
        return _base(

Many models already have bits and pieces exposed as public so the change you made may be correct. This @_spi piece is something I learned about recently and it looks like it might be useful in cases like this -- you want to expose something but not fully advertise it.

[xocialize]

@_spi is a much better fit than plain public here — thanks for the pointer, I hadn't used it. Pushed: the whole surface is now @_spi(GemmaEncoder) public, and clients opt in with @_spi(GemmaEncoder) import MLXLLM. It keeps these internals out of the advertised API while still making the tap implementable, which seems like exactly the "expose but don't advertise" case you described.

Notes on the choices, in case you'd shape any of them differently:

• Group name. The repo currently has one group, @_spi(Testing), used for test-only escape hatches — reusing it would mislabel a surface meant for production client code, so I went with a purpose-descriptive GemmaEncoder following the MaterializedModule style. Happy to rename or fold it into a single shared repo-wide group if you'd rather have one.
• layers must carry the group since its element type is SPI — compiler requirement rather than a choice.
• config is exposed whole, which hands clients the entire Gemma3TextConfiguration even though only hiddenSize/hiddenLayers are SPI-public. I kept it because the tap wants config.hiddenLayers to size the result and threading the two values separately is clumsier — but if you'd prefer to drop config and expose only those two, that's an easy trim.
• DocC fix. Verify documentation was failing on a relative callAsFunction(_:mask:cache:) link I'd written on a property (DocC resolved it as a child of embedTokens). Qualifying it wasn't possible — Gemma3Model.callAsFunction is internal, so it isn't in the public symbol graph — so it's plain code formatting now. Verified with scripts/verify-docs.sh: MLXLLM builds clean under --warnings-as-errors. (Worth noting the @_spi change alone would have turned CI green, since SPI symbols drop out of the symbol graph — I fixed the comment anyway so it's correct on its own.)

One limitation worth stating plainly, since it's the flip side of your original mask concern: Gemma3TransformerBlock.callAsFunction takes a non-optional mask, so this surface only enables a uniform-mask tap — a client can't reproduce the per-layer sliding/global selection that Gemma3Model.callAsFunction does. That's exactly what our encoder wants (it must match the LTX-2 reference), but it does mean the surface isn't a general-purpose "run the layer stack correctly" API. If you'd rather it were, exposing the mask construction would be the additional piece — though that's more surface than this PR needs, so I've left it out.