Affected by 
Documentation
¶
Index ¶
- type Distribution
- func (d Distribution) Arrays() []*mlx.Array
- func (d Distribution) LogProbs(vocab int) *mlx.Array
- func (d Distribution) Prob(tokens *mlx.Array) *mlx.Array
- func (d Distribution) ProbsForIDs(ids *mlx.Array) *mlx.Array
- func (d Distribution) ResidualAgainst(draft Distribution) Distribution
- func (d Distribution) Rows() int
- func (d Distribution) SampleWithKey(key *mlx.Array) *mlx.Array
- func (d Distribution) SliceRows(start, stop int) Distribution
- type Options
- type Result
- type Sampler
- func (s *Sampler) Add(seqID int, opts Options, priorTokens []int32)
- func (s *Sampler) Bernoulli(seqID int, p *mlx.Array) *mlx.Array
- func (s *Sampler) Commit(seqID int, tokens []int32)
- func (s *Sampler) Distribution(seqID int, logits *mlx.Array, draftTokens *mlx.Array) Distribution
- func (s *Sampler) Free()
- func (s *Sampler) Remove(seqID int)
- func (s *Sampler) Sample(seqIDs []int, logits *mlx.Array) Result
- func (s *Sampler) SampleDistribution(seqID int, dist Distribution) *mlx.Array
- func (s *Sampler) SpeculativeScores(seqID int, logits *mlx.Array, draftTokens *mlx.Array) *mlx.Array
Constants ¶
This section is empty.
Variables ¶
This section is empty.
Functions ¶
This section is empty.
Types ¶
type Distribution ¶ added in v0.24.0
type Distribution struct {
IDs *mlx.Array // sparse token ids, shape [B,K]; nil for dense distributions
Probs *mlx.Array // probabilities, shape [B,K] or [B,V]
}
Distribution is the filtered probability distribution used by the sampler. When IDs is nil, Probs is dense over the vocabulary. When IDs is set, Probs is sparse over the token ids in IDs, preserving GPU residency for the top-k-first path used by normal and speculative sampling.
func ConcatenateDistributions ¶ added in v0.24.0
func ConcatenateDistributions(dists []Distribution) Distribution
ConcatenateDistributions concatenates distribution rows. All inputs must use the same sparse/dense layout.
func (Distribution) Arrays ¶ added in v0.24.0
func (d Distribution) Arrays() []*mlx.Array
Arrays returns the tensor fields for mlx lifecycle management.
func (Distribution) LogProbs ¶ added in v0.24.0
func (d Distribution) LogProbs(vocab int) *mlx.Array
LogProbs returns dense log-probabilities, scattering sparse distributions into a full-vocabulary tensor when needed.
func (Distribution) Prob ¶ added in v0.24.0
func (d Distribution) Prob(tokens *mlx.Array) *mlx.Array
Prob returns the probability assigned to one token per row.
func (Distribution) ProbsForIDs ¶ added in v0.24.0
func (d Distribution) ProbsForIDs(ids *mlx.Array) *mlx.Array
ProbsForIDs returns probabilities for each requested token id. ids must be rank-2 [B,N], matching the distribution rows.
func (Distribution) ResidualAgainst ¶ added in v0.24.0
func (d Distribution) ResidualAgainst(draft Distribution) Distribution
ResidualAgainst returns the Leviathan/Chen rejection distribution proportional to max(target - draft, 0). Sparse target distributions stay sparse over the target support; tokens outside target support have zero mass.
func (Distribution) Rows ¶ added in v0.24.0
func (d Distribution) Rows() int
Rows returns the number of rows in the distribution.
func (Distribution) SampleWithKey ¶ added in v0.24.0
func (d Distribution) SampleWithKey(key *mlx.Array) *mlx.Array
SampleWithKey draws one token per row using key when supplied.
func (Distribution) SliceRows ¶ added in v0.24.0
func (d Distribution) SliceRows(start, stop int) Distribution
SliceRows returns a row slice while preserving sparse/dense layout.
type Options ¶ added in v0.21.1
type Options struct {
Temperature float32
TopP float32
MinP float32
TopK int
RepeatLastN int
RepeatPenalty float32
PresencePenalty float32
FrequencyPenalty float32
Seed int
UseSeed bool
// Logprobs causes Sample to populate Result.Logprob with the selected
// token's log-probability. TopLogprobs (when > 0) adds top-K pairs.
Logprobs bool
TopLogprobs int
}
type Result ¶ added in v0.21.1
type Result struct {
Token *mlx.Array // sampled token ids, shape [B]
Logprob *mlx.Array // sampled-token logprobs, shape [B,1]; nil unless any registered slot has Logprobs
TopTokens *mlx.Array // top-K token ids, shape [B,maxK]; nil unless any registered slot has TopLogprobs>0
TopLogprobs *mlx.Array // top-K logprobs, shape [B,maxK]; same
}
Result bundles the outputs of one decode step. Logprob/TopTokens/ TopLogprobs are populated whenever any registered slot has Logprobs (respectively TopLogprobs>0). Consumers need to filter by their per-slot Options.
type Sampler ¶
type Sampler struct {
// contains filtered or unexported fields
}
Sampler is a batched, slot-based sampler. Sequences are registered with Add and released with Remove. Each Sample call takes a subset of registered slots (in any order) with their [B,V] logits, samples one token per row, and appends it to that slot's ring-buffer history. Slots not named in a given call are untouched.
func New ¶
New constructs an empty sampler with no registered slots. numCtx is the runner's context window and must be positive.
func (*Sampler) Add ¶ added in v0.22.0
Add registers a sequence under seqID. The last RepeatLastN entries of priorTokens seed the ring buffer.
func (*Sampler) Bernoulli ¶ added in v0.24.0
Bernoulli samples boolean outcomes while advancing seqID's deterministic RNG stream when a seed is configured.
func (*Sampler) Commit ¶ added in v0.23.1
Commit appends already-selected tokens to seqID's repeat-penalty history. It is used after speculative sampling once the accepted continuation is known. Normal Sample calls continue to mutate history themselves.
func (*Sampler) Distribution ¶ added in v0.24.0
Distribution applies this slot's sampling transforms to logits without mutating sampler state. Rows align with the end of the draft chain: the final row is built as if every draft token had already been appended to the slot history, each earlier row with one fewer. Validation passes len(draftTokens)+1 rows, so row i sees draftTokens[:i]; a proposal step passes a single row, which sees the whole chain so far. logits must be [R,V] or [1,R,V].
func (*Sampler) Free ¶
func (s *Sampler) Free()
Free releases the pooled history tensor and resets the sampler to the New-equivalent state so it may be reused.
func (*Sampler) Remove ¶ added in v0.22.0
Remove releases the slot. The pool tensor is rebuilt to drop the row.
func (*Sampler) Sample ¶
Sample draws one token per row of logits ([B,V]); seqIDs[i] names the slot whose logits live at row i. Each sampled token is appended to its slot's ring. Slots not named in seqIDs are untouched.
func (*Sampler) SampleDistribution ¶ added in v0.24.0
func (s *Sampler) SampleDistribution(seqID int, dist Distribution) *mlx.Array
SampleDistribution draws from a precomputed distribution while advancing seqID's deterministic RNG stream when a seed is configured.
func (*Sampler) SpeculativeScores ¶ added in v0.23.1
func (s *Sampler) SpeculativeScores(seqID int, logits *mlx.Array, draftTokens *mlx.Array) *mlx.Array
SpeculativeScores applies this slot's sampling transforms to logits without mutating sampler state and returns dense log-probability scores for sampled decoding. Greedy decoding returns the penalty-adjusted logits. Rows align with the end of the draft chain as in Distribution.
Source Files