fix: clamp embedding batch size to the ONNX graph batch axis (#970)

A model registered through TextEmbedding.add_custom_model() against the
Qwen3 causal-LM ONNX weights crashed on any batch larger than one:

  onnxruntime ... Shape mismatch attempting to re-use buffer.
  {1,1,18,18} != {8,1,18,18}

CustomTextEmbedding inherits OnnxTextEmbedding.embed, whose default is
batch_size=256, while the graph is pinned to a batch of one.
Qwen3TextEmbedding avoided this by hardcoding batch_size=1 in its own
embed; CustomTextEmbedding had no such override.

Read the batch axis from the loaded session instead of hardcoding it:
_detect_static_batch_size() records the literal the graph declares (or
None when the axis is symbolic) at load time, and _embed_documents()
clamps the requested batch size to it on both the sequential and the
worker-pool path. A graph with a dynamic batch axis keeps batching at
the size the caller asked for.

The inputs alone are not authoritative -- the Qwen3 export declares
input_ids as ['batch_size', 'sequence_length'] yet emits
last_hidden_state as [1, 'sequence_length', 1024] -- so the outputs are
inspected as well.

Author: n24q02m

qwen3_embed/common/onnx_model.py

@@ -58,6 +58,7 @@ def _post_process_onnx_output(self, output: OnnxOutputContext, **kwargs: Any) ->
     def __init__(self) -> None:
         self.model: ort.InferenceSession | None = None
         self.model_input_names: set[str] | None = None
+        self.static_batch_size: int | None = None
         self.tokenizer: Tokenizer | None = None
         self.special_token_to_id: dict[str, int] = {}
 
@@ -167,6 +168,28 @@ def _instantiate_onnx_session(
                 )
         return session, input_names
 
+    @staticmethod
+    def _detect_static_batch_size(session: "ort.InferenceSession") -> int | None:
+        """Return the batch size the graph is pinned to, or None if it is dynamic.
+
+        ONNX declares each dimension either as an int (static) or as a symbolic
+        ``dim_param`` string such as ``"batch_size"`` (dynamic). Both the inputs
+        and the outputs are inspected because they can disagree: causal-LM
+        exports keep the symbolic ``batch_size`` on ``input_ids`` while the graph
+        body was traced with a fixed batch, and only the output declaration
+        carries the literal. Feeding a larger batch to such a graph fails deep
+        inside onnxruntime with a buffer shape mismatch, so the smallest literal
+        found on a batch axis wins.
+        """
+        declared = [
+            node.shape[0]
+            for node in (*session.get_inputs(), *session.get_outputs())
+            # Rank < 2 has no batch axis to speak of, e.g. a scalar side input
+            # declared as ``[1]`` is a length, not a batch of one.
+            if node.shape is not None and len(node.shape) >= 2 and isinstance(node.shape[0], int)
+        ]
+        return min(declared) if declared else None
+
     def _load_onnx_model(
         self,
         model_dir: Path,
@@ -179,6 +202,7 @@ def _load_onnx_model(
             config=config,
         )
         self.model_input_names = set(input_names)
+        self.static_batch_size = self._detect_static_batch_size(self.model)
         self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
         return self.model, input_names
 

qwen3_embed/text/onnx_embedding.py

@@ -103,7 +103,8 @@ def embed(
 
         Args:
             documents: Iterator of documents or single document to embed
-            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
+            batch_size: Batch size for encoding -- higher values will use more memory, but be faster.
+                Capped at the batch size the ONNX graph is pinned to, if it declares one.
             parallel:
                 If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
                 If 0, use all available cores.

qwen3_embed/text/onnx_text_model.py

@@ -89,6 +89,21 @@ def onnx_embed(
             input_ids=onnx_input.get("input_ids", input_ids),
         )
 
+    def _effective_batch_size(self, batch_size: int) -> int:
+        """Clamp the requested batch size to what the loaded graph accepts.
+
+        A graph whose batch axis is dynamic keeps the requested batch size; one
+        pinned to a fixed batch is capped at that value, because exceeding it
+        aborts the run inside onnxruntime instead of raising something a caller
+        can act on. The requested size is also returned untouched when no
+        session is loaded in this process -- with ``lazy_load`` plus
+        ``parallel`` the graph only exists in the workers.
+        """
+        static_batch_size = getattr(self, "static_batch_size", None)
+        if static_batch_size is None:
+            return batch_size
+        return min(batch_size, static_batch_size)
+
     def _embed_documents(
         self,
         model_name: str,
@@ -116,7 +131,7 @@ def _embed_documents(
         if parallel is None or is_small:
             if not hasattr(self, "model") or self.model is None:
                 self.load_onnx_model()
-            for batch in iter_batch(documents, batch_size):
+            for batch in iter_batch(documents, self._effective_batch_size(batch_size)):
                 yield from self._post_process_onnx_output(
                     self.onnx_embed(batch, **kwargs), **kwargs
                 )
@@ -148,7 +163,8 @@ def _embed_documents(
                     start_method=start_method,
                 ),
             )
-            for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
+            batches = iter_batch(documents, self._effective_batch_size(batch_size))
+            for batch in pool.ordered_map(batches, **params):
                 yield from self._post_process_onnx_output(batch, **kwargs)
 
     def _token_count(self, texts: str | Iterable[str], batch_size: int = 1024, **_: Any) -> int:

qwen3_embed/text/text_embedding.py

@@ -191,7 +191,8 @@ def embed(
 
         Args:
             documents: Iterator of documents or single document to embed
-            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
+            batch_size: Batch size for encoding -- higher values will use more memory, but be faster.
+                Capped at the batch size the ONNX graph is pinned to, if it declares one.
             parallel:
                 If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
                 If 0, use all available cores.

tests/test_integration.py

@@ -294,6 +294,47 @@ def test_custom_model_parallel_mixed_case(self):
             assert emb.shape == (1024,)
 
 
+# ═══════════════════════════════════════════════════════════════════════════
+# Embedding: Custom model on a static-batch graph
+# ═══════════════════════════════════════════════════════════════════════════
+
+
+@pytest.mark.integration
+class TestCustomModelStaticBatchGraph:
+    """A custom model on causal-LM weights must not crash on the default batch size.
+
+    The Qwen3 ONNX graph is pinned to a batch of one, so embedding more
+    documents than that in a single run used to abort inside onnxruntime with
+    ``Shape mismatch attempting to re-use buffer``.
+    """
+
+    def test_default_batch_size_on_static_batch_graph(self):
+        CustomTextEmbedding._SUPPORTED.clear()
+        model_id = "Org/Custom-Qwen-Static-Batch"
+        TextEmbedding.add_custom_model(
+            model_description=DenseModelDescription(
+                model=model_id,
+                sources=ModelSource(hf="n24q02m/Qwen3-Embedding-0.6B-ONNX"),
+                model_file="onnx/model_quantized.onnx",
+                dim=1024,
+            ),
+            pooling=PoolingType.LAST_TOKEN,
+            normalization=True,
+        )
+        try:
+            model = TextEmbedding(model_name=model_id)
+            docs = [f"Document number {i} about topic {chr(65 + i)}." for i in range(8)]
+            # No batch_size: the default of 256 is what the caller gets by default.
+            embeddings = list(model.embed(docs))
+        finally:
+            CustomTextEmbedding._SUPPORTED.clear()
+
+        assert len(embeddings) == 8
+        for emb in embeddings:
+            assert emb.shape == (1024,)
+            assert abs(np.linalg.norm(emb) - 1.0) < 1e-3
+
+
 # ═══════════════════════════════════════════════════════════════════════════
 # Reranker: Basic operations
 # ═══════════════════════════════════════════════════════════════════════════

tests/test_onnx_model.py

@@ -285,6 +285,87 @@ def test_add_extra_session_options():
         ConcreteOnnxModel.add_extra_session_options(session_options, {"invalid_option": True})
 
 
+def _shape_node(name: str, shape: list[Any]) -> MagicMock:
+    """A stand-in for an onnxruntime NodeArg carrying a declared shape."""
+    node = MagicMock()
+    node.name = name
+    node.shape = shape
+    return node
+
+
+class TestDetectStaticBatchSize:
+    """A literal on the batch axis pins the graph; a dim_param string does not."""
+
+    def test_fully_dynamic_returns_none(self):
+        session = MagicMock()
+        session.get_inputs.return_value = [
+            _shape_node("input_ids", ["batch_size", "sequence_length"]),
+        ]
+        session.get_outputs.return_value = [
+            _shape_node("last_hidden_state", ["batch_size", "sequence_length", 768]),
+        ]
+        assert OnnxModel._detect_static_batch_size(session) is None
+
+    def test_literal_on_output_only(self):
+        """Shapes as declared by n24q02m/Qwen3-Embedding-0.6B-ONNX model_quantized.onnx:
+        the inputs still say ``batch_size`` while the output says ``1``."""
+        session = MagicMock()
+        session.get_inputs.return_value = [
+            _shape_node("input_ids", ["batch_size", "sequence_length"]),
+            _shape_node("attention_mask", ["batch_size", "sequence_length"]),
+        ]
+        session.get_outputs.return_value = [
+            _shape_node("last_hidden_state", [1, "sequence_length", 1024]),
+        ]
+        assert OnnxModel._detect_static_batch_size(session) == 1
+
+    def test_literal_on_input(self):
+        session = MagicMock()
+        session.get_inputs.return_value = [
+            _shape_node("input_ids", [4, "sequence_length"]),
+        ]
+        session.get_outputs.return_value = [
+            _shape_node("last_hidden_state", ["batch_size", "sequence_length", 768]),
+        ]
+        assert OnnxModel._detect_static_batch_size(session) == 4
+
+    def test_smallest_literal_wins(self):
+        session = MagicMock()
+        session.get_inputs.return_value = [
+            _shape_node("input_ids", [4, "sequence_length"]),
+        ]
+        session.get_outputs.return_value = [
+            _shape_node("last_hidden_state", [1, "sequence_length", 768]),
+        ]
+        assert OnnxModel._detect_static_batch_size(session) == 1
+
+    def test_rank_one_node_is_not_a_batch_axis(self):
+        session = MagicMock()
+        session.get_inputs.return_value = [
+            _shape_node("input_ids", ["batch_size", "sequence_length"]),
+            _shape_node("max_new_tokens", [1]),
+        ]
+        session.get_outputs.return_value = [
+            _shape_node("last_hidden_state", ["batch_size", "sequence_length", 768]),
+        ]
+        assert OnnxModel._detect_static_batch_size(session) is None
+
+
+def test_load_records_static_batch_size(model: ConcreteOnnxModel, mock_ort):
+    """The batch axis is read once, when the session is created."""
+    session_mock = mock_ort.InferenceSession.return_value
+    session_mock.get_inputs.return_value = [
+        _shape_node("input_ids", ["batch_size", "sequence_length"]),
+    ]
+    session_mock.get_outputs.return_value = [
+        _shape_node("last_hidden_state", [1, "sequence_length", 1024]),
+    ]
+
+    model._load_onnx_model(Path("dummy"), "model.onnx", OnnxSessionConfig(threads=None))
+
+    assert model.static_batch_size == 1
+
+
 def test_preprocess_onnx_input(model: ConcreteOnnxModel):
     """Test _preprocess_onnx_input returns unchanged input by default."""
     data = {"input_ids": np.array([[1, 2, 3]])}

tests/test_text_onnx_text_model.py

@@ -434,6 +434,80 @@ def test_parallel_uses_forkserver_or_spawn(self) -> None:
         assert call_kwargs["config"].start_method in ("forkserver", "spawn")
 
 
+class TestOnnxTextModelStaticBatchGraph:
+    """A graph pinned to a fixed batch must never be fed a larger one.
+
+    Feeding it more rows aborts inside onnxruntime with a buffer shape
+    mismatch, so the requested batch size is clamped to what the graph
+    declares -- and left alone when the batch axis is dynamic.
+    """
+
+    def _session(self, batch_dim: int | str) -> MagicMock:
+        """Mock session shaped like a real text-embedding graph."""
+        session = MagicMock()
+        session.run.return_value = [np.ones((1, 4), dtype=np.float32)]
+        node_ids = MagicMock()
+        node_ids.name = "input_ids"
+        node_ids.shape = ["batch_size", "sequence_length"]
+        session.get_inputs.return_value = [node_ids]
+        node_out = MagicMock()
+        node_out.name = "last_hidden_state"
+        node_out.shape = [batch_dim, "sequence_length", 4]
+        session.get_outputs.return_value = [node_out]
+        return session
+
+    def _tokenizer(self) -> MagicMock:
+        """Tokenizer whose output length tracks the batch it is given."""
+        enc = MagicMock()
+        enc.ids = [1, 2, 3, 0]
+        enc.attention_mask = [1, 1, 1, 0]
+        tok = MagicMock()
+        tok.encode_batch.side_effect = lambda documents: [enc] * len(documents)
+        return tok
+
+    def _model(self, batch_dim: int | str) -> tuple[ConcreteOnnxTextModel, MagicMock]:
+        session = self._session(batch_dim)
+        m = ConcreteOnnxTextModel()
+        m.model = session
+        m.model_input_names = {"input_ids"}
+        m.tokenizer = self._tokenizer()
+        m.static_batch_size = m._detect_static_batch_size(session)
+        return m, session
+
+    @staticmethod
+    def _batch_widths(session: MagicMock) -> list[int]:
+        return [call[0][1]["input_ids"].shape[0] for call in session.run.call_args_list]
+
+    def test_static_batch_graph_is_fed_one_document_at_a_time(self) -> None:
+        m, session = self._model(batch_dim=1)
+        documents = [f"document {i}" for i in range(8)]
+
+        list(m._embed_documents("t", "/tmp", documents=documents))
+
+        assert self._batch_widths(session) == [1] * 8
+
+    def test_dynamic_batch_graph_keeps_the_requested_batch(self) -> None:
+        m, session = self._model(batch_dim="batch_size")
+        documents = [f"document {i}" for i in range(8)]
+
+        list(m._embed_documents("t", "/tmp", documents=documents))
+
+        assert self._batch_widths(session) == [8]
+
+    def test_caller_batch_size_below_the_graph_limit_is_respected(self) -> None:
+        m, session = self._model(batch_dim=4)
+        documents = [f"document {i}" for i in range(6)]
+
+        list(m._embed_documents("t", "/tmp", documents=documents, batch_size=2))
+
+        assert self._batch_widths(session) == [2, 2, 2]
+
+    def test_unloaded_session_leaves_the_batch_size_alone(self) -> None:
+        """With lazy_load plus parallel the graph lives only in the workers."""
+        m = ConcreteOnnxTextModel()
+        assert m._effective_batch_size(256) == 256
+
+
 class TestTextEmbeddingWorkerProcess:
     """Lines 183-185 — TextEmbeddingWorker.process yields (idx, OnnxOutputContext)."""