diff --git a/AGENTS.md b/AGENTS.md
index acde8562..c73f21a1 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -7,7 +7,7 @@ The main steps of adding a new model are the following:
 2. Look at some existing model implementations in Bumblebee. In case of LLMs, copying an existing LLM implementation is typically a good starting point.
 
 3. Implement the model code.
-   - Whenever possible, reuse existing primitives, most notably `Layers.Transformer.blocks/2`, which is shared for most LLM implementations. Sometimes models introduce novelties to the transformer design, in which case it may be necessary to add a new option to `Layers.Transformer.blocks/2`.
+   - Use `Layers.Transformer.blocks/3` for the block stack. It handles the cache and outputs, and the model writes the block body as plain code, in the same order as the Python model. Use `self_attention/3`, `cross_attention/4`, `basic_ffn/4` and `gated_ffn/4` for standard parts. If the model has a new attention type, write it in the model with the attention steps (`project_heads/4`, `rotary_embedding/6`, `Layers.attention/8`, `output_projection/3`). Do not add model-specific options to the shared code.
    - Include relevant options from Python model configuration as Bumblebee model options (with matching defaults).
    - Make sure the `params_mapping/1` maps to correct Python layer names. You can use `Bumblebee.load_model(..., log_params_diff: true)` to get all logs related to params loading.
 
diff --git a/README.md b/README.md
index 745e3343..5dc0e15e 100644
--- a/README.md
+++ b/README.md
@@ -160,7 +160,7 @@ The main steps of adding a new model are the following:
 2. Look at some existing model implementations in Bumblebee. In case of LLMs, copying an existing LLM implementation is typically a good starting point.
 
 3. Implement the model code.
-   - Whenever possible, reuse existing primitives, most notably `Layers.Transformer.blocks/2`, which is shared for most LLM implementations. Sometimes models introduce novelties to the transformer design, in which case it may be necessary to add a new option to `Layers.Transformer.blocks/2`.
+   - Use `Layers.Transformer.blocks/3` for the block stack. It handles the cache and outputs, and the model writes the block body as plain code, in the same order as the Python model. Use `self_attention/3`, `cross_attention/4`, `basic_ffn/4` and `gated_ffn/4` for standard parts. If the model has a new attention type, write it in the model with the attention steps (`project_heads/4`, `rotary_embedding/6`, `Layers.attention/8`, `output_projection/3`). Do not add model-specific options to the shared code.
    - Include relevant options from Python model configuration as Bumblebee model options (with matching defaults).
    - Make sure the `params_mapping/1` maps to correct Python layer names. You can use `Bumblebee.load_model(..., log_params_diff: true)` to get all logs related to params loading.
 
diff --git a/lib/bumblebee/audio/whisper.ex b/lib/bumblebee/audio/whisper.ex
index 9b78c29d..f6d23013 100644
--- a/lib/bumblebee/audio/whisper.ex
+++ b/lib/bumblebee/audio/whisper.ex
@@ -419,24 +419,47 @@ defmodule Bumblebee.Audio.Whisper do
     name = opts[:name]
 
     outputs =
-      Layers.Transformer.blocks(hidden_state,
-        attention_head_mask: attention_head_mask,
-        num_blocks: spec.encoder_num_blocks,
-        num_attention_heads: spec.encoder_num_attention_heads,
-        hidden_size: spec.hidden_size,
-        kernel_initializer: kernel_initializer(spec),
-        dropout_rate: spec.dropout_rate,
-        attention_dropout_rate: spec.attention_dropout_rate,
-        key_use_bias: false,
-        layer_norm: [
-          epsilon: 1.0e-5
-        ],
-        ffn: [
-          intermediate_size: spec.encoder_intermediate_size,
-          activation: spec.activation
+      Layers.Transformer.blocks(
+        hidden_state,
+        [
+          num_blocks: spec.encoder_num_blocks,
+          attention_head_mask: attention_head_mask,
+          name: join(name, "blocks")
         ],
-        block_type: :norm_first,
-        name: join(name, "blocks")
+        fn hidden_state, block ->
+          name = block.name
+
+          shortcut = hidden_state
+
+          {hidden_state, attention, _self_attention_cache} =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "self_attention_norm"))
+            |> Layers.Transformer.self_attention(
+              block,
+              [name: join(name, "self_attention")] ++
+                attention_opts(spec, spec.encoder_num_attention_heads)
+            )
+
+          hidden_state =
+            hidden_state
+            |> Axon.dropout(rate: spec.dropout_rate)
+            |> Axon.add(shortcut)
+
+          shortcut = hidden_state
+
+          hidden_state =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "output_norm"))
+            |> Layers.Transformer.basic_ffn(spec.encoder_intermediate_size, spec.hidden_size,
+              activation: spec.activation,
+              dropout_rate: spec.dropout_rate,
+              kernel_initializer: kernel_initializer(spec),
+              name: join(name, "ffn")
+            )
+            |> Axon.add(shortcut)
+
+          %{hidden_state: hidden_state, attention: attention}
+        end
       )
 
     hidden_state = Axon.layer_norm(outputs.hidden_state, name: join(name, "norm"))
@@ -461,29 +484,73 @@ defmodule Bumblebee.Audio.Whisper do
     name = opts[:name]
 
     outputs =
-      Layers.Transformer.blocks(hidden_state,
-        attention_mask: attention_mask,
-        attention_head_mask: attention_head_mask,
-        cross_hidden_state: encoder_hidden_state,
-        cross_attention_head_mask: cross_attention_head_mask,
-        cache: cache,
-        causal: true,
-        num_blocks: spec.decoder_num_blocks,
-        num_attention_heads: spec.decoder_num_attention_heads,
-        hidden_size: spec.hidden_size,
-        kernel_initializer: kernel_initializer(spec),
-        dropout_rate: spec.dropout_rate,
-        attention_dropout_rate: spec.attention_dropout_rate,
-        key_use_bias: false,
-        layer_norm: [
-          epsilon: 1.0e-5
+      Layers.Transformer.blocks(
+        hidden_state,
+        [
+          num_blocks: spec.decoder_num_blocks,
+          attention_mask: attention_mask,
+          attention_head_mask: attention_head_mask,
+          cross_attention_head_mask: cross_attention_head_mask,
+          cache: cache,
+          name: join(name, "blocks")
         ],
-        ffn: [
-          intermediate_size: spec.decoder_intermediate_size,
-          activation: spec.activation
-        ],
-        block_type: :norm_first,
-        name: join(name, "blocks")
+        fn hidden_state, block ->
+          name = block.name
+
+          attention_opts = attention_opts(spec, spec.decoder_num_attention_heads)
+
+          shortcut = hidden_state
+
+          {hidden_state, attention, self_attention_cache} =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "self_attention_norm"))
+            |> Layers.Transformer.self_attention(
+              block,
+              [causal: true, name: join(name, "self_attention")] ++ attention_opts
+            )
+
+          hidden_state =
+            hidden_state
+            |> Axon.dropout(rate: spec.dropout_rate)
+            |> Axon.add(shortcut)
+
+          shortcut = hidden_state
+
+          {hidden_state, cross_attention, cross_attention_cache} =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "cross_attention_norm"))
+            |> Layers.Transformer.cross_attention(
+              encoder_hidden_state,
+              block,
+              [name: join(name, "cross_attention")] ++ attention_opts
+            )
+
+          hidden_state =
+            hidden_state
+            |> Axon.dropout(rate: spec.dropout_rate)
+            |> Axon.add(shortcut)
+
+          shortcut = hidden_state
+
+          hidden_state =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "output_norm"))
+            |> Layers.Transformer.basic_ffn(spec.decoder_intermediate_size, spec.hidden_size,
+              activation: spec.activation,
+              dropout_rate: spec.dropout_rate,
+              kernel_initializer: kernel_initializer(spec),
+              name: join(name, "ffn")
+            )
+            |> Axon.add(shortcut)
+
+          %{
+            hidden_state: hidden_state,
+            attention: attention,
+            cross_attention: cross_attention,
+            self_attention_cache: self_attention_cache,
+            cross_attention_cache: cross_attention_cache
+          }
+        end
       )
 
     hidden_state = Axon.layer_norm(outputs.hidden_state, name: join(name, "norm"))
@@ -495,6 +562,16 @@ defmodule Bumblebee.Audio.Whisper do
     }
   end
 
+  defp attention_opts(spec, num_heads) do
+    [
+      num_heads: num_heads,
+      hidden_size: spec.hidden_size,
+      key_use_bias: false,
+      dropout_rate: spec.attention_dropout_rate,
+      kernel_initializer: kernel_initializer(spec)
+    ]
+  end
+
   defp kernel_initializer(spec) do
     Axon.Initializers.normal(scale: spec.initializer_scale)
   end
diff --git a/lib/bumblebee/diffusion/layers/unet.ex b/lib/bumblebee/diffusion/layers/unet.ex
index 50837a23..fce134dc 100644
--- a/lib/bumblebee/diffusion/layers/unet.ex
+++ b/lib/bumblebee/diffusion/layers/unet.ex
@@ -314,20 +314,59 @@ defmodule Bumblebee.Diffusion.Layers.UNet do
       end
     end)
     |> Layers.Transformer.blocks(
-      cross_hidden_state: cross_hidden_state,
-      num_blocks: depth,
-      num_attention_heads: num_heads,
-      hidden_size: hidden_size,
-      query_use_bias: false,
-      key_use_bias: false,
-      value_use_bias: false,
-      layer_norm: [
-        epsilon: 1.0e-5
-      ],
-      dropout_rate: dropout,
-      ffn: &ffn_geglu(&1, 4 * hidden_size, hidden_size, dropout: dropout, name: &2),
-      block_type: :norm_first,
-      name: join(name, "blocks")
+      [num_blocks: depth, name: join(name, "blocks")],
+      fn hidden_state, block ->
+        name = block.name
+
+        attention_opts = [
+          num_heads: num_heads,
+          hidden_size: hidden_size,
+          query_use_bias: false,
+          key_use_bias: false,
+          value_use_bias: false
+        ]
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, _cache} =
+          hidden_state
+          |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "self_attention_norm"))
+          |> Layers.Transformer.self_attention(
+            block,
+            [name: join(name, "self_attention")] ++ attention_opts
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: dropout)
+          |> Axon.add(shortcut)
+
+        shortcut = hidden_state
+
+        {hidden_state, cross_attention, _cache} =
+          hidden_state
+          |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "cross_attention_norm"))
+          |> Layers.Transformer.cross_attention(
+            cross_hidden_state,
+            block,
+            [name: join(name, "cross_attention")] ++ attention_opts
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: dropout)
+          |> Axon.add(shortcut)
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "output_norm"))
+          |> ffn_geglu(4 * hidden_size, hidden_size, dropout: dropout, name: join(name, "ffn"))
+          |> Axon.add(shortcut)
+
+        %{hidden_state: hidden_state, attention: attention, cross_attention: cross_attention}
+      end
     )
     |> then(fn %{hidden_state: hidden_state} ->
       if use_linear_projection do
diff --git a/lib/bumblebee/layers/transformer.ex b/lib/bumblebee/layers/transformer.ex
index 62c0bd46..6563b59b 100644
--- a/lib/bumblebee/layers/transformer.ex
+++ b/lib/bumblebee/layers/transformer.ex
@@ -6,108 +6,67 @@ defmodule Bumblebee.Layers.Transformer do
   alias Bumblebee.Layers
 
   @doc """
-  Adds a list of transformer block to the network.
+  Adds a stack of transformer blocks to the network.
 
-  Adds transformer blocks (see `block/2`), while also accumulating
-  caches and outputs. This is a convenience function, since this
-  chunk is oftentimes the same across transformer models.
+  This function handles the parts shared by most transformer models:
+  the attention cache, per-block head masks and collecting outputs.
+  The block itself is defined by `fun`.
+
+  `fun` receives the hidden state and a `block` map with the following
+  keys:
+
+    * `:index` - the zero-based block index
+
+    * `:name` - the prefix for block layer names
+
+    * `:attention_mask` - the attention mask, combined with the cache
+      mask during iterative decoding
+
+    * `:attention_head_mask` - the head mask for this block
+
+    * `:cross_attention_head_mask` - the cross-attention head mask for
+      this block
+
+    * `:offset` - offset in the input sequence during iterative decoding
+
+    * `:self_attention_cache`, `:cross_attention_cache` - attention
+      caches for this block
+
+  `fun` must return a map with `:hidden_state`. It may also include
+  `:attention`, `:cross_attention`, `:self_attention_cache` and
+  `:cross_attention_cache`.
 
   ## Options
 
-    * `:num_blocks` (required) - the number of consecutive transformer
-      blocks to add
+    * `:num_blocks` (required) - the number of blocks
 
-    * `:block_output` - an optional 2-arity function receiving each block's
-      hidden state and zero-based index. The returned hidden state is
-      collected in the outputs and passed to the next block
+    * `:attention_mask` - a mask indicating which positions to attend to
 
-    * `:share_attention_relative_bias` - when attention relative bias
-      is configured, this option controls whether the bias from the
-      first block is used for all other blocks. Defaults to `false`
+    * `:attention_head_mask` - a mask to nullify selected attention heads,
+      with a leading axis for blocks
 
-    * `:rotary_embedding` - configuration of rotary embedding. Can be:
-      - a keyword list (applied to all blocks)
-      - a function that takes the block index and returns the configuration
+    * `:cross_attention_head_mask` - same as `:attention_head_mask`, but
+      for cross-attention
 
-    * `:attention_window_size` - sliding window attention configuration. Can be:
-      - `nil` for global attention (default)
-      - a `{left, right}` tuple (applied to all blocks)
-      - a function that takes the block index and returns `nil` or `{left, right}`.
-        This enables per-layer attention patterns like Gemma 3's alternating
-        local/global attention (5 local layers followed by 1 global layer)
+    * `:cache` - cache for iterative decoding
 
     * `:name` - the prefix for layer names
 
-  For all other options (including required options) see `block/2`.
-  Note that `:attention_head_mask`, `:cross_attention_head_mask` should
-  have an additional leading axis corresponding to the number of blocks,
-  since each block gets its own mask. Similarly `:cache` is a tuple of
-  block caches.
   """
-  def blocks(hidden_state, opts) do
-    validate_required_keys!(opts, [:num_blocks, :num_attention_heads, :hidden_size, :ffn])
-
-    # Note: :attention_window_size is NOT in block_opts_keys because it's handled
-    # specially (supports per-layer function) and passed explicitly to block/2
-    block_opts_keys = [
-      :num_attention_heads,
-      :num_key_value_heads,
-      :causal,
-      :hidden_size,
-      :ffn,
-      :kernel_initializer,
-      :attention_head_size,
-      :dropout_rate,
-      :attention_dropout_rate,
-      :query_use_bias,
-      :key_use_bias,
-      :value_use_bias,
-      :output_use_bias,
-      :layer_norm,
-      :block_type,
-      :attention_scale,
-      :query_norm,
-      :key_norm,
-      :attention
-    ]
-
+  def blocks(hidden_state, opts, fun) do
     opts =
-      Keyword.validate!(
-        opts,
-        block_opts_keys ++
-          [
-            :name,
-            :num_blocks,
-            :rotary_embedding,
-            :attention_window_size,
-            block_output: fn hidden_state, _index -> hidden_state end,
-            attention_mask: Layers.none(),
-            attention_head_mask: Layers.none(),
-            attention_relative_bias: nil,
-            share_attention_relative_bias: false,
-            cross_hidden_state: nil,
-            cross_attention_mask: Layers.none(),
-            cross_attention_head_mask: Layers.none(),
-            cache: Layers.none()
-          ]
-      )
-
-    name = opts[:name]
-    num_blocks = opts[:num_blocks]
-
-    attention_mask = opts[:attention_mask]
-    attention_head_mask = opts[:attention_head_mask]
-    cross_hidden_state = opts[:cross_hidden_state]
-    cross_attention_mask = opts[:cross_attention_mask]
-    cross_attention_head_mask = opts[:cross_attention_head_mask]
-    cache = opts[:cache]
-    rotary_embedding = opts[:rotary_embedding]
-    attention_window_size = opts[:attention_window_size]
-    block_output = opts[:block_output]
+      Keyword.validate!(opts, [
+        :name,
+        :num_blocks,
+        attention_mask: Layers.none(),
+        attention_head_mask: Layers.none(),
+        cross_attention_head_mask: Layers.none(),
+        cache: Layers.none()
+      ])
 
-    block_opts = Keyword.take(opts, block_opts_keys)
+    {attention_mask, cache} =
+      Layers.Decoder.cached_attention_mask(opts[:attention_mask], opts[:cache])
 
-    {attention_mask, cache} = Layers.Decoder.cached_attention_mask(attention_mask, cache)
     offset = Layers.Decoder.get_cache_offset(cache)
 
     state = %{
@@ -115,69 +74,52 @@ defmodule Bumblebee.Layers.Transformer do
       hidden_states: Axon.container({hidden_state}),
       attentions: Axon.container({}),
       cross_attentions: Axon.container({}),
-      cache: cache,
-      attention_relative_bias: Layers.none()
+      cache: cache
     }
 
     outputs =
-      for idx <- 0..(num_blocks - 1), reduce: state do
+      for index <- 0..(opts[:num_blocks] - 1), reduce: state do
         state ->
-          block_attention_head_mask = Axon.nx(attention_head_mask, & &1[idx])
-          block_cross_attention_head_mask = Axon.nx(cross_attention_head_mask, & &1[idx])
-          block_cache = Layers.Decoder.get_block_cache(state.cache, idx)
-
-          attention_relative_bias =
-            if opts[:share_attention_relative_bias] and idx > 0 do
-              state.attention_relative_bias
-            else
-              opts[:attention_relative_bias] || Layers.none()
-            end
-
-          block_rotary_embedding =
-            case rotary_embedding do
-              nil -> nil
-              fun when is_function(fun, 1) -> fun.(idx)
-              config when is_list(config) -> config
-            end
-
-          # Support per-layer attention window size for models like Gemma 3
-          # that alternate between local (sliding window) and global attention
-          block_attention_window_size =
-            case attention_window_size do
-              nil -> nil
-              fun when is_function(fun, 1) -> fun.(idx)
-              size -> size
-            end
-
-          {hidden_state, attention, cross_attention, block_cache, attention_relative_bias} =
-            block(
-              state.hidden_state,
-              [
-                attention_mask: attention_mask,
-                attention_head_mask: block_attention_head_mask,
-                attention_relative_bias: attention_relative_bias,
-                cross_hidden_state: cross_hidden_state,
-                cross_attention_mask: cross_attention_mask,
-                cross_attention_head_mask: block_cross_attention_head_mask,
-                block_cache: block_cache,
-                offset: offset,
-                rotary_embedding: block_rotary_embedding,
-                attention_window_size: block_attention_window_size,
-                name: join(name, idx)
-              ] ++ block_opts
-            )
+          block_cache = Layers.Decoder.get_block_cache(state.cache, index)
+
+          {self_attention_cache, cross_attention_cache} =
+            Layers.Decoder.get_attention_caches(block_cache)
+
+          block = %{
+            index: index,
+            name: join(opts[:name], index),
+            attention_mask: attention_mask,
+            attention_head_mask: Axon.nx(opts[:attention_head_mask], & &1[index]),
+            cross_attention_head_mask: Axon.nx(opts[:cross_attention_head_mask], & &1[index]),
+            offset: offset,
+            self_attention_cache: self_attention_cache,
+            cross_attention_cache: cross_attention_cache
+          }
 
-          cache = Layers.Decoder.put_block_cache(state.cache, idx, block_cache)
+          result =
+            Map.merge(
+              %{
+                attention: Layers.none(),
+                cross_attention: Layers.none(),
+                self_attention_cache: self_attention_cache,
+                cross_attention_cache: cross_attention_cache
+              },
+              fun.(state.hidden_state, block)
+            )
 
-          hidden_state = block_output.(hidden_state, idx)
+          block_cache =
+            Layers.Decoder.put_attention_caches(
+              block_cache,
+              result.self_attention_cache,
+              result.cross_attention_cache
+            )
 
           %{
-            hidden_state: hidden_state,
-            hidden_states: Layers.append(state.hidden_states, hidden_state),
-            attentions: Layers.append(state.attentions, attention),
-            cross_attentions: Layers.append(state.cross_attentions, cross_attention),
-            attention_relative_bias: attention_relative_bias,
-            cache: cache
+            hidden_state: result.hidden_state,
+            hidden_states: Layers.append(state.hidden_states, result.hidden_state),
+            attentions: Layers.append(state.attentions, result.attention),
+            cross_attentions: Layers.append(state.cross_attentions, result.cross_attention),
+            cache: Layers.Decoder.put_block_cache(state.cache, index, block_cache)
           }
       end
 
@@ -185,489 +127,137 @@ defmodule Bumblebee.Layers.Transformer do
   end
 
   @doc """
-  Adds a transformer block to the network
-
-  Depending on the configuration, this can implement both the encoder
-  and the decoder block.
-
-  A transformer block consists of self-attention, shortcut connection,
-  normalization and dropout layers. It may also include a cross-attention
-  block.
-
-  ## Options
-
-    * `:num_attention_heads` (required) - the number of attention heads
-
-    * `:hidden_size` (required) - the dimensionality of the attention
-      block and the last layer in the output feed-forward network
-
-    * `:ffn` (required) - configuration of the feed-forward network at
-      the end of the transformer block. By default the network has two
-      dense layers with an activation in-between and is configured with
-      the following options:
-
-      * `:intermediate_size` (required) - the dimensionality of the
-        first dense layer
-
-      * `:activation` - the activation used in-between the layers.
-        Defaults to `:gelu`
+  Adds a self-attention layer to a block in `blocks/3`.
 
-      Alternatively a custom 2-arity function may be given. The function
-      should add FFN nodes to the given Axon node. The function also
-      receives layer name prefix as the second argument.
+  Takes the mask, cache and offset from `block`. See
+  `multi_head_attention/4` for options.
 
-    * `:attention_mask` - a mask indicating which positions to attend to
-
-    * `:attention_head_mask` - a mask to nullify selected attention heads
-
-    * `:attention_relative_bias` - configuration of relative bias. If set,
-      will apply relative attention bias with the given options. Valid
-      options are:
-
-        * `:num_buckets` (required) - number of relative attention buckets
+  Returns `{output, attention_weights, attention_cache}`.
+  """
+  def self_attention(hidden_state, block, opts) do
+    {output, weights, cache, _attention_relative_bias} =
+      multi_head_attention(
+        hidden_state,
+        hidden_state,
+        hidden_state,
+        [
+          attention_mask: block.attention_mask,
+          attention_head_mask: block.attention_head_mask,
+          attention_cache: block.self_attention_cache,
+          offset: block.offset
+        ] ++ opts
+      )
 
-        * `:max_distance` (required) - maximum distance of the relative attention
-          bias
+    {output, weights, cache}
+  end
 
-        * `:bidirectional` (required) - whether to apply the relative attention
-          bias bidirectionally
+  @doc """
+  Adds a cross-attention layer to a block in `blocks/3`.
 
-      Alternatively an `Axon` node may be given with the computed bias.
+  Takes the head mask, cache and offset from `block`. The cross
+  attention mask should be given as `:attention_mask`. See
+  `multi_head_attention/4` for other options.
 
-    * `:cross_hidden_state` - the second input for the the cross-attention
-      block. The cross-attention block is added only when this is specified.
-      Defaults to `nil`
+  Returns `{output, attention_weights, attention_cache}`.
+  """
+  def cross_attention(hidden_state, cross_hidden_state, block, opts) do
+    {output, weights, cache, _attention_relative_bias} =
+      multi_head_attention(
+        hidden_state,
+        cross_hidden_state,
+        cross_hidden_state,
+        [
+          attention_head_mask: block.cross_attention_head_mask,
+          attention_cache: block.cross_attention_cache,
+          offset: block.offset
+        ] ++ opts
+      )
 
-    * `:cross_attention_mask` - a mask indicating which positions to
-      attend to in the cross-attention block
+    {output, weights, cache}
+  end
 
-    * `:cross_attention_head_mask` - a mask to nullify selected attention
-      heads in the cross-attention block
+  @doc """
+  Adds a feed-forward network with two dense layers and an activation
+  in-between.
 
-    * `:block_cache` - cache for all attention blocks in this block,
-      used in iterative decoding
+  ## Options
 
-    * `:offset` - offset in the input sequence during iterative decoding
+    * `:activation` - the activation. Defaults to `:gelu`
 
-    * `:causal` - whether the self-attention block should be causal.
-      Defaults to `false`
+    * `:dropout_rate` - the dropout rate at the end. Defaults to `0.0`
 
     * `:kernel_initializer` - initializer for kernel weights. Defaults
       to `:glorot_uniform`
 
-    * `:attention_head_size` - the projection size for key, value,
-      and query states per-head. Defaults to `div(hidden_size, num_attention_heads)`
-
-    * `:dropout_rate` - the dropout rate for dropout layers. Defaults
-      to `0.0`
-
-    * `:attention_dropout_rate` - the dropout rate for attention weights
-      dropout. Defaults to `0.0`
-
-    * `:query_use_bias` - whether to use bias in the query projection.
-      Defaults to `true`
-
-    * `:key_use_bias` - whether to use bias in the key projection.
-      Defaults to `true`
-
-    * `:value_use_bias` - whether to use bias in the value projection.
-      Defaults to `true`
-
-    * `:output_use_bias` - whether to use bias in the output projection.
-      Defaults to `true`
-
-    * `:layer_norm` - configuration of the layer norm operation
-      at the end of the transformer block. By default the network uses regular
-      layer normalization configured with the following options:
-
-        * `:epsilon` - the epsilon used by the layer normalization
-          layers. Defaults to `1.0e-5`
-
-      Alternatively a custom 2-arity function may be given. The function
-      should add a normalization node to the given Axon node. The function
-      also receives the layer name prefix as the second argument.
-
-    * `:block_type` - controls which configuration of the block to use,
-      one of:
-
-        * `:standard` (default) - the original transformer block
-
-        * `:norm_first` - same as `:standard`, but with normalization layers
-          placed before each group of layers, rather than after
-
-        * `:parallel` - block with attention and FFN independently (in parallel).
-          This type doesn't support cross-attention
-
-      Alternatively a custom 3-arity function may be given. The function
-      receives the input hidden state, a map with block steps and a
-      name to prefix any additional layers.
-
-    * `:attention_window_size` - when set, enables sliding window attention.
-      Should be a `{left, right}` tuple with window size on each side
-
-    * `:attention_scale` - the scaling factor applied to the attention weights.
-      Defaults to $\frac{1}{\sqrt{d}}$.
-
-    * `:attention` - a custom 4-arity attention function, defaulting to
-      `multi_head_attention/4`. It receives query, key and value nodes,
-      followed by the attention options, and must return the same tuple
-      as `multi_head_attention/4`. It applies to self- and cross-attention
-
-    * `:rotary_embedding` - configuration of rotary embedding. If set,
-      will apply rotary position embedding with the given options. Valid
-      options are:
-
-        * `:position_ids` (required) - input position ids used for the
-          embedding
-
-        * `:max_positions` - the maximum number of distinct positions
-
-        * `:base` - base for computing rotary embedding frequency. Defaults
-        to `10_000`.
-
-        * `:percentage` - percentage of hidden dimensions to allocate to rotary embeddings.
-        Defaults to `1.0`.
-
     * `:name` - the prefix for layer names
 
-  ## References
-
-    * [Attention Is All You Need](https://arxiv.org/abs/1706.03762), Figure 1
-
   """
-  def block(hidden_state, opts) do
-    validate_required_keys!(opts, [:num_attention_heads, :hidden_size, :ffn])
-
+  def basic_ffn(x, intermediate_size, output_size, opts) do
     opts =
       Keyword.validate!(opts, [
         :name,
-        :num_attention_heads,
-        :hidden_size,
-        :ffn,
-        :num_key_value_heads,
-        attention_mask: Layers.none(),
-        attention_head_mask: Layers.none(),
-        attention_relative_bias: Layers.none(),
-        cross_hidden_state: nil,
-        cross_attention_mask: Layers.none(),
-        cross_attention_head_mask: Layers.none(),
-        block_cache: Layers.none(),
-        offset: Layers.none(),
-        causal: false,
-        kernel_initializer: :glorot_uniform,
-        attention_head_size: nil,
+        activation: :gelu,
         dropout_rate: 0.0,
-        attention_dropout_rate: 0.0,
-        query_use_bias: true,
-        key_use_bias: true,
-        value_use_bias: true,
-        output_use_bias: true,
-        block_type: :standard,
-        layer_norm: [],
-        attention_window_size: nil,
-        attention_scale: nil,
-        rotary_embedding: nil,
-        query_norm: nil,
-        key_norm: nil,
-        attention: &multi_head_attention/4
+        kernel_initializer: :glorot_uniform
       ])
 
     name = opts[:name]
-    num_attention_heads = opts[:num_attention_heads]
-    num_key_value_heads = opts[:num_key_value_heads] || num_attention_heads
-    hidden_size = opts[:hidden_size]
-    ffn = opts[:ffn]
-    causal = opts[:causal]
-    kernel_initializer = opts[:kernel_initializer]
-    attention_head_size = opts[:attention_head_size]
-    dropout_rate = opts[:dropout_rate]
-    attention_dropout_rate = opts[:attention_dropout_rate]
-    query_use_bias = opts[:query_use_bias]
-    key_use_bias = opts[:key_use_bias]
-    value_use_bias = opts[:value_use_bias]
-    output_use_bias = opts[:output_use_bias]
-    attention_mask = opts[:attention_mask]
-    attention_head_mask = opts[:attention_head_mask]
-    attention_relative_bias = opts[:attention_relative_bias]
-    cross_hidden_state = opts[:cross_hidden_state]
-    cross_attention_mask = opts[:cross_attention_mask]
-    cross_attention_head_mask = opts[:cross_attention_head_mask]
-    block_cache = opts[:block_cache]
-    offset = opts[:offset]
-    layer_norm = opts[:layer_norm]
-    block_type = opts[:block_type]
-    attention_window_size = opts[:attention_window_size]
-    attention_scale = opts[:attention_scale]
-    rotary_embedding = opts[:rotary_embedding]
-    query_norm = opts[:query_norm]
-    key_norm = opts[:key_norm]
-
-    attention_fun = opts[:attention]
 
-    ffn_fun =
-      case ffn do
-        opts when is_list(opts) ->
-          validate_required_keys!(opts, [:intermediate_size])
-          opts = Keyword.validate!(opts, [:intermediate_size, activation: :gelu])
-
-          &basic_ffn(&1, opts[:intermediate_size], hidden_size,
-            activation: opts[:activation],
-            kernel_initializer: kernel_initializer,
-            dropout_rate: dropout_rate,
-            name: &2
-          )
-
-        fun when is_function(fun) ->
-          fun
-      end
-
-    layer_norm_fun =
-      case layer_norm do
-        opts when is_list(opts) ->
-          opts = Keyword.validate!(opts, epsilon: 1.0e-5)
-
-          &Axon.layer_norm(&1, epsilon: opts[:epsilon], name: &2)
-
-        fun when is_function(fun) ->
-          fun
-      end
-
-    {self_attention_cache, cross_attention_cache} =
-      Layers.Decoder.get_attention_caches(block_cache)
-
-    # Self-attention, shortcut connection, normalization and dropout
-
-    self_attention_norm = &layer_norm_fun.(&1, join(name, "self_attention_norm"))
-
-    self_attention = fn hidden_state ->
-      {hidden_state, attention, self_attention_cache, attention_relative_bias} =
-        attention_fun.(hidden_state, hidden_state, hidden_state,
-          attention_mask: attention_mask,
-          attention_head_mask: attention_head_mask,
-          attention_relative_bias: attention_relative_bias,
-          attention_cache: self_attention_cache,
-          offset: offset,
-          causal: causal,
-          num_heads: num_attention_heads,
-          num_key_value_heads: num_key_value_heads,
-          hidden_size: hidden_size,
-          kernel_initializer: kernel_initializer,
-          attention_head_size: attention_head_size,
-          dropout_rate: attention_dropout_rate,
-          query_use_bias: query_use_bias,
-          key_use_bias: key_use_bias,
-          value_use_bias: value_use_bias,
-          output_use_bias: output_use_bias,
-          attention_window_size: attention_window_size,
-          attention_scale: attention_scale,
-          rotary_embedding: rotary_embedding,
-          query_norm: query_norm,
-          key_norm: key_norm,
-          name: join(name, "self_attention")
-        )
-
-      hidden_state =
-        Axon.dropout(hidden_state, rate: dropout_rate, name: join(name, "self_attention_dropout"))
-
-      {hidden_state, {attention, self_attention_cache, attention_relative_bias}}
-    end
-
-    # Cross-attention, shortcut connection, normalization and dropout
-
-    cross_attention_maybe = fn hidden_state, fun ->
-      if cross_hidden_state do
-        Layers.if_present cross_hidden_state do
-          fun.(hidden_state)
-        else
-          {hidden_state, {Layers.none(), cross_attention_cache}}
-        end
-      else
-        {hidden_state, {Layers.none(), cross_attention_cache}}
-      end
-    end
-
-    cross_attention_norm = &layer_norm_fun.(&1, join(name, "cross_attention_norm"))
-
-    cross_attention = fn hidden_state ->
-      {hidden_state, cross_attention, cross_attention_cache, _cross_attention_relative_bias} =
-        attention_fun.(hidden_state, cross_hidden_state, cross_hidden_state,
-          attention_mask: cross_attention_mask,
-          attention_head_mask: cross_attention_head_mask,
-          attention_cache: cross_attention_cache,
-          offset: offset,
-          num_heads: num_attention_heads,
-          num_key_value_heads: num_key_value_heads,
-          hidden_size: hidden_size,
-          kernel_initializer: kernel_initializer,
-          attention_head_size: attention_head_size,
-          dropout_rate: attention_dropout_rate,
-          query_use_bias: query_use_bias,
-          key_use_bias: key_use_bias,
-          value_use_bias: value_use_bias,
-          output_use_bias: output_use_bias,
-          attention_window_size: attention_window_size,
-          attention_scale: attention_scale,
-          rotary_embedding: rotary_embedding,
-          name: join(name, "cross_attention")
-        )
-
-      hidden_state =
-        Axon.dropout(
-          hidden_state,
-          rate: dropout_rate,
-          name: join(name, "cross_attention_dropout")
-        )
-
-      {hidden_state, {cross_attention, cross_attention_cache}}
-    end
-
-    # Output feed-forward network, shortcut connection, normalization and dropout
-
-    output_norm = &layer_norm_fun.(&1, join(name, "output_norm"))
-
-    ffn = &ffn_fun.(&1, join(name, "ffn"))
-
-    block_impl =
-      case block_type do
-        type when is_atom(type) -> &block_impl(type, &1, &2, &3)
-        fun when is_function(fun) -> fun
-      end
-
-    {hidden_state, attention_info, cross_attention_info} =
-      block_impl.(
-        hidden_state,
-        %{
-          self_attention_norm: self_attention_norm,
-          self_attention: self_attention,
-          cross_attention_maybe: cross_attention_maybe,
-          cross_attention_norm: cross_attention_norm,
-          cross_attention: cross_attention,
-          output_norm: output_norm,
-          ffn: ffn
-        },
-        name
-      )
-
-    {attention, self_attention_cache, attention_relative_bias} = attention_info
-    {cross_attention, cross_attention_cache} = cross_attention_info
-
-    block_cache =
-      Layers.Decoder.put_attention_caches(
-        block_cache,
-        self_attention_cache,
-        cross_attention_cache
-      )
-
-    {hidden_state, attention, cross_attention, block_cache, attention_relative_bias}
-  end
-
-  defp block_impl(:standard, hidden_state, steps, _name) do
-    shortcut = hidden_state
-
-    {hidden_state, attention_info} = steps.self_attention.(hidden_state)
-
-    hidden_state =
-      hidden_state
-      |> Axon.add(shortcut)
-      |> steps.self_attention_norm.()
-
-    {hidden_state, cross_attention_info} =
-      steps.cross_attention_maybe.(hidden_state, fn hidden_state ->
-        shortcut = hidden_state
-
-        {hidden_state, cross_attention_info} = steps.cross_attention.(hidden_state)
-
-        hidden_state =
-          hidden_state
-          |> Axon.add(shortcut)
-          |> steps.cross_attention_norm.()
-
-        {hidden_state, cross_attention_info}
-      end)
-
-    shortcut = hidden_state
-
-    hidden_state =
-      hidden_state
-      |> steps.ffn.()
-      |> Axon.add(shortcut)
-      |> steps.output_norm.()
-
-    {hidden_state, attention_info, cross_attention_info}
+    x
+    |> Axon.dense(intermediate_size,
+      kernel_initializer: opts[:kernel_initializer],
+      name: join(name, "intermediate")
+    )
+    |> Layers.activation(opts[:activation])
+    |> Axon.dense(output_size,
+      kernel_initializer: opts[:kernel_initializer],
+      name: join(name, "output")
+    )
+    |> Axon.dropout(rate: opts[:dropout_rate])
   end
 
-  defp block_impl(:norm_first, hidden_state, steps, _name) do
-    shortcut = hidden_state
-
-    {hidden_state, attention_info} =
-      hidden_state
-      |> steps.self_attention_norm.()
-      |> steps.self_attention.()
-
-    hidden_state = Axon.add(hidden_state, shortcut)
-
-    {hidden_state, cross_attention_info} =
-      steps.cross_attention_maybe.(hidden_state, fn hidden_state ->
-        shortcut = hidden_state
-
-        {hidden_state, cross_attention_info} =
-          hidden_state
-          |> steps.cross_attention_norm.()
-          |> steps.cross_attention.()
-
-        hidden_state = Axon.add(hidden_state, shortcut)
-
-        {hidden_state, cross_attention_info}
-      end)
-
-    shortcut = hidden_state
+  @doc """
+  Adds a gated feed-forward network, as used in most recent LLMs.
 
-    hidden_state =
-      hidden_state
-      |> steps.output_norm.()
-      |> steps.ffn.()
-      |> Axon.add(shortcut)
+  ## Options
 
-    {hidden_state, attention_info, cross_attention_info}
-  end
+    * `:activation` - the activation applied to the gate. Defaults
+      to `:silu`
 
-  defp block_impl(:parallel, hidden_state, steps, _name) do
-    shortcut = hidden_state
+    * `:use_bias` - whether to use bias in the dense layers. Defaults
+      to `false`
 
-    {attention_hidden_state, attention_info} =
-      hidden_state
-      |> steps.self_attention_norm.()
-      |> steps.self_attention.()
+    * `:kernel_initializer` - initializer for kernel weights. Defaults
+      to `:glorot_uniform`
 
-    {_hidden_state, cross_attention_info} =
-      steps.cross_attention_maybe.(hidden_state, fn _hidden_state ->
-        raise "cross attention not supported"
-      end)
+    * `:name` - the prefix for layer names
 
-    ffn_hidden_state =
-      hidden_state
-      |> steps.output_norm.()
-      |> steps.ffn.()
+  """
+  def gated_ffn(x, intermediate_size, output_size, opts) do
+    opts =
+      Keyword.validate!(opts, [
+        :name,
+        activation: :silu,
+        use_bias: false,
+        kernel_initializer: :glorot_uniform
+      ])
 
-    hidden_state = Axon.add([shortcut, attention_hidden_state, ffn_hidden_state])
+    name = opts[:name]
 
-    {hidden_state, attention_info, cross_attention_info}
-  end
+    dense = fn x, size, name ->
+      Axon.dense(x, size,
+        use_bias: opts[:use_bias],
+        kernel_initializer: opts[:kernel_initializer],
+        name: name
+      )
+    end
 
-  defp basic_ffn(x, intermediate_size, output_size, opts) do
-    name = opts[:name]
+    intermediate = dense.(x, intermediate_size, join(name, "intermediate"))
+    gate = dense.(x, intermediate_size, join(name, "gate"))
 
-    x
-    |> Axon.dense(intermediate_size,
-      kernel_initializer: opts[:kernel_initializer],
-      name: join(name, "intermediate")
-    )
-    |> Layers.activation(opts[:activation])
-    |> Axon.dense(output_size,
-      kernel_initializer: opts[:kernel_initializer],
-      name: join(name, "output")
-    )
-    |> Axon.dropout(rate: opts[:dropout_rate])
+    intermediate
+    |> Axon.multiply(Layers.activation(gate, opts[:activation]))
+    |> dense.(output_size, join(name, "output"))
   end
 
   @doc """
@@ -818,102 +408,34 @@ defmodule Bumblebee.Layers.Transformer do
     attention_relative_bias = opts[:attention_relative_bias]
 
     attention_head_size = opts[:attention_head_size] || div(hidden_size, num_heads)
-    inner_size = num_heads * attention_head_size
-    inner_kv_size = num_key_value_heads * attention_head_size
-
-    query =
-      query
-      |> Axon.dense(inner_size,
-        kernel_initializer: kernel_initializer,
-        name: join(name, "query"),
-        use_bias: query_use_bias
-      )
-      |> Layers.split_heads(num_heads)
 
-    key =
-      key
-      |> Axon.dense(inner_kv_size,
+    project = fn input, num_heads, use_bias, suffix ->
+      project_heads(input, num_heads, attention_head_size,
+        use_bias: use_bias,
         kernel_initializer: kernel_initializer,
-        name: join(name, "key"),
-        use_bias: key_use_bias
+        name: join(name, suffix)
       )
-      |> Layers.split_heads(num_key_value_heads)
-
-    value =
-      value
-      |> Axon.dense(inner_kv_size,
-        kernel_initializer: kernel_initializer,
-        name: join(name, "value"),
-        use_bias: value_use_bias
-      )
-      |> Layers.split_heads(num_key_value_heads)
+    end
 
-    # Apply query and key normalization if configured (before rotary embedding)
-    query =
-      if query_norm do
-        query_norm.(query, join(name, "query_norm"))
-      else
-        query
-      end
+    query = project.(query, num_heads, query_use_bias, "query")
+    key = project.(key, num_key_value_heads, key_use_bias, "key")
+    value = project.(value, num_key_value_heads, value_use_bias, "value")
 
-    key =
-      if key_norm do
-        key_norm.(key, join(name, "key_norm"))
-      else
-        key
-      end
+    query = if query_norm, do: query_norm.(query, join(name, "query_norm")), else: query
+    key = if key_norm, do: key_norm.(key, join(name, "key_norm")), else: key
 
     {query, key} =
-      case rotary_embedding do
-        opts when is_list(opts) ->
-          validate_required_keys!(opts, [:position_ids])
-
-          opts =
-            Keyword.validate!(opts, [
-              :position_ids,
-              :max_positions,
-              :scaling_strategy,
-              base: 10_000,
-              percentage: 1.0
-            ])
-
-          {position_ids, opts} = Keyword.pop(opts, :position_ids)
-          {percentage, opts} = Keyword.pop(opts, :percentage)
-
-          size = trunc(attention_head_size * percentage)
-
-          rotary_opts = [name: join(name, "rotary_embedding")] ++ opts
-
-          if size == attention_head_size do
-            Layers.rotary_embedding(query, key, position_ids, attention_mask, size, rotary_opts)
-          else
-            query_rotary = Axon.nx(query, & &1[[.., .., .., 0..(size - 1)//1]])
-            query_pass = Axon.nx(query, & &1[[.., .., .., size..-1//1]])
-
-            key_rotary = Axon.nx(key, & &1[[.., .., .., 0..(size - 1)//1]])
-            key_pass = Axon.nx(key, & &1[[.., .., .., size..-1//1]])
-
-            {query_rotary, key_rotary} =
-              Layers.rotary_embedding(
-                query_rotary,
-                key_rotary,
-                position_ids,
-                attention_mask,
-                size,
-                rotary_opts
-              )
-
-            {Axon.concatenate([query_rotary, query_pass], axis: -1),
-             Axon.concatenate([key_rotary, key_pass], axis: -1)}
-          end
-
-        nil ->
-          {query, key}
+      if rotary_embedding do
+        rotary_embedding(query, key, attention_mask, attention_head_size, rotary_embedding,
+          name: join(name, "rotary_embedding")
+        )
+      else
+        {query, key}
       end
 
     num_key_value_groups = div(num_heads, num_key_value_heads)
-    key = repeat_states(key, num_key_value_groups)
-    value = repeat_states(value, num_key_value_groups)
+    key = repeat_heads(key, num_key_value_groups)
+    value = repeat_heads(value, num_key_value_groups)
 
     {key, value, attention_cache} =
       Layers.Decoder.cached_attention_key_values(key, value, attention_cache, offset)
@@ -952,20 +474,133 @@ defmodule Bumblebee.Layers.Transformer do
       )
 
     attention_output =
-      attention_output
-      |> Layers.flatten_trailing()
-      |> Axon.dense(hidden_size,
+      output_projection(attention_output, hidden_size,
+        use_bias: output_use_bias,
         kernel_initializer: kernel_initializer,
-        name: join(name, "output"),
-        use_bias: output_use_bias
+        name: join(name, "output")
       )
 
     {attention_output, attention_weights, attention_cache, attention_relative_bias}
   end
 
-  defp repeat_states(state, 1), do: state
+  @doc """
+  Projects the input to attention heads.
+
+  Returns a node with shape `{batch_size, sequence_length, num_heads, head_size}`.
+
+  ## Options
+
+    * `:use_bias` - whether to use bias in the projection. Defaults
+      to `true`
+
+    * `:kernel_initializer` - initializer for kernel weights. Defaults
+      to `:glorot_uniform`
+
+    * `:name` - the layer name
+
+  """
+  def project_heads(input, num_heads, head_size, opts) do
+    opts = Keyword.validate!(opts, [:name, use_bias: true, kernel_initializer: :glorot_uniform])
+
+    input
+    |> Axon.dense(num_heads * head_size,
+      use_bias: opts[:use_bias],
+      kernel_initializer: opts[:kernel_initializer],
+      name: opts[:name]
+    )
+    |> Layers.split_heads(num_heads)
+  end
+
+  @doc """
+  Merges attention heads and projects them to `hidden_size`.
+
+  Accepts the same options as `project_heads/4`.
+  """
+  def output_projection(attention_output, hidden_size, opts) do
+    opts = Keyword.validate!(opts, [:name, use_bias: true, kernel_initializer: :glorot_uniform])
+
+    attention_output
+    |> Layers.flatten_trailing()
+    |> Axon.dense(hidden_size,
+      use_bias: opts[:use_bias],
+      kernel_initializer: opts[:kernel_initializer],
+      name: opts[:name]
+    )
+  end
+
+  @doc """
+  Applies rotary embedding to query and key heads.
+
+  ## Rotary options
+
+    * `:position_ids` (required) - input position ids used for the
+      embedding
+
+    * `:max_positions` - the maximum number of distinct positions
+
+    * `:base` - base for computing rotary embedding frequency. Defaults
+      to `10_000`
+
+    * `:scaling_strategy` - see `Bumblebee.Layers.rotary_embedding/6`
+
+    * `:percentage` - percentage of head dimensions to apply rotary
+      embedding to. Defaults to `1.0`
+
+  ## Options
+
+    * `:name` - the prefix for layer names
+
+  """
+  def rotary_embedding(query, key, attention_mask, head_size, rotary_opts, opts) do
+    validate_required_keys!(rotary_opts, [:position_ids])
+
+    rotary_opts =
+      Keyword.validate!(rotary_opts, [
+        :position_ids,
+        :max_positions,
+        :scaling_strategy,
+        base: 10_000,
+        percentage: 1.0
+      ])
+
+    {position_ids, rotary_opts} = Keyword.pop(rotary_opts, :position_ids)
+    {percentage, rotary_opts} = Keyword.pop(rotary_opts, :percentage)
+
+    size = trunc(head_size * percentage)
+
+    rotary_opts = [name: opts[:name]] ++ rotary_opts
+
+    if size == head_size do
+      Layers.rotary_embedding(query, key, position_ids, attention_mask, size, rotary_opts)
+    else
+      query_rotary = Axon.nx(query, & &1[[.., .., .., 0..(size - 1)//1]])
+      query_pass = Axon.nx(query, & &1[[.., .., .., size..-1//1]])
+
+      key_rotary = Axon.nx(key, & &1[[.., .., .., 0..(size - 1)//1]])
+      key_pass = Axon.nx(key, & &1[[.., .., .., size..-1//1]])
+
+      {query_rotary, key_rotary} =
+        Layers.rotary_embedding(
+          query_rotary,
+          key_rotary,
+          position_ids,
+          attention_mask,
+          size,
+          rotary_opts
+        )
+
+      {Axon.concatenate([query_rotary, query_pass], axis: -1),
+       Axon.concatenate([key_rotary, key_pass], axis: -1)}
+    end
+  end
+
+  @doc """
+  Repeats key or value heads to match the number of query heads, as
+  in grouped-query attention.
+  """
+  def repeat_heads(state, 1), do: state
 
-  defp repeat_states(state, times) do
+  def repeat_heads(state, times) do
     Layers.repeat_interleave(state, times, axis: 2)
   end
 
diff --git a/lib/bumblebee/multimodal/layout_lm.ex b/lib/bumblebee/multimodal/layout_lm.ex
index 9c44854b..d0555242 100644
--- a/lib/bumblebee/multimodal/layout_lm.ex
+++ b/lib/bumblebee/multimodal/layout_lm.ex
@@ -408,23 +408,52 @@ defmodule Bumblebee.Multimodal.LayoutLm do
   defp encoder(hidden_state, attention_mask, attention_head_mask, spec, opts) do
     name = opts[:name]
 
-    Layers.Transformer.blocks(hidden_state,
-      attention_mask: attention_mask,
-      attention_head_mask: attention_head_mask,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: kernel_initializer(spec),
-      dropout_rate: spec.dropout_rate,
-      attention_dropout_rate: spec.attention_dropout_rate,
-      layer_norm: [
-        epsilon: spec.layer_norm_epsilon
-      ],
-      ffn: [
-        intermediate_size: spec.intermediate_size,
-        activation: spec.activation
+    Layers.Transformer.blocks(
+      hidden_state,
+      [
+        num_blocks: spec.num_blocks,
+        attention_mask: attention_mask,
+        attention_head_mask: attention_head_mask,
+        name: join(name, "blocks")
       ],
-      name: join(name, "blocks")
+      fn hidden_state, block ->
+        name = block.name
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, _cache} =
+          Layers.Transformer.self_attention(hidden_state, block,
+            num_heads: spec.num_attention_heads,
+            hidden_size: spec.hidden_size,
+            kernel_initializer: kernel_initializer(spec),
+            dropout_rate: spec.attention_dropout_rate,
+            name: join(name, "self_attention")
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: spec.dropout_rate)
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Layers.Transformer.basic_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            dropout_rate: spec.dropout_rate,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+
+        %{hidden_state: hidden_state, attention: attention}
+      end
     )
   end
 
diff --git a/lib/bumblebee/text/albert.ex b/lib/bumblebee/text/albert.ex
index 31a48ec1..2175ab02 100644
--- a/lib/bumblebee/text/albert.ex
+++ b/lib/bumblebee/text/albert.ex
@@ -329,16 +329,16 @@ defmodule Bumblebee.Text.Albert do
         name: "embedder"
       )
 
-    {hidden_state, hidden_states, attentions} =
+    encoder_outputs =
       encoder(embeddings, inputs["attention_mask"], spec, name: join(name, "encoder"))
 
-    pooled_state = pooler(hidden_state, spec, name: join(name, "pooler"))
+    pooled_state = pooler(encoder_outputs.hidden_state, spec, name: join(name, "pooler"))
 
     %{
-      hidden_state: hidden_state,
+      hidden_state: encoder_outputs.hidden_state,
       pooled_state: pooled_state,
-      hidden_states: hidden_states,
-      attentions: attentions
+      hidden_states: encoder_outputs.hidden_states,
+      attentions: encoder_outputs.attentions
     }
   end
 
@@ -387,42 +387,58 @@ defmodule Bumblebee.Text.Albert do
         name: join(name, "embedding_projection")
       )
 
-    hidden_states = Axon.container({hidden_state})
-    attentions = Axon.container({})
-
-    for block_idx <- 0..(spec.num_blocks - 1),
-        inner_idx <- 0..(spec.block_depth - 1),
-        reduce: {hidden_state, hidden_states, attentions} do
-      {hidden_state, hidden_states, attentions} ->
+    # Blocks within a group share parameters
+    Layers.Transformer.blocks(
+      hidden_state,
+      [
+        num_blocks: spec.num_blocks * spec.block_depth,
+        attention_mask: attention_mask
+      ],
+      fn hidden_state, block ->
+        block_idx = div(block.index, spec.block_depth)
+        inner_idx = rem(block.index, spec.block_depth)
         group_idx = div(block_idx, div(spec.num_blocks, spec.num_groups))
 
         name = name |> join("groups") |> join(group_idx) |> join("blocks") |> join(inner_idx)
 
         # TODO: wrap encoder block in a layer_drop combinator
-        {hidden_state, attention, _cross_attention, _block_cache, _position_bias} =
-          Layers.Transformer.block(hidden_state,
-            attention_mask: attention_mask,
-            num_attention_heads: spec.num_attention_heads,
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, _cache} =
+          Layers.Transformer.self_attention(hidden_state, block,
+            num_heads: spec.num_attention_heads,
             hidden_size: spec.hidden_size,
             kernel_initializer: kernel_initializer(spec),
+            dropout_rate: spec.attention_dropout_rate,
+            name: join(name, "self_attention")
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: spec.dropout_rate)
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Layers.Transformer.basic_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
             dropout_rate: spec.dropout_rate,
-            attention_dropout_rate: spec.attention_dropout_rate,
-            layer_norm: [
-              epsilon: spec.layer_norm_epsilon
-            ],
-            ffn: [
-              intermediate_size: spec.intermediate_size,
-              activation: spec.activation
-            ],
-            name: name
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "ffn")
           )
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
 
-        {
-          hidden_state,
-          Layers.append(hidden_states, hidden_state),
-          Layers.append(attentions, attention)
-        }
-    end
+        %{hidden_state: hidden_state, attention: attention}
+      end
+    )
   end
 
   defp pooler(hidden_state, spec, opts) do
diff --git a/lib/bumblebee/text/bart.ex b/lib/bumblebee/text/bart.ex
index b0e4e272..87182482 100644
--- a/lib/bumblebee/text/bart.ex
+++ b/lib/bumblebee/text/bart.ex
@@ -546,23 +546,52 @@ defmodule Bumblebee.Text.Bart do
   defp encoder(hidden_state, attention_mask, attention_head_mask, spec, opts) do
     name = opts[:name]
 
-    Layers.Transformer.blocks(hidden_state,
-      attention_mask: attention_mask,
-      attention_head_mask: attention_head_mask,
-      num_blocks: spec.encoder_num_blocks,
-      num_attention_heads: spec.encoder_num_attention_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: kernel_initializer(spec),
-      dropout_rate: spec.dropout_rate,
-      attention_dropout_rate: spec.attention_dropout_rate,
-      layer_norm: [
-        epsilon: 1.0e-5
-      ],
-      ffn: [
-        intermediate_size: spec.encoder_intermediate_size,
-        activation: spec.activation
+    Layers.Transformer.blocks(
+      hidden_state,
+      [
+        num_blocks: spec.encoder_num_blocks,
+        attention_mask: attention_mask,
+        attention_head_mask: attention_head_mask,
+        name: join(name, "blocks")
       ],
-      name: join(name, "blocks")
+      fn hidden_state, block ->
+        name = block.name
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, _cache} =
+          Layers.Transformer.self_attention(hidden_state, block,
+            num_heads: spec.encoder_num_attention_heads,
+            hidden_size: spec.hidden_size,
+            kernel_initializer: kernel_initializer(spec),
+            dropout_rate: spec.attention_dropout_rate,
+            name: join(name, "self_attention")
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: spec.dropout_rate)
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "self_attention_norm"))
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Layers.Transformer.basic_ffn(spec.encoder_intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            dropout_rate: spec.dropout_rate,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "output_norm"))
+
+        %{
+          hidden_state: hidden_state,
+          attention: attention
+        }
+      end
     )
   end
 
@@ -579,28 +608,88 @@ defmodule Bumblebee.Text.Bart do
        ) do
     name = opts[:name]
 
-    Layers.Transformer.blocks(hidden_state,
-      attention_mask: attention_mask,
-      attention_head_mask: attention_head_mask,
-      cross_hidden_state: encoder_hidden_state,
-      cross_attention_mask: encoder_attention_mask,
-      cross_attention_head_mask: cross_attention_head_mask,
-      cache: cache,
-      causal: true,
-      num_blocks: spec.decoder_num_blocks,
-      num_attention_heads: spec.decoder_num_attention_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: kernel_initializer(spec),
-      dropout_rate: spec.dropout_rate,
-      attention_dropout_rate: spec.attention_dropout_rate,
-      layer_norm: [
-        epsilon: 1.0e-5
-      ],
-      ffn: [
-        intermediate_size: spec.decoder_intermediate_size,
-        activation: spec.activation
+    Layers.Transformer.blocks(
+      hidden_state,
+      [
+        num_blocks: spec.decoder_num_blocks,
+        attention_mask: attention_mask,
+        attention_head_mask: attention_head_mask,
+        cross_attention_head_mask: cross_attention_head_mask,
+        cache: cache,
+        name: join(name, "blocks")
       ],
-      name: join(name, "blocks")
+      fn hidden_state, block ->
+        name = block.name
+
+        attention_opts = [
+          num_heads: spec.decoder_num_attention_heads,
+          hidden_size: spec.hidden_size,
+          kernel_initializer: kernel_initializer(spec),
+          dropout_rate: spec.attention_dropout_rate
+        ]
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, self_attention_cache} =
+          Layers.Transformer.self_attention(
+            hidden_state,
+            block,
+            [causal: true, name: join(name, "self_attention")] ++ attention_opts
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: spec.dropout_rate)
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "self_attention_norm"))
+
+        {hidden_state, cross_attention, cross_attention_cache} =
+          Layers.if_present encoder_hidden_state do
+            shortcut = hidden_state
+
+            {hidden_state, cross_attention, cross_attention_cache} =
+              Layers.Transformer.cross_attention(
+                hidden_state,
+                encoder_hidden_state,
+                block,
+                [
+                  attention_mask: encoder_attention_mask,
+                  name: join(name, "cross_attention")
+                ] ++ attention_opts
+              )
+
+            hidden_state =
+              hidden_state
+              |> Axon.dropout(rate: spec.dropout_rate)
+              |> Axon.add(shortcut)
+              |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "cross_attention_norm"))
+
+            {hidden_state, cross_attention, cross_attention_cache}
+          else
+            {hidden_state, Layers.none(), block.cross_attention_cache}
+          end
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Layers.Transformer.basic_ffn(spec.decoder_intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            dropout_rate: spec.dropout_rate,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "output_norm"))
+
+        %{
+          hidden_state: hidden_state,
+          attention: attention,
+          cross_attention: cross_attention,
+          self_attention_cache: self_attention_cache,
+          cross_attention_cache: cross_attention_cache
+        }
+      end
     )
   end
 
diff --git a/lib/bumblebee/text/bert.ex b/lib/bumblebee/text/bert.ex
index bf8323c6..dc35dfbe 100644
--- a/lib/bumblebee/text/bert.ex
+++ b/lib/bumblebee/text/bert.ex
@@ -521,33 +521,95 @@ defmodule Bumblebee.Text.Bert do
     Layers.Transformer.blocks(
       hidden_state,
       [
+        num_blocks: spec.num_blocks,
         attention_mask: attention_mask,
         attention_head_mask: attention_head_mask,
+        cross_attention_head_mask: cross_attention_head_mask,
         cache: cache,
-        causal: decoder?,
-        num_blocks: spec.num_blocks,
-        num_attention_heads: spec.num_attention_heads,
-        hidden_size: spec.hidden_size,
-        kernel_initializer: kernel_initializer(spec),
-        dropout_rate: spec.dropout_rate,
-        attention_dropout_rate: spec.attention_dropout_rate,
-        layer_norm: [
-          epsilon: spec.layer_norm_epsilon
-        ],
-        ffn: [
-          intermediate_size: spec.intermediate_size,
-          activation: spec.activation
-        ],
         name: join(name, "blocks")
-      ] ++
-        if(cross_attention?,
-          do: [
-            cross_hidden_state: encoder_hidden_state,
-            cross_attention_mask: encoder_attention_mask,
-            cross_attention_head_mask: cross_attention_head_mask
-          ],
-          else: []
-        )
+      ],
+      fn hidden_state, block ->
+        name = block.name
+
+        attention_opts = [
+          num_heads: spec.num_attention_heads,
+          hidden_size: spec.hidden_size,
+          kernel_initializer: kernel_initializer(spec),
+          dropout_rate: spec.attention_dropout_rate
+        ]
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, self_attention_cache} =
+          Layers.Transformer.self_attention(
+            hidden_state,
+            block,
+            [causal: decoder?, name: join(name, "self_attention")] ++ attention_opts
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: spec.dropout_rate)
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+
+        {hidden_state, cross_attention, cross_attention_cache} =
+          if cross_attention? do
+            Layers.if_present encoder_hidden_state do
+              shortcut = hidden_state
+
+              {hidden_state, cross_attention, cross_attention_cache} =
+                Layers.Transformer.cross_attention(
+                  hidden_state,
+                  encoder_hidden_state,
+                  block,
+                  [
+                    attention_mask: encoder_attention_mask,
+                    name: join(name, "cross_attention")
+                  ] ++ attention_opts
+                )
+
+              hidden_state =
+                hidden_state
+                |> Axon.dropout(rate: spec.dropout_rate)
+                |> Axon.add(shortcut)
+                |> Axon.layer_norm(
+                  epsilon: spec.layer_norm_epsilon,
+                  name: join(name, "cross_attention_norm")
+                )
+
+              {hidden_state, cross_attention, cross_attention_cache}
+            else
+              {hidden_state, Layers.none(), block.cross_attention_cache}
+            end
+          else
+            {hidden_state, Layers.none(), block.cross_attention_cache}
+          end
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Layers.Transformer.basic_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            dropout_rate: spec.dropout_rate,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+
+        %{
+          hidden_state: hidden_state,
+          attention: attention,
+          cross_attention: cross_attention,
+          self_attention_cache: self_attention_cache,
+          cross_attention_cache: cross_attention_cache
+        }
+      end
     )
   end
 
diff --git a/lib/bumblebee/text/blenderbot.ex b/lib/bumblebee/text/blenderbot.ex
index 18898e8c..491e8881 100644
--- a/lib/bumblebee/text/blenderbot.ex
+++ b/lib/bumblebee/text/blenderbot.ex
@@ -392,24 +392,53 @@ defmodule Bumblebee.Text.Blenderbot do
     name = opts[:name]
 
     encoder_outputs =
-      Layers.Transformer.blocks(hidden_state,
-        attention_mask: attention_mask,
-        attention_head_mask: attention_head_mask,
-        num_blocks: spec.encoder_num_blocks,
-        num_attention_heads: spec.encoder_num_attention_heads,
-        hidden_size: spec.hidden_size,
-        kernel_initializer: kernel_initializer(spec),
-        dropout_rate: spec.dropout_rate,
-        attention_dropout_rate: spec.attention_dropout_rate,
-        layer_norm: [
-          epsilon: 1.0e-5
-        ],
-        ffn: [
-          intermediate_size: spec.encoder_intermediate_size,
-          activation: spec.activation
+      Layers.Transformer.blocks(
+        hidden_state,
+        [
+          num_blocks: spec.encoder_num_blocks,
+          attention_mask: attention_mask,
+          attention_head_mask: attention_head_mask,
+          name: join(name, "blocks")
         ],
-        block_type: :norm_first,
-        name: join(name, "blocks")
+        fn hidden_state, block ->
+          name = block.name
+
+          shortcut = hidden_state
+
+          {hidden_state, attention, _cache} =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "self_attention_norm"))
+            |> Layers.Transformer.self_attention(block,
+              num_heads: spec.encoder_num_attention_heads,
+              hidden_size: spec.hidden_size,
+              kernel_initializer: kernel_initializer(spec),
+              dropout_rate: spec.attention_dropout_rate,
+              name: join(name, "self_attention")
+            )
+
+          hidden_state =
+            hidden_state
+            |> Axon.dropout(rate: spec.dropout_rate)
+            |> Axon.add(shortcut)
+
+          shortcut = hidden_state
+
+          hidden_state =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "output_norm"))
+            |> Layers.Transformer.basic_ffn(spec.encoder_intermediate_size, spec.hidden_size,
+              activation: spec.activation,
+              dropout_rate: spec.dropout_rate,
+              kernel_initializer: kernel_initializer(spec),
+              name: join(name, "ffn")
+            )
+            |> Axon.add(shortcut)
+
+          %{
+            hidden_state: hidden_state,
+            attention: attention
+          }
+        end
       )
 
     hidden_state = Axon.layer_norm(encoder_outputs.hidden_state, name: join(name, "norm"))
@@ -435,29 +464,88 @@ defmodule Bumblebee.Text.Blenderbot do
     name = opts[:name]
 
     decoder_outputs =
-      Layers.Transformer.blocks(hidden_state,
-        attention_mask: attention_mask,
-        attention_head_mask: attention_head_mask,
-        cross_hidden_state: encoder_hidden_state,
-        cross_attention_mask: encoder_attention_mask,
-        cross_attention_head_mask: cross_attention_head_mask,
-        cache: cache,
-        causal: true,
-        num_blocks: spec.decoder_num_blocks,
-        num_attention_heads: spec.decoder_num_attention_heads,
-        hidden_size: spec.hidden_size,
-        kernel_initializer: kernel_initializer(spec),
-        dropout_rate: spec.dropout_rate,
-        attention_dropout_rate: spec.attention_dropout_rate,
-        layer_norm: [
-          epsilon: 1.0e-5
-        ],
-        ffn: [
-          intermediate_size: spec.decoder_intermediate_size,
-          activation: spec.activation
+      Layers.Transformer.blocks(
+        hidden_state,
+        [
+          num_blocks: spec.decoder_num_blocks,
+          attention_mask: attention_mask,
+          attention_head_mask: attention_head_mask,
+          cross_attention_head_mask: cross_attention_head_mask,
+          cache: cache,
+          name: join(name, "blocks")
         ],
-        block_type: :norm_first,
-        name: join(name, "blocks")
+        fn hidden_state, block ->
+          name = block.name
+
+          attention_opts = [
+            num_heads: spec.decoder_num_attention_heads,
+            hidden_size: spec.hidden_size,
+            kernel_initializer: kernel_initializer(spec),
+            dropout_rate: spec.attention_dropout_rate
+          ]
+
+          shortcut = hidden_state
+
+          {hidden_state, attention, self_attention_cache} =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "self_attention_norm"))
+            |> Layers.Transformer.self_attention(
+              block,
+              [causal: true, name: join(name, "self_attention")] ++ attention_opts
+            )
+
+          hidden_state =
+            hidden_state
+            |> Axon.dropout(rate: spec.dropout_rate)
+            |> Axon.add(shortcut)
+
+          {hidden_state, cross_attention, cross_attention_cache} =
+            Layers.if_present encoder_hidden_state do
+              shortcut = hidden_state
+
+              {hidden_state, cross_attention, cross_attention_cache} =
+                hidden_state
+                |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "cross_attention_norm"))
+                |> Layers.Transformer.cross_attention(
+                  encoder_hidden_state,
+                  block,
+                  [
+                    attention_mask: encoder_attention_mask,
+                    name: join(name, "cross_attention")
+                  ] ++ attention_opts
+                )
+
+              hidden_state =
+                hidden_state
+                |> Axon.dropout(rate: spec.dropout_rate)
+                |> Axon.add(shortcut)
+
+              {hidden_state, cross_attention, cross_attention_cache}
+            else
+              {hidden_state, Layers.none(), block.cross_attention_cache}
+            end
+
+          shortcut = hidden_state
+
+          hidden_state =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "output_norm"))
+            |> Layers.Transformer.basic_ffn(spec.decoder_intermediate_size, spec.hidden_size,
+              activation: spec.activation,
+              dropout_rate: spec.dropout_rate,
+              kernel_initializer: kernel_initializer(spec),
+              name: join(name, "ffn")
+            )
+            |> Axon.add(shortcut)
+
+          %{
+            hidden_state: hidden_state,
+            attention: attention,
+            cross_attention: cross_attention,
+            self_attention_cache: self_attention_cache,
+            cross_attention_cache: cross_attention_cache
+          }
+        end
       )
 
     hidden_state = Axon.layer_norm(decoder_outputs.hidden_state, name: join(name, "norm"))
diff --git a/lib/bumblebee/text/blip_text.ex b/lib/bumblebee/text/blip_text.ex
index e21fa191..3a6487f7 100644
--- a/lib/bumblebee/text/blip_text.ex
+++ b/lib/bumblebee/text/blip_text.ex
@@ -312,38 +312,98 @@ defmodule Bumblebee.Text.BlipText do
     name = opts[:name]
     decoder? = opts[:decoder?]
 
-    cross_attention? = decoder?
-
     Layers.Transformer.blocks(
       hidden_state,
       [
+        num_blocks: spec.num_blocks,
         attention_mask: attention_mask,
         attention_head_mask: attention_head_mask,
+        cross_attention_head_mask: cross_attention_head_mask,
         cache: cache,
-        causal: decoder?,
-        num_blocks: spec.num_blocks,
-        num_attention_heads: spec.num_attention_heads,
-        hidden_size: spec.hidden_size,
-        kernel_initializer: kernel_initializer(spec),
-        dropout_rate: spec.dropout_rate,
-        attention_dropout_rate: spec.attention_dropout_rate,
-        layer_norm: [
-          epsilon: spec.layer_norm_epsilon
-        ],
-        ffn: [
-          intermediate_size: spec.intermediate_size,
-          activation: spec.activation
-        ],
         name: join(name, "blocks")
-      ] ++
-        if(cross_attention?,
-          do: [
-            cross_hidden_state: encoder_hidden_state,
-            cross_attention_mask: encoder_attention_mask,
-            cross_attention_head_mask: cross_attention_head_mask
-          ],
-          else: []
-        )
+      ],
+      fn hidden_state, block ->
+        name = block.name
+
+        attention_opts = [
+          num_heads: spec.num_attention_heads,
+          hidden_size: spec.hidden_size,
+          kernel_initializer: kernel_initializer(spec),
+          dropout_rate: spec.attention_dropout_rate
+        ]
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, self_attention_cache} =
+          Layers.Transformer.self_attention(
+            hidden_state,
+            block,
+            [causal: decoder?, name: join(name, "self_attention")] ++ attention_opts
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: spec.dropout_rate)
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+
+        {hidden_state, cross_attention, cross_attention_cache} =
+          if decoder? do
+            Layers.if_present encoder_hidden_state do
+              shortcut = hidden_state
+
+              {hidden_state, cross_attention, cross_attention_cache} =
+                Layers.Transformer.cross_attention(
+                  hidden_state,
+                  encoder_hidden_state,
+                  block,
+                  [
+                    attention_mask: encoder_attention_mask,
+                    name: join(name, "cross_attention")
+                  ] ++ attention_opts
+                )
+
+              hidden_state =
+                hidden_state
+                |> Axon.dropout(rate: spec.dropout_rate)
+                |> Axon.add(shortcut)
+                |> Axon.layer_norm(
+                  epsilon: spec.layer_norm_epsilon,
+                  name: join(name, "cross_attention_norm")
+                )
+
+              {hidden_state, cross_attention, cross_attention_cache}
+            else
+              {hidden_state, Layers.none(), block.cross_attention_cache}
+            end
+          else
+            {hidden_state, Layers.none(), block.cross_attention_cache}
+          end
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Layers.Transformer.basic_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            dropout_rate: spec.dropout_rate,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+
+        %{
+          hidden_state: hidden_state,
+          attention: attention,
+          cross_attention: cross_attention,
+          self_attention_cache: self_attention_cache,
+          cross_attention_cache: cross_attention_cache
+        }
+      end
     )
   end
 
diff --git a/lib/bumblebee/text/clip_text.ex b/lib/bumblebee/text/clip_text.ex
index 91abbe81..7f01de13 100644
--- a/lib/bumblebee/text/clip_text.ex
+++ b/lib/bumblebee/text/clip_text.ex
@@ -208,24 +208,51 @@ defmodule Bumblebee.Text.ClipText do
   defp encoder(embeddings, attention_mask, spec, opts) do
     name = opts[:name]
 
-    Layers.Transformer.blocks(embeddings,
-      attention_mask: attention_mask,
-      causal: true,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: Axon.Initializers.normal(scale: 0.01),
-      dropout_rate: 0.0,
-      attention_dropout_rate: spec.attention_dropout_rate,
-      layer_norm: [
-        epsilon: spec.layer_norm_epsilon
+    kernel_initializer = Axon.Initializers.normal(scale: 0.01)
+
+    Layers.Transformer.blocks(
+      embeddings,
+      [
+        num_blocks: spec.num_blocks,
+        attention_mask: attention_mask,
+        name: join(name, "blocks")
       ],
-      ffn: [
-        intermediate_size: spec.intermediate_size,
-        activation: spec.activation
-      ],
-      block_type: :norm_first,
-      name: join(name, "blocks")
+      fn hidden_state, block ->
+        name = block.name
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, _self_attention_cache} =
+          hidden_state
+          |> Axon.layer_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+          |> Layers.Transformer.self_attention(block,
+            num_heads: spec.num_attention_heads,
+            hidden_size: spec.hidden_size,
+            causal: true,
+            dropout_rate: spec.attention_dropout_rate,
+            kernel_initializer: kernel_initializer,
+            name: join(name, "self_attention")
+          )
+
+        hidden_state = Axon.add(hidden_state, shortcut)
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Axon.layer_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+          |> Layers.Transformer.basic_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            kernel_initializer: kernel_initializer,
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+
+        %{hidden_state: hidden_state, attention: attention}
+      end
     )
   end
 
diff --git a/lib/bumblebee/text/deberta_v2.ex b/lib/bumblebee/text/deberta_v2.ex
index 3c192f33..88eac201 100644
--- a/lib/bumblebee/text/deberta_v2.ex
+++ b/lib/bumblebee/text/deberta_v2.ex
@@ -151,28 +151,7 @@ defmodule Bumblebee.Text.DebertaV2 do
 
     relative = relative_embeddings(spec)
 
-    outputs =
-      Transformer.blocks(embedding,
-        num_blocks: spec.num_blocks,
-        num_attention_heads: spec.num_attention_heads,
-        hidden_size: spec.hidden_size,
-        attention_mask: mask,
-        attention: fn query, _key, _value, opts ->
-          {output, weights} = attention(query, relative, mask, spec, opts[:name])
-          {output, weights, opts[:attention_cache], Layers.none()}
-        end,
-        block_output: fn output, index ->
-          if index == 0 and spec.conv_kernel_size > 0,
-            do: convolution(embedding, output, mask, spec),
-            else: output
-        end,
-        ffn: [intermediate_size: spec.intermediate_size, activation: spec.activation],
-        dropout_rate: spec.dropout_rate,
-        attention_dropout_rate: spec.attention_dropout_rate,
-        layer_norm: [epsilon: spec.layer_norm_epsilon],
-        kernel_initializer: initializer(spec),
-        name: "encoder.blocks"
-      )
+    outputs = encoder(embedding, relative, mask, spec)
 
     Layers.output(%{
       hidden_state: outputs.hidden_state,
@@ -181,53 +160,124 @@ defmodule Bumblebee.Text.DebertaV2 do
     })
   end
 
-  defp attention(hidden, relative, mask, spec, name) do
+  defp encoder(embedding, relative, mask, spec) do
+    # Both query and key padding positions are masked
+    attention_mask =
+      Axon.nx(mask, fn mask ->
+        Nx.multiply(Nx.new_axis(Nx.new_axis(mask, 1), 2), Nx.new_axis(Nx.new_axis(mask, 1), 3))
+      end)
+
+    Transformer.blocks(
+      embedding,
+      [num_blocks: spec.num_blocks, name: "encoder.blocks"],
+      fn hidden_state, block ->
+        name = block.name
+
+        shortcut = hidden_state
+
+        {hidden_state, attention} =
+          attention(hidden_state, relative, attention_mask, spec, name <> ".self_attention")
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: spec.dropout_rate)
+          |> Axon.add(shortcut)
+          |> norm(spec, name <> ".self_attention_norm")
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Transformer.basic_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            dropout_rate: spec.dropout_rate,
+            kernel_initializer: initializer(spec),
+            name: name <> ".ffn"
+          )
+          |> Axon.add(shortcut)
+          |> norm(spec, name <> ".output_norm")
+
+        hidden_state =
+          if block.index == 0 and spec.conv_kernel_size > 0,
+            do: convolution(embedding, hidden_state, mask, spec),
+            else: hidden_state
+
+        %{hidden_state: hidden_state, attention: attention}
+      end
+    )
+  end
+
+  defp attention(hidden, relative, attention_mask, spec, name) do
+    heads = spec.num_attention_heads
+    head_size = div(spec.hidden_size, heads)
+    scale_factor = 1 + length(spec.position_attention_types)
+
     relative = Axon.dropout(relative, rate: spec.dropout_rate)
     {query, position_query} = attention_projection(hidden, relative, spec, name <> ".query")
     {key, position_key} = attention_projection(hidden, relative, spec, name <> ".key")
 
-    values =
-      hidden
-      |> dense(spec.hidden_size, spec, name <> ".value")
-      |> Axon.nx(&split_heads(&1, spec.num_attention_heads))
+    value =
+      Transformer.project_heads(hidden, heads, head_size,
+        kernel_initializer: initializer(spec),
+        name: name <> ".value"
+      )
+
+    position_projection = fn suffix ->
+      relative
+      |> Axon.nx(&Nx.new_axis(&1, 0))
+      |> Transformer.project_heads(heads, head_size,
+        kernel_initializer: initializer(spec),
+        name: name <> suffix
+      )
+    end
 
     position_key =
       if spec.use_relative_attention and not spec.share_attention_key and
-           :content_to_position in spec.position_attention_types do
-        relative
-        |> Axon.nx(&Nx.new_axis(&1, 0))
-        |> dense(spec.hidden_size, spec, name <> ".position_key")
-        |> Axon.nx(&split_heads(&1, spec.num_attention_heads))
-      else
-        position_key
-      end
+           :content_to_position in spec.position_attention_types,
+         do: position_projection.(".position_key"),
+         else: position_key
 
     position_query =
       if spec.use_relative_attention and not spec.share_attention_key and
-           :position_to_content in spec.position_attention_types do
-        relative
-        |> Axon.nx(&Nx.new_axis(&1, 0))
-        |> dense(spec.hidden_size, spec, name <> ".position_query")
-        |> Axon.nx(&split_heads(&1, spec.num_attention_heads))
+           :position_to_content in spec.position_attention_types,
+         do: position_projection.(".position_query"),
+         else: position_query
+
+    relative_bias =
+      if spec.use_relative_attention and spec.position_attention_types != [] do
+        Axon.layer(
+          &relative_attention_bias/5,
+          [query, key, position_query, position_key],
+          buckets: spec.position_buckets,
+          max_position: relative_limit(spec),
+          span: position_span(spec),
+          terms: spec.position_attention_types
+        )
       else
-        position_query
+        Layers.none()
       end
 
-    weights =
-      Axon.layer(
-        &attention_weights/6,
-        [query, key, position_query, position_key, mask],
-        buckets: spec.position_buckets,
-        max_position: relative_limit(spec),
-        span: position_span(spec),
-        relative: spec.use_relative_attention,
-        terms: spec.position_attention_types
+    # The key is scaled before the dot product, as in the Python model
+    key = Axon.nx(key, &Nx.divide(&1, Nx.sqrt(head_size * scale_factor)))
+
+    {output, weights} =
+      Layers.attention(
+        query,
+        key,
+        value,
+        attention_mask,
+        Layers.none(),
+        relative_bias,
+        Layers.none(),
+        scale: 1,
+        dropout_rate: spec.attention_dropout_rate
       )
-      |> Axon.dropout(rate: spec.attention_dropout_rate)
 
     output =
-      Axon.layer(&attention_context/3, [weights, values])
-      |> dense(spec.hidden_size, spec, name <> ".output")
+      Transformer.output_projection(output, spec.hidden_size,
+        kernel_initializer: initializer(spec),
+        name: name <> ".output"
+      )
 
     {output, weights}
   end
@@ -279,7 +329,12 @@ defmodule Bumblebee.Text.DebertaV2 do
 
     embedding =
       if width != spec.hidden_size,
-        do: dense(embedding, spec.hidden_size, spec, "embedder.projection", use_bias: false),
+        do:
+          Axon.dense(embedding, spec.hidden_size,
+            use_bias: false,
+            kernel_initializer: initializer(spec),
+            name: "embedder.projection"
+          ),
         else: embedding
 
     embedding
@@ -354,9 +409,6 @@ defmodule Bumblebee.Text.DebertaV2 do
   defp norm(hidden, spec, name),
     do: Axon.layer_norm(hidden, epsilon: spec.layer_norm_epsilon, name: name)
 
-  defp dense(hidden, units, spec, name, opts \\ []),
-    do: Axon.dense(hidden, units, [name: name, kernel_initializer: initializer(spec)] ++ opts)
-
   defp initializer(spec), do: Axon.Initializers.normal(scale: spec.initializer_scale)
 
   defp relative_limit(spec),
@@ -383,60 +435,32 @@ defmodule Bumblebee.Text.DebertaV2 do
     {content, position}
   end
 
-  defnp attention_weights(query, key, position_query, position_key, mask, opts) do
-    opts = keyword!(opts, [:buckets, :max_position, :span, :relative, :terms, :mode])
+  defnp relative_attention_bias(query, key, position_query, position_key, opts) do
+    opts = keyword!(opts, [:buckets, :max_position, :span, :terms, :mode])
+    query = Nx.transpose(query, axes: [0, 2, 1, 3])
+    key = Nx.transpose(key, axes: [0, 2, 1, 3])
     scale = Nx.sqrt(Nx.axis_size(query, -1) * (1 + term_count(opts[:terms])))
-    scores = product(query, key / scale)
-
-    scores =
-      if opts[:relative] do
-        positions =
-          relative_positions(Nx.axis_size(query, 2), opts[:buckets], opts[:max_position])
-
-        bias = Nx.broadcast(0, Nx.shape(scores))
-
-        bias =
-          if has_term?(opts[:terms], :content_to_position) do
-            indices = Nx.clip(positions + opts[:span], 0, 2 * opts[:span] - 1)
-            bias + gather(product(query, position_key), indices) / scale
-          else
-            bias
-          end
-
-        bias =
-          if has_term?(opts[:terms], :position_to_content) do
-            indices = Nx.clip(-positions + opts[:span], 0, 2 * opts[:span] - 1)
+    positions = relative_positions(Nx.axis_size(query, 2), opts[:buckets], opts[:max_position])
 
-            bias +
-              Nx.transpose(gather(product(key, position_query), indices), axes: [0, 1, 3, 2]) /
-                scale
-          else
-            bias
-          end
-
-        scores + bias
+    content_to_position =
+      if has_term?(opts[:terms], :content_to_position) do
+        position_key = Nx.transpose(position_key, axes: [0, 2, 1, 3])
+        indices = Nx.clip(positions + opts[:span], 0, 2 * opts[:span] - 1)
+        gather(product(query, position_key), indices) / scale
       else
-        scores
+        0.0
       end
 
-    allowed = Nx.new_axis(Nx.new_axis(mask, 1), 2) * Nx.new_axis(Nx.new_axis(mask, 1), 3)
-
-    scores =
-      Nx.select(
-        Nx.broadcast(allowed, Nx.shape(scores)),
-        scores,
-        Nx.Constants.min_finite(Nx.type(scores))
-      )
-
-    Axon.Activations.softmax(scores, axis: -1)
-  end
-
-  defnp attention_context(weights, value, _opts) do
-    {batch, heads, length, width} = Nx.shape(value)
+    position_to_content =
+      if has_term?(opts[:terms], :position_to_content) do
+        position_query = Nx.transpose(position_query, axes: [0, 2, 1, 3])
+        indices = Nx.clip(-positions + opts[:span], 0, 2 * opts[:span] - 1)
+        Nx.transpose(gather(product(key, position_query), indices), axes: [0, 1, 3, 2]) / scale
+      else
+        0.0
+      end
 
-    Nx.dot(weights, [3], [0, 1], value, [2], [0, 1])
-    |> Nx.transpose(axes: [0, 2, 1, 3])
-    |> Nx.reshape({batch, length, heads * width})
+    content_to_position + position_to_content
   end
 
   deftransformp(has_term?(terms, term), do: term in terms)
@@ -446,10 +470,7 @@ defmodule Bumblebee.Text.DebertaV2 do
 
   defnp split_heads(input, heads) do
     {batch, length, hidden} = Nx.shape(input)
-
-    input
-    |> Nx.reshape({batch, length, heads, div(hidden, heads)})
-    |> Nx.transpose(axes: [0, 2, 1, 3])
+    Nx.reshape(input, {batch, length, heads, div(hidden, heads)})
   end
 
   defnp product(left, right) do
diff --git a/lib/bumblebee/text/distilbert.ex b/lib/bumblebee/text/distilbert.ex
index 2a5b7c2d..11a5256c 100644
--- a/lib/bumblebee/text/distilbert.ex
+++ b/lib/bumblebee/text/distilbert.ex
@@ -369,22 +369,47 @@ defmodule Bumblebee.Text.Distilbert do
 
     Layers.Transformer.blocks(
       hidden_state,
-      attention_mask: attention_mask,
-      attention_head_mask: attention_head_mask,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: kernel_initializer(spec),
-      dropout_rate: spec.dropout_rate,
-      attention_dropout_rate: spec.attention_dropout_rate,
-      layer_norm: [
-        epsilon: 1.0e-12
-      ],
-      ffn: [
-        intermediate_size: spec.intermediate_size,
-        activation: spec.activation
+      [
+        num_blocks: spec.num_blocks,
+        attention_mask: attention_mask,
+        attention_head_mask: attention_head_mask,
+        name: join(name, "blocks")
       ],
-      name: join(name, "blocks")
+      fn hidden_state, block ->
+        name = block.name
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, _cache} =
+          Layers.Transformer.self_attention(hidden_state, block,
+            num_heads: spec.num_attention_heads,
+            hidden_size: spec.hidden_size,
+            kernel_initializer: kernel_initializer(spec),
+            dropout_rate: spec.attention_dropout_rate,
+            name: join(name, "self_attention")
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: spec.dropout_rate)
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(epsilon: 1.0e-12, name: join(name, "self_attention_norm"))
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Layers.Transformer.basic_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            dropout_rate: spec.dropout_rate,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(epsilon: 1.0e-12, name: join(name, "output_norm"))
+
+        %{hidden_state: hidden_state, attention: attention}
+      end
     )
   end
 
diff --git a/lib/bumblebee/text/gemma.ex b/lib/bumblebee/text/gemma.ex
index cc04731a..969890e1 100644
--- a/lib/bumblebee/text/gemma.ex
+++ b/lib/bumblebee/text/gemma.ex
@@ -336,54 +336,72 @@ defmodule Bumblebee.Text.Gemma do
        ) do
     name = opts[:name]
 
-    Layers.Transformer.blocks(hidden_state,
-      attention_mask: attention_mask,
-      attention_head_mask: attention_head_mask,
-      attention_head_size: spec.attention_head_size,
-      cache: cache,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      num_key_value_heads: spec.num_key_value_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: kernel_initializer(spec),
-      layer_norm:
-        &Layers.rms_norm(&1, shift: 1.0, name: &2, epsilon: spec.layer_norm_epsilon, upcast: :all),
-      ffn:
-        &gated_ffn(&1, spec.intermediate_size, spec.hidden_size,
-          name: &2,
-          activation: spec.activation
-        ),
-      block_type: :norm_first,
-      causal: true,
-      rotary_embedding: [
-        position_ids: position_ids,
-        max_positions: spec.max_positions,
-        base: spec.rotary_embedding_base,
-        scaling_strategy: spec.rotary_embedding_scaling_strategy
-      ],
-      query_use_bias: spec.use_attention_bias,
-      key_use_bias: spec.use_attention_bias,
-      value_use_bias: spec.use_attention_bias,
-      output_use_bias: spec.use_attention_bias,
-      name: join(name, "blocks")
-    )
-  end
-
-  defp gated_ffn(hidden_state, intermediate_size, output_size, opts) do
-    name = opts[:name]
-    activation = opts[:activation]
-
-    intermediate =
-      Axon.dense(hidden_state, intermediate_size,
-        name: join(name, "intermediate"),
-        use_bias: false
+    norm = fn hidden_state, name ->
+      Layers.rms_norm(hidden_state,
+        shift: 1.0,
+        epsilon: spec.layer_norm_epsilon,
+        upcast: :all,
+        name: name
       )
+    end
 
-    gate = Axon.dense(hidden_state, intermediate_size, name: join(name, "gate"), use_bias: false)
-
-    hidden_state = Axon.multiply(intermediate, Axon.activation(gate, activation))
-
-    Axon.dense(hidden_state, output_size, name: join(name, "output"), use_bias: false)
+    Layers.Transformer.blocks(
+      hidden_state,
+      [
+        num_blocks: spec.num_blocks,
+        attention_mask: attention_mask,
+        attention_head_mask: attention_head_mask,
+        cache: cache,
+        name: join(name, "blocks")
+      ],
+      fn hidden_state, block ->
+        name = block.name
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, self_attention_cache} =
+          hidden_state
+          |> norm.(join(name, "self_attention_norm"))
+          |> Layers.Transformer.self_attention(block,
+            num_heads: spec.num_attention_heads,
+            num_key_value_heads: spec.num_key_value_heads,
+            hidden_size: spec.hidden_size,
+            attention_head_size: spec.attention_head_size,
+            causal: true,
+            rotary_embedding: [
+              position_ids: position_ids,
+              max_positions: spec.max_positions,
+              base: spec.rotary_embedding_base,
+              scaling_strategy: spec.rotary_embedding_scaling_strategy
+            ],
+            query_use_bias: spec.use_attention_bias,
+            key_use_bias: spec.use_attention_bias,
+            value_use_bias: spec.use_attention_bias,
+            output_use_bias: spec.use_attention_bias,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "self_attention")
+          )
+
+        hidden_state = Axon.add(hidden_state, shortcut)
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> norm.(join(name, "output_norm"))
+          |> Layers.Transformer.gated_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+
+        %{
+          hidden_state: hidden_state,
+          attention: attention,
+          self_attention_cache: self_attention_cache
+        }
+      end
+    )
   end
 
   defp language_modeling_head(hidden_state, spec, opts) do
diff --git a/lib/bumblebee/text/gemma3_text.ex b/lib/bumblebee/text/gemma3_text.ex
index afe1502a..f1797926 100644
--- a/lib/bumblebee/text/gemma3_text.ex
+++ b/lib/bumblebee/text/gemma3_text.ex
@@ -373,148 +373,95 @@ defmodule Bumblebee.Text.Gemma3Text do
        ) do
     name = opts[:name]
 
-    # QK-norm functions for Gemma 3 (uses shift: 1.0 for (1+weight) formula)
-    query_norm = &Layers.rms_norm(&1, shift: 1.0, epsilon: spec.layer_norm_epsilon, name: &2)
-    key_norm = &Layers.rms_norm(&1, shift: 1.0, epsilon: spec.layer_norm_epsilon, name: &2)
-
-    # Per-layer attention window size based on layer_types
-    # :attention_window_size uses local (sliding window) attention
-    # :full_attention uses global attention (nil window size)
     layer_types = spec.layer_types || generate_layer_types(spec.num_blocks)
 
-    attention_window_size = fn idx ->
-      case Enum.at(layer_types, idx, :sliding_attention) do
-        :full_attention -> nil
-        :sliding_attention -> {spec.attention_window_size, spec.attention_window_size}
-      end
-    end
-
-    # Per-layer rotary embedding base: local layers use rotary_embedding_base_local,
-    # global layers use rotary_embedding_base
-    rotary_embedding = fn idx ->
-      base =
-        case Enum.at(layer_types, idx, :sliding_attention) do
-          :full_attention -> spec.rotary_embedding_base
-          :sliding_attention -> spec.rotary_embedding_base_local
-        end
-
-      [
-        position_ids: position_ids,
-        max_positions: spec.max_positions,
-        base: base,
-        scaling_strategy: spec.rotary_embedding_scaling_strategy
-      ]
-    end
-
-    attention_scale = :math.pow(spec.attention_scale_base, -0.5)
-
-    Layers.Transformer.blocks(hidden_state,
-      attention_mask: attention_mask,
-      attention_head_mask: attention_head_mask,
-      cache: cache,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      num_key_value_heads: spec.num_key_value_heads,
-      hidden_size: spec.hidden_size,
-      attention_head_size: spec.attention_head_size,
-      attention_scale: attention_scale,
-      kernel_initializer: kernel_initializer(spec),
-      layer_norm:
-        &Layers.rms_norm(&1,
-          shift: 1.0,
-          name: &2,
-          epsilon: spec.layer_norm_epsilon,
-          upcast: :all
-        ),
-      ffn:
-        &gated_ffn(&1, spec.intermediate_size, spec.hidden_size,
-          name: &2,
-          activation: spec.activation
-        ),
-      block_type: &gemma3_block_impl(&1, &2, &3, spec),
-      causal: true,
-      rotary_embedding: rotary_embedding,
-      attention_window_size: attention_window_size,
-      query_norm: query_norm,
-      key_norm: key_norm,
-      query_use_bias: spec.use_attention_bias,
-      key_use_bias: spec.use_attention_bias,
-      value_use_bias: spec.use_attention_bias,
-      output_use_bias: spec.use_attention_bias,
-      name: join(name, "blocks")
-    )
-  end
-
-  # Custom block implementation for Gemma 3's unique normalization structure:
-  # - Post-attention norm BEFORE residual add
-  # - Pre/post FFN norms
-  defp gemma3_block_impl(hidden_state, steps, name, spec) do
-    # Pre-attention norm + attention (using provided steps)
-    shortcut = hidden_state
-
-    {hidden_state, attention_info} =
-      hidden_state
-      |> steps.self_attention_norm.()
-      |> steps.self_attention.()
-
-    # Post-attention norm BEFORE residual (Gemma 3 specific)
-    hidden_state =
-      Layers.rms_norm(hidden_state,
-        shift: 1.0,
-        name: join(name, "post_attention_norm"),
-        epsilon: spec.layer_norm_epsilon,
-        upcast: :all
-      )
-
-    hidden_state = Axon.add(shortcut, hidden_state)
-
-    # FFN with pre/post norms (Gemma 3 specific)
-    shortcut = hidden_state
-
-    hidden_state =
-      Layers.rms_norm(hidden_state,
-        shift: 1.0,
-        name: join(name, "pre_ffn_norm"),
-        epsilon: spec.layer_norm_epsilon,
-        upcast: :all
-      )
-
-    hidden_state = steps.ffn.(hidden_state)
-
-    hidden_state =
+    norm = fn hidden_state, name ->
       Layers.rms_norm(hidden_state,
         shift: 1.0,
-        name: join(name, "post_ffn_norm"),
         epsilon: spec.layer_norm_epsilon,
-        upcast: :all
+        upcast: :all,
+        name: name
       )
+    end
 
-    hidden_state = Axon.add(shortcut, hidden_state)
-
-    # Handle cross-attention (required by block interface but not used by Gemma 3)
-    {_hidden_state, cross_attention_info} =
-      steps.cross_attention_maybe.(hidden_state, fn _ ->
-        raise "cross attention not supported"
-      end)
-
-    {hidden_state, attention_info, cross_attention_info}
-  end
-
-  defp gated_ffn(hidden_state, intermediate_size, output_size, opts) do
-    name = opts[:name]
-    activation = opts[:activation]
-
-    intermediate =
-      Axon.dense(hidden_state, intermediate_size,
-        name: join(name, "intermediate"),
-        use_bias: false
-      )
-
-    gate = Axon.dense(hidden_state, intermediate_size, name: join(name, "gate"), use_bias: false)
-
-    hidden_state = Axon.multiply(intermediate, Layers.activation(gate, activation))
-
-    Axon.dense(hidden_state, output_size, name: join(name, "output"), use_bias: false)
+    Layers.Transformer.blocks(
+      hidden_state,
+      [
+        num_blocks: spec.num_blocks,
+        attention_mask: attention_mask,
+        attention_head_mask: attention_head_mask,
+        cache: cache,
+        name: join(name, "blocks")
+      ],
+      fn hidden_state, block ->
+        name = block.name
+
+        # Local layers use sliding window attention and a separate rotary base
+        {attention_window_size, rotary_embedding_base} =
+          case Enum.at(layer_types, block.index, :sliding_attention) do
+            :full_attention ->
+              {nil, spec.rotary_embedding_base}
+
+            :sliding_attention ->
+              {{spec.attention_window_size, spec.attention_window_size},
+               spec.rotary_embedding_base_local}
+          end
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, self_attention_cache} =
+          hidden_state
+          |> norm.(join(name, "self_attention_norm"))
+          |> Layers.Transformer.self_attention(block,
+            num_heads: spec.num_attention_heads,
+            num_key_value_heads: spec.num_key_value_heads,
+            hidden_size: spec.hidden_size,
+            attention_head_size: spec.attention_head_size,
+            attention_scale: :math.pow(spec.attention_scale_base, -0.5),
+            attention_window_size: attention_window_size,
+            causal: true,
+            query_norm:
+              &Layers.rms_norm(&1, shift: 1.0, epsilon: spec.layer_norm_epsilon, name: &2),
+            key_norm:
+              &Layers.rms_norm(&1, shift: 1.0, epsilon: spec.layer_norm_epsilon, name: &2),
+            rotary_embedding: [
+              position_ids: position_ids,
+              max_positions: spec.max_positions,
+              base: rotary_embedding_base,
+              scaling_strategy: spec.rotary_embedding_scaling_strategy
+            ],
+            query_use_bias: spec.use_attention_bias,
+            key_use_bias: spec.use_attention_bias,
+            value_use_bias: spec.use_attention_bias,
+            output_use_bias: spec.use_attention_bias,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "self_attention")
+          )
+
+        hidden_state =
+          hidden_state
+          |> norm.(join(name, "post_attention_norm"))
+          |> Axon.add(shortcut)
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> norm.(join(name, "pre_ffn_norm"))
+          |> Layers.Transformer.gated_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            name: join(name, "ffn")
+          )
+          |> norm.(join(name, "post_ffn_norm"))
+          |> Axon.add(shortcut)
+
+        %{
+          hidden_state: hidden_state,
+          attention: attention,
+          self_attention_cache: self_attention_cache
+        }
+      end
+    )
   end
 
   defp language_modeling_head(hidden_state, spec, opts) do
diff --git a/lib/bumblebee/text/gpt2.ex b/lib/bumblebee/text/gpt2.ex
index f23aa032..05e6bd06 100644
--- a/lib/bumblebee/text/gpt2.ex
+++ b/lib/bumblebee/text/gpt2.ex
@@ -394,35 +394,98 @@ defmodule Bumblebee.Text.Gpt2 do
     Layers.Transformer.blocks(
       hidden_state,
       [
+        num_blocks: spec.num_blocks,
         attention_mask: attention_mask,
         attention_head_mask: attention_head_mask,
+        cross_attention_head_mask: cross_attention_head_mask,
         cache: cache,
-        causal: true,
-        num_blocks: spec.num_blocks,
-        num_attention_heads: spec.num_attention_heads,
-        hidden_size: spec.hidden_size,
-        kernel_initializer: kernel_initializer(spec),
-        dropout_rate: spec.dropout_rate,
-        attention_dropout_rate: spec.attention_dropout_rate,
-        layer_norm: [
-          epsilon: spec.layer_norm_epsilon
-        ],
-        ffn: [
-          intermediate_size: spec.intermediate_size || 4 * spec.hidden_size,
-          activation: spec.activation
-        ],
-        block_type: :norm_first,
-        attention_scale: if(not spec.scale_attention_weights, do: 1),
         name: join(name, "blocks")
-      ] ++
-        if(spec.use_cross_attention,
-          do: [
-            cross_hidden_state: encoder_hidden_state,
-            cross_attention_mask: encoder_attention_mask,
-            cross_attention_head_mask: cross_attention_head_mask
-          ],
-          else: []
-        )
+      ],
+      fn hidden_state, block ->
+        name = block.name
+
+        attention_opts = [
+          num_heads: spec.num_attention_heads,
+          hidden_size: spec.hidden_size,
+          attention_scale: if(not spec.scale_attention_weights, do: 1),
+          kernel_initializer: kernel_initializer(spec),
+          dropout_rate: spec.attention_dropout_rate
+        ]
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, self_attention_cache} =
+          hidden_state
+          |> Axon.layer_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+          |> Layers.Transformer.self_attention(
+            block,
+            [causal: true, name: join(name, "self_attention")] ++ attention_opts
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: spec.dropout_rate)
+          |> Axon.add(shortcut)
+
+        {hidden_state, cross_attention, cross_attention_cache} =
+          if spec.use_cross_attention do
+            Layers.if_present encoder_hidden_state do
+              shortcut = hidden_state
+
+              {hidden_state, cross_attention, cross_attention_cache} =
+                hidden_state
+                |> Axon.layer_norm(
+                  epsilon: spec.layer_norm_epsilon,
+                  name: join(name, "cross_attention_norm")
+                )
+                |> Layers.Transformer.cross_attention(
+                  encoder_hidden_state,
+                  block,
+                  [
+                    attention_mask: encoder_attention_mask,
+                    name: join(name, "cross_attention")
+                  ] ++ attention_opts
+                )
+
+              hidden_state =
+                hidden_state
+                |> Axon.dropout(rate: spec.dropout_rate)
+                |> Axon.add(shortcut)
+
+              {hidden_state, cross_attention, cross_attention_cache}
+            else
+              {hidden_state, Layers.none(), block.cross_attention_cache}
+            end
+          else
+            {hidden_state, Layers.none(), block.cross_attention_cache}
+          end
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Axon.layer_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+          |> Layers.Transformer.basic_ffn(
+            spec.intermediate_size || 4 * spec.hidden_size,
+            spec.hidden_size,
+            activation: spec.activation,
+            dropout_rate: spec.dropout_rate,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+
+        %{
+          hidden_state: hidden_state,
+          attention: attention,
+          cross_attention: cross_attention,
+          self_attention_cache: self_attention_cache,
+          cross_attention_cache: cross_attention_cache
+        }
+      end
     )
   end
 
diff --git a/lib/bumblebee/text/gpt_big_code.ex b/lib/bumblebee/text/gpt_big_code.ex
index db4dd558..76684f5b 100644
--- a/lib/bumblebee/text/gpt_big_code.ex
+++ b/lib/bumblebee/text/gpt_big_code.ex
@@ -398,36 +398,99 @@ defmodule Bumblebee.Text.GptBigCode do
     Layers.Transformer.blocks(
       hidden_state,
       [
+        num_blocks: spec.num_blocks,
         attention_mask: attention_mask,
         attention_head_mask: attention_head_mask,
+        cross_attention_head_mask: cross_attention_head_mask,
         cache: cache,
-        causal: true,
-        num_blocks: spec.num_blocks,
-        num_attention_heads: spec.num_attention_heads,
-        num_key_value_heads: spec.num_key_value_heads,
-        hidden_size: spec.hidden_size,
-        kernel_initializer: kernel_initializer(spec),
-        dropout_rate: spec.dropout_rate,
-        attention_dropout_rate: spec.attention_dropout_rate,
-        layer_norm: [
-          epsilon: spec.layer_norm_epsilon
-        ],
-        ffn: [
-          intermediate_size: spec.intermediate_size || 4 * spec.hidden_size,
-          activation: spec.activation
-        ],
-        block_type: :norm_first,
-        attention_scale: if(not spec.scale_attention_weights, do: 1),
         name: join(name, "blocks")
-      ] ++
-        if(spec.use_cross_attention,
-          do: [
-            cross_hidden_state: encoder_hidden_state,
-            cross_attention_mask: encoder_attention_mask,
-            cross_attention_head_mask: cross_attention_head_mask
-          ],
-          else: []
-        )
+      ],
+      fn hidden_state, block ->
+        name = block.name
+
+        attention_opts = [
+          num_heads: spec.num_attention_heads,
+          num_key_value_heads: spec.num_key_value_heads,
+          hidden_size: spec.hidden_size,
+          attention_scale: if(not spec.scale_attention_weights, do: 1),
+          kernel_initializer: kernel_initializer(spec),
+          dropout_rate: spec.attention_dropout_rate
+        ]
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, self_attention_cache} =
+          hidden_state
+          |> Axon.layer_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+          |> Layers.Transformer.self_attention(
+            block,
+            [causal: true, name: join(name, "self_attention")] ++ attention_opts
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: spec.dropout_rate)
+          |> Axon.add(shortcut)
+
+        {hidden_state, cross_attention, cross_attention_cache} =
+          if spec.use_cross_attention do
+            Layers.if_present encoder_hidden_state do
+              shortcut = hidden_state
+
+              {hidden_state, cross_attention, cross_attention_cache} =
+                hidden_state
+                |> Axon.layer_norm(
+                  epsilon: spec.layer_norm_epsilon,
+                  name: join(name, "cross_attention_norm")
+                )
+                |> Layers.Transformer.cross_attention(
+                  encoder_hidden_state,
+                  block,
+                  [
+                    attention_mask: encoder_attention_mask,
+                    name: join(name, "cross_attention")
+                  ] ++ attention_opts
+                )
+
+              hidden_state =
+                hidden_state
+                |> Axon.dropout(rate: spec.dropout_rate)
+                |> Axon.add(shortcut)
+
+              {hidden_state, cross_attention, cross_attention_cache}
+            else
+              {hidden_state, Layers.none(), block.cross_attention_cache}
+            end
+          else
+            {hidden_state, Layers.none(), block.cross_attention_cache}
+          end
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Axon.layer_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+          |> Layers.Transformer.basic_ffn(
+            spec.intermediate_size || 4 * spec.hidden_size,
+            spec.hidden_size,
+            activation: spec.activation,
+            dropout_rate: spec.dropout_rate,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+
+        %{
+          hidden_state: hidden_state,
+          attention: attention,
+          cross_attention: cross_attention,
+          self_attention_cache: self_attention_cache,
+          cross_attention_cache: cross_attention_cache
+        }
+      end
     )
   end
 
diff --git a/lib/bumblebee/text/gpt_neo_x.ex b/lib/bumblebee/text/gpt_neo_x.ex
index 1f99a09c..71513942 100644
--- a/lib/bumblebee/text/gpt_neo_x.ex
+++ b/lib/bumblebee/text/gpt_neo_x.ex
@@ -334,32 +334,74 @@ defmodule Bumblebee.Text.GptNeoX do
        ) do
     name = opts[:name]
 
-    Layers.Transformer.blocks(hidden_state,
-      attention_mask: attention_mask,
-      attention_head_mask: attention_head_mask,
-      cache: cache,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: kernel_initializer(spec),
-      layer_norm: [
-        epsilon: spec.layer_norm_epsilon
+    Layers.Transformer.blocks(
+      hidden_state,
+      [
+        num_blocks: spec.num_blocks,
+        attention_mask: attention_mask,
+        attention_head_mask: attention_head_mask,
+        cache: cache,
+        name: join(name, "blocks")
       ],
-      ffn: [
-        intermediate_size: spec.intermediate_size
-      ],
-      block_type: if(spec.use_parallel_transformer_block, do: :parallel, else: :norm_first),
-      causal: true,
-      rotary_embedding: [
-        position_ids: position_ids,
-        percentage: spec.rotary_embedding_percentage,
-        base: spec.rotary_embedding_base
-      ],
-      query_use_bias: true,
-      key_use_bias: true,
-      value_use_bias: true,
-      output_use_bias: true,
-      name: join(name, "blocks")
+      fn hidden_state, block ->
+        name = block.name
+
+        {attention_output, attention, self_attention_cache} =
+          hidden_state
+          |> Axon.layer_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+          |> Layers.Transformer.self_attention(block,
+            num_heads: spec.num_attention_heads,
+            hidden_size: spec.hidden_size,
+            causal: true,
+            rotary_embedding: [
+              position_ids: position_ids,
+              percentage: spec.rotary_embedding_percentage,
+              base: spec.rotary_embedding_base
+            ],
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "self_attention")
+          )
+
+        hidden_state =
+          if spec.use_parallel_transformer_block do
+            ffn_output =
+              hidden_state
+              |> Axon.layer_norm(
+                epsilon: spec.layer_norm_epsilon,
+                name: join(name, "output_norm")
+              )
+              |> Layers.Transformer.basic_ffn(spec.intermediate_size, spec.hidden_size,
+                kernel_initializer: kernel_initializer(spec),
+                name: join(name, "ffn")
+              )
+
+            Axon.add([hidden_state, attention_output, ffn_output])
+          else
+            hidden_state = Axon.add(attention_output, hidden_state)
+
+            ffn_output =
+              hidden_state
+              |> Axon.layer_norm(
+                epsilon: spec.layer_norm_epsilon,
+                name: join(name, "output_norm")
+              )
+              |> Layers.Transformer.basic_ffn(spec.intermediate_size, spec.hidden_size,
+                kernel_initializer: kernel_initializer(spec),
+                name: join(name, "ffn")
+              )
+
+            Axon.add(ffn_output, hidden_state)
+          end
+
+        %{
+          hidden_state: hidden_state,
+          attention: attention,
+          self_attention_cache: self_attention_cache
+        }
+      end
     )
   end
 
diff --git a/lib/bumblebee/text/llama.ex b/lib/bumblebee/text/llama.ex
index 18141523..d6126d74 100644
--- a/lib/bumblebee/text/llama.ex
+++ b/lib/bumblebee/text/llama.ex
@@ -330,55 +330,68 @@ defmodule Bumblebee.Text.Llama do
        ) do
     name = opts[:name]
 
-    Layers.Transformer.blocks(hidden_state,
-      attention_mask: attention_mask,
-      attention_head_mask: attention_head_mask,
-      attention_head_size: spec.attention_head_size,
-      cache: cache,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      num_key_value_heads: spec.num_key_value_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: kernel_initializer(spec),
-      layer_norm: &Layers.rms_norm(&1, name: &2, epsilon: spec.layer_norm_epsilon),
-      ffn:
-        &gated_ffn(&1, spec.intermediate_size, spec.hidden_size,
-          name: &2,
-          activation: spec.activation
-        ),
-      block_type: :norm_first,
-      causal: true,
-      rotary_embedding: [
-        position_ids: position_ids,
-        max_positions: spec.max_positions,
-        base: spec.rotary_embedding_base,
-        scaling_strategy: spec.rotary_embedding_scaling_strategy
+    Layers.Transformer.blocks(
+      hidden_state,
+      [
+        num_blocks: spec.num_blocks,
+        attention_mask: attention_mask,
+        attention_head_mask: attention_head_mask,
+        cache: cache,
+        name: join(name, "blocks")
       ],
-      query_use_bias: false,
-      key_use_bias: false,
-      value_use_bias: false,
-      output_use_bias: false,
-      name: join(name, "blocks")
+      fn hidden_state, block ->
+        name = block.name
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, self_attention_cache} =
+          hidden_state
+          |> Layers.rms_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+          |> Layers.Transformer.self_attention(block,
+            num_heads: spec.num_attention_heads,
+            num_key_value_heads: spec.num_key_value_heads,
+            hidden_size: spec.hidden_size,
+            attention_head_size: spec.attention_head_size,
+            causal: true,
+            rotary_embedding: [
+              position_ids: position_ids,
+              max_positions: spec.max_positions,
+              base: spec.rotary_embedding_base,
+              scaling_strategy: spec.rotary_embedding_scaling_strategy
+            ],
+            query_use_bias: false,
+            key_use_bias: false,
+            value_use_bias: false,
+            output_use_bias: false,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "self_attention")
+          )
+
+        hidden_state = Axon.add(hidden_state, shortcut)
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Layers.rms_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+          |> Layers.Transformer.gated_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+
+        %{
+          hidden_state: hidden_state,
+          attention: attention,
+          self_attention_cache: self_attention_cache
+        }
+      end
     )
   end
 
-  defp gated_ffn(hidden_state, intermediate_size, output_size, opts) do
-    name = opts[:name]
-    activation = opts[:activation]
-
-    intermediate =
-      Axon.dense(hidden_state, intermediate_size,
-        name: join(name, "intermediate"),
-        use_bias: false
-      )
-
-    gate = Axon.dense(hidden_state, intermediate_size, name: join(name, "gate"), use_bias: false)
-
-    hidden_state = Axon.multiply(intermediate, Axon.activation(gate, activation))
-
-    Axon.dense(hidden_state, output_size, name: join(name, "output"), use_bias: false)
-  end
-
   defp language_modeling_head(hidden_state, spec, opts) do
     name = opts[:name]
 
diff --git a/lib/bumblebee/text/m2m100.ex b/lib/bumblebee/text/m2m100.ex
index 5f064741..0fe1ecde 100644
--- a/lib/bumblebee/text/m2m100.ex
+++ b/lib/bumblebee/text/m2m100.ex
@@ -350,24 +350,53 @@ defmodule Bumblebee.Text.M2m100 do
     name = opts[:name]
 
     encoder_outputs =
-      Layers.Transformer.blocks(hidden_state,
-        attention_mask: attention_mask,
-        attention_head_mask: attention_head_mask,
-        num_blocks: spec.encoder_num_blocks,
-        num_attention_heads: spec.encoder_num_attention_heads,
-        hidden_size: spec.hidden_size,
-        kernel_initializer: kernel_initializer(spec),
-        dropout_rate: spec.dropout_rate,
-        attention_dropout_rate: spec.attention_dropout_rate,
-        layer_norm: [
-          epsilon: 1.0e-5
-        ],
-        block_type: :norm_first,
-        ffn: [
-          intermediate_size: spec.encoder_intermediate_size,
-          activation: spec.activation
+      Layers.Transformer.blocks(
+        hidden_state,
+        [
+          num_blocks: spec.encoder_num_blocks,
+          attention_mask: attention_mask,
+          attention_head_mask: attention_head_mask,
+          name: join(name, "blocks")
         ],
-        name: join(name, "blocks")
+        fn hidden_state, block ->
+          name = block.name
+
+          shortcut = hidden_state
+
+          {hidden_state, attention, _cache} =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "self_attention_norm"))
+            |> Layers.Transformer.self_attention(block,
+              num_heads: spec.encoder_num_attention_heads,
+              hidden_size: spec.hidden_size,
+              kernel_initializer: kernel_initializer(spec),
+              dropout_rate: spec.attention_dropout_rate,
+              name: join(name, "self_attention")
+            )
+
+          hidden_state =
+            hidden_state
+            |> Axon.dropout(rate: spec.dropout_rate)
+            |> Axon.add(shortcut)
+
+          shortcut = hidden_state
+
+          hidden_state =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "output_norm"))
+            |> Layers.Transformer.basic_ffn(spec.encoder_intermediate_size, spec.hidden_size,
+              activation: spec.activation,
+              dropout_rate: spec.dropout_rate,
+              kernel_initializer: kernel_initializer(spec),
+              name: join(name, "ffn")
+            )
+            |> Axon.add(shortcut)
+
+          %{
+            hidden_state: hidden_state,
+            attention: attention
+          }
+        end
       )
 
     hidden_state = Axon.layer_norm(encoder_outputs.hidden_state, name: join(name, "norm"))
@@ -465,29 +494,88 @@ defmodule Bumblebee.Text.M2m100 do
     name = opts[:name]
 
     decoder_outputs =
-      Layers.Transformer.blocks(hidden_state,
-        attention_mask: attention_mask,
-        attention_head_mask: attention_head_mask,
-        cross_hidden_state: encoder_hidden_state,
-        cross_attention_mask: encoder_attention_mask,
-        cross_attention_head_mask: cross_attention_head_mask,
-        cache: cache,
-        causal: true,
-        num_blocks: spec.decoder_num_blocks,
-        num_attention_heads: spec.decoder_num_attention_heads,
-        hidden_size: spec.hidden_size,
-        kernel_initializer: kernel_initializer(spec),
-        dropout_rate: spec.dropout_rate,
-        attention_dropout_rate: spec.attention_dropout_rate,
-        layer_norm: [
-          epsilon: 1.0e-5
-        ],
-        block_type: :norm_first,
-        ffn: [
-          intermediate_size: spec.decoder_intermediate_size,
-          activation: spec.activation
+      Layers.Transformer.blocks(
+        hidden_state,
+        [
+          num_blocks: spec.decoder_num_blocks,
+          attention_mask: attention_mask,
+          attention_head_mask: attention_head_mask,
+          cross_attention_head_mask: cross_attention_head_mask,
+          cache: cache,
+          name: join(name, "blocks")
         ],
-        name: join(name, "blocks")
+        fn hidden_state, block ->
+          name = block.name
+
+          attention_opts = [
+            num_heads: spec.decoder_num_attention_heads,
+            hidden_size: spec.hidden_size,
+            kernel_initializer: kernel_initializer(spec),
+            dropout_rate: spec.attention_dropout_rate
+          ]
+
+          shortcut = hidden_state
+
+          {hidden_state, attention, self_attention_cache} =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "self_attention_norm"))
+            |> Layers.Transformer.self_attention(
+              block,
+              [causal: true, name: join(name, "self_attention")] ++ attention_opts
+            )
+
+          hidden_state =
+            hidden_state
+            |> Axon.dropout(rate: spec.dropout_rate)
+            |> Axon.add(shortcut)
+
+          {hidden_state, cross_attention, cross_attention_cache} =
+            Layers.if_present encoder_hidden_state do
+              shortcut = hidden_state
+
+              {hidden_state, cross_attention, cross_attention_cache} =
+                hidden_state
+                |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "cross_attention_norm"))
+                |> Layers.Transformer.cross_attention(
+                  encoder_hidden_state,
+                  block,
+                  [
+                    attention_mask: encoder_attention_mask,
+                    name: join(name, "cross_attention")
+                  ] ++ attention_opts
+                )
+
+              hidden_state =
+                hidden_state
+                |> Axon.dropout(rate: spec.dropout_rate)
+                |> Axon.add(shortcut)
+
+              {hidden_state, cross_attention, cross_attention_cache}
+            else
+              {hidden_state, Layers.none(), block.cross_attention_cache}
+            end
+
+          shortcut = hidden_state
+
+          hidden_state =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "output_norm"))
+            |> Layers.Transformer.basic_ffn(spec.decoder_intermediate_size, spec.hidden_size,
+              activation: spec.activation,
+              dropout_rate: spec.dropout_rate,
+              kernel_initializer: kernel_initializer(spec),
+              name: join(name, "ffn")
+            )
+            |> Axon.add(shortcut)
+
+          %{
+            hidden_state: hidden_state,
+            attention: attention,
+            cross_attention: cross_attention,
+            self_attention_cache: self_attention_cache,
+            cross_attention_cache: cross_attention_cache
+          }
+        end
       )
 
     hidden_state = Axon.layer_norm(decoder_outputs.hidden_state, name: join(name, "norm"))
diff --git a/lib/bumblebee/text/mbart.ex b/lib/bumblebee/text/mbart.ex
index 3b93a0a7..4120c786 100644
--- a/lib/bumblebee/text/mbart.ex
+++ b/lib/bumblebee/text/mbart.ex
@@ -515,24 +515,53 @@ defmodule Bumblebee.Text.Mbart do
     name = opts[:name]
 
     encoder_outputs =
-      Layers.Transformer.blocks(hidden_state,
-        attention_mask: attention_mask,
-        attention_head_mask: attention_head_mask,
-        num_blocks: spec.encoder_num_blocks,
-        num_attention_heads: spec.encoder_num_attention_heads,
-        hidden_size: spec.hidden_size,
-        kernel_initializer: kernel_initializer(spec),
-        dropout_rate: spec.dropout_rate,
-        attention_dropout_rate: spec.attention_dropout_rate,
-        layer_norm: [
-          epsilon: 1.0e-5
-        ],
-        block_type: :norm_first,
-        ffn: [
-          intermediate_size: spec.encoder_intermediate_size,
-          activation: spec.activation
+      Layers.Transformer.blocks(
+        hidden_state,
+        [
+          num_blocks: spec.encoder_num_blocks,
+          attention_mask: attention_mask,
+          attention_head_mask: attention_head_mask,
+          name: join(name, "blocks")
         ],
-        name: join(name, "blocks")
+        fn hidden_state, block ->
+          name = block.name
+
+          shortcut = hidden_state
+
+          {hidden_state, attention, _cache} =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "self_attention_norm"))
+            |> Layers.Transformer.self_attention(block,
+              num_heads: spec.encoder_num_attention_heads,
+              hidden_size: spec.hidden_size,
+              kernel_initializer: kernel_initializer(spec),
+              dropout_rate: spec.attention_dropout_rate,
+              name: join(name, "self_attention")
+            )
+
+          hidden_state =
+            hidden_state
+            |> Axon.dropout(rate: spec.dropout_rate)
+            |> Axon.add(shortcut)
+
+          shortcut = hidden_state
+
+          hidden_state =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "output_norm"))
+            |> Layers.Transformer.basic_ffn(spec.encoder_intermediate_size, spec.hidden_size,
+              activation: spec.activation,
+              dropout_rate: spec.dropout_rate,
+              kernel_initializer: kernel_initializer(spec),
+              name: join(name, "ffn")
+            )
+            |> Axon.add(shortcut)
+
+          %{
+            hidden_state: hidden_state,
+            attention: attention
+          }
+        end
       )
 
     hidden_state = Axon.layer_norm(encoder_outputs.hidden_state, name: join(name, "norm"))
@@ -606,29 +635,88 @@ defmodule Bumblebee.Text.Mbart do
     name = opts[:name]
 
     decoder_outputs =
-      Layers.Transformer.blocks(hidden_state,
-        attention_mask: attention_mask,
-        attention_head_mask: attention_head_mask,
-        cross_hidden_state: encoder_hidden_state,
-        cross_attention_mask: encoder_attention_mask,
-        cross_attention_head_mask: cross_attention_head_mask,
-        cache: cache,
-        causal: true,
-        num_blocks: spec.decoder_num_blocks,
-        num_attention_heads: spec.decoder_num_attention_heads,
-        hidden_size: spec.hidden_size,
-        kernel_initializer: kernel_initializer(spec),
-        dropout_rate: spec.dropout_rate,
-        attention_dropout_rate: spec.attention_dropout_rate,
-        layer_norm: [
-          epsilon: 1.0e-5
-        ],
-        block_type: :norm_first,
-        ffn: [
-          intermediate_size: spec.decoder_intermediate_size,
-          activation: spec.activation
+      Layers.Transformer.blocks(
+        hidden_state,
+        [
+          num_blocks: spec.decoder_num_blocks,
+          attention_mask: attention_mask,
+          attention_head_mask: attention_head_mask,
+          cross_attention_head_mask: cross_attention_head_mask,
+          cache: cache,
+          name: join(name, "blocks")
         ],
-        name: join(name, "blocks")
+        fn hidden_state, block ->
+          name = block.name
+
+          attention_opts = [
+            num_heads: spec.decoder_num_attention_heads,
+            hidden_size: spec.hidden_size,
+            kernel_initializer: kernel_initializer(spec),
+            dropout_rate: spec.attention_dropout_rate
+          ]
+
+          shortcut = hidden_state
+
+          {hidden_state, attention, self_attention_cache} =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "self_attention_norm"))
+            |> Layers.Transformer.self_attention(
+              block,
+              [causal: true, name: join(name, "self_attention")] ++ attention_opts
+            )
+
+          hidden_state =
+            hidden_state
+            |> Axon.dropout(rate: spec.dropout_rate)
+            |> Axon.add(shortcut)
+
+          {hidden_state, cross_attention, cross_attention_cache} =
+            Layers.if_present encoder_hidden_state do
+              shortcut = hidden_state
+
+              {hidden_state, cross_attention, cross_attention_cache} =
+                hidden_state
+                |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "cross_attention_norm"))
+                |> Layers.Transformer.cross_attention(
+                  encoder_hidden_state,
+                  block,
+                  [
+                    attention_mask: encoder_attention_mask,
+                    name: join(name, "cross_attention")
+                  ] ++ attention_opts
+                )
+
+              hidden_state =
+                hidden_state
+                |> Axon.dropout(rate: spec.dropout_rate)
+                |> Axon.add(shortcut)
+
+              {hidden_state, cross_attention, cross_attention_cache}
+            else
+              {hidden_state, Layers.none(), block.cross_attention_cache}
+            end
+
+          shortcut = hidden_state
+
+          hidden_state =
+            hidden_state
+            |> Axon.layer_norm(epsilon: 1.0e-5, name: join(name, "output_norm"))
+            |> Layers.Transformer.basic_ffn(spec.decoder_intermediate_size, spec.hidden_size,
+              activation: spec.activation,
+              dropout_rate: spec.dropout_rate,
+              kernel_initializer: kernel_initializer(spec),
+              name: join(name, "ffn")
+            )
+            |> Axon.add(shortcut)
+
+          %{
+            hidden_state: hidden_state,
+            attention: attention,
+            cross_attention: cross_attention,
+            self_attention_cache: self_attention_cache,
+            cross_attention_cache: cross_attention_cache
+          }
+        end
       )
 
     hidden_state = Axon.layer_norm(decoder_outputs.hidden_state, name: join(name, "norm"))
diff --git a/lib/bumblebee/text/mistral.ex b/lib/bumblebee/text/mistral.ex
index afbe26d9..e14e6b9f 100644
--- a/lib/bumblebee/text/mistral.ex
+++ b/lib/bumblebee/text/mistral.ex
@@ -315,55 +315,69 @@ defmodule Bumblebee.Text.Mistral do
        ) do
     name = opts[:name]
 
-    Layers.Transformer.blocks(hidden_state,
-      attention_mask: attention_mask,
-      attention_head_mask: attention_head_mask,
-      cache: cache,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      num_key_value_heads: spec.num_key_value_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: kernel_initializer(spec),
-      layer_norm: &Layers.rms_norm(&1, name: &2, epsilon: spec.layer_norm_epsilon),
-      ffn:
-        &gated_ffn(&1, spec.intermediate_size, spec.hidden_size,
-          name: &2,
-          activation: spec.activation
-        ),
-      block_type: :norm_first,
-      causal: true,
-      attention_window_size:
-        spec.attention_window_size && {spec.attention_window_size, spec.attention_window_size},
-      rotary_embedding: [
-        position_ids: position_ids,
-        max_positions: spec.max_positions,
-        base: spec.rotary_embedding_base
+    Layers.Transformer.blocks(
+      hidden_state,
+      [
+        num_blocks: spec.num_blocks,
+        attention_mask: attention_mask,
+        attention_head_mask: attention_head_mask,
+        cache: cache,
+        name: join(name, "blocks")
       ],
-      query_use_bias: false,
-      key_use_bias: false,
-      value_use_bias: false,
-      output_use_bias: false,
-      name: join(name, "blocks")
+      fn hidden_state, block ->
+        name = block.name
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, self_attention_cache} =
+          hidden_state
+          |> Layers.rms_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+          |> Layers.Transformer.self_attention(block,
+            num_heads: spec.num_attention_heads,
+            num_key_value_heads: spec.num_key_value_heads,
+            hidden_size: spec.hidden_size,
+            causal: true,
+            attention_window_size:
+              spec.attention_window_size &&
+                {spec.attention_window_size, spec.attention_window_size},
+            rotary_embedding: [
+              position_ids: position_ids,
+              max_positions: spec.max_positions,
+              base: spec.rotary_embedding_base
+            ],
+            query_use_bias: false,
+            key_use_bias: false,
+            value_use_bias: false,
+            output_use_bias: false,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "self_attention")
+          )
+
+        hidden_state = Axon.add(hidden_state, shortcut)
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Layers.rms_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+          |> Layers.Transformer.gated_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+
+        %{
+          hidden_state: hidden_state,
+          attention: attention,
+          self_attention_cache: self_attention_cache
+        }
+      end
     )
   end
 
-  defp gated_ffn(hidden_state, intermediate_size, output_size, opts) do
-    name = opts[:name]
-    activation = opts[:activation]
-
-    intermediate =
-      Axon.dense(hidden_state, intermediate_size,
-        name: join(name, "intermediate"),
-        use_bias: false
-      )
-
-    gate = Axon.dense(hidden_state, intermediate_size, name: join(name, "gate"), use_bias: false)
-
-    hidden_state = Axon.multiply(intermediate, Layers.activation(gate, activation))
-
-    Axon.dense(hidden_state, output_size, name: join(name, "output"), use_bias: false)
-  end
-
   defp language_modeling_head(hidden_state, spec, opts) do
     name = opts[:name]
 
diff --git a/lib/bumblebee/text/modernbert.ex b/lib/bumblebee/text/modernbert.ex
index bf621108..01048ccc 100644
--- a/lib/bumblebee/text/modernbert.ex
+++ b/lib/bumblebee/text/modernbert.ex
@@ -315,64 +315,79 @@ defmodule Bumblebee.Text.ModernBert do
 
     layer_types = spec.layer_types || generate_layer_types(spec.num_blocks)
 
-    attention_window_size = fn idx ->
-      case Enum.at(layer_types, idx, :sliding_attention) do
-        :full_attention ->
-          nil
-
-        :sliding_attention ->
-          half_window = div(spec.local_attention_window, 2)
-          {half_window, half_window}
-      end
-    end
-
-    rotary_embedding = fn idx ->
-      base =
-        case Enum.at(layer_types, idx, :sliding_attention) do
-          :full_attention -> spec.rotary_embedding_base
-          :sliding_attention -> spec.rotary_embedding_base_local
-        end
-
-      [
-        position_ids: position_ids,
-        max_positions: spec.max_positions,
-        base: base
-      ]
-    end
-
-    layer_norm = fn input, name ->
-      if String.ends_with?(name, "encoder.blocks.0.self_attention_norm") do
-        # The first self-attention norm is skipped.
-        input
-      else
-        layer_norm(input, epsilon: spec.layer_norm_epsilon, name: name)
-      end
-    end
-
     outputs =
-      Layers.Transformer.blocks(hidden_state,
-        attention_mask: attention_mask,
-        attention_head_mask: attention_head_mask,
-        num_blocks: spec.num_blocks,
-        num_attention_heads: spec.num_attention_heads,
-        hidden_size: spec.hidden_size,
-        kernel_initializer: kernel_initializer(spec),
-        dropout_rate: spec.dropout_rate,
-        attention_dropout_rate: spec.attention_dropout_rate,
-        layer_norm: layer_norm,
-        ffn:
-          &gated_ffn(&1, spec.intermediate_size, spec.hidden_size,
-            activation: spec.activation,
-            name: &2
-          ),
-        block_type: :norm_first,
-        rotary_embedding: rotary_embedding,
-        attention_window_size: attention_window_size,
-        query_use_bias: false,
-        key_use_bias: false,
-        value_use_bias: false,
-        output_use_bias: false,
-        name: join(name, "blocks")
+      Layers.Transformer.blocks(
+        hidden_state,
+        [
+          num_blocks: spec.num_blocks,
+          attention_mask: attention_mask,
+          attention_head_mask: attention_head_mask,
+          name: join(name, "blocks")
+        ],
+        fn hidden_state, block ->
+          name = block.name
+
+          # Local layers use sliding window attention and a separate rotary base
+          {attention_window_size, rotary_embedding_base} =
+            case Enum.at(layer_types, block.index, :sliding_attention) do
+              :full_attention ->
+                {nil, spec.rotary_embedding_base}
+
+              :sliding_attention ->
+                half_window = div(spec.local_attention_window, 2)
+                {{half_window, half_window}, spec.rotary_embedding_base_local}
+            end
+
+          shortcut = hidden_state
+
+          # The first block has no attention norm
+          hidden_state =
+            if block.index == 0 do
+              hidden_state
+            else
+              layer_norm(hidden_state,
+                epsilon: spec.layer_norm_epsilon,
+                name: join(name, "self_attention_norm")
+              )
+            end
+
+          {hidden_state, attention, _self_attention_cache} =
+            Layers.Transformer.self_attention(hidden_state, block,
+              num_heads: spec.num_attention_heads,
+              hidden_size: spec.hidden_size,
+              attention_window_size: attention_window_size,
+              rotary_embedding: [
+                position_ids: position_ids,
+                max_positions: spec.max_positions,
+                base: rotary_embedding_base
+              ],
+              query_use_bias: false,
+              key_use_bias: false,
+              value_use_bias: false,
+              output_use_bias: false,
+              kernel_initializer: kernel_initializer(spec),
+              dropout_rate: spec.attention_dropout_rate,
+              name: join(name, "self_attention")
+            )
+
+          hidden_state =
+            hidden_state
+            |> Axon.dropout(rate: spec.dropout_rate)
+            |> Axon.add(shortcut)
+
+          shortcut = hidden_state
+
+          hidden_state =
+            hidden_state
+            |> layer_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+            |> gated_ffn(spec.intermediate_size, spec.hidden_size,
+              activation: spec.activation,
+              name: join(name, "ffn")
+            )
+            |> Axon.add(shortcut)
+
+          %{hidden_state: hidden_state, attention: attention}
+        end
       )
 
     hidden_state =
@@ -388,6 +403,8 @@ defmodule Bumblebee.Text.ModernBert do
     }
   end
 
+  # Unlike the shared gated FFN, the activation is applied to the
+  # intermediate projection, not the gate
   defp gated_ffn(hidden_state, intermediate_size, output_size, opts) do
     name = opts[:name]
     activation = opts[:activation]
diff --git a/lib/bumblebee/text/modernbert_decoder.ex b/lib/bumblebee/text/modernbert_decoder.ex
index b061bccb..dcc549b7 100644
--- a/lib/bumblebee/text/modernbert_decoder.ex
+++ b/lib/bumblebee/text/modernbert_decoder.ex
@@ -287,46 +287,77 @@ defmodule Bumblebee.Text.ModernBertDecoder do
        ) do
     name = opts[:name]
 
-    layer_norm = fn input, name ->
-      if String.ends_with?(name, "decoder.blocks.0.self_attention_norm") do
-        # The first self-attention norm is skipped.
-        input
-      else
-        layer_norm(input, epsilon: spec.layer_norm_epsilon, name: name)
-      end
-    end
-
-    Layers.Transformer.blocks(hidden_state,
-      attention_mask: attention_mask,
-      attention_head_mask: attention_head_mask,
-      cache: cache,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: kernel_initializer(spec),
-      dropout_rate: spec.dropout_rate,
-      attention_dropout_rate: spec.attention_dropout_rate,
-      layer_norm: layer_norm,
-      ffn:
-        &gated_ffn(&1, spec.intermediate_size, spec.hidden_size,
-          activation: spec.activation,
-          name: &2
-        ),
-      block_type: :norm_first,
-      causal: true,
-      rotary_embedding: [
-        position_ids: position_ids,
-        max_positions: spec.max_positions,
-        base: spec.rotary_embedding_base
+    Layers.Transformer.blocks(
+      hidden_state,
+      [
+        num_blocks: spec.num_blocks,
+        attention_mask: attention_mask,
+        attention_head_mask: attention_head_mask,
+        cache: cache,
+        name: join(name, "blocks")
       ],
-      query_use_bias: false,
-      key_use_bias: false,
-      value_use_bias: false,
-      output_use_bias: false,
-      name: join(name, "blocks")
+      fn hidden_state, block ->
+        name = block.name
+
+        shortcut = hidden_state
+
+        # The first block has no attention norm
+        hidden_state =
+          if block.index == 0 do
+            hidden_state
+          else
+            layer_norm(hidden_state,
+              epsilon: spec.layer_norm_epsilon,
+              name: join(name, "self_attention_norm")
+            )
+          end
+
+        {hidden_state, attention, self_attention_cache} =
+          Layers.Transformer.self_attention(hidden_state, block,
+            num_heads: spec.num_attention_heads,
+            hidden_size: spec.hidden_size,
+            causal: true,
+            rotary_embedding: [
+              position_ids: position_ids,
+              max_positions: spec.max_positions,
+              base: spec.rotary_embedding_base
+            ],
+            query_use_bias: false,
+            key_use_bias: false,
+            value_use_bias: false,
+            output_use_bias: false,
+            kernel_initializer: kernel_initializer(spec),
+            dropout_rate: spec.attention_dropout_rate,
+            name: join(name, "self_attention")
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: spec.dropout_rate)
+          |> Axon.add(shortcut)
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> layer_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+          |> gated_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+
+        %{
+          hidden_state: hidden_state,
+          attention: attention,
+          self_attention_cache: self_attention_cache
+        }
+      end
     )
   end
 
+  # Unlike the shared gated FFN, the activation is applied to the
+  # intermediate projection, not the gate
   defp gated_ffn(hidden_state, intermediate_size, output_size, opts) do
     name = opts[:name]
     activation = opts[:activation]
diff --git a/lib/bumblebee/text/mpnet.ex b/lib/bumblebee/text/mpnet.ex
index 817aee00..9751a43f 100644
--- a/lib/bumblebee/text/mpnet.ex
+++ b/lib/bumblebee/text/mpnet.ex
@@ -364,30 +364,63 @@ defmodule Bumblebee.Text.MpNet do
   defp encoder(hidden_state, attention_mask, attention_head_mask, spec, opts) do
     name = opts[:name]
 
-    Layers.Transformer.blocks(
-      hidden_state,
-      attention_mask: attention_mask,
-      attention_head_mask: attention_head_mask,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: kernel_initializer(spec),
-      dropout_rate: spec.dropout_rate,
-      attention_dropout_rate: spec.attention_dropout_rate,
-      layer_norm: [
-        epsilon: spec.layer_norm_epsilon
-      ],
-      ffn: [
-        intermediate_size: spec.intermediate_size,
-        activation: spec.activation
-      ],
-      attention_relative_bias: [
+    # The bias is computed once and shared by all blocks
+    attention_relative_bias =
+      Layers.relative_attention_bias(hidden_state, hidden_state, Layers.none(), Layers.none(),
         bidirectional: true,
         num_buckets: spec.relative_attention_num_buckets,
-        max_distance: 128
+        max_distance: 128,
+        num_heads: spec.num_attention_heads,
+        name: join(name, "blocks.0.self_attention.relative_attention_bias")
+      )
+
+    Layers.Transformer.blocks(
+      hidden_state,
+      [
+        num_blocks: spec.num_blocks,
+        attention_mask: attention_mask,
+        attention_head_mask: attention_head_mask,
+        name: join(name, "blocks")
       ],
-      share_attention_relative_bias: true,
-      name: join(name, "blocks")
+      fn hidden_state, block ->
+        name = block.name
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, _cache} =
+          Layers.Transformer.self_attention(hidden_state, block,
+            num_heads: spec.num_attention_heads,
+            hidden_size: spec.hidden_size,
+            kernel_initializer: kernel_initializer(spec),
+            dropout_rate: spec.attention_dropout_rate,
+            attention_relative_bias: attention_relative_bias,
+            name: join(name, "self_attention")
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: spec.dropout_rate)
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Layers.Transformer.basic_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            dropout_rate: spec.dropout_rate,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+
+        %{hidden_state: hidden_state, attention: attention}
+      end
     )
   end
 
diff --git a/lib/bumblebee/text/nomic_bert.ex b/lib/bumblebee/text/nomic_bert.ex
index 98a4ef84..a679f91c 100644
--- a/lib/bumblebee/text/nomic_bert.ex
+++ b/lib/bumblebee/text/nomic_bert.ex
@@ -212,34 +212,60 @@ defmodule Bumblebee.Text.NomicBert do
   defp encoder(hidden_state, position_ids, attention_mask, attention_head_mask, spec, opts) do
     name = opts[:name]
 
-    Layers.Transformer.blocks(hidden_state,
-      attention_mask: attention_mask,
-      attention_head_mask: attention_head_mask,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: kernel_initializer(spec),
-      layer_norm: [epsilon: spec.layer_norm_epsilon],
-      ffn:
-        &gated_ffn(&1, intermediate_size(spec), spec.hidden_size,
-          name: &2,
-          activation: spec.activation,
-          gate_use_bias: spec.ffn_gate_bias,
-          output_use_bias: spec.ffn_output_bias
-        ),
-      block_type: :standard,
-      causal: false,
-      rotary_embedding: [
-        position_ids: position_ids,
-        max_positions: spec.max_positions,
-        base: spec.rotary_embedding_base,
-        percentage: spec.rotary_embedding_percentage
+    Layers.Transformer.blocks(
+      hidden_state,
+      [
+        num_blocks: spec.num_blocks,
+        attention_mask: attention_mask,
+        attention_head_mask: attention_head_mask,
+        name: join(name, "blocks")
       ],
-      query_use_bias: false,
-      key_use_bias: false,
-      value_use_bias: false,
-      output_use_bias: false,
-      name: join(name, "blocks")
+      fn hidden_state, block ->
+        name = block.name
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, _cache} =
+          Layers.Transformer.self_attention(hidden_state, block,
+            num_heads: spec.num_attention_heads,
+            hidden_size: spec.hidden_size,
+            kernel_initializer: kernel_initializer(spec),
+            query_use_bias: false,
+            key_use_bias: false,
+            value_use_bias: false,
+            output_use_bias: false,
+            rotary_embedding: [
+              position_ids: position_ids,
+              max_positions: spec.max_positions,
+              base: spec.rotary_embedding_base,
+              percentage: spec.rotary_embedding_percentage
+            ],
+            name: join(name, "self_attention")
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> gated_ffn(intermediate_size(spec), spec.hidden_size,
+            activation: spec.activation,
+            gate_use_bias: spec.ffn_gate_bias,
+            output_use_bias: spec.ffn_output_bias,
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+
+        %{hidden_state: hidden_state, attention: attention}
+      end
     )
   end
 
diff --git a/lib/bumblebee/text/phi.ex b/lib/bumblebee/text/phi.ex
index 0d7b4250..dada527a 100644
--- a/lib/bumblebee/text/phi.ex
+++ b/lib/bumblebee/text/phi.ex
@@ -342,58 +342,61 @@ defmodule Bumblebee.Text.Phi do
        ) do
     name = opts[:name]
 
-    Layers.Transformer.blocks(hidden_state,
-      attention_mask: attention_mask,
-      attention_head_mask: attention_head_mask,
-      cache: cache,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      num_key_value_heads: spec.num_key_value_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: kernel_initializer(spec),
-      layer_norm: [
-        epsilon: spec.layer_norm_epsilon
-      ],
-      ffn: [
-        intermediate_size: spec.intermediate_size,
-        activation: spec.activation
+    Layers.Transformer.blocks(
+      hidden_state,
+      [
+        num_blocks: spec.num_blocks,
+        attention_mask: attention_mask,
+        attention_head_mask: attention_head_mask,
+        cache: cache,
+        name: join(name, "blocks")
       ],
-      block_type: &block_impl/3,
-      causal: true,
-      rotary_embedding: [
-        position_ids: position_ids,
-        max_positions: spec.max_positions,
-        base: spec.rotary_embedding_base,
-        percentage: spec.rotary_embedding_percentage
-      ],
-      query_use_bias: true,
-      key_use_bias: true,
-      value_use_bias: true,
-      output_use_bias: true,
-      name: join(name, "blocks")
+      fn hidden_state, block ->
+        name = block.name
+
+        shortcut = hidden_state
+
+        hidden_state =
+          Axon.layer_norm(hidden_state,
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+
+        # Attention and FFN run in parallel on the same input
+        {attention_hidden_state, attention, self_attention_cache} =
+          Layers.Transformer.self_attention(hidden_state, block,
+            num_heads: spec.num_attention_heads,
+            num_key_value_heads: spec.num_key_value_heads,
+            hidden_size: spec.hidden_size,
+            causal: true,
+            rotary_embedding: [
+              position_ids: position_ids,
+              max_positions: spec.max_positions,
+              base: spec.rotary_embedding_base,
+              percentage: spec.rotary_embedding_percentage
+            ],
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "self_attention")
+          )
+
+        ffn_hidden_state =
+          Layers.Transformer.basic_ffn(hidden_state, spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "ffn")
+          )
+
+        hidden_state = Axon.add([shortcut, attention_hidden_state, ffn_hidden_state])
+
+        %{
+          hidden_state: hidden_state,
+          attention: attention,
+          self_attention_cache: self_attention_cache
+        }
+      end
     )
   end
 
-  # :parallel block with attention norm applied earlier and without ffn norm
-  defp block_impl(hidden_state, steps, _name) do
-    shortcut = hidden_state
-
-    hidden_state = steps.self_attention_norm.(hidden_state)
-
-    {attention_hidden_state, attention_info} = steps.self_attention.(hidden_state)
-
-    {_hidden_state, cross_attention_info} =
-      steps.cross_attention_maybe.(hidden_state, fn _hidden_state ->
-        raise "cross attention not supported"
-      end)
-
-    ffn_hidden_state = steps.ffn.(hidden_state)
-
-    hidden_state = Axon.add([shortcut, attention_hidden_state, ffn_hidden_state])
-
-    {hidden_state, attention_info, cross_attention_info}
-  end
-
   defp language_modeling_head(hidden_state, spec, opts) do
     name = opts[:name]
 
diff --git a/lib/bumblebee/text/phi3.ex b/lib/bumblebee/text/phi3.ex
index 9348ad5b..c09e083a 100644
--- a/lib/bumblebee/text/phi3.ex
+++ b/lib/bumblebee/text/phi3.ex
@@ -356,57 +356,71 @@ defmodule Bumblebee.Text.Phi3 do
        ) do
     name = opts[:name]
 
-    Layers.Transformer.blocks(hidden_state,
-      attention_mask: attention_mask,
-      attention_head_mask: attention_head_mask,
-      cache: cache,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      num_key_value_heads: spec.num_key_value_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: kernel_initializer(spec),
-      layer_norm: &Layers.rms_norm(&1, name: &2, epsilon: spec.layer_norm_epsilon),
-      ffn:
-        &gated_ffn(&1, spec.intermediate_size, spec.hidden_size,
-          name: &2,
-          activation: spec.activation
-        ),
-      block_type: :norm_first,
-      causal: true,
-      attention_window_size:
-        spec.attention_window_size && {spec.attention_window_size, spec.attention_window_size},
-      rotary_embedding: [
-        position_ids: position_ids,
-        max_positions: spec.max_positions,
-        base: spec.rotary_embedding_base,
-        percentage: spec.rotary_embedding_percentage,
-        scaling_strategy: spec.rotary_embedding_scaling_strategy
+    Layers.Transformer.blocks(
+      hidden_state,
+      [
+        num_blocks: spec.num_blocks,
+        attention_mask: attention_mask,
+        attention_head_mask: attention_head_mask,
+        cache: cache,
+        name: join(name, "blocks")
       ],
-      query_use_bias: false,
-      key_use_bias: false,
-      value_use_bias: false,
-      output_use_bias: false,
-      name: join(name, "blocks")
+      fn hidden_state, block ->
+        name = block.name
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, self_attention_cache} =
+          hidden_state
+          |> Layers.rms_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+          |> Layers.Transformer.self_attention(block,
+            num_heads: spec.num_attention_heads,
+            num_key_value_heads: spec.num_key_value_heads,
+            hidden_size: spec.hidden_size,
+            causal: true,
+            attention_window_size:
+              spec.attention_window_size &&
+                {spec.attention_window_size, spec.attention_window_size},
+            rotary_embedding: [
+              position_ids: position_ids,
+              max_positions: spec.max_positions,
+              base: spec.rotary_embedding_base,
+              percentage: spec.rotary_embedding_percentage,
+              scaling_strategy: spec.rotary_embedding_scaling_strategy
+            ],
+            query_use_bias: false,
+            key_use_bias: false,
+            value_use_bias: false,
+            output_use_bias: false,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "self_attention")
+          )
+
+        hidden_state = Axon.add(hidden_state, shortcut)
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Layers.rms_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+          |> Layers.Transformer.gated_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+
+        %{
+          hidden_state: hidden_state,
+          attention: attention,
+          self_attention_cache: self_attention_cache
+        }
+      end
     )
   end
 
-  defp gated_ffn(hidden_state, intermediate_size, output_size, opts) do
-    name = opts[:name]
-    activation = opts[:activation]
-
-    intermediate =
-      Axon.dense(hidden_state, intermediate_size,
-        name: join(name, "intermediate"),
-        use_bias: false
-      )
-
-    gate = Axon.dense(hidden_state, intermediate_size, name: join(name, "gate"), use_bias: false)
-
-    hidden_state = Axon.multiply(intermediate, Axon.activation(gate, activation))
-
-    Axon.dense(hidden_state, output_size, name: join(name, "output"), use_bias: false)
-  end
-
   defp language_modeling_head(hidden_state, spec, opts) do
     name = opts[:name]
 
diff --git a/lib/bumblebee/text/qwen3.ex b/lib/bumblebee/text/qwen3.ex
index 568dd4e6..a89e0edc 100644
--- a/lib/bumblebee/text/qwen3.ex
+++ b/lib/bumblebee/text/qwen3.ex
@@ -332,68 +332,76 @@ defmodule Bumblebee.Text.Qwen3 do
        ) do
     name = opts[:name]
 
-    # Build query and key normalization functions for Qwen3
-    query_norm =
+    # Query and key norm, applied per head
+    qk_norm =
       if spec.use_qk_norm do
         &Layers.rms_norm(&1, epsilon: spec.layer_norm_epsilon, channel_index: -1, name: &2)
       end
 
-    key_norm =
-      if spec.use_qk_norm do
-        &Layers.rms_norm(&1, epsilon: spec.layer_norm_epsilon, channel_index: -1, name: &2)
-      end
-
-    Layers.Transformer.blocks(hidden_state,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      num_key_value_heads: spec.num_key_value_heads,
-      hidden_size: spec.hidden_size,
-      attention_head_size: spec.attention_head_size,
-      kernel_initializer: kernel_initializer(spec),
-      query_use_bias: false,
-      key_use_bias: false,
-      value_use_bias: false,
-      output_use_bias: false,
-      block_type: :norm_first,
-      attention_mask: attention_mask,
-      attention_head_mask: attention_head_mask,
-      cache: cache,
-      causal: true,
-      layer_norm: &Layers.rms_norm(&1, epsilon: spec.layer_norm_epsilon, name: &2),
-      ffn:
-        &gated_ffn(&1, spec.intermediate_size, spec.hidden_size,
-          name: &2,
-          activation: spec.activation
-        ),
-      rotary_embedding: [
-        position_ids: position_ids,
-        max_positions: spec.max_positions,
-        base: spec.rotary_embedding_base,
-        scaling_strategy: spec.rotary_embedding_scaling_strategy
+    Layers.Transformer.blocks(
+      hidden_state,
+      [
+        num_blocks: spec.num_blocks,
+        attention_mask: attention_mask,
+        attention_head_mask: attention_head_mask,
+        cache: cache,
+        name: join(name, "blocks")
       ],
-      query_norm: query_norm,
-      key_norm: key_norm,
-      name: join(name, "blocks")
+      fn hidden_state, block ->
+        name = block.name
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, self_attention_cache} =
+          hidden_state
+          |> Layers.rms_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+          |> Layers.Transformer.self_attention(block,
+            num_heads: spec.num_attention_heads,
+            num_key_value_heads: spec.num_key_value_heads,
+            hidden_size: spec.hidden_size,
+            attention_head_size: spec.attention_head_size,
+            causal: true,
+            query_norm: qk_norm,
+            key_norm: qk_norm,
+            rotary_embedding: [
+              position_ids: position_ids,
+              max_positions: spec.max_positions,
+              base: spec.rotary_embedding_base,
+              scaling_strategy: spec.rotary_embedding_scaling_strategy
+            ],
+            query_use_bias: false,
+            key_use_bias: false,
+            value_use_bias: false,
+            output_use_bias: false,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "self_attention")
+          )
+
+        hidden_state = Axon.add(hidden_state, shortcut)
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Layers.rms_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+          |> Layers.Transformer.gated_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+
+        %{
+          hidden_state: hidden_state,
+          attention: attention,
+          self_attention_cache: self_attention_cache
+        }
+      end
     )
   end
 
-  defp gated_ffn(hidden_state, intermediate_size, output_size, opts) do
-    name = opts[:name]
-    activation = opts[:activation]
-
-    intermediate =
-      Axon.dense(hidden_state, intermediate_size,
-        name: join(name, "intermediate"),
-        use_bias: false
-      )
-
-    gate = Axon.dense(hidden_state, intermediate_size, name: join(name, "gate"), use_bias: false)
-
-    hidden_state = Axon.multiply(intermediate, Axon.activation(gate, activation))
-
-    Axon.dense(hidden_state, output_size, name: join(name, "output"), use_bias: false)
-  end
-
   defp language_modeling_head(hidden_state, spec, opts) do
     name = opts[:name]
 
diff --git a/lib/bumblebee/text/roberta.ex b/lib/bumblebee/text/roberta.ex
index 4db674f3..e70d7226 100644
--- a/lib/bumblebee/text/roberta.ex
+++ b/lib/bumblebee/text/roberta.ex
@@ -474,33 +474,95 @@ defmodule Bumblebee.Text.Roberta do
     Layers.Transformer.blocks(
       hidden_state,
       [
+        num_blocks: spec.num_blocks,
         attention_mask: attention_mask,
         attention_head_mask: attention_head_mask,
+        cross_attention_head_mask: cross_attention_head_mask,
         cache: cache,
-        causal: decoder?,
-        num_blocks: spec.num_blocks,
-        num_attention_heads: spec.num_attention_heads,
-        hidden_size: spec.hidden_size,
-        kernel_initializer: kernel_initializer(spec),
-        dropout_rate: spec.dropout_rate,
-        attention_dropout_rate: spec.attention_dropout_rate,
-        layer_norm: [
-          epsilon: spec.layer_norm_epsilon
-        ],
-        ffn: [
-          intermediate_size: spec.intermediate_size,
-          activation: spec.activation
-        ],
         name: join(name, "blocks")
-      ] ++
-        if(cross_attention?,
-          do: [
-            cross_hidden_state: encoder_hidden_state,
-            cross_attention_mask: encoder_attention_mask,
-            cross_attention_head_mask: cross_attention_head_mask
-          ],
-          else: []
-        )
+      ],
+      fn hidden_state, block ->
+        name = block.name
+
+        attention_opts = [
+          num_heads: spec.num_attention_heads,
+          hidden_size: spec.hidden_size,
+          kernel_initializer: kernel_initializer(spec),
+          dropout_rate: spec.attention_dropout_rate
+        ]
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, self_attention_cache} =
+          Layers.Transformer.self_attention(
+            hidden_state,
+            block,
+            [causal: decoder?, name: join(name, "self_attention")] ++ attention_opts
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: spec.dropout_rate)
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+
+        {hidden_state, cross_attention, cross_attention_cache} =
+          if cross_attention? do
+            Layers.if_present encoder_hidden_state do
+              shortcut = hidden_state
+
+              {hidden_state, cross_attention, cross_attention_cache} =
+                Layers.Transformer.cross_attention(
+                  hidden_state,
+                  encoder_hidden_state,
+                  block,
+                  [
+                    attention_mask: encoder_attention_mask,
+                    name: join(name, "cross_attention")
+                  ] ++ attention_opts
+                )
+
+              hidden_state =
+                hidden_state
+                |> Axon.dropout(rate: spec.dropout_rate)
+                |> Axon.add(shortcut)
+                |> Axon.layer_norm(
+                  epsilon: spec.layer_norm_epsilon,
+                  name: join(name, "cross_attention_norm")
+                )
+
+              {hidden_state, cross_attention, cross_attention_cache}
+            else
+              {hidden_state, Layers.none(), block.cross_attention_cache}
+            end
+          else
+            {hidden_state, Layers.none(), block.cross_attention_cache}
+          end
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Layers.Transformer.basic_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            dropout_rate: spec.dropout_rate,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+          |> Axon.layer_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+
+        %{
+          hidden_state: hidden_state,
+          attention: attention,
+          cross_attention: cross_attention,
+          self_attention_cache: self_attention_cache,
+          cross_attention_cache: cross_attention_cache
+        }
+      end
     )
   end
 
diff --git a/lib/bumblebee/text/smollm3.ex b/lib/bumblebee/text/smollm3.ex
index 646f5951..cba0d7ad 100644
--- a/lib/bumblebee/text/smollm3.ex
+++ b/lib/bumblebee/text/smollm3.ex
@@ -404,72 +404,75 @@ defmodule Bumblebee.Text.SmolLm3 do
        ) do
     name = opts[:name]
 
-    rotary_embedding_config = [
-      position_ids: position_ids,
-      max_positions: spec.max_positions,
-      base: spec.rotary_embedding_base,
-      scaling_strategy: spec.rotary_embedding_scaling_strategy
-    ]
-
-    rotary_embedding =
-      case spec.rotary_embedding_enabled do
-        nil ->
-          rotary_embedding_config
-
-        rotary_embedding_enabled ->
-          fn layer_index ->
-            if Enum.at(rotary_embedding_enabled, layer_index) do
-              rotary_embedding_config
-            else
-              nil
-            end
+    Layers.Transformer.blocks(
+      hidden_state,
+      [
+        num_blocks: spec.num_blocks,
+        attention_mask: attention_mask,
+        attention_head_mask: attention_head_mask,
+        cache: cache,
+        name: join(name, "blocks")
+      ],
+      fn hidden_state, block ->
+        name = block.name
+
+        # Some blocks may skip rotary embedding (no-rope layers)
+        rotary_embedding =
+          if spec.rotary_embedding_enabled == nil or
+               Enum.at(spec.rotary_embedding_enabled, block.index) do
+            [
+              position_ids: position_ids,
+              max_positions: spec.max_positions,
+              base: spec.rotary_embedding_base,
+              scaling_strategy: spec.rotary_embedding_scaling_strategy
+            ]
           end
-      end
 
-    Layers.Transformer.blocks(hidden_state,
-      attention_mask: attention_mask,
-      attention_head_mask: attention_head_mask,
-      attention_head_size: spec.attention_head_size,
-      cache: cache,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      num_key_value_heads: spec.num_key_value_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: kernel_initializer(spec),
-      layer_norm: &Layers.rms_norm(&1, name: &2, epsilon: spec.layer_norm_epsilon),
-      ffn:
-        &gated_ffn(&1, spec.intermediate_size, spec.hidden_size,
-          name: &2,
-          activation: spec.activation
-        ),
-      block_type: :norm_first,
-      causal: true,
-      rotary_embedding: rotary_embedding,
-      query_use_bias: false,
-      key_use_bias: false,
-      value_use_bias: false,
-      output_use_bias: false,
-      name: join(name, "blocks")
+        shortcut = hidden_state
+
+        {hidden_state, attention, self_attention_cache} =
+          hidden_state
+          |> Layers.rms_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+          |> Layers.Transformer.self_attention(block,
+            num_heads: spec.num_attention_heads,
+            num_key_value_heads: spec.num_key_value_heads,
+            hidden_size: spec.hidden_size,
+            attention_head_size: spec.attention_head_size,
+            causal: true,
+            rotary_embedding: rotary_embedding,
+            query_use_bias: false,
+            key_use_bias: false,
+            value_use_bias: false,
+            output_use_bias: false,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "self_attention")
+          )
+
+        hidden_state = Axon.add(hidden_state, shortcut)
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Layers.rms_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+          |> Layers.Transformer.gated_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+
+        %{
+          hidden_state: hidden_state,
+          attention: attention,
+          self_attention_cache: self_attention_cache
+        }
+      end
     )
   end
 
-  defp gated_ffn(hidden_state, intermediate_size, output_size, opts) do
-    name = opts[:name]
-    activation = opts[:activation]
-
-    intermediate =
-      Axon.dense(hidden_state, intermediate_size,
-        name: join(name, "intermediate"),
-        use_bias: false
-      )
-
-    gate = Axon.dense(hidden_state, intermediate_size, name: join(name, "gate"), use_bias: false)
-
-    hidden_state = Axon.multiply(intermediate, Axon.activation(gate, activation))
-
-    Axon.dense(hidden_state, output_size, name: join(name, "output"), use_bias: false)
-  end
-
   defp language_modeling_head(hidden_state, spec, opts) do
     name = opts[:name]
 
diff --git a/lib/bumblebee/text/t5.ex b/lib/bumblebee/text/t5.ex
index ba92cd3d..fb9337d4 100644
--- a/lib/bumblebee/text/t5.ex
+++ b/lib/bumblebee/text/t5.ex
@@ -389,31 +389,60 @@ defmodule Bumblebee.Text.T5 do
   defp encoder(hidden_state, attention_mask, attention_head_mask, spec, opts) do
     name = opts[:name]
 
+    # The relative bias is computed once and shared by all blocks
+    attention_relative_bias =
+      Layers.relative_attention_bias(hidden_state, hidden_state, Layers.none(), Layers.none(),
+        bidirectional: true,
+        num_heads: spec.encoder_num_attention_heads,
+        num_buckets: spec.relative_attention_num_buckets,
+        max_distance: spec.relative_attention_max_distance,
+        name: join(name, "blocks.0.self_attention.relative_attention_bias")
+      )
+
     encoder_outputs =
-      Layers.Transformer.blocks(hidden_state,
-        attention_mask: attention_mask,
-        attention_head_mask: attention_head_mask,
-        num_blocks: spec.encoder_num_blocks,
-        num_attention_heads: spec.encoder_num_attention_heads,
-        hidden_size: spec.hidden_size,
-        kernel_initializer: kernel_initializer(spec),
-        dropout_rate: spec.dropout_rate,
-        layer_norm: &Layers.rms_norm(&1, name: &2, epsilon: spec.layer_norm_epsilon),
-        ffn: &ffn(&1, spec, name: &2),
-        block_type: :norm_first,
-        attention_head_size: spec.attention_head_size,
-        query_use_bias: false,
-        key_use_bias: false,
-        value_use_bias: false,
-        output_use_bias: false,
-        attention_relative_bias: [
-          bidirectional: true,
-          num_buckets: spec.relative_attention_num_buckets,
-          max_distance: spec.relative_attention_max_distance
+      Layers.Transformer.blocks(
+        hidden_state,
+        [
+          num_blocks: spec.encoder_num_blocks,
+          attention_mask: attention_mask,
+          attention_head_mask: attention_head_mask,
+          name: join(name, "blocks")
         ],
-        share_attention_relative_bias: true,
-        attention_scale: 1,
-        name: join(name, "blocks")
+        fn hidden_state, block ->
+          name = block.name
+
+          shortcut = hidden_state
+
+          {hidden_state, attention, _self_attention_cache} =
+            hidden_state
+            |> Layers.rms_norm(
+              epsilon: spec.layer_norm_epsilon,
+              name: join(name, "self_attention_norm")
+            )
+            |> Layers.Transformer.self_attention(
+              block,
+              [
+                attention_relative_bias: attention_relative_bias,
+                name: join(name, "self_attention")
+              ] ++ attention_opts(spec, spec.encoder_num_attention_heads)
+            )
+
+          hidden_state =
+            hidden_state
+            |> Axon.dropout(rate: spec.dropout_rate)
+            |> Axon.add(shortcut)
+
+          shortcut = hidden_state
+
+          hidden_state =
+            hidden_state
+            |> Layers.rms_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+            |> ffn(spec, name: join(name, "ffn"))
+            |> Axon.dropout(rate: spec.dropout_rate)
+            |> Axon.add(shortcut)
+
+          %{hidden_state: hidden_state, attention: attention}
+        end
       )
 
     hidden_state =
@@ -441,36 +470,104 @@ defmodule Bumblebee.Text.T5 do
        ) do
     name = opts[:name]
 
+    # The relative bias is computed once and shared by all blocks. It
+    # depends on the self-attention cache length and offset, which are
+    # the same for all blocks
+    {self_attention_cache, _cross_attention_cache} =
+      cache
+      |> Layers.Decoder.get_block_cache(0)
+      |> Layers.Decoder.get_attention_caches()
+
+    attention_relative_bias =
+      Layers.relative_attention_bias(
+        hidden_state,
+        hidden_state,
+        self_attention_cache,
+        Layers.Decoder.get_cache_offset(cache),
+        bidirectional: false,
+        num_heads: spec.decoder_num_attention_heads,
+        num_buckets: spec.relative_attention_num_buckets,
+        max_distance: spec.relative_attention_max_distance,
+        name: join(name, "blocks.0.self_attention.relative_attention_bias")
+      )
+
     decoder_outputs =
-      Layers.Transformer.blocks(hidden_state,
-        attention_mask: attention_mask,
-        attention_head_mask: attention_head_mask,
-        cross_hidden_state: encoder_hidden_state,
-        cross_attention_mask: encoder_attention_mask,
-        cross_attention_head_mask: cross_attention_head_mask,
-        cache: cache,
-        causal: true,
-        num_blocks: spec.decoder_num_blocks,
-        num_attention_heads: spec.decoder_num_attention_heads,
-        hidden_size: spec.hidden_size,
-        kernel_initializer: kernel_initializer(spec),
-        dropout_rate: spec.dropout_rate,
-        attention_head_size: spec.attention_head_size,
-        layer_norm: &Layers.rms_norm(&1, name: &2, epsilon: spec.layer_norm_epsilon),
-        ffn: &ffn(&1, spec, name: &2),
-        block_type: :norm_first,
-        query_use_bias: false,
-        key_use_bias: false,
-        value_use_bias: false,
-        output_use_bias: false,
-        attention_relative_bias: [
-          bidirectional: false,
-          num_buckets: spec.relative_attention_num_buckets,
-          max_distance: spec.relative_attention_max_distance
+      Layers.Transformer.blocks(
+        hidden_state,
+        [
+          num_blocks: spec.decoder_num_blocks,
+          attention_mask: attention_mask,
+          attention_head_mask: attention_head_mask,
+          cross_attention_head_mask: cross_attention_head_mask,
+          cache: cache,
+          name: join(name, "blocks")
         ],
-        share_attention_relative_bias: true,
-        attention_scale: 1,
-        name: join(name, "blocks")
+        fn hidden_state, block ->
+          name = block.name
+
+          attention_opts = attention_opts(spec, spec.decoder_num_attention_heads)
+
+          shortcut = hidden_state
+
+          {hidden_state, attention, self_attention_cache} =
+            hidden_state
+            |> Layers.rms_norm(
+              epsilon: spec.layer_norm_epsilon,
+              name: join(name, "self_attention_norm")
+            )
+            |> Layers.Transformer.self_attention(
+              block,
+              [
+                attention_relative_bias: attention_relative_bias,
+                causal: true,
+                name: join(name, "self_attention")
+              ] ++ attention_opts
+            )
+
+          hidden_state =
+            hidden_state
+            |> Axon.dropout(rate: spec.dropout_rate)
+            |> Axon.add(shortcut)
+
+          shortcut = hidden_state
+
+          {hidden_state, cross_attention, cross_attention_cache} =
+            hidden_state
+            |> Layers.rms_norm(
+              epsilon: spec.layer_norm_epsilon,
+              name: join(name, "cross_attention_norm")
+            )
+            |> Layers.Transformer.cross_attention(
+              encoder_hidden_state,
+              block,
+              [
+                attention_mask: encoder_attention_mask,
+                name: join(name, "cross_attention")
+              ] ++ attention_opts
+            )
+
+          hidden_state =
+            hidden_state
+            |> Axon.dropout(rate: spec.dropout_rate)
+            |> Axon.add(shortcut)
+
+          shortcut = hidden_state
+
+          hidden_state =
+            hidden_state
+            |> Layers.rms_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+            |> ffn(spec, name: join(name, "ffn"))
+            |> Axon.dropout(rate: spec.dropout_rate)
+            |> Axon.add(shortcut)
+
+          %{
+            hidden_state: hidden_state,
+            attention: attention,
+            cross_attention: cross_attention,
+            self_attention_cache: self_attention_cache,
+            cross_attention_cache: cross_attention_cache
+          }
+        end
       )
 
     hidden_state =
@@ -487,6 +584,20 @@ defmodule Bumblebee.Text.T5 do
     }
   end
 
+  defp attention_opts(spec, num_heads) do
+    [
+      num_heads: num_heads,
+      hidden_size: spec.hidden_size,
+      attention_head_size: spec.attention_head_size,
+      attention_scale: 1,
+      query_use_bias: false,
+      key_use_bias: false,
+      value_use_bias: false,
+      output_use_bias: false,
+      kernel_initializer: kernel_initializer(spec)
+    ]
+  end
+
   defp ffn(hidden_state, spec, opts) do
     name = opts[:name]
 
@@ -512,7 +623,6 @@ defmodule Bumblebee.Text.T5 do
     hidden_state
     |> Axon.dropout(rate: spec.dropout_rate)
     |> Axon.dense(spec.hidden_size, name: join(name, "output"), use_bias: false)
-    |> Axon.dropout(rate: spec.dropout_rate)
   end
 
   defp language_modeling_head(hidden_state, spec, opts) do
diff --git a/lib/bumblebee/vision/blip_vision.ex b/lib/bumblebee/vision/blip_vision.ex
index 741f3298..4743fa12 100644
--- a/lib/bumblebee/vision/blip_vision.ex
+++ b/lib/bumblebee/vision/blip_vision.ex
@@ -182,22 +182,44 @@ defmodule Bumblebee.Vision.BlipVision do
   defp encoder(embeddings, spec, opts) do
     name = opts[:name]
 
-    Layers.Transformer.blocks(embeddings,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: kernel_initializer(spec),
-      dropout_rate: 0.0,
-      attention_dropout_rate: spec.attention_dropout_rate,
-      layer_norm: [
-        epsilon: spec.layer_norm_epsilon
-      ],
-      ffn: [
-        intermediate_size: spec.intermediate_size,
-        activation: spec.activation
-      ],
-      block_type: :norm_first,
-      name: join(name, "blocks")
+    Layers.Transformer.blocks(
+      embeddings,
+      [num_blocks: spec.num_blocks, name: join(name, "blocks")],
+      fn hidden_state, block ->
+        name = block.name
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, _cache} =
+          hidden_state
+          |> Axon.layer_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+          |> Layers.Transformer.self_attention(block,
+            num_heads: spec.num_attention_heads,
+            hidden_size: spec.hidden_size,
+            kernel_initializer: kernel_initializer(spec),
+            dropout_rate: spec.attention_dropout_rate,
+            name: join(name, "self_attention")
+          )
+
+        hidden_state = Axon.add(hidden_state, shortcut)
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Axon.layer_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+          |> Layers.Transformer.basic_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+
+        %{hidden_state: hidden_state, attention: attention}
+      end
     )
   end
 
diff --git a/lib/bumblebee/vision/clip_vision.ex b/lib/bumblebee/vision/clip_vision.ex
index c77c66e3..7930de79 100644
--- a/lib/bumblebee/vision/clip_vision.ex
+++ b/lib/bumblebee/vision/clip_vision.ex
@@ -206,22 +206,46 @@ defmodule Bumblebee.Vision.ClipVision do
   defp encoder(embeddings, spec, opts) do
     name = opts[:name]
 
-    Layers.Transformer.blocks(embeddings,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: Axon.Initializers.normal(scale: 0.01),
-      dropout_rate: 0.0,
-      attention_dropout_rate: spec.attention_dropout_rate,
-      layer_norm: [
-        epsilon: spec.layer_norm_epsilon
-      ],
-      ffn: [
-        intermediate_size: spec.intermediate_size,
-        activation: spec.activation
-      ],
-      block_type: :norm_first,
-      name: join(name, "blocks")
+    kernel_initializer = Axon.Initializers.normal(scale: 0.01)
+
+    Layers.Transformer.blocks(
+      embeddings,
+      [num_blocks: spec.num_blocks, name: join(name, "blocks")],
+      fn hidden_state, block ->
+        name = block.name
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, _cache} =
+          hidden_state
+          |> Axon.layer_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+          |> Layers.Transformer.self_attention(block,
+            num_heads: spec.num_attention_heads,
+            hidden_size: spec.hidden_size,
+            kernel_initializer: kernel_initializer,
+            dropout_rate: spec.attention_dropout_rate,
+            name: join(name, "self_attention")
+          )
+
+        hidden_state = Axon.add(hidden_state, shortcut)
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Axon.layer_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+          |> Layers.Transformer.basic_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            kernel_initializer: kernel_initializer,
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+
+        %{hidden_state: hidden_state, attention: attention}
+      end
     )
   end
 
diff --git a/lib/bumblebee/vision/deit.ex b/lib/bumblebee/vision/deit.ex
index 9eaf7ce3..e97accfc 100644
--- a/lib/bumblebee/vision/deit.ex
+++ b/lib/bumblebee/vision/deit.ex
@@ -303,25 +303,51 @@ defmodule Bumblebee.Vision.Deit do
   defp encoder(hidden_state, spec, opts) do
     name = opts[:name]
 
-    Layers.Transformer.blocks(hidden_state,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: kernel_initializer(spec),
-      dropout_rate: spec.dropout_rate,
-      attention_dropout_rate: spec.attention_dropout_rate,
-      query_use_bias: spec.use_attention_bias,
-      key_use_bias: spec.use_attention_bias,
-      value_use_bias: spec.use_attention_bias,
-      layer_norm: [
-        epsilon: spec.layer_norm_epsilon
-      ],
-      ffn: [
-        intermediate_size: spec.intermediate_size,
-        activation: spec.activation
-      ],
-      block_type: :norm_first,
-      name: join(name, "blocks")
+    Layers.Transformer.blocks(
+      hidden_state,
+      [num_blocks: spec.num_blocks, name: join(name, "blocks")],
+      fn hidden_state, block ->
+        name = block.name
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, _cache} =
+          hidden_state
+          |> Axon.layer_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+          |> Layers.Transformer.self_attention(block,
+            num_heads: spec.num_attention_heads,
+            hidden_size: spec.hidden_size,
+            kernel_initializer: kernel_initializer(spec),
+            dropout_rate: spec.attention_dropout_rate,
+            query_use_bias: spec.use_attention_bias,
+            key_use_bias: spec.use_attention_bias,
+            value_use_bias: spec.use_attention_bias,
+            name: join(name, "self_attention")
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: spec.dropout_rate)
+          |> Axon.add(shortcut)
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Axon.layer_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+          |> Layers.Transformer.basic_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            dropout_rate: spec.dropout_rate,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+
+        %{hidden_state: hidden_state, attention: attention}
+      end
     )
   end
 
diff --git a/lib/bumblebee/vision/dino_v2.ex b/lib/bumblebee/vision/dino_v2.ex
index 956bd6e5..db60febb 100644
--- a/lib/bumblebee/vision/dino_v2.ex
+++ b/lib/bumblebee/vision/dino_v2.ex
@@ -343,94 +343,91 @@ defmodule Bumblebee.Vision.DinoV2 do
   defp encoder(hidden_state, spec, opts) do
     name = opts[:name]
 
-    ffn =
-      if spec.ffn_swiglu_activation do
-        intermediate_size =
-          div(floor(floor(spec.hidden_size * spec.intermediate_size_ratio) * 2 / 3 + 7), 8) * 8
-
-        &ffn_swiglu(&1, intermediate_size, spec.hidden_size, name: &2)
-      else
-        intermediate_size = floor(spec.hidden_size * spec.intermediate_size_ratio)
-
-        [
-          intermediate_size: intermediate_size,
-          activation: spec.activation
-        ]
-      end
+    Layers.Transformer.blocks(
+      hidden_state,
+      [num_blocks: spec.num_blocks, name: join(name, "blocks")],
+      fn hidden_state, block ->
+        name = block.name
 
-    Layers.Transformer.blocks(hidden_state,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: kernel_initializer(spec),
-      dropout_rate: spec.dropout_rate,
-      attention_dropout_rate: spec.attention_dropout_rate,
-      query_use_bias: spec.use_attention_bias,
-      key_use_bias: spec.use_attention_bias,
-      value_use_bias: spec.use_attention_bias,
-      layer_norm: [
-        epsilon: spec.layer_norm_epsilon
-      ],
-      ffn: ffn,
-      block_type: &block_impl(&1, &2, &3, spec),
-      name: join(name, "blocks")
+        shortcut = hidden_state
+
+        {hidden_state, attention, _cache} =
+          hidden_state
+          |> Axon.layer_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+          |> Layers.Transformer.self_attention(block,
+            num_heads: spec.num_attention_heads,
+            hidden_size: spec.hidden_size,
+            kernel_initializer: kernel_initializer(spec),
+            dropout_rate: spec.attention_dropout_rate,
+            query_use_bias: spec.use_attention_bias,
+            key_use_bias: spec.use_attention_bias,
+            value_use_bias: spec.use_attention_bias,
+            name: join(name, "self_attention")
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: spec.dropout_rate)
+          |> Layers.scale(
+            scale_initializer: Axon.Initializers.full(spec.scale_initial_value),
+            name: join(name, "self_attention_scale")
+          )
+          |> Axon.add(shortcut)
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Axon.layer_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+          |> ffn(spec, name: join(name, "ffn"))
+          |> Layers.scale(
+            scale_initializer: Axon.Initializers.full(spec.scale_initial_value),
+            name: join(name, "output_scale")
+          )
+          |> Axon.add(shortcut)
+
+        %{hidden_state: hidden_state, attention: attention}
+      end
     )
   end
 
+  defp ffn(hidden_state, spec, opts) do
+    name = opts[:name]
+
+    if spec.ffn_swiglu_activation do
+      intermediate_size =
+        div(floor(floor(spec.hidden_size * spec.intermediate_size_ratio) * 2 / 3 + 7), 8) * 8
+
+      ffn_swiglu(hidden_state, intermediate_size, spec.hidden_size, name: name)
+    else
+      intermediate_size = floor(spec.hidden_size * spec.intermediate_size_ratio)
+
+      Layers.Transformer.basic_ffn(hidden_state, intermediate_size, spec.hidden_size,
+        activation: spec.activation,
+        dropout_rate: spec.dropout_rate,
+        kernel_initializer: kernel_initializer(spec),
+        name: name
+      )
+    end
+  end
+
   # A feed-forward network with SwiGLU nonlinearity as in https://arxiv.org/abs/2002.05202
   defp ffn_swiglu(x, intermediate_size, output_size, opts) do
     name = opts[:name]
-    dropout = opts[:dropout] || 0.0
 
     {gate, x} =
       x
       |> Axon.dense(intermediate_size * 2, name: join(name, "intermediate"))
       |> Axon.split(2, axis: -1)
 
-    x = Axon.multiply(x, Axon.silu(gate))
-
     x
-    |> Axon.dropout(rate: dropout, name: join(name, "dropout"))
+    |> Axon.multiply(Axon.silu(gate))
     |> Axon.dense(output_size, name: join(name, "output"))
   end
 
-  # :norm_first block with additional scaling layers
-  defp block_impl(hidden_state, steps, name, spec) do
-    shortcut = hidden_state
-
-    {hidden_state, attention_info} =
-      hidden_state
-      |> steps.self_attention_norm.()
-      |> steps.self_attention.()
-
-    hidden_state =
-      hidden_state
-      |> Bumblebee.Layers.scale(
-        scale_initializer: Axon.Initializers.full(spec.scale_initial_value),
-        name: join(name, "self_attention_scale")
-      )
-      |> Axon.add(shortcut)
-
-    {_hidden_state, cross_attention_info} =
-      steps.cross_attention_maybe.(hidden_state, fn _hidden_state ->
-        raise "cross attention not supported"
-      end)
-
-    shortcut = hidden_state
-
-    hidden_state =
-      hidden_state
-      |> steps.output_norm.()
-      |> steps.ffn.()
-      |> Bumblebee.Layers.scale(
-        scale_initializer: Axon.Initializers.full(spec.scale_initial_value),
-        name: join(name, "output_scale")
-      )
-      |> Axon.add(shortcut)
-
-    {hidden_state, attention_info, cross_attention_info}
-  end
-
   defp kernel_initializer(spec) do
     Axon.Initializers.normal(scale: spec.initializer_scale)
   end
diff --git a/lib/bumblebee/vision/vit.ex b/lib/bumblebee/vision/vit.ex
index 3f2d3b72..20768751 100644
--- a/lib/bumblebee/vision/vit.ex
+++ b/lib/bumblebee/vision/vit.ex
@@ -253,25 +253,51 @@ defmodule Bumblebee.Vision.Vit do
   defp encoder(hidden_state, spec, opts) do
     name = opts[:name]
 
-    Layers.Transformer.blocks(hidden_state,
-      num_blocks: spec.num_blocks,
-      num_attention_heads: spec.num_attention_heads,
-      hidden_size: spec.hidden_size,
-      kernel_initializer: kernel_initializer(spec),
-      dropout_rate: spec.dropout_rate,
-      attention_dropout_rate: spec.attention_dropout_rate,
-      query_use_bias: spec.use_attention_bias,
-      key_use_bias: spec.use_attention_bias,
-      value_use_bias: spec.use_attention_bias,
-      layer_norm: [
-        epsilon: spec.layer_norm_epsilon
-      ],
-      ffn: [
-        intermediate_size: spec.intermediate_size,
-        activation: spec.activation
-      ],
-      block_type: :norm_first,
-      name: join(name, "blocks")
+    Layers.Transformer.blocks(
+      hidden_state,
+      [num_blocks: spec.num_blocks, name: join(name, "blocks")],
+      fn hidden_state, block ->
+        name = block.name
+
+        shortcut = hidden_state
+
+        {hidden_state, attention, _cache} =
+          hidden_state
+          |> Axon.layer_norm(
+            epsilon: spec.layer_norm_epsilon,
+            name: join(name, "self_attention_norm")
+          )
+          |> Layers.Transformer.self_attention(block,
+            num_heads: spec.num_attention_heads,
+            hidden_size: spec.hidden_size,
+            kernel_initializer: kernel_initializer(spec),
+            dropout_rate: spec.attention_dropout_rate,
+            query_use_bias: spec.use_attention_bias,
+            key_use_bias: spec.use_attention_bias,
+            value_use_bias: spec.use_attention_bias,
+            name: join(name, "self_attention")
+          )
+
+        hidden_state =
+          hidden_state
+          |> Axon.dropout(rate: spec.dropout_rate)
+          |> Axon.add(shortcut)
+
+        shortcut = hidden_state
+
+        hidden_state =
+          hidden_state
+          |> Axon.layer_norm(epsilon: spec.layer_norm_epsilon, name: join(name, "output_norm"))
+          |> Layers.Transformer.basic_ffn(spec.intermediate_size, spec.hidden_size,
+            activation: spec.activation,
+            dropout_rate: spec.dropout_rate,
+            kernel_initializer: kernel_initializer(spec),
+            name: join(name, "ffn")
+          )
+          |> Axon.add(shortcut)
+
+        %{hidden_state: hidden_state, attention: attention}
+      end
     )
   end
 
diff --git a/test/bumblebee/layers/transformer_test.exs b/test/bumblebee/layers/transformer_test.exs
deleted file mode 100644
index 269f8f02..00000000
--- a/test/bumblebee/layers/transformer_test.exs
+++ /dev/null
@@ -1,93 +0,0 @@
-defmodule Bumblebee.Layers.TransformerTest do
-  use ExUnit.Case, async: true
-
-  import Bumblebee.TestHelpers
-
-  alias Bumblebee.Layers
-  alias Bumblebee.Layers.Transformer
-
-  test "custom attention receives self- and cross-attention options" do
-    hidden = Axon.input("hidden", shape: {nil, nil, 4})
-    context = Axon.input("context", shape: {nil, nil, 4})
-
-    opts = [
-      num_blocks: 2,
-      num_attention_heads: 2,
-      hidden_size: 4,
-      ffn: [intermediate_size: 8],
-      cross_hidden_state: context,
-      name: "blocks"
-    ]
-
-    reference = hidden |> Transformer.blocks(opts) |> model_outputs()
-    owner = self()
-
-    attention = fn query, key, value, opts ->
-      send(owner, {:attention, opts[:name]})
-      Transformer.multi_head_attention(query, key, value, opts)
-    end
-
-    model = hidden |> Transformer.blocks([attention: attention] ++ opts) |> model_outputs()
-
-    for index <- 0..1, kind <- ["self_attention", "cross_attention"] do
-      name = "blocks.#{index}.#{kind}"
-      assert_received {:attention, ^name}
-    end
-
-    inputs = %{
-      "hidden" => Nx.iota({1, 3, 4}, type: :f32),
-      "context" => Nx.iota({1, 2, 4}, type: :f32)
-    }
-
-    {init, predict} = Axon.build(model)
-    params = init.(inputs, Axon.ModelState.empty())
-    actual = predict.(params, inputs)
-    expected = Axon.predict(reference, params, inputs)
-
-    assert_all_close(actual.hidden_state, expected.hidden_state)
-
-    for {left, right} <-
-          Enum.zip(Tuple.to_list(actual.attentions), Tuple.to_list(expected.attentions)) do
-      assert_all_close(left, right)
-    end
-
-    for {left, right} <-
-          Enum.zip(
-            Tuple.to_list(actual.cross_attentions),
-            Tuple.to_list(expected.cross_attentions)
-          ) do
-      assert_all_close(left, right)
-    end
-  end
-
-  test "block output is collected and passed to the following block" do
-    hidden = Axon.input("hidden", shape: {nil, nil, 2})
-
-    outputs =
-      Transformer.blocks(hidden,
-        num_blocks: 2,
-        num_attention_heads: 1,
-        hidden_size: 2,
-        attention: fn query, _key, _value, opts ->
-          {query, Layers.none(), opts[:attention_cache], Layers.none()}
-        end,
-        ffn: fn hidden, _name -> Axon.nx(hidden, &Nx.multiply(&1, 0)) end,
-        layer_norm: fn hidden, _name -> hidden end,
-        block_output: fn hidden, index -> Axon.nx(hidden, &Nx.add(&1, index + 1)) end
-      )
-
-    model = Axon.container(Map.take(outputs, [:hidden_state, :hidden_states]))
-    inputs = Nx.tensor([[[1.0, 2.0]]])
-    {init, predict} = Axon.build(model)
-    actual = predict.(init.(inputs, Axon.ModelState.empty()), inputs)
-
-    assert_equal(elem(actual.hidden_states, 0), inputs)
-    assert_equal(elem(actual.hidden_states, 1), Nx.tensor([[[3.0, 5.0]]]))
-    assert_equal(elem(actual.hidden_states, 2), Nx.tensor([[[8.0, 12.0]]]))
-    assert_equal(actual.hidden_state, Nx.tensor([[[8.0, 12.0]]]))
-  end
-
-  defp model_outputs(outputs) do
-    Axon.container(Map.take(outputs, [:hidden_state, :attentions, :cross_attentions]))
-  end
-end