feat(multimodal): CUDA-graph bucket ladder (Dynamo-side target_bucket) + Qwen3-VL ViT example

components/src/dynamo/vllm/multimodal_utils/async_vision_encoder.py

@@ -87,9 +87,9 @@ def load(self, model_id: str) -> None:
         if self._batcher is not None or self._pool is not None:
             raise RuntimeError("AsyncVisionEncoder.load() called twice")
         # Construct the pool + batcher INSIDE the try so a constructor failure
-        # (e.g. a backend exposing a max_batch_cost the batcher rejects) still
-        # reaps the pool via shutdown() instead of leaking it. shutdown() is
-        # None-safe on the not-yet-assigned member.
+        # (e.g. a backend exposing a misconfigured buckets / max_batch_cost the
+        # batcher rejects) still reaps the pool via shutdown() instead of leaking
+        # it. shutdown() is None-safe on the not-yet-assigned member.
         try:
             self._pool = ThreadPoolExecutor(
                 max_workers=self._preprocess_concurrency,
@@ -98,6 +98,7 @@ def load(self, model_id: str) -> None:
             self._batcher = ThreadedMicroBatcher(
                 self._backend.forward_batch,
                 max_batch_cost=self._backend.max_batch_cost,
+                buckets=self._backend.buckets,
                 on_start=lambda: self._backend.build(model_id),
                 on_stop=self._backend.close,
                 name=self._name,

components/src/dynamo/vllm/multimodal_utils/threaded_micro_batcher.py

@@ -15,11 +15,17 @@
   pooled and split into batches whose summed ``cost`` stays within
   ``max_batch_cost`` (a compute/token budget, not a raw count). Packing is
   **one-dimensional** — by scalar ``cost`` alone; the batcher never inspects item
-  shape.
+  shape;
+- a **graph ladder** (optional ``buckets``): when set, the batcher rounds each
+  batch's ``sum(cost)`` **up to the nearest rung** and passes it as
+  ``target_bucket`` so ``fn`` can pad to that rung and replay its captured graph.
+  ``buckets=None`` ⇒ eager: ``target_bucket=None``.
 
 The caller speaks in opaque items plus a per-item scalar ``cost`` (int), computed
 once off-thread (see ``Preprocessed``); the batcher never interprets the items, so
-all model knowledge stays in the caller.
+all model knowledge stays in the caller. Final padding of a batch to a captured
+CUDA-graph shape is the ``fn``'s job (it owns
+the model), not the batcher's — the batcher only decides *which rung*.
 
 Coalescing window — **eager drain-on-completion, no timer** (the design default):
 whenever the worker is free it pulls everything queued and runs it, then repeats.
@@ -56,7 +62,7 @@
 import time
 from dataclasses import dataclass
 from enum import Enum, auto
-from typing import Callable, Generic, List, Optional, TypeVar
+from typing import Callable, Generic, List, Optional, Sequence, TypeVar
 
 logger = logging.getLogger(__name__)
 
@@ -113,18 +119,25 @@ class _Work(Generic[T]):
 
 
 class ThreadedMicroBatcher(Generic[T, R]):
-    """Run ``fn(list[item]) -> list[result]`` on a dedicated thread, coalescing
-    concurrent ``submit()`` calls into cost-bounded batches.
+    """Run ``fn(list[item], target_bucket) -> list[result]`` on a dedicated thread,
+    coalescing concurrent ``submit()`` calls into cost-bounded batches.
 
     Args:
-        fn: Batched work; one result per item, in order. Runs on the worker thread.
+        fn: Batched work; one result per item, in order. Called as
+            ``fn(items, target_bucket)`` on the worker thread — ``target_bucket``
+            is the ladder rung to pad to (``None`` in eager mode).
         max_batch_cost: Max summed ``cost`` of a single ``fn`` batch (>= 1).
             ``None`` (default) ⇒ **pass-through**: no cap — the whole drained set
-            runs as one ``fn`` call (``cost`` ignored).
+            runs as one ``fn`` call (``cost`` ignored). With ``buckets`` set,
+            ``None`` derives the ceiling as ``max(buckets)``.
+        buckets: Optional sorted graph ladder. When set, the batcher rounds a
+            batch's ``sum(cost)`` up to the nearest rung and passes it as
+            ``target_bucket``; ``None``/empty ⇒ eager (``target_bucket=None``).
         max_wait_ms: Opt-in coalescing hold after the first item arrives. Default
             ``0`` ⇒ eager drain-on-completion (no timer).
         on_start: Optional callable run once on the worker thread before serving
-            (model build / warmup); its failure surfaces from ``start()``.
+            (model build / warmup / graph capture); its failure surfaces from
+            ``start()``.
         on_stop: Optional callable run once on the worker thread at teardown (after
             the serving loop ends), iff ``on_start`` succeeded. Its failure is
             logged, never raised.
@@ -136,9 +149,10 @@ class ThreadedMicroBatcher(Generic[T, R]):
 
     def __init__(
         self,
-        fn: Callable[[List[T]], List[R]],
+        fn: Callable[[List[T], Optional[int]], List[R]],
         *,
         max_batch_cost: Optional[int] = None,
+        buckets: Optional[Sequence[int]] = None,
         max_wait_ms: float = 0.0,
         on_start: Optional[Callable[[], None]] = None,
         on_stop: Optional[Callable[[], None]] = None,
@@ -150,6 +164,11 @@ def __init__(
             raise ValueError("max_batch_cost must be >= 1 (or None for pass-through)")
         if max_outstanding_cost is not None and max_outstanding_cost < 1:
             raise ValueError("max_outstanding_cost must be >= 1")
+        self._buckets = self._validate_buckets(buckets, max_batch_cost)
+        # Graph mode needs a bounded ceiling, so derive it from the ladder when the
+        # author left it None (pass-through is only meaningful in eager mode).
+        if self._buckets is not None and max_batch_cost is None:
+            max_batch_cost = self._buckets[-1]
         self._fn = fn
         self._max_batch_cost = max_batch_cost
         self._max_wait_s = max_wait_ms / 1000.0
@@ -173,6 +192,37 @@ def __init__(
         self._live: set[_Request] = set()
         self._thread: Optional[threading.Thread] = None
 
+    @staticmethod
+    def _validate_buckets(
+        buckets: Optional[Sequence[int]], max_batch_cost: Optional[int]
+    ) -> Optional[tuple]:
+        """Normalise the ladder to a sorted tuple of positive ints (or None).
+
+        When ``max_batch_cost`` is set, the ladder must cover it so every packed
+        batch has a rung to round up to: ``max(buckets) >= max_batch_cost``. When
+        it is ``None``, the ceiling is derived as ``max(buckets)`` by the caller."""
+        if not buckets:
+            return None
+        rungs = tuple(sorted(int(b) for b in buckets))
+        if any(b < 1 for b in rungs):
+            raise ValueError("buckets must be positive ints")
+        if max_batch_cost is not None and rungs[-1] < max_batch_cost:
+            raise ValueError(
+                f"max(buckets)={rungs[-1]} < max_batch_cost={max_batch_cost}; the "
+                "ladder must cover the dispatch ceiling (set max_batch_cost to "
+                "max(buckets) for a graphed encoder)"
+            )
+        return rungs
+
+    def _target_bucket(self, batch_cost: int) -> Optional[int]:
+        """Round a batch's summed cost up to the nearest ladder rung (None=eager)."""
+        if self._buckets is None:
+            return None
+        for rung in self._buckets:  # sorted ascending
+            if rung >= batch_cost:
+                return rung
+        return self._buckets[-1]  # unreachable: max(buckets) >= max_batch_cost
+
     # ---- lifecycle ---------------------------------------------------------
 
     def start(self) -> None:
@@ -457,8 +507,9 @@ def _run_batch(self, batch: List[_Work]) -> None:
                 )
             return
         items = [w.item for w in runnable]
+        target_bucket = self._target_bucket(sum(w.cost for w in runnable))
         try:
-            results = self._fn(items)
+            results = self._fn(items, target_bucket)
         except (
             BaseException
         ) as exc:  # noqa: BLE001 — a bad batch must not hang awaiters

components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_async_vision_encoder.py

@@ -174,13 +174,14 @@ def test_preprocess_concurrency_must_be_positive():
 
 
 def test_load_reaps_pool_if_batcher_ctor_fails():
-    """A backend exposing a max_batch_cost the batcher rejects must not leak the pool."""
+    """A backend exposing a ladder the batcher rejects must not leak the pool."""
 
-    class _BadBudget(_FakeBackend):
-        max_batch_cost = 0  # ThreadedMicroBatcher requires >= 1 → ctor raises
+    class _BadBuckets(_FakeBackend):
+        max_batch_cost = 8
+        buckets = [2, 4]  # max(buckets) < max_batch_cost → batcher ctor raises
 
-    enc = AsyncVisionEncoder(_BadBudget())
-    with pytest.raises(ValueError, match="max_batch_cost"):
+    enc = AsyncVisionEncoder(_BadBuckets())
+    with pytest.raises(ValueError, match="ladder must cover"):
         enc.load("m")
     assert enc._pool is not None and enc._pool._shutdown is True
     assert enc._batcher is None

components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_threaded_micro_batcher.py

@@ -4,13 +4,13 @@
 """Unit tests for dynamo.vllm.multimodal_utils.threaded_micro_batcher.
 
 Pin the execution contract: on_start + every fn call (+ on_stop) run on one
-dedicated thread (so CUDA-graph capture/replay share a thread), concurrent submits
-coalesce into cost-bounded batches up to max_batch_cost (or pass-through when
-None), eager-drain pulls all queued work when free, errors reach every awaiting
-caller, and the shutdown lifecycle behaves.
+dedicated thread (so CUDA-graph capture/replay share a thread), concurrent
+submits coalesce into cost-bounded same-bucket batches, the graph ladder rounds a
+batch's cost up to a rung (target_bucket), eager-drain pulls all queued work when
+free, errors reach every awaiting caller, and the shutdown lifecycle behaves.
 
-``fn`` is ``fn(items)``; ``cost`` is a precomputed scalar that rides on
-``submit(items, costs)`` (one-dimensional packing — no bucket_key, no ladder).
+``fn`` is ``fn(items, target_bucket)``; ``cost`` is a precomputed scalar that rides
+on ``submit(items, costs)`` (one-dimensional packing — no bucket_key).
 """
 
 import asyncio
@@ -32,16 +32,17 @@
 ]
 
 
-def _echo(items):
+def _echo(items, target_bucket=None):
     return list(items)
 
 
 class _Recorder:
-    """fn that records the threads it ran on and the batches it received."""
+    """fn that records the threads it ran on and the batches / target_buckets it saw."""
 
     def __init__(self):
         self.threads: list[int] = []
         self.batches: list[list] = []
+        self.target_buckets: list = []
         self.start_thread: int | None = None
         self.stop_thread: int | None = None
 
@@ -51,9 +52,10 @@ def on_start(self):
     def on_stop(self):
         self.stop_thread = threading.get_ident()
 
-    def fn(self, items):
+    def fn(self, items, target_bucket=None):
         self.threads.append(threading.get_ident())
         self.batches.append(list(items))
+        self.target_buckets.append(target_bucket)
         return [("r", x) for x in items]
 
 
@@ -117,7 +119,7 @@ async def test_eager_drain_pulls_all_queued_when_free():
     release = threading.Event()
     batches: list[list] = []
 
-    def fn(items):
+    def fn(items, target_bucket=None):
         batches.append(list(items))
         if "block" in items:
             entered.set()
@@ -155,9 +157,48 @@ async def test_cost_budget_caps_each_batch():
         b.shutdown()
 
 
+async def test_target_bucket_rounds_packed_cost_up_to_nearest_rung():
+    """Graph mode: the batcher rounds a batch's sum(cost) up to the nearest rung
+    and passes it as target_bucket."""
+    rec = _Recorder()
+    b = ThreadedMicroBatcher(
+        rec.fn, max_batch_cost=8, buckets=[2, 4, 8], max_wait_ms=200.0
+    )
+    b.start()
+    try:
+        # One coalesced batch of cost 3 → rounds up to rung 4.
+        await b.submit(["x", "y", "z"], costs=[1, 1, 1])
+        assert rec.batches == [["x", "y", "z"]]
+        assert rec.target_buckets == [4]
+    finally:
+        b.shutdown()
+
+
+async def test_eager_mode_passes_none_target_bucket():
+    rec = _Recorder()
+    b = ThreadedMicroBatcher(rec.fn)  # buckets=None ⇒ eager
+    b.start()
+    try:
+        await b.submit(["a"])
+        assert rec.target_buckets == [None]
+    finally:
+        b.shutdown()
+
+
+def test_buckets_below_max_batch_cost_rejected():
+    with pytest.raises(ValueError, match="ladder must cover"):
+        ThreadedMicroBatcher(_echo, max_batch_cost=8, buckets=[2, 4])
+
+
+def test_buckets_derive_max_batch_cost_when_none():
+    """Graph mode with no explicit budget derives the ceiling from the ladder."""
+    b = ThreadedMicroBatcher(_echo, buckets=[2, 4, 8])  # max_batch_cost=None
+    assert b._max_batch_cost == 8
+
+
 async def test_max_batch_cost_none_is_passthrough():
     """Default (max_batch_cost=None): no cap and no per-item ceiling — the whole
-    drained set runs as ONE fn call regardless of summed cost."""
+    drained same-bucket group runs as ONE fn call regardless of summed cost."""
     rec = _Recorder()
     b = ThreadedMicroBatcher(rec.fn, max_wait_ms=200.0)  # max_batch_cost=None
     b.start()
@@ -166,12 +207,13 @@ async def test_max_batch_cost_none_is_passthrough():
         out = await b.submit([1, 2, 3, 4], costs=[1000, 1000, 1000, 1000])
         assert len(out) == 4
         assert rec.batches == [[1, 2, 3, 4]]  # one un-split batch
+        assert rec.target_buckets == [None]  # eager (no ladder)
     finally:
         b.shutdown()
 
 
 async def test_error_reaches_every_caller():
-    def boom(items):
+    def boom(items, target_bucket=None):
         raise ValueError("boom")
 
     b = ThreadedMicroBatcher(boom, max_wait_ms=50.0)
@@ -186,7 +228,7 @@ def boom(items):
 
 
 async def test_wrong_result_count_raises():
-    b = ThreadedMicroBatcher(lambda items: [], max_wait_ms=10.0)
+    b = ThreadedMicroBatcher(lambda items, target_bucket=None: [], max_wait_ms=10.0)
     b.start()
     try:
         with pytest.raises(RuntimeError, match="one result per item"):
@@ -233,7 +275,7 @@ async def test_shutdown_fails_queued_items():
     entered = threading.Event()
     release = threading.Event()
 
-    def blocking(items):
+    def blocking(items, target_bucket=None):
         entered.set()
         release.wait(timeout=5.0)
         return [("r", x) for x in items]
@@ -275,7 +317,7 @@ async def test_cancelled_submit_is_retired_and_releases_admission():
     entered = threading.Event()
     release = threading.Event()
 
-    def blocking(items):
+    def blocking(items, target_bucket=None):
         entered.set()
         release.wait(timeout=5.0)
         return [("r", x) for x in items]
@@ -309,7 +351,7 @@ async def test_max_outstanding_cost_rejects_when_full():
     entered = threading.Event()
     release = threading.Event()
 
-    def blocking(items):
+    def blocking(items, target_bucket=None):
         entered.set()
         release.wait(timeout=5.0)
         return [("r", x) for x in items]
@@ -379,7 +421,7 @@ async def test_partial_batch_failure_fails_request_once_and_releases():
     later sibling item is tombstoned — it never reaches fn."""
     seen: list = []
 
-    def fn(items):
+    def fn(items, target_bucket=None):
         seen.extend(items)
         if "bad" in items:
             raise ValueError("boom")
@@ -407,7 +449,7 @@ async def test_no_fn_after_shutdown_for_collected_items():
     release = threading.Event()
     seen: list = []
 
-    def fn(items):
+    def fn(items, target_bucket=None):
         seen.extend(items)
         if "a" in items:
             entered.set()