Skip to content

[worker]: support Gemma 4 per-layer KV cache specs#3

Open
Carlos-Marques wants to merge 1 commit into
release-0.5.0from
devin/1781041203-gemma4-e2b-support
Open

[worker]: support Gemma 4 per-layer KV cache specs#3
Carlos-Marques wants to merge 1 commit into
release-0.5.0from
devin/1781041203-gemma4-e2b-support

Conversation

@Carlos-Marques

Copy link
Copy Markdown

Summary

Adds Gemma 4-aware KV cache spec generation for the Neuron worker.

For multimodal hf_config objects with nested text_config, get_kv_cache_spec() now derives per-layer cache properties from the text config:

layer_types[i] == "sliding_attention" -> head_size=head_dim, sliding_window=sliding_window
layer_types[i] == "full_attention"    -> head_size=global_head_dim, sliding_window=None

This is needed because Gemma 4 E2B mixes local sliding attention (head_dim=256) with full/global attention (global_head_dim=512), so a uniform self.model.head_dim cache spec is incorrect. Unit coverage exercises the nested Gemma4 text config path.

Link to Devin session: https://app.devin.ai/sessions/8501976bc5364801a8d6d8b8b768a547
Requested by: @Carlos-Marques

Co-Authored-By: carlos <carlosmarques.personal@gmail.com>

Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
@devin-ai-integration

Copy link
Copy Markdown

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

  • Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
  • Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

  • Disable automatic comment, CI, and merge conflict monitoring

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant