mlxrunner: apply in-flight drafts to proposal penalty history

Sampler.Distribution built row i as if draftTokens[:i] were appended, leaving
a single-row proposal call with no draft history, so a drafter skipped the
repeat/presence penalties the target's validation applies and re-proposed
penalized tokens. Align rows with the end of the draft chain instead: the
final row sees every draft token, each earlier row one fewer.

x/mlxrunner/sample/sample.go

@@ -443,8 +443,12 @@ func (s *Sampler) Sample(seqIDs []int, logits *mlx.Array) Result {
 }
 
 // Distribution applies this slot's sampling transforms to logits without
-// mutating sampler state. Row i is built as if draftTokens[:i] had already
-// been appended to the slot history. logits must be [R,V] or [1,R,V].
+// 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 (s *Sampler) Distribution(seqID int, logits *mlx.Array, draftTokens *mlx.Array) Distribution {
 	slot, logits, draftTokens := s.speculativeInputs("Distribution", seqID, logits, draftTokens)
 	rows := logits.Dim(0)
@@ -465,7 +469,8 @@ func (s *Sampler) Distribution(seqID int, logits *mlx.Array, draftTokens *mlx.Ar
 
 // 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.
+// decoding. Greedy decoding returns the penalty-adjusted logits. Rows align
+// with the end of the draft chain as in Distribution.
 func (s *Sampler) SpeculativeScores(seqID int, logits *mlx.Array, draftTokens *mlx.Array) *mlx.Array {
 	slot, logits, draftTokens := s.speculativeInputs("SpeculativeScores", seqID, logits, draftTokens)
 	rows := logits.Dim(0)
@@ -522,6 +527,17 @@ func (s *Sampler) speculativeInputs(caller string, seqID int, logits *mlx.Array,
 	if draftTokens != nil && draftTokens.NumDims() == 1 {
 		draftTokens = draftTokens.ExpandDims(0)
 	}
+
+	// Rows align with the end of the draft chain, so the earliest row sees
+	// draftCount-rows+1 prior drafts. More rows than draftCount+1 would make
+	// that count negative and silently drop the prefix; reject it loudly.
+	draftCount := 0
+	if draftTokens != nil {
+		draftCount = draftTokens.Dim(1)
+	}
+	if logits.Dim(0) > draftCount+1 {
+		panic(fmt.Sprintf("sample.Sampler.%s: %d logit rows exceed the %d-token draft chain", caller, logits.Dim(0), draftCount))
+	}
 	return slot, logits, draftTokens
 }
 
@@ -578,7 +594,7 @@ func (s *Sampler) speculativeDistributionSerial(slot *slotState, logits *mlx.Arr
 	for i := range rows {
 		rowLogits := logits.Slice(mlx.Slice(i, i+1), mlx.Slice())
 		hist := base
-		prefixLen := min(i, draftCount)
+		prefixLen := draftCount - rows + 1 + i
 		if prefixLen > 0 {
 			prefix := draftTokens.Slice(mlx.Slice(), mlx.Slice(0, prefixLen))
 			if hist == nil {
@@ -616,7 +632,9 @@ func (s *Sampler) speculativeHistory(slot *slotState, draftTokens *mlx.Array, ro
 	sourceIdx := make([]int32, rows*width)
 	writeMask := make([]bool, rows*width)
 	for i := range rows {
-		prefixLen := min(i, draftCount)
+		// Non-positive lengths run no iterations: rows beyond the draft
+		// chain's start carry the base history unchanged.
+		prefixLen := draftCount - rows + 1 + i
 		for j := range prefixLen {
 			pos := (next + j) % width
 			sourceIdx[i*width+pos] = int32(j)

x/mlxrunner/sample/sample_test.go

@@ -336,6 +336,34 @@ func TestSpeculativeScoresUsesDraftHistoryWithoutCommit(t *testing.T) {
 	}
 }
 
+func TestDistributionSingleRowAppliesDraftPrefix(t *testing.T) {
+	skipIfNoMLX(t)
+
+	s := New(128)
+	t.Cleanup(func() {
+		s.Free()
+		mlx.Sweep()
+	})
+
+	// A proposal step passes one logits row with the chain's earlier drafts:
+	// the single row is the chain's final step, so every draft belongs to
+	// its history. Slot 0 exercises the batched history path (full ring),
+	// slot 1 the serial path (ring not yet full).
+	s.Add(0, Options{RepeatLastN: 2, RepeatPenalty: 10}, []int32{0, 1})
+	s.Add(1, Options{RepeatLastN: 8, RepeatPenalty: 10}, []int32{0, 1})
+	prefix := mlx.NewArrayInt32([]int32{3, 4}, []int32{1, 2})
+
+	for _, seqID := range []int{0, 1} {
+		// Drafts 3 and 4 are penalized, so token 2 wins over the higher raw
+		// scores; with the drafts absent from the history, token 3 would.
+		dist := s.Distribution(seqID, batchLogits([]float32{0, 0, 9, 9, 8}), prefix)
+		mlx.Eval(dist.IDs)
+		if got := dist.IDs.Ints()[0]; got != 2 {
+			t.Fatalf("seq %d token = %d, want 2 (drafts 3 and 4 penalized)", seqID, got)
+		}
+	}
+}
+
 func TestCommitBatchesRingWrites(t *testing.T) {
 	skipIfNoMLX(t)