diff --git a/src/maxtext/configs/inference/vllm.yml b/src/maxtext/configs/inference/vllm.yml index 157845c5b7..44efdcb161 100644 --- a/src/maxtext/configs/inference/vllm.yml +++ b/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) # ========================================== diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index 0ea2913265..fc570e69bb 100644 --- a/src/maxtext/layers/decoders.py +++ b/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,7 +1368,13 @@ 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, @@ -1354,8 +1382,23 @@ def _apply_gemma3_scanned_blocks( 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 diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 1a1ed65924..488fe04f73 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -1111,38 +1111,22 @@ def get_deepseek(): DecoderBlockType.GEMMA: [gemma.GemmaDecoderLayer], DecoderBlockType.GEMMA2: [gemma2.Gemma2DecoderLayer], DecoderBlockType.GEMMA3: [gemma3.Gemma3DecoderLayer], - DecoderBlockType.GEMMA4: get_scannable( - gemma4.Gemma4DecoderLayer, gemma4.Gemma4ScannableBlock - ), + DecoderBlockType.GEMMA4: get_scannable(gemma4.Gemma4DecoderLayer, gemma4.Gemma4ScannableBlock), DecoderBlockType.GEMMA4_SMALL: [gemma4_small.Gemma4SmallDecoderLayer], DecoderBlockType.GPT3: [gpt3.Gpt3DecoderLayer], DecoderBlockType.QWEN2: [qwen2.Qwen2DecoderLayer], DecoderBlockType.QWEN3: [qwen3.Qwen3DecoderLayer], DecoderBlockType.QWEN3_MOE: [qwen3.Qwen3MoeDecoderLayer], - DecoderBlockType.QWEN3_CUSTOM_MOE: [ - qwen3_custom.Qwen3CustomMoeDecoderLayer - ], + DecoderBlockType.QWEN3_CUSTOM_MOE: [qwen3_custom.Qwen3CustomMoeDecoderLayer], DecoderBlockType.SIMPLE: [simple_layer.SimpleDecoderLayer], DecoderBlockType.SIMPLE_MLP: [simple_layer.SimpleMlpDecoderLayer], DecoderBlockType.DEEPSEEK: get_deepseek(), - DecoderBlockType.DEEPSEEK4: get_scannable( - deepseek4.DeepSeek4DecoderLayer, deepseek4.DeepSeek4ScannableBlock - ), - DecoderBlockType.GPT_OSS: get_scannable( - gpt_oss.GptOssDecoderLayer, gpt_oss.GptOssScannableBlock - ), - DecoderBlockType.QWEN3_NEXT: get_scannable( - qwen3.Qwen3NextDecoderLayer, qwen3.Qwen3NextScannableBlock - ), - DecoderBlockType.QWEN3_5: get_scannable( - qwen3_5.Qwen3_5DecoderLayer, qwen3_5.Qwen3_5ScannableBlock - ), - DecoderBlockType.LLAMA4: get_scannable( - llama4.Llama4DecoderLayer, llama4.Llama4ScannableBlock - ), - DecoderBlockType.OLMO3: get_scannable( - olmo3.Olmo3DecoderLayer, olmo3.Olmo3ScannableBlock - ), + DecoderBlockType.DEEPSEEK4: get_scannable(deepseek4.DeepSeek4DecoderLayer, deepseek4.DeepSeek4ScannableBlock), + DecoderBlockType.GPT_OSS: get_scannable(gpt_oss.GptOssDecoderLayer, gpt_oss.GptOssScannableBlock), + DecoderBlockType.QWEN3_NEXT: get_scannable(qwen3.Qwen3NextDecoderLayer, qwen3.Qwen3NextScannableBlock), + DecoderBlockType.QWEN3_5: get_scannable(qwen3_5.Qwen3_5DecoderLayer, qwen3_5.Qwen3_5ScannableBlock), + DecoderBlockType.LLAMA4: get_scannable(llama4.Llama4DecoderLayer, llama4.Llama4ScannableBlock), + DecoderBlockType.OLMO3: get_scannable(olmo3.Olmo3DecoderLayer, olmo3.Olmo3ScannableBlock), } if cfg.decoder_block not in layer_map: @@ -1664,6 +1648,19 @@ def __call__( model_mode, logical_partition_spec=logical_partition_spec, ) + elif self.is_gemma4: + y = self._apply_gemma4_scanned_blocks( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + bidirectional_mask, + previous_chunk, + slot, + kv_caches=kv_caches, + attention_metadata=attention_metadata, + ) else: # Standard pipeline run (non-DeepSeek, incl. Gemma4 — matches Linen decoders.py). # Gemma4 routes through the pipeline here; _apply_gemma4_scanned_blocks is @@ -1932,6 +1929,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.""" @@ -1943,10 +1942,20 @@ def _apply_gemma3_scanned_blocks( layer_args = (decoder_segment_ids, decoder_positions, deterministic, model_mode) layer_kwargs = {"bidirectional_mask": bidirectional_mask} + if attention_metadata is not None: + layer_kwargs["attention_metadata"] = attention_metadata # Apply the main scan over the full blocks if scan_length > 0: - y, self.layers, _ = self._apply_layers_sequentially(self.layers, y, *layer_args, length=scan_length, **layer_kwargs) + grouped_kv_caches = maxtext_utils.prepare_kv_caches_for_scan( + kv_caches, scan_length, attention_pattern_length, stack=False + ) + y, self.layers, _ = self._apply_layers_sequentially( + self.layers, y, *layer_args, length=scan_length, kv_caches_stacked=grouped_kv_caches, **layer_kwargs + ) + maxtext_utils.update_kv_caches_after_scan( + kv_caches, grouped_kv_caches, scan_length, attention_pattern_length, stacked=False + ) # Apply any remaining layers that did not fit into a full scanned block num_remaining_layers = cfg.num_decoder_layers % attention_pattern_length @@ -1954,23 +1963,38 @@ def _apply_gemma3_scanned_blocks( policy = self.get_remat_policy() prevent_cse = maxtext_utils.should_prevent_cse_in_remat(cfg) - def pure_gemma_fn(graphdef, state_in, y_in): + 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]) + + def pure_gemma_fn(graphdef, state_in, y_in, kv_in): merged_layer = nnx.merge(graphdef, state_in) - out_y, _ = merged_layer( - y_in, - *layer_args, - previous_chunk=previous_chunk, - slot=slot, - **layer_kwargs, - ) - return out_y, nnx.state(merged_layer) + call_kwargs = dict(layer_kwargs) + if kv_in is not None: + call_kwargs["kv_cache"] = kv_in + call_kwargs["previous_chunk"] = previous_chunk + call_kwargs["slot"] = slot + out_res = merged_layer(y_in, *layer_args, **call_kwargs) + if isinstance(out_res, tuple): + out_y = out_res[0] + out_kv = out_res[1] if len(out_res) > 1 else None + else: + out_y = out_res + out_kv = None + return out_y, out_kv, nnx.state(merged_layer) checkpointed_gemma_fn = jax.checkpoint(pure_gemma_fn, policy=policy, prevent_cse=prevent_cse) graphdef, state = nnx.split(self.layers_remainder) - y, new_state = checkpointed_gemma_fn(graphdef, state, y) + y, updated_remainder_kv, new_state = checkpointed_gemma_fn(graphdef, state, y, remainder_kv) nnx.update(self.layers_remainder, new_state) + 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( @@ -1983,6 +2007,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.""" @@ -1994,10 +2020,20 @@ def _apply_gemma4_scanned_blocks( layer_args = (decoder_segment_ids, decoder_positions, deterministic, model_mode) layer_kwargs = {"bidirectional_mask": bidirectional_mask, "slot": slot, "previous_chunk": previous_chunk} + if attention_metadata is not None: + layer_kwargs["attention_metadata"] = attention_metadata # Apply the main scan over the full blocks if scan_length > 0: - y, self.layers, _ = self._apply_layers_sequentially(self.layers, y, *layer_args, length=scan_length, **layer_kwargs) + grouped_kv_caches = maxtext_utils.prepare_kv_caches_for_scan( + kv_caches, scan_length, attention_pattern_length, stack=False + ) + y, self.layers, _ = self._apply_layers_sequentially( + self.layers, y, *layer_args, length=scan_length, kv_caches_stacked=grouped_kv_caches, **layer_kwargs + ) + maxtext_utils.update_kv_caches_after_scan( + kv_caches, grouped_kv_caches, scan_length, attention_pattern_length, stacked=False + ) # Apply any remaining layers that did not fit into a full scanned block num_remaining_layers = cfg.num_decoder_layers % attention_pattern_length @@ -2005,21 +2041,38 @@ def _apply_gemma4_scanned_blocks( policy = self.get_remat_policy() prevent_cse = maxtext_utils.should_prevent_cse_in_remat(cfg) - def pure_gemma_fn(graphdef, state_in, y_in): + 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]) + + def pure_gemma_fn(graphdef, state_in, y_in, kv_in): merged_layer = nnx.merge(graphdef, state_in) - out_y, _ = merged_layer( - y_in, - *layer_args, - **layer_kwargs, - ) - return out_y, nnx.state(merged_layer) + call_kwargs = dict(layer_kwargs) + if kv_in is not None: + call_kwargs["kv_cache"] = kv_in + call_kwargs["previous_chunk"] = previous_chunk + call_kwargs["slot"] = slot + out_res = merged_layer(y_in, *layer_args, **call_kwargs) + if isinstance(out_res, tuple): + out_y = out_res[0] + out_kv = out_res[1] if len(out_res) > 1 else None + else: + out_y = out_res + out_kv = None + return out_y, out_kv, nnx.state(merged_layer) checkpointed_gemma_fn = jax.checkpoint(pure_gemma_fn, policy=policy, prevent_cse=prevent_cse) graphdef, state = nnx.split(self.layers_remainder) - y, new_state = checkpointed_gemma_fn(graphdef, state, y) + y, updated_remainder_kv, new_state = checkpointed_gemma_fn(graphdef, state, y, remainder_kv) nnx.update(self.layers_remainder, new_state) + 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_small_layers( diff --git a/src/maxtext/models/gemma3.py b/src/maxtext/models/gemma3.py index 2d2b36a431..c33272d277 100644 --- a/src/maxtext/models/gemma3.py +++ b/src/maxtext/models/gemma3.py @@ -259,8 +259,6 @@ def update_cache(cache, val): 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 @@ -324,6 +322,8 @@ def __call__( page_state=None, previous_chunk=None, bidirectional_mask=None, + kv_cache=None, + attention_metadata=None, ): cfg = self.config @@ -331,8 +331,10 @@ def __call__( inputs = checkpoint_name(inputs, "decoder_layer_input") y = inputs + updated_kvs = [] for layer_id in range(self.num_of_layers): - y = getattr(self, f"layers_{layer_id}")( + current_kv = kv_cache[layer_id] if kv_cache is not None else None + y, new_kv = getattr(self, f"layers_{layer_id}")( y, decoder_segment_ids, decoder_positions, @@ -341,10 +343,15 @@ def __call__( previous_chunk=previous_chunk, slot=slot, bidirectional_mask=bidirectional_mask, + kv_cache=current_kv, + attention_metadata=attention_metadata, ) - if cfg.scan_layers: - y = y[0] - if cfg.scan_layers: + if kv_cache is not None: + updated_kvs.append(new_kv) + + if kv_cache is not None: + return y, tuple(updated_kvs) + elif cfg.scan_layers: return y, None else: return y diff --git a/src/maxtext/models/gemma4.py b/src/maxtext/models/gemma4.py index 8cf927629c..dbf8efc4de 100644 --- a/src/maxtext/models/gemma4.py +++ b/src/maxtext/models/gemma4.py @@ -398,8 +398,6 @@ def update_cache(cache, val): 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 @@ -464,13 +462,18 @@ def __call__( page_state=None, previous_chunk=None, bidirectional_mask=None, + kv_cache=None, + attention_metadata=None, ): + cfg = self.config inputs = nn.with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) inputs = checkpoint_name(inputs, "decoder_layer_input") y = inputs + updated_kvs = [] for layer_id in range(self.num_of_layers): - y, _ = getattr(self, f"layers_{layer_id}")( + current_kv = kv_cache[layer_id] if kv_cache is not None else None + y, new_kv = getattr(self, f"layers_{layer_id}")( y, decoder_segment_ids, decoder_positions, @@ -479,9 +482,18 @@ def __call__( previous_chunk=previous_chunk, slot=slot, bidirectional_mask=bidirectional_mask, + kv_cache=current_kv, + attention_metadata=attention_metadata, ) + if kv_cache is not None: + updated_kvs.append(new_kv) - return y, None + if kv_cache is not None: + return y, tuple(updated_kvs) + elif cfg.scan_layers: + return y, None + else: + return y Gemma4ScannableBlockToLinen = nnx_wrappers.to_linen_class( diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index 3451ec824d..47c2c7e587 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -2067,3 +2067,43 @@ def get_mesh_from_config( axis_types = tuple([AxisType.Auto] * len(config.mesh_axes)) return Mesh(devices_array, config.mesh_axes, axis_types=axis_types) + + +def prepare_kv_caches_for_scan(kv_caches, scan_length, block_len, stack=False): + """Groups the flat list of KV caches into block-sized tuples and optionally stacks them.""" + if kv_caches is None: + return None + + if not isinstance(kv_caches, list): + raise TypeError(f"kv_caches must be a list, got {type(kv_caches)}") + + grouped = [tuple(kv_caches[i * block_len : (i + 1) * block_len]) for i in range(scan_length)] + + if stack: + # Stack the list of tuples into a tuple of stacked PyTrees along a new leading axis + return jax.tree_util.tree_map(lambda *args: jnp.stack(args, axis=0), *grouped) + else: + return grouped + + +def update_kv_caches_after_scan(kv_caches, returned_kv_cache, scan_length, block_len, stacked=False): + """Updates the original flat list of KV caches from the scanned outputs.""" + if kv_caches is None or returned_kv_cache is None: + return + + if not isinstance(kv_caches, list): + raise TypeError(f"kv_caches must be a list, got {type(kv_caches)}") + + if stacked: + # Unstack the returned tuple of stacked PyTrees back into a list of tuples + unstacked = [jax.tree_util.tree_map(lambda x, i=i: x[i], returned_kv_cache) for i in range(scan_length)] + for i in range(scan_length): + start_idx = i * block_len + for offset, updated_item in enumerate(unstacked[i]): + kv_caches[start_idx + offset] = updated_item + else: + # returned_kv_cache is already a list of tuples + for i in range(scan_length): + start_idx = i * block_len + for offset, updated_item in enumerate(returned_kv_cache[i]): + kv_caches[start_idx + offset] = updated_item diff --git a/tests/unit/maxtext_utils_test.py b/tests/unit/maxtext_utils_test.py index 0d51e1e1c5..670fcce5e1 100644 --- a/tests/unit/maxtext_utils_test.py +++ b/tests/unit/maxtext_utils_test.py @@ -352,9 +352,7 @@ def setUp(self): self.mesh = Mesh(devices_array, self.config.mesh_axes) quant = quantizations.configure_quantization(self.config) if self.config.pure_nnx: - self._create_model_partial, self.model = ( - model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) - ) + self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) else: self.model = models.transformer_as_linen(self.config, mesh=self.mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) @@ -951,7 +949,21 @@ class MockConfig: def setUp(self): # Setup a dummy device array for the mock to return - self.devices_array = np.array(jax.devices()) + devices = jax.devices() + if len(devices) < 2: + mock_devices = [] + for i in range(2): + d = MagicMock(spec=jax.Device) + d.id = i + d.device_kind = "cpu" + d.platform = "cpu" + # make it hashable + d.__hash__ = lambda self, idx=i: idx + d.__eq__ = lambda self, other, idx=i: isinstance(other, MagicMock) and other.id == idx + mock_devices.append(d) + self.devices_array = np.array(mock_devices) + else: + self.devices_array = np.array(devices) @patch("maxtext.utils.maxtext_utils.create_device_mesh") def test_get_mesh_explicit_mode(self, mock_create_device_mesh): @@ -1377,13 +1389,9 @@ def setUp(self): self.mesh = Mesh(devices_array, self.config.mesh_axes) quant = quantizations.configure_quantization(self.config) if self.config.pure_nnx: - self._create_model_partial, self.model = ( - model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) - ) + self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) else: - self.model = Transformer( - self.config, mesh=self.mesh, quant=quant, model_mode=MODEL_MODE_TRAIN - ) + self.model = Transformer(self.config, mesh=self.mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) def test_setup_training_state_returns_train_state(self): rng = jax.random.PRNGKey(0) @@ -1422,13 +1430,9 @@ def setUp(self): self.mesh = Mesh(devices_array, self.config.mesh_axes) quant = quantizations.configure_quantization(self.config) if self.config.pure_nnx: - self._create_model_partial, self.model = ( - model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) - ) + self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) else: - self.model = Transformer( - self.config, mesh=self.mesh, quant=quant, model_mode=MODEL_MODE_TRAIN - ) + self.model = Transformer(self.config, mesh=self.mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) self.rng = jax.random.PRNGKey(0) self.tx = optax.adam(learning_rate=0.001) @@ -1442,9 +1446,7 @@ def create_train_state_fn(): init_state_fn = create_train_state_fn annotations = maxtext_utils_nnx.get_partition_spec_nnx( - maxtext_utils.get_abstract_state( - self.config, self.mesh, init_state_fn, True - )[2] + maxtext_utils.get_abstract_state(self.config, self.mesh, init_state_fn, True)[2] ) else: init_state_fn = functools.partial( @@ -1455,9 +1457,7 @@ def create_train_state_fn(): True, self.rng, ) - annotations = maxtext_utils.get_logical_annotations( - self.config, self.mesh, init_state_fn - ) + annotations = maxtext_utils.get_logical_annotations(self.config, self.mesh, init_state_fn) # Result should be a pytree with PartitionSpec leaves leaves = jax.tree_util.tree_leaves(annotations) self.assertGreater(len(leaves), 0) @@ -1593,9 +1593,7 @@ class MockConfig: optimizer_memory_host_offload: bool = False parameter_memory_host_offload: bool = False param_scan_axis: int = 0 - logical_axis_rules: list = field( - default_factory=lambda: [["data", ["data"]], ["model", ["model"]]] - ) + logical_axis_rules: list = field(default_factory=lambda: [["data", ["data"]], ["model", ["model"]]]) class MockTrainState(nnx.Module): """Simulates a TrainState with params and optimizer state.""" @@ -1692,5 +1690,43 @@ def test_invalid_init_fn(self): maxtext_utils.get_abstract_state_nnx(self.config, self.mesh, None) +class TestKVCacheScanHelpers(unittest.TestCase): + """Tests for KV cache scan helper functions.""" + + def test_prepare_kv_caches_for_scan_valid(self): + kv_caches = [jnp.array([1.0]), jnp.array([2.0]), jnp.array([3.0]), jnp.array([4.0])] + # scan_length = 2, block_len = 2 + grouped = maxtext_utils.prepare_kv_caches_for_scan(kv_caches, scan_length=2, block_len=2, stack=False) + self.assertEqual(len(grouped), 2) + self.assertEqual(grouped[0], (kv_caches[0], kv_caches[1])) + self.assertEqual(grouped[1], (kv_caches[2], kv_caches[3])) + + def test_prepare_kv_caches_for_scan_invalid_type(self): + kv_caches_tuple = (jnp.array([1.0]), jnp.array([2.0])) + with self.assertRaises(TypeError): + maxtext_utils.prepare_kv_caches_for_scan(kv_caches_tuple, scan_length=1, block_len=2) + + kv_caches_array = jnp.array([1.0, 2.0]) + with self.assertRaises(TypeError): + maxtext_utils.prepare_kv_caches_for_scan(kv_caches_array, scan_length=1, block_len=2) + + def test_update_kv_caches_after_scan_valid(self): + kv_caches = [jnp.array([1.0]), jnp.array([2.0]), jnp.array([3.0]), jnp.array([4.0])] + returned_kv_cache = [(jnp.array([10.0]), jnp.array([20.0])), (jnp.array([30.0]), jnp.array([40.0]))] + + maxtext_utils.update_kv_caches_after_scan(kv_caches, returned_kv_cache, scan_length=2, block_len=2, stacked=False) + + self.assertTrue(jnp.array_equal(kv_caches[0], jnp.array([10.0]))) + self.assertTrue(jnp.array_equal(kv_caches[1], jnp.array([20.0]))) + self.assertTrue(jnp.array_equal(kv_caches[2], jnp.array([30.0]))) + self.assertTrue(jnp.array_equal(kv_caches[3], jnp.array([40.0]))) + + def test_update_kv_caches_after_scan_invalid_type(self): + kv_caches_tuple = (jnp.array([1.0]), jnp.array([2.0])) + returned_kv_cache = [(jnp.array([10.0]), jnp.array([20.0]))] + with self.assertRaises(TypeError): + maxtext_utils.update_kv_caches_after_scan(kv_caches_tuple, returned_kv_cache, scan_length=1, block_len=2) + + if __name__ == "__main__": unittest.main()