| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Per-layer config resolution for Gemma4 and its variants.""" |
| 4 | |
| 5 | from copy import copy |
| 6 | |
| 7 | from transformers import PreTrainedConfig |
| 8 | |
| 9 | |
| 10 | def gemma4_layer_config( |
| 11 | text_config: PreTrainedConfig, layer_idx: int |
| 12 | ) -> PreTrainedConfig: |
| 13 | """The Gemma4 text config as it applies to one layer. |
| 14 | |
| 15 | Gemma4 uses a larger head dimension on its full attention layers than on its |
| 16 | sliding ones, and with `attention_k_eq_v` it uses more KV heads there too. |
| 17 | Transformers >= 5.15.0 says so in the config itself; before that the values |
| 18 | are flat attributes that `layer_types` picks between, so resolve them here. |
| 19 | The result is homogeneous either way, so callers read `head_dim` and |
| 20 | `num_key_value_heads` off it without caring which version they are on. |
| 21 | """ |
| 22 | if getattr(text_config, "is_heterogeneous", False): |
| 23 | return text_config.per_layer_config[layer_idx] |
| 24 | |
| 25 | layer = copy(text_config) |
| 26 | if text_config.layer_types[layer_idx] == "full_attention": |
| 27 | global_head_dim = getattr(text_config, "global_head_dim", None) |
| 28 | layer.head_dim = global_head_dim or text_config.head_dim |
| 29 | if getattr(text_config, "attention_k_eq_v", False): |
| 30 | global_kv_heads = getattr(text_config, "num_global_key_value_heads", None) |
| 31 | layer.num_key_value_heads = ( |
| 32 | global_kv_heads or text_config.num_key_value_heads |
| 33 | ) |
| 34 | return layer |