fix: bound multi-slot backend allocations

.agents/skills/mblt-npu/SKILL.md

@@ -14,6 +14,7 @@ description: >-
 2. Preserve the `MobilintNPUBackend` compatibility contract: lifecycle methods, callable MXQ
    inference, input dtype inspection, serialization, and board-selected subclasses.
    For multi-slot backends, `max_batch_size` is aggregate capacity; keep slot-zero handles writable,
+   accept only non-boolean integers from 1 through 1024, cap each backend at 64 model slots,
    preserve `infer_slot()`, and roll back every slot after allocation or launch failure.
 3. Normalize legacy `aries`/`regulus` target values to board identifiers. Accepted boards are
    `aries-rb`, `regulus-ra`, `regulus-rb`, `regulus-ra-usb`, and `regulus-rb-usb`. Forward the

AGENTS.md

@@ -26,8 +26,9 @@ aligned with that floor before using newer language syntax or standard-library A
   selection behavior.
 - `max_batch_size` is aggregate capacity: probe the compiled per-slot batch axis
   `K`, launch `ceil(max_batch_size / K)` slots, share accelerators per device,
-  and preserve slot-zero compatibility handles. Keep allocation rollback and
-  `infer_slot()` deterministic across all slots.
+  and preserve slot-zero compatibility handles. Accept only non-boolean integers
+  from 1 through 1024 and cap the derived allocation at 64 model slots. Keep
+  allocation rollback and `infer_slot()` deterministic across all slots.
 - Normalize legacy `aries` to `aries-rb` and `regulus` to `regulus-ra`. Supported
   board identifiers are `aries-rb`, `regulus-ra`, `regulus-rb`, `regulus-ra-usb`,
   and `regulus-rb-usb`. The board name is passed as the first argument to

README.md

@@ -68,6 +68,10 @@ need to inspect the supported targets.
 
 `max_batch_size` is aggregate capacity. At `create()`, the backend probes the
 compiled per-model capacity `K` and loads `ceil(max_batch_size / K)` model slots.
+For predictable resource use, `max_batch_size` must be a non-boolean integer
+from 1 through 1024, and one backend may load at most 64 model slots. If the
+probed `K` would require more slots, creation disposes slot zero and fails before
+allocating any additional slots.
 Slots are distributed round-robin over the devices named by canonical target
 strings and reuse one accelerator per device. `mxq_model` and `acc` continue to
 refer to slot zero for compatibility; concurrent callers can use

mblt_npu/__init__.py

@@ -12,6 +12,8 @@
 from .npu_backend import (
     BACKEND_CLASSES,
     DEFAULT_TARGET_DEVICE,
+    MAX_BATCH_SIZE,
+    MAX_MODEL_SLOTS,
     MobilintBackendAllocError,
     MobilintAriesBackend,
     MobilintNPUBackend,
@@ -27,6 +29,8 @@
 __all__ = [
     "BACKEND_CLASSES",
     "DEFAULT_TARGET_DEVICE",
+    "MAX_BATCH_SIZE",
+    "MAX_MODEL_SLOTS",
     "MobilintBackendAllocError",
     "MobilintAriesBackend",
     "MobilintNPUBackend",

mblt_npu/npu_backend.py

@@ -27,7 +27,6 @@
 """
 
 import logging
-import math
 import os
 import re
 import sys
@@ -62,9 +61,26 @@
 DEFAULT_TARGET_DEVICE = "aries-rb"
 """Default supported Mobilint NPU board."""
 
+MAX_BATCH_SIZE = 1024
+"""Largest aggregate batch capacity accepted by the public configuration API."""
+
+MAX_MODEL_SLOTS = 64
+"""Largest number of native model slots a backend instance may allocate."""
+
 _TARGET_DEVICE_ALIASES = {"aries": "aries-rb", "regulus": "regulus-ra"}
 
 
+def _validate_max_batch_size(value: object) -> int:
+    """Return a safe aggregate capacity or reject it before native allocation."""
+    if isinstance(value, bool) or not isinstance(value, int):
+        raise TypeError("max_batch_size must be a non-boolean integer.")
+    if not 1 <= value <= MAX_BATCH_SIZE:
+        raise ValueError(
+            f"max_batch_size must be between 1 and {MAX_BATCH_SIZE}, got {value}."
+        )
+    return value
+
+
 def normalize_target_device(target_device: str) -> str:
     """Normalize a board identifier while accepting legacy product aliases."""
     if not isinstance(target_device, str):
@@ -242,9 +258,12 @@ def __init__(
                 expanded into ``target_cores`` / ``target_clusters`` when
                 those lists are empty, and prepends the device prefix to
                 legacy 2-part items.
-            max_batch_size: Requested aggregate batch capacity. The backend
+            max_batch_size: Requested aggregate batch capacity in the inclusive
+                range ``1..MAX_BATCH_SIZE``. Booleans and non-integers are rejected.
+                The backend
                 launches enough slots so that ``N * K >= max_batch_size``,
-                where ``K`` is the compiled batch axis of the MXQ artifact.
+                where ``K`` is the compiled batch axis of the MXQ artifact, up
+                to ``MAX_MODEL_SLOTS`` slots.
             core_mode: Execution mode that determines how NPU cores are
                 allocated. One of ``"single"``, ``"multi"``, ``"global4"``,
                 or ``"global8"``.
@@ -278,7 +297,7 @@ def __init__(
         self.revision = revision
         self._commit_hash = commit_hash
         self.mxq_path = mxq_path
-        self.max_batch_size = max(1, max_batch_size)
+        self.max_batch_size = _validate_max_batch_size(max_batch_size)
 
         # Multi-slot backing state; populated in create()/launch().
         # ``self.acc`` and ``self.mxq_model`` remain accessible as
@@ -951,15 +970,22 @@ def create(self) -> None:
         user-config or artifact bug.
 
         Raises:
+            TypeError: If the current ``max_batch_size`` is not a non-boolean
+                integer.
             MobilintBackendAllocError: If any slot hits a device-memory
                 ``BadAlloc`` or the slot 0 K probe hits a ``BadAlloc``.
             QbRuntimeError: If any slot or the slot 0 K probe fails for a
                 non-alloc reason (after partial-state rollback).
-            ValueError: If ``self.core_mode`` is not one of the supported
-                values.
+            ValueError: If the current ``max_batch_size`` is outside the
+                supported range or ``self.core_mode`` is unsupported.
             AssertionError: If ``"global8"`` mode is requested but a device
                 does not cover both clusters.
         """
+        # ``max_batch_size`` remains writable for compatibility. Revalidate
+        # its current value before opening accelerators or creating slot 0 so
+        # post-construction mutation cannot bypass the public capacity bounds.
+        self.max_batch_size = _validate_max_batch_size(self.max_batch_size)
+
         unique_devs = self._unique_devs_from_targets()
         if not unique_devs:
             unique_devs = [self._fallback_dev()]
@@ -1053,9 +1079,20 @@ def _spawn_slot(slot_idx: int, dev: int) -> "Model":
                 file=sys.stderr,
             )
             raise
-        self.n_models = max(
-            1, math.ceil(self.max_batch_size / max(1, self.k_per_model))
-        )
+        # Integer ceiling avoids converting attacker-controlled integers to a
+        # float. Keep this post-probe guard even though construction validates
+        # max_batch_size: the attribute remains writable for compatibility.
+        k_per_model = max(1, self.k_per_model)
+        self.n_models = (self.max_batch_size + k_per_model - 1) // k_per_model
+        if self.n_models > MAX_MODEL_SLOTS:
+            self._dispose_all_slots()
+            self.accs = {}
+            self.n_models = 0
+            raise ValueError(
+                f"max_batch_size={self.max_batch_size} with "
+                f"k_per_model={k_per_model} requires more than the supported "
+                f"maximum of {MAX_MODEL_SLOTS} model slots."
+            )
 
         for slot_idx in range(1, self.n_models):
             d = int(unique_devs[slot_idx % len(unique_devs)])

tests/test_npu_backend_multi_slot.py

@@ -21,6 +21,8 @@
 
 from mblt_npu import npu_backend as npu_backend_module
 from mblt_npu.npu_backend import (
+    MAX_BATCH_SIZE,
+    MAX_MODEL_SLOTS,
     MobilintBackendAllocError,
     MobilintNPUBackend,
     _is_qbruntime_bad_alloc,
@@ -153,6 +155,65 @@ def _make_backend_at(tmp_path, **kwargs) -> MobilintNPUBackend:
     return MobilintNPUBackend(**kwargs)
 
 
+@pytest.mark.parametrize("value", [True, False, 1.5, "2", None])
+def test_max_batch_size_rejects_non_integer_values(value) -> None:
+    """Public construction rejects values that are not strict integers."""
+    with pytest.raises(TypeError, match="non-boolean integer"):
+        MobilintNPUBackend(max_batch_size=value)
+
+
+@pytest.mark.parametrize("value", [0, -1, MAX_BATCH_SIZE + 1])
+def test_max_batch_size_rejects_values_outside_safe_range(value: int) -> None:
+    """Public construction bounds aggregate capacity before native allocation."""
+    with pytest.raises(ValueError, match=f"between 1 and {MAX_BATCH_SIZE}"):
+        MobilintNPUBackend(max_batch_size=value)
+
+
+def test_from_dict_validates_max_batch_size() -> None:
+    """Deserialized configuration uses the same strict admission checks."""
+    with pytest.raises(TypeError, match="non-boolean integer"):
+        MobilintNPUBackend.from_dict({"max_batch_size": 2.5})
+
+
+@pytest.mark.parametrize(
+    ("value", "error", "message"),
+    [
+        (True, TypeError, "non-boolean integer"),
+        (0, ValueError, f"between 1 and {MAX_BATCH_SIZE}"),
+        (MAX_BATCH_SIZE + 1, ValueError, f"between 1 and {MAX_BATCH_SIZE}"),
+    ],
+)
+def test_max_batch_size_revalidates_writable_assignments(
+    tmp_path, stub_qbruntime, value, error, message
+) -> None:
+    """Compatibility writes are revalidated before any native allocation."""
+    backend = _make_backend_at(tmp_path, max_batch_size=1)
+    backend.max_batch_size = value
+
+    with pytest.raises(error, match=message):
+        backend.create()
+
+    assert stub_qbruntime.models == []
+    assert backend.accs == {}
+
+
+def test_create_rejects_excessive_derived_slots_and_rolls_back(
+    tmp_path, stub_qbruntime
+) -> None:
+    """The post-probe slot cap disposes slot zero before further allocations."""
+    backend = _make_backend_at(tmp_path, max_batch_size=MAX_MODEL_SLOTS)
+    backend.max_batch_size = MAX_MODEL_SLOTS + 1
+
+    with pytest.raises(ValueError, match=f"maximum of {MAX_MODEL_SLOTS} model slots"):
+        backend.create()
+
+    assert len(stub_qbruntime.models) == 1
+    assert stub_qbruntime.models[0].disposed is True
+    assert backend.mxq_models == []
+    assert backend.accs == {}
+    assert backend.n_models == 0
+
+
 def test_backend_create_n1_single_slot_single_device(tmp_path, stub_qbruntime) -> None:
     """max_batch_size=1 with a K=1 model yields exactly one slot on the target device."""
     backend = _make_backend_at(