Expose MLX sampler parameters and honor GenerationOptions.sampling - #231
Merged
Merged
Conversation
…onOptions The MLX backend hardcoded sampling parameters in toGenerateParameters / toStructuredGenerateParameters (topP: 1.0, repetitionPenalty: nil, topK/minP at defaults) and never read GenerationOptions.sampling, so callers could only tune temperature and maximumResponseTokens. MLXLMCommon.GenerateParameters already supports the full set. Add topP / topK / minP / repetitionPenalty / repetitionContextSize to MLXLanguageModel.CustomGenerationOptions (all optional, default nil → existing behavior unchanged) and forward them in both parameter mappers, preserving each path's prior defaults via `custom?.field ?? <previous default>`. Implements #165.
Reads the core GenerationOptions.sampling (SamplingMode) in toGenerateParameters / toStructuredGenerateParameters so top-p/top-k/greedy set via the standard sampling surface reach MLX, not only the custom block. Precedence: custom block wins, then sampling-derived, then existing default. Seed is not forwarded (no per-call seed in MLXLMCommon.GenerateParameters). Adds derivation + precedence tests.
There was a problem hiding this comment.
🟡 Changes recommended
Explicit temperature currently overrides greedy sampling, causing stochastic generation contrary to the sampling contract.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Exposes MLX sampler controls and maps core sampling options into MLX generation parameters.
Changes:
- Adds MLX-specific sampler and repetition options.
- Maps greedy, top-k, and nucleus sampling with precedence handling.
- Adds parameter-mapping tests.
File summaries
| File | Description |
|---|---|
Sources/AnyLanguageModel/Models/MLXLanguageModel.swift |
Adds and maps MLX sampling controls. |
Tests/AnyLanguageModelTests/CustomGenerationOptionsTests.swift |
Tests sampling derivation and precedence. |
Review details
Suppressed comments (3)
Sources/AnyLanguageModel/Models/MLXLanguageModel.swift:299
nildoes not disable top-k whenGenerationOptions.samplingsupplies a top-k value; only an explicit0does. Documentnilas deferring to core sampling and update the repeated initializer documentation.
/// Set this to `nil` or `0` to disable top-k sampling.
Sources/AnyLanguageModel/Models/MLXLanguageModel.swift:309
- For structured generation,
nilresolves to the existing1.1penalty rather than disabling repetition. Document the path-dependent default and explain that1.0explicitly neutralizes the penalty; update the repeated initializer documentation too.
/// Set this to `nil` to disable the repetition penalty.
Sources/AnyLanguageModel/Models/MLXLanguageModel.swift:1396
- The structured path has the same precedence inversion: an explicit nonzero temperature defeats
.greedy, so this no longer implements the sampling mode's argmax contract. ResolvegreedyTemperaturefirst here as well.
temperature: Float(options.temperature ?? derived.greedyTemperature.map(Double.init) ?? 0.2),
- Files reviewed: 2/2 changed files
- Comments generated: 4
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
There was a problem hiding this comment.
🟡 Changes recommended
MLX discards supported sampling seeds, leaving seeded generation nondeterministic.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (2)
Previously missed (1) — in code that hasn't changed since the last review.
Sources/AnyLanguageModel/Models/MLXLanguageModel.swift:1360
- This now drops the caller's sampling seed even though the compatible
mlx-swift-lm3.31.4 release exposesGenerateParameters.seed; the package requirement atPackage.swift:45can resolve that release, and the Ollama/Llama mappings already honor this field. As a result,.random(..., seed:)remains nondeterministic on MLX. Preserve the seed while destructuring.topK/.nucleusand pass it to both parameter constructors (with coverage for each path).
Sources/AnyLanguageModel/Models/MLXLanguageModel.swift:302
- This says every non-
nilvalue disables top-k, but positive values enable it. Document the fallback and positive-value behavior so callers are not instructed to do the opposite of the API's purpose.
/// Set this to `nil` to inherit top-k sampling from `GenerationOptions.sampling`,
/// otherwise disable top-k sampling. Set this to `0` to disable it explicitly.
- Files reviewed: 2/2 changed files
- Comments generated: 1
- Review effort level: Balanced
This was referenced Sep 11, 2026
mattt
added a commit
that referenced
this pull request
Sep 14, 2026
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.
@SpiraMira’s first two commits from #168, cherry-picked onto
mainwith their authorship intact, plus test-formatting fixes and review follow-ups for greedy sampling, sampling seeds, structured-generation coverage, and option documentation.MLX now reads
GenerationOptions.samplingfor top-k, top-p, and greedy sampling and exposes MLX-specific sampler controls throughCustomGenerationOptions. Precedence is custom options, then sampling, then the existing defaults; structured generation retains its own defaults. Greedy sampling forces temperature zero in both paths, even when an explicit temperature is supplied. Top-k and nucleus seeds are forwarded in both paths; this requiresmlx-swift-lm3.31.4 or later, where per-generation seed support was added.The speculative OS 27
toolCallingModecommit is excluded; #206 owns that SystemLanguageModel work. This replaces #168 because pushes to the contributor’s fork are blocked. Original description and review discussion: #168.