Add compile-once-and-pad path for Qwen3-VL multimodal generation

Lets one compiled generation graph serve repeated calls with images
of varying sizes. The first call JIT-compiles for the configured
upper bounds; subsequent calls hit the EXLA cache and run ~2-3x
faster on CPU.

Featurizer changes:
- Two new options, :max_patches and :max_num_images. When set,
  pixel_values is right-padded along the patches axis with zeros and
  image_grid_thw is padded with [0, 0, 0] rows.
- Guard rails: :max_patches must be a multiple of merge_size**2 and
  must accommodate the largest image the user plans to send.

Vision encoder changes:
- patch_metadata derives a patch_valid mask (i < total_real_patches)
  from the padded grid_thw. Padded image_id values are clipped to a
  valid index so gather operations succeed, and the safe_grid_w guard
  prevents division by zero in the row/col derivation for padded
  positions.
- The block-diagonal attention mask is ANDed with patch_valid so
  padded patches neither attend nor are attended to. Their embedding
  contributions therefore drop out of the output entirely, which is
  what makes the padding correctness-preserving.

API:
- Bumblebee.Multimodal.ImageTextToText.compile/5 configures the
  featurizer + tokenizer at the upper-bound shapes and returns a
  state struct.
- run/3 takes that state plus a prompt+image and runs generation,
  hitting the cached compiled graph.

Validation:
- 4 new featurizer tests covering the padding shape, padding values,
  the merge_size**2 constraint, and the "too many patches" error.
- Full fast suite: 289 passed, 0 regressions.
- Real Qwen3-VL-2B-Instruct on COCO 39769 with greedy decode and
  max_patches=1024: cold call 27.3s, warm call 10.1s (2.7x speedup),
  both produce "A group of cats are lying on a pink blanket with
  remote controls."

Refs #442.

Author: nyo16

lib/bumblebee/multimodal/image_text_to_text.ex

@@ -1,13 +1,19 @@
 defmodule Bumblebee.Multimodal.ImageTextToText do
   @moduledoc """
-  Generation helper for vision-language models like Qwen3-VL.
-
-  This wraps featurization, prompt expansion, and `Bumblebee.Text.Generation`
-  in a single call. Each call recompiles the generation graph if the
-  image or prompt produces a different total patch count or sequence
-  length, which makes this best suited for interactive or one-shot use.
-  For high-throughput serving with batched, varying image sizes, see
-  the static-shape padding follow-up.
+  Generation helpers for vision-language models like Qwen3-VL.
+
+  Two entry points:
+
+    * `generate/6` — one-shot call. Featurizes, expands the prompt
+      placeholder, and runs generation. Each call recompiles the graph
+      when the image or sequence length changes, so it suits
+      interactive use.
+
+    * `compile/5` + `run/3` — compile the generation graph **once** for
+      upper-bound shapes, then run repeatedly with images of varying
+      sizes. The featurizer pads `pixel_values` and `image_grid_thw` to
+      the configured maxima, and the vision encoder excludes padded
+      patches from attention via `patch_valid`.
   """
 
   alias Bumblebee.Text
@@ -97,6 +103,123 @@ defmodule Bumblebee.Multimodal.ImageTextToText do
     %{text: decoded, token_ids: token_ids}
   end
 
+  @doc """
+  Compiles the generation graph once for the given upper-bound shapes.
+
+  The returned struct can be passed to `run/3` repeatedly. Calls with
+  images that produce fewer than `:max_patches` real patches or
+  shorter than `:sequence_length` prompts are padded; the vision
+  encoder masks the padded positions out of attention.
+
+  ## Options
+
+    * `:max_patches` (required) — upper bound on total patches across
+      all images in one call. Must be a multiple of `merge_size ** 2`.
+    * `:max_num_images` (required) — upper bound on number of images
+      per call.
+    * `:sequence_length` (required) — upper bound on token count
+      (prompt + generated).
+  """
+  def compile(
+        model_info,
+        featurizer,
+        tokenizer,
+        %Text.GenerationConfig{} = generation_config,
+        opts
+      ) do
+    opts = Keyword.validate!(opts, [:max_patches, :max_num_images, :sequence_length])
+    max_patches = Keyword.fetch!(opts, :max_patches)
+    max_num_images = Keyword.fetch!(opts, :max_num_images)
+    sequence_length = Keyword.fetch!(opts, :sequence_length)
+
+    %{model: model, params: params, spec: spec} = model_info
+
+    unless Map.has_key?(spec, :image_token_id) do
+      raise ArgumentError,
+            "expected a multimodal model with :image_token_id, got #{inspect(spec.__struct__)}"
+    end
+
+    merge_size = spec.vision_spec.spatial_merge_size
+
+    featurizer =
+      Bumblebee.configure(featurizer,
+        max_patches: max_patches,
+        max_num_images: max_num_images
+      )
+
+    tokenizer =
+      Bumblebee.configure(tokenizer,
+        length: sequence_length,
+        pad_direction: :left,
+        return_token_type_ids: false
+      )
+
+    generate_fun = Text.Generation.build_generate(model, spec, generation_config)
+
+    %{
+      generate_fun: generate_fun,
+      params: params,
+      spec: spec,
+      featurizer: featurizer,
+      tokenizer: tokenizer,
+      merge_size: merge_size,
+      max_patches: max_patches,
+      max_num_images: max_num_images,
+      sequence_length: sequence_length
+    }
+  end
+
+  @doc """
+  Runs a prompt + image through a pre-compiled generator from `compile/5`.
+
+  EXLA caches the compiled graph by input shape; since the featurizer
+  pads to the upper bounds configured in `compile/5`, every call hits
+  the same cached graph.
+  """
+  def run(compiled, text, image) do
+    %{
+      generate_fun: generate_fun,
+      params: params,
+      featurizer: featurizer,
+      tokenizer: tokenizer,
+      merge_size: merge_size
+    } = compiled
+
+    image_inputs = Bumblebee.apply_featurizer(featurizer, image)
+    grid_thw_real = unpad_grid_thw(image_inputs["image_grid_thw"])
+    visual_tokens = visual_tokens_for(grid_thw_real, merge_size)
+    expanded_text = expand_marker(text, visual_tokens)
+
+    text_inputs = Bumblebee.apply_tokenizer(tokenizer, expanded_text)
+
+    inputs =
+      text_inputs
+      |> Map.merge(image_inputs)
+      |> Map.put("seed", Nx.tensor([:erlang.system_time()], type: :s64))
+
+    %{token_ids: token_ids} = generate_fun.(params, inputs)
+
+    decoded =
+      token_ids
+      |> Nx.to_batched(1)
+      |> Enum.map(&Bumblebee.Tokenizer.decode(tokenizer, Nx.to_flat_list(&1)))
+      |> hd()
+
+    %{text: decoded, token_ids: token_ids}
+  end
+
+  # Drops padding rows ([0, 0, 0]) so visual_tokens_for matches the
+  # actual prompt expansion length.
+  defp unpad_grid_thw(grid_thw) do
+    grid_thw
+    |> Nx.to_list()
+    |> Enum.reject(fn [t, h, w] -> t == 0 and h == 0 and w == 0 end)
+    |> case do
+      [] -> Nx.tensor([[0, 0, 0]], type: :s64)
+      rows -> Nx.tensor(rows, type: :s64)
+    end
+  end
+
   defp expand_marker(text, visual_tokens) do
     case String.split(text, @placeholder) do
       [_only] ->

lib/bumblebee/vision/qwen3_vl_featurizer.ex

@@ -57,6 +57,21 @@ defmodule Bumblebee.Vision.Qwen3VLFeaturizer do
       explicit maximum total pixels after smart-resize. Overrides the `:quality`
       preset when set.
       """
+    ],
+    max_patches: [
+      default: nil,
+      doc: """
+      when set, pads `pixel_values` along the patches axis to this size with
+      zeros. Required for compile-once-and-pad serving of variable-size
+      images. Must be a multiple of `merge_size ** 2`.
+      """
+    ],
+    max_num_images: [
+      default: nil,
+      doc: """
+      when set, pads `image_grid_thw` to this many rows with `[0, 0, 0]`.
+      Required alongside `:max_patches` for compile-once-and-pad serving.
+      """
     ]
   ]
 
@@ -117,12 +132,71 @@ defmodule Bumblebee.Vision.Qwen3VLFeaturizer do
       |> Enum.map(& &1.grid_thw)
       |> Nx.stack()
 
+    {pixel_values, image_grid_thw} =
+      maybe_pad_to_max(pixel_values, image_grid_thw, featurizer)
+
     %{
       "pixel_values" => pixel_values,
       "image_grid_thw" => image_grid_thw
     }
   end
 
+  defp maybe_pad_to_max(pixel_values, image_grid_thw, featurizer) do
+    pixel_values = maybe_pad_patches(pixel_values, featurizer)
+    image_grid_thw = maybe_pad_grid_thw(image_grid_thw, featurizer)
+    {pixel_values, image_grid_thw}
+  end
+
+  defp maybe_pad_patches(pixel_values, %{max_patches: nil}), do: pixel_values
+
+  defp maybe_pad_patches(pixel_values, featurizer) do
+    {num_patches, flat} = Nx.shape(pixel_values)
+    max_patches = featurizer.max_patches
+    merge_sq = featurizer.merge_size * featurizer.merge_size
+
+    unless rem(max_patches, merge_sq) == 0 do
+      raise ArgumentError,
+            ":max_patches (#{max_patches}) must be a multiple of merge_size**2 " <>
+              "(= #{merge_sq})"
+    end
+
+    if num_patches > max_patches do
+      raise ArgumentError,
+            "featurizer produced #{num_patches} patches but :max_patches is " <>
+              "#{max_patches}; raise :max_patches or lower :quality / :max_pixels"
+    end
+
+    pad_rows = max_patches - num_patches
+
+    if pad_rows == 0 do
+      pixel_values
+    else
+      padding = Nx.broadcast(Nx.tensor(0.0, type: Nx.type(pixel_values)), {pad_rows, flat})
+      Nx.concatenate([pixel_values, padding], axis: 0)
+    end
+  end
+
+  defp maybe_pad_grid_thw(image_grid_thw, %{max_num_images: nil}), do: image_grid_thw
+
+  defp maybe_pad_grid_thw(image_grid_thw, featurizer) do
+    {num_images, 3} = Nx.shape(image_grid_thw)
+    max_num_images = featurizer.max_num_images
+
+    if num_images > max_num_images do
+      raise ArgumentError,
+            "got #{num_images} images but :max_num_images is #{max_num_images}"
+    end
+
+    pad_rows = max_num_images - num_images
+
+    if pad_rows == 0 do
+      image_grid_thw
+    else
+      padding = Nx.broadcast(Nx.tensor(0, type: Nx.type(image_grid_thw)), {pad_rows, 3})
+      Nx.concatenate([image_grid_thw, padding], axis: 0)
+    end
+  end
+
   defp normalize_input(input) when is_list(input), do: input
   defp normalize_input(%{image: _} = input), do: [input]
   defp normalize_input(%{video: _} = input), do: [input]

lib/bumblebee/vision/qwen3_vl_vision.ex

@@ -250,7 +250,7 @@ defmodule Bumblebee.Vision.Qwen3VLVision do
     src_grid_size = trunc(:math.sqrt(spec.num_position_embeddings))
     merge_size = spec.spatial_merge_size
 
-    {row_in_image, col_in_image, grid_h_per_patch, grid_w_per_patch, _image_id} =
+    {row_in_image, col_in_image, grid_h_per_patch, grid_w_per_patch, _image_id, _patch_valid} =
       patch_metadata(grid_thw, total_patches, merge_size)
 
     src_max_f = Nx.tensor(src_grid_size - 1, type: :f32)
@@ -317,24 +317,41 @@ defmodule Bumblebee.Vision.Qwen3VLVision do
 
     cumulative = Nx.cumulative_sum(patches_per_image)
     exclusive_cumulative = Nx.subtract(cumulative, patches_per_image)
+    total_real_patches = Nx.sum(patches_per_image)
 
     patch_indices = Nx.iota({total_patches}, type: :s64)
 
-    image_id_per_patch =
+    # Patches beyond total_real_patches are padding slots (when the
+    # featurizer was configured with :max_patches). Mark them invalid so
+    # downstream attention masking can exclude them entirely.
+    patch_valid = Nx.less(patch_indices, total_real_patches)
+
+    image_id_raw =
       patch_indices
       |> Nx.new_axis(-1)
       |> Nx.greater_equal(Nx.new_axis(cumulative, 0))
       |> Nx.sum(axes: [-1])
       |> Nx.as_type(:s64)
 
+    n_images = Nx.axis_size(grid_thw, 0)
+    # Padded patches map to image_id == n_images (out of bounds). Clip so
+    # gather operations succeed. Their derived row/col/grid values are
+    # garbage but get masked out via `patch_valid` in the attention step.
+    image_id_per_patch = Nx.clip(image_id_raw, 0, n_images - 1)
+
     offset_per_patch = Nx.take(exclusive_cumulative, image_id_per_patch)
     local_index = Nx.subtract(patch_indices, offset_per_patch)
 
     grid_h_per_patch = Nx.take(grid_h, image_id_per_patch)
     grid_w_per_patch = Nx.take(grid_w, image_id_per_patch)
 
+    # Padded images have grid_w == 0; guard the divisions so we don't
+    # divide by zero. The resulting coordinates for padded patches are
+    # arbitrary and are masked out downstream.
+    safe_grid_w = Nx.max(grid_w_per_patch, merge_size)
+
     merge_sq = merge_size * merge_size
-    merged_w_per_patch = Nx.quotient(grid_w_per_patch, merge_size)
+    merged_w_per_patch = Nx.quotient(safe_grid_w, merge_size)
 
     block_idx = Nx.quotient(local_index, merge_sq)
     within = Nx.remainder(local_index, merge_sq)
@@ -346,7 +363,8 @@ defmodule Bumblebee.Vision.Qwen3VLVision do
     row_in_image = block_row |> Nx.multiply(merge_size) |> Nx.add(within_h)
     col_in_image = block_col |> Nx.multiply(merge_size) |> Nx.add(within_w)
 
-    {row_in_image, col_in_image, grid_h_per_patch, grid_w_per_patch, image_id_per_patch}
+    {row_in_image, col_in_image, grid_h_per_patch, grid_w_per_patch, image_id_per_patch,
+     patch_valid}
   end
 
   defp encoder(embeddings, grid_thw, spec, opts) do
@@ -365,7 +383,7 @@ defmodule Bumblebee.Vision.Qwen3VLVision do
         fn embed, grid_thw_t, _opts ->
           {_batch, total_patches, _hidden} = Nx.shape(embed)
 
-          {row_in_image, col_in_image, _, _, _} =
+          {row_in_image, col_in_image, _, _, _, _} =
             patch_metadata(grid_thw_t, total_patches, spec.spatial_merge_size)
 
           compute_2d_rotary_from_positions(
@@ -384,10 +402,10 @@ defmodule Bumblebee.Vision.Qwen3VLVision do
         fn embed, grid_thw_t, _opts ->
           {_batch, total_patches, _hidden} = Nx.shape(embed)
 
-          {_, _, _, _, image_id_per_patch} =
+          {_, _, _, _, image_id_per_patch, patch_valid} =
             patch_metadata(grid_thw_t, total_patches, spec.spatial_merge_size)
 
-          block_diagonal_attention_mask(image_id_per_patch)
+          block_diagonal_attention_mask(image_id_per_patch, patch_valid)
         end,
         [embeddings, grid_thw],
         op_name: :attention_mask
@@ -418,11 +436,13 @@ defmodule Bumblebee.Vision.Qwen3VLVision do
   end
 
   # Returns {total_patches, total_patches} boolean tensor where True means
-  # the two patches share an image (and are therefore allowed to attend).
-  defnp block_diagonal_attention_mask(image_id_per_patch) do
+  # the two patches share an image AND both are valid (not padding).
+  defnp block_diagonal_attention_mask(image_id_per_patch, patch_valid) do
     a = Nx.new_axis(image_id_per_patch, -1)
     b = Nx.new_axis(image_id_per_patch, 0)
-    Nx.equal(a, b)
+    same_image = Nx.equal(a, b)
+    valid_pair = Nx.multiply(Nx.new_axis(patch_valid, -1), Nx.new_axis(patch_valid, 0))
+    Nx.logical_and(same_image, valid_pair)
   end
 
   defp vision_transformer_blocks(

notebooks/qwen3_vl.livemd

@@ -138,10 +138,41 @@ Bumblebee.Multimodal.ImageTextToText.generate(
 #=> %{text: "A group of cats lying on a pink blanket with remote controls.", token_ids: ...}
 ```
 
-> Note: this is a single-call helper — each call recompiles the
-> generation graph if the image size or sequence length changes. The
-> follow-up static-shape padding work lets one compiled graph serve
-> repeated calls with varying image sizes.
+> Note: each `generate/6` call recompiles the generation graph when
+> the image size or sequence length changes. For repeated calls, use
+> `compile/5` + `run/3` (see below).
+
+## Compile Once, Run Many
+
+For serving-style use where many images of varying sizes share one
+compiled graph, configure upper bounds with `compile/5`, then call
+`run/3` repeatedly. The featurizer pads `pixel_values` and
+`image_grid_thw` to the maxima you set, and the vision encoder
+excludes the padded patches from attention.
+
+```elixir
+compiled =
+  Bumblebee.Multimodal.ImageTextToText.compile(
+    model_info,
+    featurizer,
+    tokenizer,
+    generation_config,
+    max_patches: 1024,
+    max_num_images: 1,
+    sequence_length: 384
+  )
+
+# First call: JIT-compiles for these upper-bound shapes
+Bumblebee.Multimodal.ImageTextToText.run(compiled, prompt, image)
+
+# Subsequent calls reuse the same compiled graph, even if the new
+# image produces fewer real patches — padding makes the shapes match.
+Bumblebee.Multimodal.ImageTextToText.run(compiled, prompt, another_image)
+```
+
+On `Qwen3-VL-2B-Instruct` + CPU + a 640×480 COCO image, the warm
+call runs in ~10s while the cold (JIT-compiling) call takes ~27s — a
+2.7x speedup that scales with the number of repeated calls.
 
 ## Multiple Images in One Prompt