sample: define multi-row distributions without a draft chain

Distribution aligns its rows with the end of the draft chain, so when
the caller passes no chain, every row sees the slot history unchanged.
That case already worked; only the row-count guard rejected it. The
guard now applies only when a chain is present, which is where more
rows than chain positions would silently drop history. A block drafter
needs the chainless case to sample its whole proposal batch in one
call.

x/mlxrunner/sample/sample.go

@@ -446,8 +446,8 @@ func (s *Sampler) Sample(seqIDs []int, logits *mlx.Array) Result {
 // 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
+// len(draftTokens)+1 rows, so row i sees draftTokens[:i]; a proposal passes
+// no chain and every row sees the slot history as it stands. 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)
@@ -529,14 +529,10 @@ func (s *Sampler) speculativeInputs(caller string, seqID int, logits *mlx.Array,
 	}
 
 	// 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))
+	// draftCount-rows+1 prior drafts; more rows than that would silently drop
+	// prefix — reject loudly. Without a chain, rows see the history unchanged.
+	if draftTokens != nil && logits.Dim(0) > draftTokens.Dim(1)+1 {
+		panic(fmt.Sprintf("sample.Sampler.%s: %d logit rows exceed the %d-token draft chain", caller, logits.Dim(0), draftTokens.Dim(1)))
 	}
 	return slot, logits, draftTokens
 }

x/mlxrunner/sample/sample_test.go

@@ -364,6 +364,39 @@ func TestDistributionSingleRowAppliesDraftPrefix(t *testing.T) {
 	}
 }
 
+func TestDistributionMultiRowWithoutChain(t *testing.T) {
+	skipIfNoMLX(t)
+
+	s := New(128)
+	t.Cleanup(func() {
+		s.Free()
+		mlx.Sweep()
+	})
+
+	// A block drafter's proposal batch samples every row from one call with
+	// no draft chain: each row sees the slot history unchanged. 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{3, 4})
+	s.Add(1, Options{RepeatLastN: 8, RepeatPenalty: 10}, []int32{3, 4})
+
+	for _, seqID := range []int{0, 1} {
+		// Tokens 3 and 4 are penalized in every row alike; rows 1 and 3
+		// share logits, so a chain alignment leaking between rows would
+		// split their winners.
+		dist := s.Distribution(seqID, batchLogits(
+			[]float32{0, 0, 8, 9, 9},
+			[]float32{0, 8, 0, 9, 9},
+			[]float32{0, 0, 8, 9, 9},
+		), nil)
+		top := dist.IDs.Slice(mlx.Slice(), mlx.Slice(0, 1))
+		mlx.Eval(top)
+		if got, want := top.Ints(), []int{2, 1, 2}; !slices.Equal(got, want) {
+			t.Fatalf("seq %d top tokens = %v, want %v", seqID, got, want)
+		}
+	}
+}
+
 func TestCommitBatchesRingWrites(t *testing.T) {
 	skipIfNoMLX(t)