From fc9b5d223195e83dcdc53637d623d4140565f121 Mon Sep 17 00:00:00 2001 From: Qi Wang Date: Tue, 30 Jun 2026 00:12:09 -0700 Subject: [PATCH] feat(multimodal): CUDA-graph bucket ladder (Dynamo-side target_bucket) + Qwen3-VL ViT example MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR4 of the custom vision-encoder series, stacked on PR3. Introduces bucket-wise graph support on the Dynamo side and a real graphed encoder example. - ThreadedMicroBatcher: optional `buckets` ladder. When set, the batcher rounds a batch's packed sum(cost) UP to the nearest rung and passes it as `target_bucket` to fn(items, target_bucket). With max_batch_cost=None it derives the ceiling as max(buckets); buckets=None stays eager / pass-through (PR3 behavior unchanged). - AsyncVisionEncoder passes backend.buckets to the batcher. - Qwen3VLViTEncoder example: loads the real Qwen3-VL vision tower (build(model_id), picks its own device), captures one CUDA graph per rung via torch.compile(reduce-overhead), and in forward_batch pads sum(cost) up to target_bucket, replays, slices the real images back out (CPU). Hardcodes image_token_id via the Qwen base. The author owns padding; Dynamo only picks the rung. `buckets`/`target_bucket` were forward-compat in the contract since PR1; this PR makes them live. Tests: target_bucket rounding (boundary + eager None + buckets-derive + ladder-covers-budget). Graph smoke validated on the integration branch. Deepstack limitation: Qwen3-VL deepstack features are not carried by the one-tensor-per-image contract — this exercises the mechanism, not accuracy. Note: based on the series' validated merge-base; rebase onto main before merge. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../multimodal_utils/async_vision_encoder.py | 7 +- .../threaded_micro_batcher.py | 71 ++++- .../test_vllm_async_vision_encoder.py | 11 +- .../test_vllm_threaded_micro_batcher.py | 78 ++++-- .../custom_encoder/qwen3vl_vit_encoder.py | 246 ++++++++++++++++++ 5 files changed, 377 insertions(+), 36 deletions(-) create mode 100644 examples/custom_encoder/qwen3vl_vit_encoder.py diff --git a/components/src/dynamo/vllm/multimodal_utils/async_vision_encoder.py b/components/src/dynamo/vllm/multimodal_utils/async_vision_encoder.py index 836c3bc65b2b..5c02e2644190 100644 --- a/components/src/dynamo/vllm/multimodal_utils/async_vision_encoder.py +++ b/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, diff --git a/components/src/dynamo/vllm/multimodal_utils/threaded_micro_batcher.py b/components/src/dynamo/vllm/multimodal_utils/threaded_micro_batcher.py index 222d35b43bfc..9c2fe5b23259 100644 --- a/components/src/dynamo/vllm/multimodal_utils/threaded_micro_batcher.py +++ b/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 diff --git a/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_async_vision_encoder.py b/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_async_vision_encoder.py index 0c0689c48919..f9d1ad40e83f 100644 --- a/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_async_vision_encoder.py +++ b/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 diff --git a/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_threaded_micro_batcher.py b/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_threaded_micro_batcher.py index 59be48f6c2f0..13199341ec39 100644 --- a/components/src/dynamo/vllm/tests/multimodal_utils/test_vllm_threaded_micro_batcher.py +++ b/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() diff --git a/examples/custom_encoder/qwen3vl_vit_encoder.py b/examples/custom_encoder/qwen3vl_vit_encoder.py new file mode 100644 index 000000000000..86b1fe83f6bb --- /dev/null +++ b/examples/custom_encoder/qwen3vl_vit_encoder.py @@ -0,0 +1,246 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Real Qwen3-VL vision-tower ``VisionEncoderBackend`` (in-process, CUDA-graphed). + +Loads the **actual** Qwen3-VL vision tower (ViT patch-embed + transformer blocks + +spatial merger) and runs it on the ``AsyncVisionEncoder`` dedicated actor thread. +It demonstrates the **bucket-ladder** graph scheme (design theme D2, the vLLM +``EncoderCudaGraphManager`` pattern) end to end: + +- ``buckets`` exposes a sorted ladder of **merged-visual-token** rungs. +- The ``ThreadedMicroBatcher`` packs images by ``cost`` (merged tokens) and rounds + the packed ``sum(cost)`` **up to the nearest rung**, passing it as + ``target_bucket``. +- ``forward_batch`` **pads** the packed batch up to ``target_bucket`` (appends + dummy images), replays the graph captured for that rung's shape, and slices the + real images' embeds back out. + +Padding the input to a rung quantises the forward's shape to the (bounded) ladder, +so ``torch.compile(mode="reduce-overhead")`` captures **one CUDA graph per rung** +(not one per arbitrary batch size — that would be SGLang's unbounded per-exact-S +scheme, which the design rejects). Capture + replay both happen on the actor +thread, the affinity the batcher guarantees. + +Stable per-image shape: every image is resized to one fixed square, so each image +is exactly ``TOKENS_PER_SIDE**2`` merged tokens — every rung is a whole number of +images and padding is always whole dummy images. ``forward_batch`` copies its +output to CPU, which also detaches it from the reused CUDA-graph output buffer. + +Limitation (intentional for this harness): Qwen3-VL's vision tower also emits +``deepstack_features`` that the LM injects at specific layers. The contract +returns one embed tensor per image (the merged ``pooler_output``); deepstack +features are not plumbed through the mixed-embeds splice path, so image grounding +is approximate. This exercises the in-process encoder + batcher + CUDA-graph +bucket-ladder path, **not** Qwen3-VL accuracy. + +Usage (via agg_custom.sh): + DYN_MODEL=Qwen/Qwen3-VL-2B-Instruct + DYN_ENCODER_CLASS=examples.custom_encoder.qwen3vl_vit_encoder.Qwen3VLViTEncoder + DYN_WORKER_GPU=2 ./agg_custom.sh + +Env knobs: + DYN_VIT_COMPILE 1 (default) → torch.compile + bucket ladder; 0 → eager (no graphs) + DYN_VIT_TOKENS_PER_SIDE merged visual tokens per image side (default 16 → 256/img) + DYN_VIT_MAX_IMAGES max images per forward → top rung (default 16) +""" + +from __future__ import annotations + +import base64 +import io +import logging +import os +from typing import Any, Dict, List, Optional + +import requests +import torch +from PIL import Image +from transformers import AutoModelForImageTextToText, AutoProcessor + +from dynamo.vllm.multimodal_utils.vision_encoder_backend import Preprocessed +from examples.custom_encoder.qwen_vision_encoder import QwenVisionEncoderBackend + +logger = logging.getLogger(__name__) + +_COMPILE = os.environ.get("DYN_VIT_COMPILE", "1") == "1" +_TOKENS_PER_SIDE = int(os.environ.get("DYN_VIT_TOKENS_PER_SIDE", "16")) +_MAX_IMAGES = int(os.environ.get("DYN_VIT_MAX_IMAGES", "16")) + + +def _load_image(image_url: str) -> Image.Image: + """Load a PIL RGB image from a data: URI, http(s) URL, or local path.""" + if image_url.startswith("data:"): + b64 = image_url.split(",", 1)[1] + return Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB") + if image_url.startswith(("http://", "https://")): + resp = requests.get(image_url, timeout=30) + resp.raise_for_status() + return Image.open(io.BytesIO(resp.content)).convert("RGB") + return Image.open(image_url).convert("RGB") + + +def _power_of_two_ladder(tokens_per_img: int, max_images: int) -> List[int]: + """Merged-token rungs at 1, 2, 4, ... images up to ``max_images``. + + Power-of-two image counts (not every count) so the ladder — and thus the + captured-graph count — stays bounded; sub-rung batches pad up to the next.""" + mults: List[int] = [] + m = 1 + while m < max_images: + mults.append(m) + m *= 2 + mults.append(max_images) + return sorted({tokens_per_img * x for x in mults}) + + +class Qwen3VLViTEncoder(QwenVisionEncoderBackend): + """In-process Qwen3-VL vision-tower backend with a CUDA-graphed bucket ladder.""" + + def __init__(self) -> None: + self._tokens_per_img = _TOKENS_PER_SIDE**2 + if _COMPILE: + # Graphed: expose the ladder; the top rung is the dispatch ceiling. + self.buckets = _power_of_two_ladder(self._tokens_per_img, _MAX_IMAGES) + self.max_batch_cost = self.buckets[-1] + else: + # Eager: no ladder; pack up to max_images worth of tokens, no padding. + self.buckets = None + self.max_batch_cost = self._tokens_per_img * _MAX_IMAGES + + def build(self, model_id: str) -> None: + """Load tokenizer (Qwen base) + processor + ViT; compile and warm up so one + CUDA graph per rung is captured on this (the actor) thread.""" + super().build(model_id) # self.tokenizer + # The worker pins the GPU via CUDA_VISIBLE_DEVICES, so the current device + # ("cuda") is correct — the backend picks its own device (no device arg). + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.processor = AutoProcessor.from_pretrained(model_id) + + model = AutoModelForImageTextToText.from_pretrained( + model_id, dtype=torch.bfloat16 + ) + # Qwen3VLForConditionalGeneration.model is Qwen3VLModel(.visual, .language_model). + inner = getattr(model, "model", model) + visual = getattr(inner, "visual", None) or getattr(model, "visual") + self.visual = visual.to(self.device).eval() + del model # drop the LM half; the vLLM worker owns the LM + torch.cuda.empty_cache() + + vc = self.visual.config + self.merge = int(vc.spatial_merge_size) + patch = int(vc.patch_size) + # Fixed square so grid_thw (→ cost) is constant: every image is + # exactly _tokens_per_img merged tokens. + self.side = patch * self.merge * _TOKENS_PER_SIDE + self._fixed_hw = (self.side, self.side) + + # One dummy (gray) image's processed tensors, reused for padding to a rung. + self._dummy = self._process(Image.new("RGB", self._fixed_hw, (127, 127, 127))) + + self._eager_visual = self.visual + self._compiled = False + if _COMPILE: + self.visual = torch.compile(self.visual, mode="reduce-overhead") + self._compiled = True + + self._warmup() + + def _warmup(self) -> None: + """Forward once at **each rung** so torch.compile captures that rung's CUDA + graph here; fall back to eager if compile/capture fails.""" + try: + rungs = self.buckets if self.buckets else [self._tokens_per_img] + for _ in range(2 if self._compiled else 1): + for rung in rungs: + # One real image padded up to the rung (the forward_batch path). + self.forward_batch( + [self._dummy], target_bucket=rung if self._compiled else None + ) + torch.cuda.synchronize() + logger.info( + "[Qwen3VLViTEncoder] ready: side=%d merge=%d compile=%s " + "tokens/img=%d buckets=%s max_batch_cost=%d", + self.side, + self.merge, + self._compiled, + self._tokens_per_img, + list(self.buckets) if self.buckets else None, + self.max_batch_cost, + ) + except Exception as exc: # noqa: BLE001 — compile/capture is best-effort + if self._compiled: + logger.warning( + "[Qwen3VLViTEncoder] compile/warmup failed (%s); using eager", exc + ) + self.visual = self._eager_visual + self._compiled = False + self.buckets = None + self.forward_batch([self._dummy]) + torch.cuda.synchronize() + else: + raise + + # ---- preprocess (off the actor thread) --------------------------------- + + def preprocess(self, image_url: str) -> Preprocessed[Dict[str, Any]]: + """Off-thread: fetch + resize to the fixed square + HF patchify. + + ``cost`` = merged visual tokens for this image (the scalar Dynamo packs + by). Every image is the same fixed square, so the ViT shapes are uniform — + the author owns any shape/padding concerns inside ``forward_batch``.""" + img = _load_image(image_url).resize(self._fixed_hw) + item = self._process(img) + t, h, w = item["grid_thw"][0].tolist() + cost = (t * h * w) // (self.merge**2) + return Preprocessed(item=item, cost=cost) + + def _process(self, img: Image.Image) -> Dict[str, Any]: + out = self.processor.image_processor(images=[img], return_tensors="pt") + return {"pixel_values": out["pixel_values"], "grid_thw": out["image_grid_thw"]} + + # ---- forward (on the actor thread) ------------------------------------- + + @torch.inference_mode() + def forward_batch( + self, items: List[Dict[str, Any]], target_bucket: Optional[int] = None + ) -> List[torch.Tensor]: + """Pad the packed batch up to ``target_bucket``, replay that rung's graph, + and slice the real images' embeds back out (one tensor per item, in order).""" + if os.environ.get("DYN_VIT_LOG_BATCH") == "1": + logger.info( + "[Qwen3VLViTEncoder] forward_batch images=%d target_bucket=%s", + len(items), + target_bucket, + ) + n_real = len(items) + pix_parts = [it["pixel_values"] for it in items] + grid_parts = [it["grid_thw"] for it in items] + + # Pad up to the rung (whole dummy images) so the forward shape is one of the + # captured rungs. target_bucket is None in eager mode → no padding. + if target_bucket is not None: + real_tokens = sum( + (t * h * w) // (self.merge**2) + for t, h, w in torch.cat(grid_parts, dim=0).tolist() + ) + n_dummy = (target_bucket - real_tokens) // self._tokens_per_img + for _ in range(max(0, n_dummy)): + pix_parts.append(self._dummy["pixel_values"]) + grid_parts.append(self._dummy["grid_thw"]) + + pix = torch.cat(pix_parts, dim=0).to(self.device, dtype=torch.bfloat16) + grid = torch.cat(grid_parts, dim=0).to(self.device) + embeds = self.visual(pix, grid).pooler_output # (total_merged_tokens, hidden) + sizes = [(t * h * w) // (self.merge**2) for t, h, w in grid.tolist()] + parts = torch.split(embeds, sizes, dim=0) + # Copy to CPU: detaches from any reused CUDA-graph output buffer and matches + # the assembler's CPU prompt_embeds layout (it preserves dtype). Drop the + # padding dummies — return only the real images, in input order. + return [parts[i].detach().to("cpu", copy=True) for i in range(n_real)] + + def close(self) -> None: + """Drop the ViT + compiled graphs on the actor thread.""" + self.visual = None + self._eager_visual = None + torch.cuda.empty_cache()