Something went wrong. Try again.
Routing embedding stack for Qwen3-Embedding-0.6B on Apple Silicon: orchestrator + MLX + llama.cpp backends
Something went wrong. Try again.
1.3 kB · 33 lines
Swift
at main
12345678910111213141516171819202122232425262728293031323334import Foundationimport MLXimport MLXEmbedders
/// Runs dummy forward passes at every (batch, length) bucket combination to/// pre-compile the MLX graph. MLX recompiles on shape changes, so we eat the/// cost once per shape up front instead of on the user's first request.enum WarmUp { static func warm( container: MLXEmbedders.ModelContainer, maxBatch: Int, maxLen: Int ) async { let lengths = Buckets.lengths.filter { $0 <= maxLen } let batches = Buckets.batches.filter { $0 <= maxBatch }
await container.perform { model, tokenizer, pooling in let padId: Int32 = { if let eos = tokenizer.eosTokenId { return Int32(eos) } return 0 }() for b in batches { for l in lengths { let ids = MLXArray([Int32](repeating: padId, count: b * l), [b, l]) let mask = MLXArray([Int32](repeating: 1, count: b * l), [b, l]) let out = model(ids, positionIds: nil, tokenTypeIds: nil, attentionMask: mask) let pooled = pooling(out, mask: mask, normalize: true) eval(pooled) } } } }}