Merge pull request #4066 from AI-Hypercomputer:jackyf/gemma3_4-base-decoding-fixes

PiperOrigin-RevId: 941501455

Author: Google-ML-Automation

src/maxtext/configs/inference/vllm.yml

@@ -88,6 +88,9 @@ logical_axis_rules: [
                       ['mlp', ['attn_dp', 'model']],
                       ['embed', []],
                       ['norm', []],
+                      ['layers', []],
+                      ['dense_layers', []],
+                      ['moe_layers', []],
                       # ==========================================
                       # Inference(Prefill, Decode, Cache)
                       # ==========================================

src/maxtext/layers/decoders.py

@@ -1056,6 +1056,8 @@ def __call__(
               bidirectional_mask_value,
               previous_chunk,
               slot,
+              kv_caches=kv_caches,
+              attention_metadata=attention_metadata,
           )
         elif cfg.decoder_block == DecoderBlockType.GEMMA4:
           bidirectional_mask_value = multimodal_input.bidirectional_mask if multimodal_input is not None else None
@@ -1068,6 +1070,8 @@ def __call__(
               bidirectional_mask_value,
               previous_chunk,
               slot,
+              kv_caches=kv_caches,
+              attention_metadata=attention_metadata,
           )
         elif cfg.decoder_block == DecoderBlockType.DEEPSEEK4:
           y = self._apply_deepseek4_scanned_blocks(
@@ -1303,6 +1307,8 @@ def _apply_gemma3_scanned_blocks(
       bidirectional_mask,
       previous_chunk,
       slot,
+      kv_caches=None,
+      attention_metadata=None,
   ):
     """Applies Gemma3 scanned decoder blocks, handling main scan and remainders."""
 
@@ -1316,27 +1322,43 @@ def _apply_gemma3_scanned_blocks(
     policy = self.get_remat_policy()
     RemattedGemma3Block = self.set_remat_policy([gemma3.Gemma3ScannableBlockToLinen], policy)[0]
 
-    layer_call_kwargs = {"bidirectional_mask": bidirectional_mask}
     layer_kwargs = {"num_of_layers": attention_pattern_length}
 
     # Apply the main scan over the full blocks
     if scan_length > 0:
-      broadcast_args = (
-          decoder_segment_ids,
-          decoder_positions,
-          deterministic,
-          model_mode,
+      kv_cache_scanned = maxtext_utils.prepare_kv_caches_for_scan(
+          kv_caches, scan_length, attention_pattern_length, stack=True
       )
-      y, _ = self.scan_decoder_layers(
+
+      broadcast_args_spec = [
+          (decoder_segment_ids, nn.broadcast),
+          (decoder_positions, nn.broadcast),
+          (deterministic, nn.broadcast),
+          (model_mode, nn.broadcast),
+          (slot, nn.broadcast),
+          (None, nn.broadcast),  # page_state
+          (previous_chunk, nn.broadcast),
+          (bidirectional_mask, nn.broadcast),
+          (kv_cache_scanned, 0 if kv_caches is not None else nn.broadcast),
+          (attention_metadata, nn.broadcast),
+      ]
+      broadcast_args = tuple(arg for arg, _ in broadcast_args_spec)
+      in_axes_tuple = tuple(axis for _, axis in broadcast_args_spec)
+
+      y, returned_kv_cache = self.scan_decoder_layers(
           cfg,
           RemattedGemma3Block,
           scan_length,
           "layers",
           mesh,
-          in_axes_tuple=(nn.broadcast,) * len(broadcast_args),
+          in_axes_tuple=in_axes_tuple,
           model_mode=self.model_mode,
           **layer_kwargs,
-      )(y, *broadcast_args, **layer_call_kwargs)
+      )(y, *broadcast_args)
+
+      maxtext_utils.update_kv_caches_after_scan(
+          kv_caches, returned_kv_cache, scan_length, attention_pattern_length, stacked=True
+      )
 
     # Apply any remaining layers that did not fit into a full scanned block
     num_remaining_layers = cfg.num_decoder_layers % attention_pattern_length
@@ -1346,16 +1368,37 @@ def _apply_gemma3_scanned_blocks(
       layer = RemattedGemma3Block(
           config=cfg, mesh=mesh, quant=self.quant, model_mode=self.model_mode, name="layers_remainder", **rem_layer_kwargs
       )  # pytype: disable=wrong-keyword-args
-      y, _ = layer(
+
+      remainder_kv = None
+      if kv_caches is not None:
+        start_idx = scan_length * attention_pattern_length
+        remainder_kv = tuple(kv_caches[start_idx : start_idx + num_remaining_layers])
+
+      y_and_kv = layer(
           y,
           decoder_segment_ids,
           decoder_positions,
           deterministic,
           model_mode,
           previous_chunk=previous_chunk,
           slot=slot,
-          **layer_call_kwargs,
+          bidirectional_mask=bidirectional_mask,
+          kv_cache=remainder_kv,
+          attention_metadata=attention_metadata,
       )
+
+      if isinstance(y_and_kv, tuple):
+        y = y_and_kv[0]
+        updated_remainder_kv = y_and_kv[1]
+      else:
+        y = y_and_kv
+        updated_remainder_kv = None
+
+      if kv_caches is not None and updated_remainder_kv is not None:
+        start_idx = scan_length * attention_pattern_length
+        for offset, updated_item in enumerate(updated_remainder_kv):
+          kv_caches[start_idx + offset] = updated_item
+
     return y
 
   def _apply_gemma4_scanned_blocks(
@@ -1368,6 +1411,8 @@ def _apply_gemma4_scanned_blocks(
       bidirectional_mask,
       previous_chunk,
       slot,
+      kv_caches=None,
+      attention_metadata=None,
   ):
     """Applies Gemma4 scanned decoder blocks, handling main scan and remainders."""
 
@@ -1379,28 +1424,32 @@ def _apply_gemma4_scanned_blocks(
     num_full_blocks = cfg.num_decoder_layers // block_pattern_len
     remainder_layers = cfg.num_decoder_layers % block_pattern_len
 
-    broadcast_args = (
-        decoder_segment_ids,
-        decoder_positions,
-        deterministic,
-        model_mode,
-    )
-
-    # Pass slot/previous_chunk/bidirectional_mask by keyword only (via layer_call_kwargs),
-    # never positionally in broadcast_args: Gemma4DecoderLayer and Gemma4ScannableBlock
-    # declare slot and previous_chunk in swapped order, so positional passing misroutes them.
-    layer_call_kwargs = {
-        "slot": slot,
-        "previous_chunk": previous_chunk,
-        "bidirectional_mask": bidirectional_mask,
-    }
-
     if num_full_blocks > 0:
       ScannableBlockToLinen = gemma4.Gemma4ScannableBlockToLinen
       policy = self.get_remat_policy()
       RemattedGemma4Block = self.set_remat_policy([ScannableBlockToLinen], policy)[0]
+
+      kv_cache_scanned = maxtext_utils.prepare_kv_caches_for_scan(
+          kv_caches, num_full_blocks, block_pattern_len, stack=True
+      )
+
+      broadcast_args_spec = [
+          (decoder_segment_ids, nn.broadcast),
+          (decoder_positions, nn.broadcast),
+          (deterministic, nn.broadcast),
+          (model_mode, nn.broadcast),
+          (slot, nn.broadcast),
+          (None, nn.broadcast),  # page_state
+          (previous_chunk, nn.broadcast),
+          (bidirectional_mask, nn.broadcast),
+          (kv_cache_scanned, 0 if kv_caches is not None else nn.broadcast),
+          (attention_metadata, nn.broadcast),
+      ]
+      broadcast_args = tuple(arg for arg, _ in broadcast_args_spec)
+      in_axes_tuple = tuple(axis for _, axis in broadcast_args_spec)
+
       # For a fully scanned block, apply it inside a nn.scan over the calculated number of full blocks
-      y, _ = nn.scan(
+      y, returned_kv_cache = nn.scan(
           RemattedGemma4Block,
           variable_axes={
               "params": cfg.param_scan_axis,
@@ -1410,7 +1459,7 @@ def _apply_gemma4_scanned_blocks(
               "_overwrite_with_gradient": 0,
           },
           split_rngs={"params": True, "dropout": cfg.enable_dropout},
-          in_axes=(nn.broadcast,) * len(broadcast_args),
+          in_axes=in_axes_tuple,
           length=num_full_blocks,
           metadata_params={
               nn.PARTITION_NAME: "layers",
@@ -1424,7 +1473,11 @@ def _apply_gemma4_scanned_blocks(
           num_of_layers=block_pattern_len,
           name="scanned_blocks",
       )(
-          y, *broadcast_args, **layer_call_kwargs
+          y, *broadcast_args
+      )
+
+      maxtext_utils.update_kv_caches_after_scan(
+          kv_caches, returned_kv_cache, num_full_blocks, block_pattern_len, stacked=True
       )
 
     # Process any remaining layers that don't fit into a full scanned block
@@ -1438,9 +1491,33 @@ def _apply_gemma4_scanned_blocks(
           attention_type=attention_type,
           layer_idx=layer_id,
       )
-      y = layer(y, *broadcast_args, **layer_call_kwargs)
-      if cfg.scan_layers:
-        y = y[0]
+      kv_cache = kv_caches[layer_id] if kv_caches is not None else None
+
+      # inputs, decoder_segment_ids, decoder_positions, deterministic, model_mode,
+      # previous_chunk, page_state, slot, bidirectional_mask, kv_cache, attention_metadata
+      remainder_args = (
+          decoder_segment_ids,
+          decoder_positions,
+          deterministic,
+          model_mode,
+          previous_chunk,
+          None,  # page_state
+          slot,
+          bidirectional_mask,
+          kv_cache,
+          attention_metadata,
+      )
+
+      y_and_kv = layer(y, *remainder_args)
+      if isinstance(y_and_kv, tuple):
+        y = y_and_kv[0]
+        new_kv = y_and_kv[1]
+      else:
+        y = y_and_kv
+        new_kv = None
+
+      if kv_caches is not None and new_kv is not None:
+        kv_caches[layer_id] = new_kv
 
     return y