refactor: pass type options for init_cache

Author: polvalente

lib/bumblebee/audio/whisper.ex

@@ -227,7 +227,7 @@ defmodule Bumblebee.Audio.Whisper do
   end
 
   @impl true
-  def init_cache(spec, batch_size, max_length, inputs) do
+  def init_cache(spec, batch_size, max_length, inputs, opts \\ []) do
     encoder_sequence_length =
       if encoder_hidden_state = inputs["encoder_hidden_state"] do
         Nx.axis_size(encoder_hidden_state, 1)
@@ -238,7 +238,8 @@ defmodule Bumblebee.Audio.Whisper do
       decoder_num_attention_heads: spec.decoder_num_attention_heads,
       encoder_num_attention_heads: spec.encoder_num_attention_heads,
       decoder_num_blocks: spec.decoder_num_blocks,
-      encoder_sequence_length: encoder_sequence_length
+      encoder_sequence_length: encoder_sequence_length,
+      attention_cache_type: opts[:cache_type]
     )
   end
 

lib/bumblebee/multimodal/blip.ex

@@ -178,7 +178,8 @@ defmodule Bumblebee.Multimodal.Blip do
         %{vision_spec: vision_spec, text_spec: text_spec},
         batch_size,
         max_length,
-        inputs
+        inputs,
+        opts \\ []
       ) do
     num_patches = div(vision_spec.image_size, vision_spec.patch_size) ** 2
     encoder_sequence_length = num_patches + 1
@@ -193,7 +194,7 @@ defmodule Bumblebee.Multimodal.Blip do
       }
       |> Map.reject(&match?({_, nil}, &1))
 
-    text_spec.__struct__.init_cache(text_spec, batch_size, max_length, inputs)
+    text_spec.__struct__.init_cache(text_spec, batch_size, max_length, inputs, opts)
   end
 
   @impl true

lib/bumblebee/text/bart.ex

@@ -417,7 +417,7 @@ defmodule Bumblebee.Text.Bart do
   end
 
   @impl true
-  def init_cache(spec, batch_size, max_length, inputs) do
+  def init_cache(spec, batch_size, max_length, inputs, opts \\ []) do
     encoder_sequence_length =
       if encoder_hidden_state = inputs["encoder_hidden_state"] do
         Nx.axis_size(encoder_hidden_state, 1)
@@ -428,7 +428,8 @@ defmodule Bumblebee.Text.Bart do
       decoder_num_attention_heads: spec.decoder_num_attention_heads,
       encoder_num_attention_heads: spec.encoder_num_attention_heads,
       decoder_num_blocks: spec.decoder_num_blocks,
-      encoder_sequence_length: encoder_sequence_length
+      encoder_sequence_length: encoder_sequence_length,
+      attention_cache_type: opts[:cache_type]
     )
   end
 

lib/bumblebee/text/bert.ex

@@ -374,7 +374,7 @@ defmodule Bumblebee.Text.Bert do
   end
 
   @impl true
-  def init_cache(spec, batch_size, max_length, inputs) do
+  def init_cache(spec, batch_size, max_length, inputs, opts \\ []) do
     encoder_sequence_length =
       if encoder_hidden_state = inputs["encoder_hidden_state"] do
         Nx.axis_size(encoder_hidden_state, 1)
@@ -385,7 +385,8 @@ defmodule Bumblebee.Text.Bert do
       decoder_num_attention_heads: spec.num_attention_heads,
       encoder_num_attention_heads: spec.num_attention_heads,
       decoder_num_blocks: spec.num_blocks,
-      encoder_sequence_length: encoder_sequence_length
+      encoder_sequence_length: encoder_sequence_length,
+      attention_cache_type: opts[:cache_type]
     )
   end
 

lib/bumblebee/text/blenderbot.ex

@@ -269,7 +269,7 @@ defmodule Bumblebee.Text.Blenderbot do
   end
 
   @impl true
-  def init_cache(spec, batch_size, max_length, inputs) do
+  def init_cache(spec, batch_size, max_length, inputs, opts \\ []) do
     encoder_sequence_length =
       if encoder_hidden_state = inputs["encoder_hidden_state"] do
         Nx.axis_size(encoder_hidden_state, 1)
@@ -280,7 +280,8 @@ defmodule Bumblebee.Text.Blenderbot do
       decoder_num_attention_heads: spec.decoder_num_attention_heads,
       encoder_num_attention_heads: spec.encoder_num_attention_heads,
       decoder_num_blocks: spec.decoder_num_blocks,
-      encoder_sequence_length: encoder_sequence_length
+      encoder_sequence_length: encoder_sequence_length,
+      attention_cache_type: opts[:cache_type]
     )
   end
 

lib/bumblebee/text/blip_text.ex

@@ -182,7 +182,7 @@ defmodule Bumblebee.Text.BlipText do
   end
 
   @impl true
-  def init_cache(spec, batch_size, max_length, inputs) do
+  def init_cache(spec, batch_size, max_length, inputs, opts \\ []) do
     encoder_sequence_length =
       if encoder_hidden_state = inputs["encoder_hidden_state"] do
         Nx.axis_size(encoder_hidden_state, 1)
@@ -193,7 +193,8 @@ defmodule Bumblebee.Text.BlipText do
       decoder_num_attention_heads: spec.num_attention_heads,
       encoder_num_attention_heads: spec.num_attention_heads,
       decoder_num_blocks: spec.num_blocks,
-      encoder_sequence_length: encoder_sequence_length
+      encoder_sequence_length: encoder_sequence_length,
+      attention_cache_type: opts[:cache_type]
     )
   end
 

lib/bumblebee/text/gemma.ex

@@ -173,13 +173,14 @@ defmodule Bumblebee.Text.Gemma do
   end
 
   @impl true
-  def init_cache(spec, batch_size, max_length, _inputs) do
+  def init_cache(spec, batch_size, max_length, _inputs, opts \\ []) do
     Layers.Decoder.init_cache(batch_size, max_length,
       hidden_size: spec.hidden_size,
       attention_head_size: spec.attention_head_size,
       decoder_num_attention_heads: spec.num_attention_heads,
       decoder_num_key_value_heads: spec.num_key_value_heads,
-      decoder_num_blocks: spec.num_blocks
+      decoder_num_blocks: spec.num_blocks,
+      attention_cache_type: opts[:cache_type]
     )
   end
 

lib/bumblebee/text/gemma3_text.ex

@@ -209,13 +209,14 @@ defmodule Bumblebee.Text.Gemma3Text do
   end
 
   @impl true
-  def init_cache(spec, batch_size, max_length, _inputs) do
+  def init_cache(spec, batch_size, max_length, _inputs, opts \\ []) do
     Layers.Decoder.init_cache(batch_size, max_length,
       hidden_size: spec.hidden_size,
       attention_head_size: spec.attention_head_size,
       decoder_num_attention_heads: spec.num_attention_heads,
       decoder_num_key_value_heads: spec.num_key_value_heads,
-      decoder_num_blocks: spec.num_blocks
+      decoder_num_blocks: spec.num_blocks,
+      attention_cache_type: opts[:cache_type]
     )
   end
 

lib/bumblebee/text/generation.ex

@@ -12,7 +12,8 @@ defmodule Bumblebee.Text.Generation do
               spec :: Bumblebee.ModelSpec.t(),
               batch_size :: pos_integer(),
               max_length :: pos_integer(),
-              inputs :: map()
+              inputs :: map(),
+              opts :: keyword()
             ) :: cache()
 
   @doc """
@@ -42,9 +43,10 @@ defmodule Bumblebee.Text.Generation do
   @doc """
   Initializes an opaque cache input for iterative inference.
   """
-  @spec init_cache(Bumblebee.ModelSpec.t(), pos_integer(), pos_integer(), map()) :: cache()
-  def init_cache(%module{} = spec, batch_size, max_length, inputs) do
-    module.init_cache(spec, batch_size, max_length, inputs)
+  @spec init_cache(Bumblebee.ModelSpec.t(), pos_integer(), pos_integer(), map(), keyword()) ::
+          cache()
+  def init_cache(%module{} = spec, batch_size, max_length, inputs, opts \\ []) do
+    module.init_cache(spec, batch_size, max_length, inputs, opts)
   end
 
   @doc """
@@ -313,17 +315,13 @@ defmodule Bumblebee.Text.Generation do
       |> Map.put(prefix <> "position_ids", position_ids)
 
     batch_size = Nx.axis_size(input_ids, 0)
-    cache = init_cache(spec, batch_size, max_length, inputs)
 
     output_policy = model_output_policy(model)
-
-    # Cast all float cache tensors to match the model output. This way
-    # we make sure the cache we pass as input has the same types as
-    # the updated cache returned from the model
-    cache =
-      Bumblebee.Utils.Nx.map(cache, fn tensor ->
-        Axon.MixedPrecision.cast(output_policy, tensor, :output)
-      end)
+    # Use the compute precision as the cache type. The key/value tensors are
+    # produced by projection layers running in compute precision, so this
+    # matches what the model will actually return for the cache.
+    cache_type = output_policy.compute || {:f, 32}
+    cache = init_cache(spec, batch_size, max_length, inputs, cache_type: cache_type)
 
     Map.put(inputs, "cache", cache)
   end

lib/bumblebee/text/gpt2.ex

@@ -278,7 +278,7 @@ defmodule Bumblebee.Text.Gpt2 do
   end
 
   @impl true
-  def init_cache(spec, batch_size, max_length, inputs) do
+  def init_cache(spec, batch_size, max_length, inputs, opts \\ []) do
     encoder_sequence_length =
       if encoder_hidden_state = inputs["encoder_hidden_state"] do
         Nx.axis_size(encoder_hidden_state, 1)
@@ -289,7 +289,8 @@ defmodule Bumblebee.Text.Gpt2 do
       decoder_num_attention_heads: spec.num_attention_heads,
       encoder_num_attention_heads: spec.num_attention_heads,
       decoder_num_blocks: spec.num_blocks,
-      encoder_sequence_length: encoder_sequence_length
+      encoder_sequence_length: encoder_sequence_length,
+      attention_cache_type: opts[:cache_type]
     )
   end