feat(backend): auto-size batch to context for embedding and rerank models

Embedding and rerank models pool over the whole input in a single physical batch (n_ubatch). With batch left at the 512 default, the backend rejects longer inputs with "input is too large to process", silently capping a large-context embedder (e.g. 8k/32k) at 512 tokens. Size n_batch to the context for these single-pass usecases, mirroring the existing FLAG_SCORE behaviour; an explicit batch: still wins.

Extracts EffectiveContextSize/EffectiveBatchSize from grpcModelOpts so the effective decode window has one home for other callers to reuse.

Adds an e2e-aio regression test that embeds a >512-token input. The AIO embedding model is switched to nomic-embed-text-v1.5 (2048 context) because the previous granite model was capped at 512 tokens and could not exercise the larger batch.

Assisted-by: claude-code:claude-opus-4-8 [Claude Code]
Signed-off-by: Richard Palethorpe <io@richiejp.com>

Author: richiejp

core/backend/options.go

@@ -87,25 +87,47 @@ func getSeed(c config.ModelConfig) int32 {
 	return seed
 }
 
-func grpcModelOpts(c config.ModelConfig, modelPath string) *pb.ModelOptions {
-	// Resolve context size first — the backend defaults to 4096 when unset,
-	// and batch sizing below has to match that effective value or the
-	// FLAG_SCORE guard misses the n_batch < n_ctx GGML_ASSERT crash.
-	ctxSize := 4096
+// DefaultContextSize and DefaultBatchSize are the backend's fallbacks when a
+// model config leaves them unset. Exported so callers that must respect the
+// effective decode window — notably the router's prompt trimmer — resolve the
+// same numbers grpcModelOpts does instead of guessing.
+const (
+	DefaultContextSize = 4096
+	DefaultBatchSize   = 512
+)
+
+// EffectiveContextSize is the context window the backend will run with: the
+// configured value, or DefaultContextSize when unset.
+func EffectiveContextSize(c config.ModelConfig) int {
 	if c.ContextSize != nil {
-		ctxSize = *c.ContextSize
+		return *c.ContextSize
 	}
+	return DefaultContextSize
+}
 
-	b := 512
+// EffectiveBatchSize is the single-decode batch the backend will run with.
+// Score, embedding and rerank all process the whole input in one pass: score
+// decodes prompt+candidate (asserts n_tokens <= n_batch), and embedding/rerank
+// pool over the full sequence in one physical batch (n_ubatch). So the batch
+// is sized to the context — anything that fits the context fits one pass,
+// avoiding both the GGML_ASSERT crash and the "input is too large to process"
+// error. Explicit `batch:` always wins.
+func EffectiveBatchSize(c config.ModelConfig) int {
 	if c.Batch != 0 {
-		b = c.Batch
-	} else if c.HasUsecases(config.FLAG_SCORE) && ctxSize > b {
-		// Score models decode prompt+candidate in one llama_decode which
-		// asserts n_tokens <= n_batch and aborts on failure. Sizing the
-		// batch to n_ctx means anything that fits the context fits one
-		// decode. Explicit `batch:` in the config still wins.
-		b = ctxSize
+		return c.Batch
 	}
+	singlePass := c.HasUsecases(config.FLAG_SCORE) ||
+		c.HasUsecases(config.FLAG_EMBEDDINGS) ||
+		c.HasUsecases(config.FLAG_RERANK)
+	if ctx := EffectiveContextSize(c); singlePass && ctx > DefaultBatchSize {
+		return ctx
+	}
+	return DefaultBatchSize
+}
+
+func grpcModelOpts(c config.ModelConfig, modelPath string) *pb.ModelOptions {
+	ctxSize := EffectiveContextSize(c)
+	b := EffectiveBatchSize(c)
 
 	flashAttention := "auto"
 

core/backend/options_internal_test.go

@@ -102,6 +102,25 @@ var _ = Describe("grpcModelOpts NBatch", func() {
 		Expect(opts.NBatch).To(BeEquivalentTo(1024))
 	})
 
+	It("sizes the batch to the context window for embedding models", func() {
+		// Embedding/rerank pool over the whole sequence in one physical batch
+		// (n_ubatch); without this the input is capped at the 512 default and
+		// the backend returns "input is too large to process".
+		embeddings := true
+		cfg := config.ModelConfig{Threads: &threads, LLMConfig: config.LLMConfig{ContextSize: &ctx}}
+		cfg.Embeddings = &embeddings
+		opts := grpcModelOpts(cfg, "/tmp/models")
+		Expect(opts.NBatch).To(BeEquivalentTo(4096))
+	})
+
+	It("sizes the batch to the context window for rerank models", func() {
+		reranking := true
+		cfg := config.ModelConfig{Threads: &threads, LLMConfig: config.LLMConfig{ContextSize: &ctx}}
+		cfg.Reranking = &reranking
+		opts := grpcModelOpts(cfg, "/tmp/models")
+		Expect(opts.NBatch).To(BeEquivalentTo(4096))
+	})
+
 	It("does not raise the batch when a score model's context is below the default", func() {
 		small := 256
 		cfg := config.ModelConfig{Threads: &threads, LLMConfig: config.LLMConfig{ContextSize: &small}, KnownUsecases: &scoreUsecase}

tests/e2e-aio/e2e_test.go

@@ -222,6 +222,37 @@ var _ = Describe("E2E test", func() {
 				Expect(resp3.Data[1].Embedding).To(Equal(resp2.Data[0].Embedding))
 				Expect(resp3.Data[0].Embedding).ToNot(Equal(resp3.Data[1].Embedding))
 			})
+
+			// Regression guard for the auto-batch fix (core/backend/options.go
+			// EffectiveBatchSize). Embeddings pool over the whole sequence in a
+			// single physical batch (n_ubatch == n_batch), so an input longer
+			// than n_batch is rejected by the backend with "input is too large
+			// to process". Before the fix n_batch defaulted to 512 regardless of
+			// the model's context, so any prompt over ~512 tokens failed here.
+			// The embedding model is configured with a 2048 context (see
+			// models/embeddings.yaml); this input is comfortably over 512 tokens
+			// and under that context, so it must embed in one pass.
+			It("embeds an input larger than the default 512 batch", func() {
+				var b bytes.Buffer
+				// ~100 short sentences ≈ 1000+ tokens: well past the old 512
+				// batch ceiling, well within the 2048 context.
+				for i := range 100 {
+					fmt.Fprintf(&b, "This is sentence number %d discussing organic skincare and machine learning. ", i)
+				}
+				longInput := b.String()
+
+				resp, err := client.Embeddings.New(context.TODO(),
+					openai.EmbeddingNewParams{
+						Input: openai.EmbeddingNewParamsInputUnion{
+							OfArrayOfStrings: []string{longInput},
+						},
+						Model: openai.EmbeddingModelTextEmbeddingAda002,
+					},
+				)
+				Expect(err).ToNot(HaveOccurred(), "a >512-token input must embed in a single batch (auto-batch sizing)")
+				Expect(len(resp.Data)).To(Equal(1), fmt.Sprint(resp))
+				Expect(resp.Data[0].Embedding).ToNot(BeEmpty())
+			})
 		})
 
 		Context("vision", func() {

tests/e2e-aio/models/embeddings.yaml

@@ -1,5 +1,12 @@
 embeddings: true
 name: text-embedding-ada-002
 backend: llama-cpp
+# nomic-embed-text-v1.5 has a 2048-token context, unlike the previous 512-token
+# granite model. The larger context is what makes the long-input embedding test
+# (e2e_test.go) meaningful: it exercises the auto-batch fix where n_batch is
+# sized up to the context window (core/backend/options.go EffectiveBatchSize) so
+# a >512-token input embeds in a single pass instead of failing with "input is
+# too large to process" against the default 512 batch.
+context_size: 2048
 parameters:
-  model: huggingface://bartowski/granite-embedding-107m-multilingual-GGUF/granite-embedding-107m-multilingual-f16.gguf
+  model: huggingface://nomic-ai/nomic-embed-text-v1.5-GGUF/nomic-embed-text-v1.5.f16.gguf