feat: add init cache type opts

[polvalente]

This PR adds init_cache types and num_heads for better flexibility of text generation models

[jonatanklosko]

Do we need these new options? We prune these by default, and in order to actually return it in the model output, the user needs to opt-in by configuring global layer options.

[polvalente]

I think we can drop these. I must've missed them in my self-review. My focus was on the cache typing

[jonatanklosko]

Why do we need a separate function? If cross attention is not enabled then get_attention_caches already returns none in the second element (or rather a model that compiles to none).

[polvalente]

bumblebee/lib/bumblebee/layers/decoder.ex

Lines 218 to 224 in ccd6f4a

	@doc """
	Retrieves self-attention and cross-attention caches from a block
	cache.
	"""
	def get_attention_caches(block_cache) do
	{Axon.nx(block_cache, & &1.self_attention), Axon.nx(block_cache, & &1.cross_attention)}
	end

It always returns cross-attention. Do you mean that cross attention is always Axon.None?

Also, being eager means fewer Axon.nx calls, which reduces the overall graph in Nx.Defn.Evaluator

[jonatanklosko]

Do we get anything from passing it downstream instead of casting as above?

[polvalente]

It means we can fit in a smaller memory footprint, if we allocate bf16 instead of f32 and then downcast to bf16

[jonatanklosko]

IIRC the cache value returned from attention layers is cast using :output precision (since it's the layer output). That's why we cast as output here.

I'm not really sure how to model this with mixed precision policy. It may be that we don't want to cast cache at any point, but then we don't have granularity to specify that, since it's a specific input/output.

[polvalente]

I went with using the policy, but originally what I wanted was to introduce a new explicit parameter for the cache type. When I just used the output, at least in my use-case, things ended up using f32 for the cache instead of bf16 like I wanted