Support scanned layers KV cache carry and fix vLLM integration

Author: khatwanimohit

src/maxtext/configs/inference/vllm.yml

@@ -78,6 +78,7 @@ logical_axis_rules: [
                       ['cache_heads', ['model']],
                       ['exp', ['expert', 'attn_dp_expert']],
                       ['paged_kv_heads', ['model']],
+                      ['layers', []],
                     ]
 data_sharding: [['data', 'attn_dp', 'model', 'expert', 'attn_dp_expert']]
 input_data_sharding_logical_axes: ['activation_embed_and_logits_batch']

src/maxtext/inference/vllm_decode.py

@@ -48,6 +48,8 @@
 from vllm import LLM
 from vllm.sampling_params import SamplingParams
 from maxtext.configs import pyconfig
+import maxtext.integration.vllm.maxtext_vllm_adapter as adapter
+adapter.register()
 
 os.environ["SKIP_JAX_PRECOMPILE"] = "1"
 os.environ["NEW_MODEL_DESIGN"] = "1"
@@ -82,6 +84,7 @@ def decode_with_vllm(config: Config) -> None:
               "weight_dtype": "bfloat16",
               "allow_split_physical_axes": True,
               "debug_sharding": config.debug_sharding,
+              "scan_layers": config.scan_layers,
           },
           "sharding": {
               "sharding_strategy": {
@@ -140,11 +143,15 @@ def decode_with_vllm(config: Config) -> None:
         f"max_target_length ({config.max_target_length}) must be greater than max_prompt_length ({max_prompt_length})"
     )
 
+  # MaxText uses -1 to mean "disabled"; vLLM requires top_p in (0, 1].
+  top_p = config.decode_sampling_nucleus_p if config.decode_sampling_nucleus_p > 0 else 1.0
+  top_k = config.decode_sampling_top_k if config.decode_sampling_top_k > 0 else -1
+
   sampling_params = SamplingParams(
       temperature=config.decode_sampling_temperature,
       max_tokens=max_tokens_to_generate,
-      top_k=config.decode_sampling_top_k,
-      top_p=config.decode_sampling_nucleus_p,
+      top_k=top_k,
+      top_p=top_p,
   )
 
   outputs = llm.generate(prompts, sampling_params)

src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py

@@ -104,6 +104,11 @@ class MaxTextForCausalLM(nnx.Module):
   of the decoding step.
   """
 
+  # Signal to tpu-inference model_loader that this class manages its own
+  # JIT-sharded initialization (via create_nnx_model with out_shardings).
+  # When True, model_loader skips wrapping __init__ in an outer bare @jax.jit,
+  _self_manages_sharding: bool = True
+
   def __init__(self, vllm_config: VllmConfig, rng_key: jax.Array, mesh: Mesh):
     """Initializes the MaxTextForCausalLM model.
 
@@ -250,7 +255,7 @@ def load_weights(self, rng_key: jax.Array) -> None:
     if self.model is not None:
       return
 
-    with self.mesh, nn.logical_axis_rules(""):
+    with self.mesh, nn.logical_axis_rules(self.maxtext_config.logical_axis_rules):
       model, _ = model_creation_utils.create_nnx_model(
           self.maxtext_config, mesh=self.mesh, model_mode=self.model_mode, rng_key=rng_key
       )

src/maxtext/layers/attentions.py

@@ -981,7 +981,7 @@ def forward_serve_vllm(
       value: Array,
       rpa_kv_cache: list[Array] | None = None,
       rpa_metadata: dict[str, Any] | None = None,
-  ) -> tuple[list[Array], Array]:
+  ) -> tuple[Array, list[Array]]:
     """Forward function for vLLM serving with RPA attention."""
     try:
       # pylint: disable=import-outside-toplevel
@@ -992,13 +992,14 @@ def forward_serve_vllm(
           "vLLM RPA attention ops require the vllm-tpu package. Please install it with `pip install vllm-tpu`."
       ) from e
 
-    if rpa_kv_cache is None or rpa_metadata is None:
-      raise ValueError("kv_cache and attention_metadata must be provided when using vLLM.")
-
     query = query.reshape(-1, query.shape[2], query.shape[3])
     key = key.reshape(-1, key.shape[2], key.shape[3])
     value = value.reshape(-1, value.shape[2], value.shape[3])
 
+    if rpa_kv_cache is None or rpa_metadata is None:
+      # Return dummy values for dry runs (e.g. during model initialization or JIT tracing)
+      return query, []
+
     if self.config.sliding_window_size > 0:
       attention_chunk_size = self.config.sliding_window_size
     else:
@@ -1026,7 +1027,7 @@ def forward_serve_vllm(
         k_scale,
         v_scale,
     )
-    return kv_cache, output
+    return output, kv_cache
 
   def __call__(
       self,
@@ -1169,7 +1170,7 @@ def __call__(
 
     elif self.config.attention == "vllm_rpa" and model_mode != MODEL_MODE_TRAIN:
       batch, seq_len, num_heads, head_dim = query.shape
-      updated_kv, attn_out = self.forward_serve_vllm(
+      attn_out, updated_kv = self.forward_serve_vllm(
           query, key, value, rpa_kv_cache=kv_cache, rpa_metadata=attention_metadata
       )
       out = attn_out.reshape(batch, seq_len, num_heads, head_dim)

src/maxtext/layers/decoders.py

@@ -981,16 +981,58 @@ def __call__(
                 "nope_layer_interval": self.config.nope_layer_interval,
                 "interleave_moe_layer_step": self.config.interleave_moe_layer_step,
             }
-          y, _ = self.scan_decoder_layers(
-              cfg,
-              RemattedBlockLayer,
-              scan_length,
-              "layers",
-              mesh,
-              in_axes_tuple=(nn.broadcast,) * len(broadcast_args),
-              model_mode=model_mode,
-              **layer_kwargs,
-          )(y, *broadcast_args)
+          # Update broadcast_args and in_axes_tuple for vLLM RPA
+          in_axes_tuple = (nn.broadcast,) * len(broadcast_args)
+          current_broadcast_args = list(broadcast_args)
+          current_in_axes_tuple = list(in_axes_tuple)
+
+          if kv_caches is not None:
+            # Stack kv_caches for scan: [num_layers, ...]
+            stacked_kv_cache = jnp.stack(kv_caches, axis=0)
+            
+            # We pass (y, stacked_kv_cache, 0) as the carry
+            carry = (y, stacked_kv_cache, 0)
+            
+            # We don't pass kv_cache as a scanned argument anymore
+            
+            # Pass None for previous_chunk, slot, page_state, kv_cache to align with __call__ signature
+            current_broadcast_args.extend([None, None, None, None, attention_metadata])
+            current_in_axes_tuple.extend([nn.broadcast] * 5)
+            
+            max_logging.info(f"DEBUG: len(current_broadcast_args)={len(current_broadcast_args)}")
+            max_logging.info(f"DEBUG: current_broadcast_args={[type(a) for a in current_broadcast_args]}")
+            
+            final_carry, _ = self.scan_decoder_layers(
+                cfg,
+                RemattedBlockLayer,
+                scan_length,
+                "layers",
+                mesh,
+                in_axes_tuple=tuple(current_in_axes_tuple),
+                model_mode=model_mode,
+                **layer_kwargs,
+            )(carry, *current_broadcast_args)
+            
+            y, returned_kv_cache, _ = final_carry
+            
+            # Update the list of KV caches from the scanned results
+            for i in range(cfg.num_decoder_layers):
+              kv_caches[i] = returned_kv_cache[i]
+          else:
+            # Fallback to old behavior if kv_caches is None (not vLLM RPA)
+            current_broadcast_args.append(None)
+            current_in_axes_tuple.append(nn.broadcast)
+            
+            y, _ = self.scan_decoder_layers(
+                cfg,
+                RemattedBlockLayer,
+                scan_length,
+                "layers",
+                mesh,
+                in_axes_tuple=tuple(current_in_axes_tuple),
+                model_mode=model_mode,
+                **layer_kwargs,
+            )(y, *current_broadcast_args)
       else:
         if cfg.decoder_block == DecoderBlockType.DEEPSEEK:
           assert len(RemattedBlockLayers) == 2, "Unscanned layers must have a length of 2 using deepseek."

src/maxtext/layers/nnx_decoders.py

@@ -428,8 +428,23 @@ def pure_layer_fn(state_in, y_in):
 
     return out
 
-  def _apply_layers_sequentially(self, layers, x_in, *args, length: int, **kwargs):
-    """Runs the layer stack using nnx.scan."""
+  def _apply_layers_sequentially(self, layers, x_in, *args, length: int, kv_caches_stacked=None, **kwargs):
+    """Runs the layer stack using nnx.scan.
+
+    Args:
+      layers: The stacked NNX module whose params are scanned over.
+      x_in: The carry (hidden state) fed into the first layer.
+      *args: Positional args broadcast to every layer call.
+      length: Number of scan iterations (= number of layers).
+      kv_caches_stacked: Optional pytree whose leaves have shape [num_layers, ...].
+        When provided, the i-th slice is passed as `kv_cache=` to layer i and the
+        updated caches are returned as a third element of the tuple.
+      **kwargs: Keyword args forwarded to the layer (filtered by the layer signature).
+
+    Returns:
+      (final_carry, updated_layers) when kv_caches_stacked is None.
+      (final_carry, updated_layers, returned_kv_stacked) otherwise.
+    """
     policy = self.get_remat_policy()
     prevent_cse = maxtext_utils.should_prevent_cse_in_remat(self.config)
     graphdef, params, state = nnx.split(
@@ -450,35 +465,83 @@ def _apply_layers_sequentially(self, layers, x_in, *args, length: int, **kwargs)
     # Filter kwargs to only include keys that exist in the layer's signature
     valid_kwargs = {k: v for k, v in kwargs.items() if k in sig.parameters or "kwargs" in sig.parameters}
 
+    use_kv = kv_caches_stacked is not None
+
     def layer_fn(carry, scanned_vars):
       # Unpack the sliced variables for THIS layer
-      current_params, current_state = scanned_vars
+      if use_kv:
+        current_params, current_state, kv_cache_layer = scanned_vars
+      else:
+        current_params, current_state = scanned_vars
+        kv_cache_layer = None
 
       if self.config.parameter_memory_host_offload:
         current_params = jax.tree.map(lambda x: jax.device_put(x, max_utils.device_space()), current_params)
 
       # Merge using the SLICED state
       layer = nnx.merge(graphdef, current_params, current_state)
 
-      # Run the layer (Filter kwargs if using the solution from previous turn)
-      layer_out = layer(carry, *args, **valid_kwargs)
+      # Build call kwargs, injecting per-layer kv_cache when available
+      call_kwargs = dict(valid_kwargs)
+      if kv_cache_layer is not None:
+        call_kwargs["kv_cache"] = kv_cache_layer
 
-      new_carry = layer_out[0] if isinstance(layer_out, tuple) else layer_out
+      layer_out = layer(carry, *args, **call_kwargs)
+
+      if isinstance(layer_out, tuple):
+        new_carry = layer_out[0]
+        updated_kv = layer_out[1] if len(layer_out) > 1 else None
+      else:
+        new_carry = layer_out
+        updated_kv = None
 
       # Extract the updated state to return it
-      # _, new_current_state = nnx.split(layer, nnx.Param, ...)
       new_current_state = nnx.state(layer)
+
+      if use_kv:
+        return new_carry, (new_current_state, updated_kv)
       return new_carry, new_current_state
 
     layer_fn = jax.checkpoint(layer_fn, policy=policy, prevent_cse=prevent_cse)
 
-    final_carry, scanned_state = jax.lax.scan(layer_fn, x_in, (params, state))
+    if use_kv:
+      # If kv_caches is provided (e.g., from vLLM), we CANNOT use jax.lax.scan
+      # because scanning requires stacking the kv_caches list, which creates a copy
+      # and breaks the in-place memory updates required by vLLM's PagedAttention.
+      # Therefore, we must unroll the loop statically when kv_caches is provided.
+      
+      # kv_caches_stacked is actually the original kv_caches list in this new flow
+      kv_caches_list = kv_caches_stacked
+      
+      current_carry = x_in
+      
+      for i in range(length):
+        # Statically slice the parameters and state for this layer
+        current_params = jax.tree.map(lambda x: x[i], params)
+        current_state = jax.tree.map(lambda x: x[i], state)
+        
+        # Call the layer
+        current_carry, (new_current_state, updated_kv) = layer_fn(
+            current_carry, (current_params, current_state, kv_caches_list[i])
+        )
+        
+        # Update the list in-place (mutates the list passed by reference)
+        kv_caches_list[i] = updated_kv
+        
+      # We don't need to rebuild scanned_state or return it because during
+      # inference with vLLM, parameters do not change and we don't need intermediates.
+      return current_carry, layers, None
+    else:
+      final_carry, scanned_state = jax.lax.scan(layer_fn, x_in, (params, state))
+      returned_kv_stacked = None
 
     if scan_axis != 0:
       scanned_params, scanned_other = scanned_state.split(nnx.Param, ...)
       scanned_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), scanned_params)
       scanned_state = nnx.State.merge(scanned_params, scanned_other)
 
+    if use_kv:
+      return final_carry, nnx.merge(graphdef, scanned_state), returned_kv_stacked
     return final_carry, nnx.merge(graphdef, scanned_state)
 
   def get_decoder_layers(self):
@@ -1001,7 +1064,19 @@ def __call__(
         )
       else:
         scan_length = int(cfg.num_decoder_layers / cfg.inhomogeneous_layer_cycle_interval)
-        y, self.layers = self._apply_layers_sequentially(self.layers, y, *layer_args, length=scan_length, **layer_kwargs)
+        if kv_caches is not None:
+          # Pass the kv_caches list directly to avoid copying in jnp.stack,
+          # which breaks vLLM PagedAttention in-place memory updates.
+          # The _apply_layers_sequentially function will handle it by statically unrolling.
+          y, self.layers, returned_kv = self._apply_layers_sequentially(
+              self.layers, y, *layer_args, length=scan_length,
+              kv_caches_stacked=kv_caches, **layer_kwargs
+          )
+          # kv_caches list is updated in-place inside _apply_layers_sequentially
+        else:
+          y, self.layers = self._apply_layers_sequentially(
+              self.layers, y, *layer_args, length=scan_length, **layer_kwargs
+          )
     else:
       prevent_cse = maxtext_utils.should_prevent_cse_in_remat(cfg)
 

src/maxtext/models/gemma.py

@@ -30,6 +30,7 @@
 from maxtext.layers.linears import Dropout, MlpBlock
 from maxtext.layers.normalizations import RMSNorm
 from maxtext.layers.quantizations import AqtQuantization as Quant
+from maxtext.inference import page_manager
 from maxtext.utils import max_utils
 
 
@@ -126,8 +127,7 @@ def __call__(
       deterministic,
       model_mode,
       previous_chunk=None,
-      page_manager=None,
-      page_state=None,
+      page_state: None | page_manager.PageState = None,
       slot=None,
       kv_cache=None,
       attention_metadata=None,

src/maxtext/models/gemma3.py

@@ -194,7 +194,13 @@ def __call__(
   ):
     cfg = self.config
     # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache))
-    if isinstance(inputs, tuple):
+    is_scan_carry = False
+    if isinstance(inputs, tuple) and len(inputs) == 3:
+      hidden_states, stacked_kv_cache, layer_idx = inputs
+      kv_cache = stacked_kv_cache[layer_idx]
+      inputs = hidden_states
+      is_scan_carry = True
+    elif isinstance(inputs, tuple):
       inputs = inputs[0]
     inputs = nn.with_logical_constraint(inputs, self.activation_axis_names)
     inputs = checkpoint_name(inputs, "decoder_layer_input")
@@ -244,7 +250,15 @@ def __call__(
           jnp.sum(layer_output == 0) / jnp.size(layer_output),
       )
 
-    if cfg.scan_layers:
+    if is_scan_carry:
+      def update_cache(cache, val):
+        if jnp.size(val) > 0:
+          return cache.at[layer_idx].set(val)
+        return cache
+
+      stacked_kv_cache = jax.tree_util.tree_map(update_cache, stacked_kv_cache, kv_cache)
+      return (layer_output, stacked_kv_cache, layer_idx + 1), None
+    elif cfg.scan_layers:
       return layer_output, None
     else:
       return layer_output, kv_cache

src/maxtext/models/gemma4.py

@@ -322,7 +322,13 @@ def __call__(
   ):
     cfg = self.config
     # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache))
-    if isinstance(inputs, tuple):
+    is_scan_carry = False
+    if isinstance(inputs, tuple) and len(inputs) == 3:
+      hidden_states, stacked_kv_cache, layer_idx = inputs
+      kv_cache = stacked_kv_cache[layer_idx]
+      inputs = hidden_states
+      is_scan_carry = True
+    elif isinstance(inputs, tuple):
       inputs = inputs[0]
     inputs = nn.with_logical_constraint(inputs, self.activation_axis_names)
     inputs = checkpoint_name(inputs, "decoder_layer_input")
@@ -383,7 +389,15 @@ def __call__(
           jnp.sum(layer_output == 0) / jnp.size(layer_output),
       )
 
-    if cfg.scan_layers:
+    if is_scan_carry:
+      def update_cache(cache, val):
+        if jnp.size(val) > 0:
+          return cache.at[layer_idx].set(val)
+        return cache
+
+      stacked_kv_cache = jax.tree_util.tree_map(update_cache, stacked_kv_cache, kv_cache)
+      return (layer_output, stacked_kv_cache, layer_idx + 1), None
+    elif cfg.scan_layers:
       return layer_output, None
     else:
       return layer_output, kv_cache