A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 17:00 UTC
Fix hard-coded parameter name in gemma3n.py (#27946)
Signed-off-by: Seungduk Kim <seungduk.kim@yanolja.com> Signed-off-by: Biswa Panda <biswa.panda@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Biswa Panda <biswa.panda@gmail.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
Seungduk Kim committed
Nov 6, 2025 at 16:07 UTC
201dc98acc5f9be105281e674b46c95ad68d9fe9
1 file changed
+20
-1
vllm/model_executor/models/gemma3n.py
+20
-1
@@ -357,8 +357,27 @@ class Gemma3nAttention(nn.Module):
357
offset = 2 if self.sliding_window is not None else 1
358
kv_shared_layer_index = first_kv_shared_layer_idx - offset
359
if kv_shared_layer_index >= 0:
360
+ # Different model wrappers expose layer parameters under
361
+ # different parent attributes.
362
+ # For example:
363
+ # - Gemma3nForCausalLM → parameters live under "model.layers"
364
+ # - Gemma3nForConditionalGeneration →
365
+ # under "language_model.model.layers"
366
+ # This logic extracts the portion of the parameter name
367
+ # *before* ".layers."
368
+ # so downstream code can consistently reference the correct
369
+ # model root regardless of which wrapper class was used.
370
+ if ".layers." in prefix:
371
+ param_name_before_layers = prefix.split(".layers.")[0]
372
+ else:
373
+ raise ValueError(
374
+ "Unexpected prefix format for Gemma3nAttention: "
375
+ f"'{prefix}'. The prefix is expected to contain "
376
+ "'.layers.' to correctly determine the KV sharing "
377
+ "target layer."
378
+ )
379
# Only the greater layer is required to specify sharing.
361
- kv_sharing_target_layer_name = f"language_model.model.layers.{kv_shared_layer_index}.self_attn.attn" # noqa: E501
380
+ kv_sharing_target_layer_name = f"{param_name_before_layers}.layers.{kv_shared_layer_index}.self_attn.attn" # noqa: E501
381
382
self.rotary_emb = get_rope(
383
self.head_dim,