Affected by 
Documentation
¶
Index ¶
- type Options
- type Result
- type Sampler
- func (s *Sampler) Add(seqID int, opts Options, priorTokens []int32)
- func (s *Sampler) Commit(seqID int, tokens []int32)
- func (s *Sampler) Free()
- func (s *Sampler) Remove(seqID int)
- func (s *Sampler) Sample(seqIDs []int, logits *mlx.Array) Result
- 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 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
// 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) 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) 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) 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 non-sampling transforms to logits without mutating sampler state. Row i is scored as if draftTokens[:i] had already been appended to the slot history. logits must be [R,V] or [1,R,V].
Source Files