Repository navigation
Sync upstream ml-explore/mlx-swift-lm (7eb77ef), manual conflict resolution - #82
Merged
Merged
Conversation
…re#654) * MLXEmbedders: add createBidirectionalAttentionMask Encoder models need a bidirectional attention mask with an optional symmetric window and key padding. The existing helpers do not fit: createAttentionMask is causal, and createBidirectionalMask and createBidirectionalSlidingWindowMask in MLXLMCommon build additive masks with no padding and a window on the absolute key position. createBidirectionalAttentionMask(length:halfWindow:paddingMask:) returns .none when halfWindow and paddingMask are both nil. Otherwise it returns a boolean [batch or 1, 1, length, length] mask where query i may attend key j when |i - j| <= halfWindow and key j is not padding, or when i == j. - Padding hides key columns only, as in transformers, so a padded query row still sees the real keys. The boolean key mask follows Qwen3. - The diagonal is always kept, so no row is fully masked. For a fully masked row, MLX's attention returns a uniform average, zeros or another finite value depending on the kernel path. - A window that covers the whole sequence skips the band, so a huge halfWindow never reaches Int32. Tests cover the inclusive band bound, a zero window, padding on columns, padding dtypes, a single token, and agreement on real rows with the transformers 5.17 Gemma 3 sliding-window rule for windows 7, 8 and 9. Co-Authored-By: Claude <noreply@anthropic.com> * MLXEmbedders: use bidirectional attention in EmbeddingGemma EmbeddingGemma is a bidirectional encoder and its config.json says so ("use_bidirectional_attention": true), but the Swift backbone ignored the key and built causal masks, so every token saw only the tokens before it. Against transformers 5.17 with the original google/embeddinggemma-300m weights, the final embeddings had a cosine of only 0.45 to 0.70. Gemma3Configuration now decodes use_bidirectional_attention; absent or null means false. When it is true: - Sliding layers see sliding_window / 2 tokens on each side. This is the band transformers uses: it stores sliding_window // 2 + 1 and tests |i - j| < it, so 512 gives |i - j| <= 256. - Global layers see every token. - The padding mask passed to EmbeddingGemma.callAsFunction reaches the backbone and hides padded keys on every layer, so a padded batch gives the same states as its texts run one at a time. A 1-D mask is accepted with a 1-D input, and the backbone checks that the mask and the inputs have the same shape. Without a mask, attention sees no padding, as in transformers; padded batches need attentionMask. When the flag is false or absent the masks are the causal ones as before, so gemma3 and gemma3_text checkpoints that are causal by design do not change. Tests cover the flag-false path (padding ignored, causal, causal window), config decoding, bidirectionality in both directions, the window bound for sliding_window 8 and 9, right, left and mixed-length padded batches against single texts, 1-D inputs, and bfloat16 weights with padding. Fixes ml-explore#638 Co-Authored-By: Claude <noreply@anthropic.com>
Qwen3VL delegates to the shared QwenVL.splitPreparedInput with its own image and video token ids and spatial merge size, so ChatSession can reuse the KV cache on append-only media turns as it does for Qwen2.5-VL. The shared split holds for this model because its vision attention is block-diagonal per media item and deepstack features are rebuilt from the pixels on every forward. Qwen3VLMoE does not conform because it has no offset-aware continuation path.
* Parse GLM-4-0414 markerless tool calls (ml-explore#491) GLM-4-9B-0414 emits `name\n{json}` with no delimiters, so the marker driven pipeline never invoked the `.glm4` parser and the call leaked into the response text. Add `NamedJSONCallScanner`, a schema-anchored scanner (byte trie over the declared tool names + balanced JSON object), shared by the recovery lexer for streaming detection and by `GLM4ToolCallParser` as a fallback when no `<arg_key>` markers are present. Detection is line-anchored and only fires for a declared tool name, so ordinary `word\n{...}` prose is untouched. The integration harness now passes `tools:` to `generate`, as `ChatSession` already does; markerless detection needs the declared names. Re-enable `glm4EndToEnd`. * Fix GLM4 protocol markers in JSON arguments
Qwen3.5 and Qwen3-Next normalize q and k in the gated delta net with `invScale * rmsNorm(x, eps: 1e-6)`, where invScale = D^-0.5 and D = linear_key_head_dim. rmsNorm adds eps to mean(x^2), so the result is x / sqrt(sum(x^2) + D * 1e-6). The reference `l2norm` in transformers (qwen3_5 and qwen3_next) is x / sqrt(sum(x^2) + 1e-6), so the effective eps was D times too large (128x for Qwen/Qwen3.5-0.8B). Pass eps = 1e-6 / D to rmsNorm. Then sqrt(mean(x^2) + 1e-6 / D) = sqrt(sum(x^2) + 1e-6) / sqrt(D), which matches the reference. The invScale factors do not change. This affects the MLXLLM qwen3_5, qwen3_5_moe, qwen3_5_text and qwen3_next models and the MLXVLM qwen3_5 and qwen3_5_moe models. GatedDeltaQKNormTests compares each of the three layers with an output built from the same modules and the reference l2norm, on tokens scaled from 1 down to 1e-6. The largest per-token relative error goes from 0.984 to 2.4e-6 or less.
) Add `ChatSession.fork(history:)` to reuse prefix KV caches across conversations Starting every `ChatSession` with an empty cache causes expensive, redundant prefills for shared prompts (e.g. system instructions, tool schemas). `fork(history:)` creates a new session sharing the existing session's cache, configuration, and draft model, prefilling only the diff. Key changes: - `ChatSession.fork(history:)`: Returns a session with a sliced copy of the KV and draft caches; shares memory until write. - Cache pre-flight checks: Validates media presence before copying to avoid wasteful deep copies (e.g., in `VarianceNormalizedKVCache`). - Transcript validation: Requires exact token-by-token prefix matches and falls back to an empty cache if media is present. - Cache support: Implemented `copy()` for `TurboQuantKVCache`. - Concurrency: Made `fork` `nonisolated(nonsending)` so caller actors can fork without hop/executor overhead.
Fix Qwen reasoning history, Swift 6.3 isolation, and test sync - Qwen: Preserve prompt-priming provenance, pass reasoning via reasoning_content, and add 3.8/3.6 regression tests for cache reuse and fallback. - ChatSession: Fix Swift 6.3 isolation checker failure in Task closure. - Tests: Fix cache-clearing synchronization race.
* scripted Tokenizer - part 1 of scripted Tokenizers, Tools and LanguageModels - create a new Tokenizer from a simple corpus - behaves like a real Tokenizer (BPE style) and can do bidirectional encode/decode
…lore#685) * MLXEmbedders: build the README example mask from text lengths The "Generate embeddings" example pads every text with the EOS token and builds the attention mask as `padded .!= eosTokenId`. Tokenizers that end every text with a real EOS token, such as EmbeddingGemma's, then lose that token as well: the mask hides it from pooling and, since ml-explore#654, from attention. Build the mask from the real lengths instead, in the same shape as the padding above it. EmbeddingGemma (mlx-community/embeddinggemma-300m-bf16) in float32, an EOS-padded batch of an 11-token query and a 548-token text, cosine against sentence-transformers: 0.776 and 0.969 with the old mask, 1.00000 for both rows with the new one.
Add compiled decode for Qwen3.5 VLM models - Compile eligible single-token decode for plain, rotating, and quantized caches (per-layer GDN/MoE), keeping dynamic M-RoPE and cache mutation eager. - Correctly invalidate dependent traces on submodule updates using weak references to preserve deallocation. - Synchronize fused-projection trace registry to support concurrent registration. - Clean up unused test helper and add regression coverage (exactness, lifecycle, concurrency, submodule updates). AI assistance: Codex implemented the fix and regression tests.
Manual conflict resolution: - Package.swift: keep both MLXLMTestSupport and MLXScriptedLM targets - LanguageModel.swift: keep fork MTP/TrailingLogits/DualModelMTP + upstream HiddenStateLanguageModel - Gemma4Text.swift: keep fork optional finalLogitSoftcapping via gemma4LogitSoftcap - Evaluate.swift: keep SSDStreamingErrorLatch wrapper around upstream nextTokenLogits - VLM Qwen35.swift: keep fork LayerPartitionable/StreamableMoE/MathRMSNorm with upstream compiled decode segments Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 5.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
7 participants
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.
Real merge (merge commit) of ml-explore/mlx-swift-lm main @ 7eb77ef, 12 commits, with 5 conflicting files resolved by hand:
Package.swift: keep bothMLXLMTestSupportandMLXScriptedLMtargetsLanguageModel.swift: forkMTPLanguageModel/TrailingLogitsLanguageModel/DualModelMTP+ upstreamHiddenStateLanguageModelGemma4Text.swift: fork's optionalfinalLogitSoftcappingviagemma4LogitSoftcapEvaluate.swift: forkSSDStreamingErrorLatch.withActivewrapping upstreamnextTokenLogitsMLXVLM/Models/Qwen35.swift: forkLayerPartitionable/StreamableMoE/MathRMSNorm+ upstream compiled decode segments (layer.forward(... allowCompiledDecode:)), wrapped inpartitionedLayerCallMust be merged with a merge commit (not squash). Not built locally (toolchain mismatch); CI is the check.
🤖 Generated with Claude Code