diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py
index 98c5839c515..51d414f8094 100644
--- a/backends/qualcomm/tests/test_qnn_delegate.py
+++ b/backends/qualcomm/tests/test_qnn_delegate.py
@@ -8144,6 +8144,14 @@ def setUp(self):
                 hellaswag_acc_norm=None,
                 sqnr=10,
             ),
+            "gemma4-e2b": TestExampleLLMScript.LlmSpecs(
+                SM8650=20,
+                SM8750=30,
+                pte_size=4_500_000_000,  # 4.5 GB
+                wikitext_ppl=120,
+                hellaswag_acc_norm=None,
+                sqnr=10,
+            ),
             "glm-1_5b": TestExampleLLMScript.LlmSpecs(
                 SM8650=42,
                 SM8750=52,
@@ -8345,6 +8353,14 @@ def test_static_llm_model(self):  # noqa: C901
                 ]
             )
 
+        if self.model_name == "gemma4-e2b":
+            cmds.extend(
+                [
+                    "--embedding-quantize",
+                    "4,32",
+                ]
+            )
+
         self.add_default_cmds(cmds)
 
         p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL)
diff --git a/examples/models/gemma4/text_decoder/convert_weights.py b/examples/models/gemma4/text_decoder/convert_weights.py
index e6ba1488a79..14628d52f20 100644
--- a/examples/models/gemma4/text_decoder/convert_weights.py
+++ b/examples/models/gemma4/text_decoder/convert_weights.py
@@ -17,7 +17,7 @@
 
 import logging
 from pathlib import Path
-from typing import Dict, Optional
+from typing import Callable, Dict, Optional
 
 import torch
 
@@ -227,6 +227,7 @@ def convert_hf_to_custom(
     checkpoint_path: str,
     config: Gemma4Config,
     dtype: Optional[torch.dtype] = None,
+    get_weight_mapping: Optional[Callable] = False,
 ) -> Dict[str, torch.Tensor]:
     """Convert HuggingFace checkpoint to custom format.
 
@@ -242,7 +243,10 @@ def convert_hf_to_custom(
 
     hf_state_dict = _load_safetensors_weights(checkpoint_path)
 
-    mapping = _get_weight_mapping(config)
+    if not get_weight_mapping:
+        mapping = _get_weight_mapping(config)
+    else:
+        mapping = get_weight_mapping(config)
 
     converted_state_dict = {}
     unmapped_keys = []
diff --git a/examples/models/llama/evaluate/eager_eval.py b/examples/models/llama/evaluate/eager_eval.py
index 5c129e1c250..cefb951281c 100644
--- a/examples/models/llama/evaluate/eager_eval.py
+++ b/examples/models/llama/evaluate/eager_eval.py
@@ -50,6 +50,8 @@ def eot_token_id(self):
 
     @property
     def prefix_token_id(self):
+        if hasattr(self._tokenizer, "bos_id"):
+            return self._tokenizer.bos_id
         return self.eot_token_id
 
     @property
diff --git a/examples/qualcomm/oss_scripts/gemma4/model_wrapper.py b/examples/qualcomm/oss_scripts/gemma4/model_wrapper.py
new file mode 100644
index 00000000000..f555a2480ad
--- /dev/null
+++ b/examples/qualcomm/oss_scripts/gemma4/model_wrapper.py
@@ -0,0 +1,232 @@
+# Copyright (c) Qualcomm Innovation Center, Inc.
+# All rights reserved
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import List, Optional, Tuple
+
+import torch
+
+from executorch.examples.models.gemma4.text_decoder.gemma4_config import Gemma4Config
+from executorch.examples.qualcomm.oss_scripts.gemma4.text_decoder.text_model import (
+    Gemma4TextModel,
+)
+from executorch.examples.qualcomm.oss_scripts.llama.masking_utils import (
+    AttentionMask,
+    CausalAttentionMask,
+    SlidingWindowAttentionMask,
+)
+from torch import nn
+
+
+def _precompute_freqs(
+    head_dim: int,
+    max_context_len: int,
+    rope_theta: float,
+    partial_rotary_factor: float,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """HuggingFace-style RoPE cos/sin with zero-padding for partial RoPE.
+
+    Returns (freqs_cos, freqs_sin) each of shape [max_context_len, head_dim // 2].
+    Non-rotated dims are zero-padded so ROTARY_EMB_REGISTRY["partial"] handles
+    the half-and-half split correctly.
+    """
+    rotary_dim = int(head_dim * partial_rotary_factor)
+    half_rot = rotary_dim // 2
+    inv_freq_rot = 1.0 / (
+        rope_theta ** (torch.arange(0, rotary_dim, 2).float() / head_dim)
+    )
+    total_half = head_dim // 2
+    nope = total_half - half_rot
+    inv_freq = (
+        torch.cat([inv_freq_rot, torch.zeros(nope)]) if nope > 0 else inv_freq_rot
+    )
+    t = torch.arange(max_context_len, dtype=torch.float32)
+    freqs = torch.outer(t, inv_freq)
+    return torch.cos(freqs), torch.sin(freqs)
+
+
+class Gemma4TextModelWrapper(nn.Module):
+    """QNN-export-ready Gemma 4 text decoder (static_llama.py convention).
+
+    KV-cache I/O covers only the 15 self-decoder layers.  Each of these layers
+    has its own head_dim (full-attention=512, sliding=256), so cache shapes
+    are heterogeneous.
+
+    forward(tokens, atten_mask, window_atten_mask, input_pos,
+            *k_caches[0..14], *v_caches[0..14])
+        -> (logits, out_k_caches, out_v_caches)
+    """
+
+    def __init__(
+        self,
+        config: Gemma4Config,
+        ar_len: int = 1,
+        output_new_cache_only: bool = True,
+        output_cache: bool = True,
+        use_i64_token: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+        self.config = config
+        self.ar_len = ar_len
+        self.use_kv_cache = config.use_kv_cache
+        self.output_cache = output_cache
+        self.use_i64_token = use_i64_token
+        self.n_self_layers = config.num_self_decoder_layers
+        self.max_batch_size = config.max_batch_size
+        self.vocab_size = config.vocab_size
+        self.kv_io_bit_width = kwargs.get("kv_io_bit_width", 32)
+
+        self.model = Gemma4TextModel(
+            config,
+            output_new_cache_only,
+            kwargs.get("enable_masked_softmax", False),
+        )
+
+        # Precompute RoPE buffers.
+        # Global: full-attention layers, partial_rotary_factor=0.25, head_dim=512.
+        # Local:  sliding-attention layers, full rotation, head_dim=256.
+        global_cos, global_sin = _precompute_freqs(
+            config.global_head_dim,
+            config.max_context_len,
+            config.rope_theta,
+            config.partial_rotary_factor,
+        )
+        local_cos, local_sin = _precompute_freqs(
+            config.head_dim,
+            config.max_context_len,
+            config.rope_local_base_freq,
+            1.0,
+        )
+        self.register_buffer("global_freqs_cos", global_cos, persistent=False)
+        self.register_buffer("global_freqs_sin", global_sin, persistent=False)
+        self.register_buffer("local_freqs_cos", local_cos, persistent=False)
+        self.register_buffer("local_freqs_sin", local_sin, persistent=False)
+
+    @property
+    def layers(self) -> List[nn.Module]:
+        """Flat view of all 35 decoder layers (self + cross).
+
+        Satisfies callers that expect decoder.layers to iterate all layers,
+        matching the LlamaModel convention. The actual nn.ModuleLists live in
+        self.model.self_decoder.layers / self.model.cross_decoder.layers
+        to preserve weight-path naming for checkpoint loading.
+        """
+        return list(self.model.self_decoder.layers) + list(
+            self.model.cross_decoder.layers
+        )
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        atten_mask: torch.Tensor,
+        window_atten_mask: torch.Tensor,
+        input_pos: Optional[torch.Tensor] = None,
+        *args,
+    ) -> Tuple[
+        torch.Tensor, List[Optional[torch.Tensor]], List[Optional[torch.Tensor]]
+    ]:
+        if self.use_kv_cache and input_pos is not None:
+            global_cos = self.global_freqs_cos[input_pos][0]
+            global_sin = self.global_freqs_sin[input_pos][0]
+            local_cos = self.local_freqs_cos[input_pos][0]
+            local_sin = self.local_freqs_sin[input_pos][0]
+        else:
+            global_cos = self.global_freqs_cos
+            global_sin = self.global_freqs_sin
+            local_cos = self.local_freqs_cos
+            local_sin = self.local_freqs_sin
+
+        k_caches = list(args[: self.n_self_layers]) if self.use_kv_cache else []
+        v_caches = (
+            list(args[self.n_self_layers : 2 * self.n_self_layers])
+            if self.use_kv_cache
+            else []
+        )
+
+        logits, output_k_cache, output_v_cache = self.model(
+            tokens,
+            atten_mask,
+            window_atten_mask,
+            input_pos,
+            global_cos,
+            global_sin,
+            local_cos,
+            local_sin,
+            k_caches,
+            v_caches,
+        )
+
+        if self.output_cache:
+            return logits, output_k_cache, output_v_cache
+        return logits
+
+    def get_example_inputs(self):
+        dtype = torch.int64 if self.use_i64_token else torch.int32
+        tokens = torch.randint(
+            self.config.vocab_size,
+            (self.config.max_batch_size, self.ar_len),
+            dtype=dtype,
+        )
+        causal_mask = CausalAttentionMask(
+            self.config.max_batch_size, self.ar_len, self.config.max_context_len
+        )
+        sliding_mask = SlidingWindowAttentionMask(
+            self.config.max_batch_size,
+            self.ar_len,
+            self.config.max_context_len,
+            sliding_window=self.config.sliding_window,
+        )
+        attn_mask = AttentionMask([causal_mask, sliding_mask])
+
+        if self.use_kv_cache:
+            pos_ids = torch.zeros(
+                (self.config.max_batch_size, self.ar_len), dtype=torch.int32
+            )
+            k_caches, v_caches = [], []
+            for i in range(self.n_self_layers):
+                hd = self.config.get_head_dim(i)
+                k_caches.append(
+                    torch.zeros(
+                        self.config.max_batch_size,
+                        self.config.num_key_value_heads,
+                        hd,
+                        self.config.max_context_len - self.ar_len,
+                    )
+                )
+                v_caches.append(
+                    torch.zeros(
+                        self.config.max_batch_size,
+                        self.config.num_key_value_heads,
+                        self.config.max_context_len - self.ar_len,
+                        hd,
+                    )
+                )
+            return (tokens, attn_mask, pos_ids, k_caches, v_caches)
+
+        return (tokens, attn_mask)
+
+    def get_metadata(self):
+        return {
+            "get_ar_len": self.ar_len,
+            "get_bos_id": 2,
+            "get_eos_id": 1,
+            "get_dim": self.config.hidden_size,
+            "get_head_dim": self.config.head_dim,
+            "get_global_head_dim": self.config.global_head_dim,
+            "get_max_batch_size": self.config.max_batch_size,
+            "get_max_seq_len": self.config.max_seq_len,
+            "get_max_context_len": self.config.max_context_len,
+            "get_n_bos": 1,
+            "get_n_eos": 1,
+            "get_n_kv_heads": self.config.num_key_value_heads,
+            "get_n_layers": self.config.num_hidden_layers,
+            "get_n_self_layers": self.n_self_layers,
+            "get_vocab_size": self.config.vocab_size,
+            "get_use_kv_cache": self.use_kv_cache,
+            "get_sliding_window": self.config.sliding_window,
+            "get_num_kv_shared_layers": self.config.num_kv_shared_layers,
+            "get_kv_io_bit_width": self.kv_io_bit_width,
+        }
diff --git a/examples/qualcomm/oss_scripts/gemma4/text_decoder/__init__.py b/examples/qualcomm/oss_scripts/gemma4/text_decoder/__init__.py
new file mode 100644
index 00000000000..c8942323171
--- /dev/null
+++ b/examples/qualcomm/oss_scripts/gemma4/text_decoder/__init__.py
@@ -0,0 +1,9 @@
+# Copyright (c) Qualcomm Innovation Center, Inc.
+# All rights reserved
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .text_model import Gemma4TextModel  # noqa: F401
+
+__all__ = ["Gemma4TextModel"]
diff --git a/examples/qualcomm/oss_scripts/gemma4/text_decoder/attention.py b/examples/qualcomm/oss_scripts/gemma4/text_decoder/attention.py
new file mode 100644
index 00000000000..b309663fc8b
--- /dev/null
+++ b/examples/qualcomm/oss_scripts/gemma4/text_decoder/attention.py
@@ -0,0 +1,253 @@
+# Copyright (c) Qualcomm Innovation Center, Inc.
+# All rights reserved
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Optional, Tuple
+
+import torch
+
+from executorch.examples.models.gemma4.text_decoder.gemma4_config import Gemma4Config
+from executorch.examples.qualcomm.oss_scripts.llama.model import (
+    NORM_REGISTRY,
+    repeat_kv,
+    ROTARY_EMB_REGISTRY,
+)
+from torch import nn
+
+
+class Gemma4Attention(nn.Module):
+    """Gemma 4 attention in static (flat KV I/O) style.
+
+    Handles per-layer head_dim, QK-norm (before RoPE), scaleless V-norm,
+    partial/full RoPE, MQA (n_kv_heads=1), and YOCO donor/shared roles.
+    """
+
+    def __init__(
+        self,
+        config: Gemma4Config,
+        layer_idx: int,
+        output_new_cache_only: bool = True,
+        enable_masked_softmax: bool = False,
+    ):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.get_head_dim(layer_idx)
+        self.num_key_value_groups = self.num_heads // self.num_kv_heads
+
+        self.is_sliding = config.is_sliding_attention(layer_idx)
+        self.is_kv_shared_layer = config.is_kv_shared_layer(layer_idx)
+        self.is_kv_donor_layer = config.is_kv_donor_layer(layer_idx)
+        self.output_new_cache_only = output_new_cache_only
+        self.enable_masked_softmax = enable_masked_softmax
+
+        # Gemma 4 uses scale=1.0; QK-norm handles normalisation
+        self.scaling = 1.0
+
+        # Q projection is present on every layer
+        self.q_proj = nn.Linear(
+            self.hidden_size, self.num_heads * self.head_dim, bias=False
+        )
+        self.q_norm = NORM_REGISTRY["rmsnorm"](self.head_dim, eps=config.rms_norm_eps)
+
+        # KV projections and norms are always created (checkpoint compatibility).
+        # For YOCO-shared layers these weights are loaded but not used in forward.
+        self.k_proj = nn.Linear(
+            self.hidden_size, self.num_kv_heads * self.head_dim, bias=False
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size, self.num_kv_heads * self.head_dim, bias=False
+        )
+        self.k_norm = NORM_REGISTRY["rmsnorm"](self.head_dim, eps=config.rms_norm_eps)
+        self.v_norm = nn.RMSNorm(
+            self.head_dim, eps=config.rms_norm_eps, elementwise_affine=False
+        )
+
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim, self.hidden_size, bias=False
+        )
+
+        self.apply_rope = (
+            ROTARY_EMB_REGISTRY["default"]
+            if self.is_sliding
+            else ROTARY_EMB_REGISTRY["partial"]
+        )
+
+        self.attn_softmax = nn.Softmax(dim=-1)
+
+    def prepare_attention_conv(self):
+        self.q_proj_conv = nn.Conv2d(
+            self.hidden_size, self.num_heads * self.head_dim, 1, bias=False
+        )
+        self.k_proj_conv = nn.Conv2d(
+            self.hidden_size, self.num_kv_heads * self.head_dim, 1, bias=False
+        )
+        self.v_proj_conv = nn.Conv2d(
+            self.hidden_size, self.num_kv_heads * self.head_dim, 1, bias=False
+        )
+        self.o_proj_conv = nn.Conv2d(
+            self.num_heads * self.head_dim, self.hidden_size, 1, bias=False
+        )
+
+        self.forward_no_conv = self.forward
+        self.forward = self.forward_attention_conv
+
+        self.q_proj_conv.weight.data.copy_(self.q_proj.weight[:, :, None, None])
+        self.k_proj_conv.weight.data.copy_(self.k_proj.weight[:, :, None, None])
+        self.v_proj_conv.weight.data.copy_(self.v_proj.weight[:, :, None, None])
+        self.o_proj_conv.weight.data.copy_(self.o_proj.weight[:, :, None, None])
+
+        del self.q_proj
+        del self.k_proj
+        del self.v_proj
+        del self.o_proj
+
+    def forward_attention_conv(
+        self,
+        hidden_states: torch.Tensor,
+        freqs_cos: torch.Tensor,
+        freqs_sin: torch.Tensor,
+        atten_mask: torch.Tensor,
+        k_cache: Optional[torch.Tensor] = None,
+        v_cache: Optional[torch.Tensor] = None,
+        donor_k: Optional[torch.Tensor] = None,
+        donor_v: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+        bsz, seq_len, _ = hidden_states.shape
+        hs = torch.reshape(
+            hidden_states, (bsz, seq_len, 1, self.hidden_size)
+        ).transpose(1, 3)
+
+        q = self.q_proj_conv(hs)
+        q = q.permute(0, 3, 1, 2).squeeze(-1)
+        q = q.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        q = self.q_norm(q)
+
+        if self.is_kv_shared_layer:
+            q = self.apply_rope(q, freqs_cos, freqs_sin)
+            kh = donor_k
+            vh = donor_v
+            new_k, new_v = None, None
+        else:
+            k = self.k_proj_conv(hs)
+            v = self.v_proj_conv(hs)
+            k = k.permute(0, 3, 1, 2).squeeze(-1)
+            v = v.permute(0, 3, 1, 2).squeeze(-1)
+            k = k.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+            v = v.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+            k = self.k_norm(k)
+            v = self.v_norm(v)
+
+            q = self.apply_rope(q, freqs_cos, freqs_sin)
+            k = self.apply_rope(k, freqs_cos, freqs_sin)
+
+            k = k.transpose(2, 3)
+
+            if k_cache is not None:
+                kh = torch.cat([k_cache, k], dim=-1)
+                vh = torch.cat([v_cache, v], dim=2)
+            else:
+                kh = k
+                vh = v
+
+            new_k = k if self.output_new_cache_only else kh
+            new_v = v if self.output_new_cache_only else vh
+
+        kh = repeat_kv(kh, self.num_key_value_groups)
+        vh = repeat_kv(vh, self.num_key_value_groups)
+
+        attn = q @ kh
+        attn = attn * self.scaling
+        if self.enable_masked_softmax:
+            attn_min = torch.amin(attn, dim=-1, keepdim=True)
+            minus_value = -20
+            attn = torch.where(atten_mask == 0, attn, attn_min + minus_value)
+        else:
+            attn = attn + atten_mask
+        attn = self.attn_softmax(attn)
+
+        y = attn @ vh
+        y = y.transpose(1, 2)
+        y = y.reshape(bsz, seq_len, 1, -1).transpose(1, 3)
+        y = self.o_proj_conv(y)
+        y = y.transpose(1, 3)
+        y = y.reshape(bsz, seq_len, -1)
+
+        return y, new_k, new_v
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        freqs_cos: torch.Tensor,
+        freqs_sin: torch.Tensor,
+        atten_mask: torch.Tensor,
+        k_cache: Optional[torch.Tensor] = None,
+        v_cache: Optional[torch.Tensor] = None,
+        donor_k: Optional[torch.Tensor] = None,
+        donor_v: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+        """
+        K cache layout (static_llama convention): [b, n_kv, head_dim, ctx]
+        V cache layout:                           [b, n_kv, ctx, head_dim]
+
+        Returns (output, new_k, new_v):
+          - self-decoder / donor layers: new_k/new_v hold the new cache slice
+            (or full cache if output_new_cache_only=False).
+          - non-donor shared layers: (None, None) — no KV I/O.
+        """
+        bsz, seq_len, _ = hidden_states.shape
+
+        q = self.q_proj(hidden_states)
+        q = q.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        q = self.q_norm(q)
+
+        if self.is_kv_shared_layer:
+            q = self.apply_rope(q, freqs_cos, freqs_sin)
+            kh = donor_k
+            vh = donor_v
+            new_k, new_v = None, None
+        else:
+            k = self.k_proj(hidden_states)
+            v = self.v_proj(hidden_states)
+            k = k.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+            v = v.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+            k = self.k_norm(k)
+            v = self.v_norm(v)
+
+            q = self.apply_rope(q, freqs_cos, freqs_sin)
+            k = self.apply_rope(k, freqs_cos, freqs_sin)
+
+            k = k.transpose(2, 3)
+
+            if k_cache is not None:
+                kh = torch.cat([k_cache, k], dim=-1)
+                vh = torch.cat([v_cache, v], dim=2)
+            else:
+                kh = k
+                vh = v
+
+            new_k = k if self.output_new_cache_only else kh
+            new_v = v if self.output_new_cache_only else vh
+
+        kh = repeat_kv(kh, self.num_key_value_groups)
+        vh = repeat_kv(vh, self.num_key_value_groups)
+
+        attn = q @ kh
+        attn = attn * self.scaling
+        if self.enable_masked_softmax:
+            attn_min = torch.amin(attn, dim=-1, keepdim=True)
+            minus_value = -20
+            attn = torch.where(atten_mask == 0, attn, attn_min + minus_value)
+        else:
+            attn = attn + atten_mask
+        attn = self.attn_softmax(attn)
+
+        y = attn @ vh
+        y = y.transpose(1, 2).reshape(bsz, seq_len, -1)
+        y = self.o_proj(y)
+
+        return y, new_k, new_v
diff --git a/examples/qualcomm/oss_scripts/gemma4/text_decoder/convert_weights.py b/examples/qualcomm/oss_scripts/gemma4/text_decoder/convert_weights.py
new file mode 100644
index 00000000000..b61d7a1b84d
--- /dev/null
+++ b/examples/qualcomm/oss_scripts/gemma4/text_decoder/convert_weights.py
@@ -0,0 +1,26 @@
+# Copyright (c) Qualcomm Innovation Center, Inc.
+# All rights reserved
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Re-key weights from upstream convert_hf_to_custom format to QNN decoder format.
+
+Upstream uses:  *.self_attn.*  and  *.mlp.*
+QNN uses:       *.attention.*  and  *.feed_forward.*
+"""
+
+from typing import Dict
+
+import torch
+
+
+def remap_keys(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    remapped = {}
+    for key, tensor in state_dict.items():
+        new_key = key.replace(".self_attn.", ".attention.").replace(
+            ".mlp.", ".feed_forward."
+        )
+        remapped[new_key] = tensor
+    return remapped
diff --git a/examples/qualcomm/oss_scripts/gemma4/text_decoder/cross_decoder.py b/examples/qualcomm/oss_scripts/gemma4/text_decoder/cross_decoder.py
new file mode 100644
index 00000000000..1bdf773b01b
--- /dev/null
+++ b/examples/qualcomm/oss_scripts/gemma4/text_decoder/cross_decoder.py
@@ -0,0 +1,80 @@
+# Copyright (c) Qualcomm Innovation Center, Inc.
+# All rights reserved
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Dict, List, Tuple
+
+import torch
+
+from executorch.examples.models.gemma4.text_decoder.gemma4_config import Gemma4Config
+from executorch.examples.qualcomm.oss_scripts.gemma4.text_decoder.decoder_layer import (
+    Gemma4DecoderLayer,
+)
+from torch import nn
+
+
+class Gemma4CrossDecoder(nn.Module):
+    """Cross-decoder: last 20 layers with KV-sharing"""
+
+    def __init__(
+        self,
+        config: Gemma4Config,
+        output_new_cache_only: bool = True,
+        enable_masked_softmax: bool = False,
+    ):
+        super().__init__()
+        self.config = config
+        self.start_layer_idx = config.num_self_decoder_layers  # 35 - 20 - 15
+        self.num_layers = config.num_cross_decoder_layers
+
+        self.layers = nn.ModuleList(
+            [
+                Gemma4DecoderLayer(
+                    config,
+                    self.start_layer_idx + i,
+                    output_new_cache_only,
+                    enable_masked_softmax,
+                    use_cache=False,
+                )
+                for i in range(self.num_layers)
+            ]
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        per_layer_inputs: torch.Tensor,
+        global_freqs_cos: torch.Tensor,
+        global_freqs_sin: torch.Tensor,
+        local_freqs_cos: torch.Tensor,
+        local_freqs_sin: torch.Tensor,
+        atten_mask: torch.Tensor,
+        window_atten_mask: torch.Tensor,
+        layer_types: List[str],
+        shared_kv: Dict[int, Tuple[torch.Tensor, torch.Tensor]],
+    ) -> torch.Tensor:
+        for i, layer in enumerate(self.layers):
+            layer_idx = self.start_layer_idx + i
+            kv_src = self.config.get_kv_shared_layer_index(layer_idx)
+            donor_k, donor_v = None, None
+            if kv_src is not None and kv_src in shared_kv:
+                donor_k, donor_v = shared_kv[kv_src]
+
+            if layer_types[layer_idx] == "sliding_attention":
+                fc, fs, mask = local_freqs_cos, local_freqs_sin, window_atten_mask
+            else:
+                fc, fs, mask = global_freqs_cos, global_freqs_sin, atten_mask
+
+            hidden_states = layer(
+                hidden_states,
+                per_layer_inputs[i],
+                fc,
+                fs,
+                mask,
+                donor_k=donor_k,
+                donor_v=donor_v,
+            )
+
+        return hidden_states
diff --git a/examples/qualcomm/oss_scripts/gemma4/text_decoder/decoder_layer.py b/examples/qualcomm/oss_scripts/gemma4/text_decoder/decoder_layer.py
new file mode 100644
index 00000000000..85e04dee2a3
--- /dev/null
+++ b/examples/qualcomm/oss_scripts/gemma4/text_decoder/decoder_layer.py
@@ -0,0 +1,114 @@
+# Copyright (c) Qualcomm Innovation Center, Inc.
+# All rights reserved
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+
+from executorch.examples.models.gemma4.text_decoder.gemma4_config import Gemma4Config
+from executorch.examples.qualcomm.oss_scripts.gemma4.text_decoder.attention import (
+    Gemma4Attention,
+)
+from executorch.examples.qualcomm.oss_scripts.gemma4.text_decoder.feed_forward import (
+    Gemma4MLP,
+)
+from executorch.examples.qualcomm.oss_scripts.llama.model import NORM_REGISTRY
+from torch import nn
+
+
+class Gemma4DecoderLayer(nn.Module):
+    def __init__(
+        self,
+        config: Gemma4Config,
+        layer_idx: int,
+        output_new_cache_only: bool = True,
+        enable_masked_softmax: bool = False,
+        use_cache: bool = True,
+    ):
+        super().__init__()
+        self.use_res_clamp = config.use_res_clamp
+
+        self.attention = Gemma4Attention(
+            config, layer_idx, output_new_cache_only, enable_masked_softmax
+        )
+        self.feed_forward = Gemma4MLP(
+            config.hidden_size, config.get_intermediate_size(layer_idx)
+        )
+        self.layer_scalar = nn.Parameter(torch.ones(1))
+
+        # Per-Layer Embedding projections
+        self.per_layer_input_gate = nn.Linear(
+            config.hidden_size, config.hidden_size_per_layer_input, bias=False
+        )
+        self.per_layer_projection = nn.Linear(
+            config.hidden_size_per_layer_input, config.hidden_size, bias=False
+        )
+
+        RMSNorm = NORM_REGISTRY["rmsnorm"]
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.pre_feedforward_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.post_feedforward_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.post_per_layer_input_norm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.use_cache = use_cache
+
+    def _res_clamp(self, x: torch.Tensor) -> torch.Tensor:
+        if not self.use_res_clamp:
+            return x
+        finfo = torch.finfo(torch.float16)
+        return torch.clamp(x, finfo.min, finfo.max)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        per_layer_input: torch.Tensor,
+        freqs_cos: torch.Tensor,
+        freqs_sin: torch.Tensor,
+        atten_mask: torch.Tensor,
+        k_cache: Optional[torch.Tensor] = None,
+        v_cache: Optional[torch.Tensor] = None,
+        donor_k: Optional[torch.Tensor] = None,
+        donor_v: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+        """Returns (hidden_states, new_k, new_v); new_k/new_v are None for shared non-donor layers."""
+        residual = hidden_states
+        h, new_k, new_v = self.attention(
+            self.input_layernorm(hidden_states),
+            freqs_cos,
+            freqs_sin,
+            atten_mask,
+            k_cache,
+            v_cache,
+            donor_k,
+            donor_v,
+        )
+        hidden_states = self._res_clamp(residual + self.post_attention_layernorm(h))
+
+        residual = hidden_states
+        h = self.feed_forward(self.pre_feedforward_layernorm(hidden_states))
+        hidden_states = self._res_clamp(residual + self.post_feedforward_layernorm(h))
+
+        # PLE block
+        residual = hidden_states
+        gated = F.gelu(self.per_layer_input_gate(hidden_states), approximate="tanh")
+        projected = self.per_layer_projection(gated * per_layer_input)
+        hidden_states = (
+            residual + self.post_per_layer_input_norm(projected)
+        ) * self.layer_scalar
+
+        if self.use_cache:
+            return hidden_states, new_k, new_v
+        else:
+            return hidden_states
diff --git a/examples/qualcomm/oss_scripts/gemma4/text_decoder/feed_forward.py b/examples/qualcomm/oss_scripts/gemma4/text_decoder/feed_forward.py
new file mode 100644
index 00000000000..c3742d21a99
--- /dev/null
+++ b/examples/qualcomm/oss_scripts/gemma4/text_decoder/feed_forward.py
@@ -0,0 +1,56 @@
+# Copyright (c) Qualcomm Innovation Center, Inc.
+# All rights reserved
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class Gemma4MLP(nn.Module):
+    def __init__(self, hidden_size: int, intermediate_size: int):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def prepare_feedforward_conv(self):
+        self.gate_proj_conv = nn.Conv2d(
+            self.hidden_size, self.intermediate_size, 1, bias=False
+        )
+        self.up_proj_conv = nn.Conv2d(
+            self.hidden_size, self.intermediate_size, 1, bias=False
+        )
+        self.down_proj_conv = nn.Conv2d(
+            self.intermediate_size, self.hidden_size, 1, bias=False
+        )
+
+        self.forward_no_conv = self.forward
+        self.forward = self.forward_feedforward_conv
+
+        self.gate_proj_conv.weight.data.copy_(self.gate_proj.weight[:, :, None, None])
+        self.up_proj_conv.weight.data.copy_(self.up_proj.weight[:, :, None, None])
+        self.down_proj_conv.weight.data.copy_(self.down_proj.weight[:, :, None, None])
+
+        del self.gate_proj
+        del self.up_proj
+        del self.down_proj
+
+    def forward_feedforward_conv(self, x: torch.Tensor) -> torch.Tensor:
+        bsz, seq_len, _ = x.size()
+        x = torch.reshape(x, (bsz, seq_len, 1, self.hidden_size)).transpose(1, 3)
+        x = self.down_proj_conv(
+            F.gelu(self.gate_proj_conv(x), approximate="tanh") * self.up_proj_conv(x)
+        )
+        x = x.transpose(1, 3)
+        x = torch.reshape(x, (bsz, seq_len, -1))
+        return x
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(
+            F.gelu(self.gate_proj(x), approximate="tanh") * self.up_proj(x)
+        )
diff --git a/examples/qualcomm/oss_scripts/gemma4/text_decoder/self_decoder.py b/examples/qualcomm/oss_scripts/gemma4/text_decoder/self_decoder.py
new file mode 100644
index 00000000000..ba9d874e9d6
--- /dev/null
+++ b/examples/qualcomm/oss_scripts/gemma4/text_decoder/self_decoder.py
@@ -0,0 +1,80 @@
+# Copyright (c) Qualcomm Innovation Center, Inc.
+# All rights reserved
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+
+import torch
+
+from executorch.examples.models.gemma4.text_decoder.gemma4_config import Gemma4Config
+from executorch.examples.qualcomm.oss_scripts.gemma4.text_decoder.decoder_layer import (
+    Gemma4DecoderLayer,
+)
+from executorch.examples.qualcomm.oss_scripts.llama.model import NORM_REGISTRY
+from torch import nn
+
+
+class Gemma4SelfDecoder(nn.Module):
+    def __init__(
+        self,
+        config: Gemma4Config,
+        output_new_cache_only: bool = True,
+        enable_masked_softmax: bool = False,
+    ):
+        super().__init__()
+        self.config = config
+        self.num_layers = config.num_self_decoder_layers
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.embed_scale = math.sqrt(config.hidden_size)
+
+        self.embed_tokens_per_layer = nn.Embedding(
+            config.vocab_size_per_layer_input,
+            config.num_hidden_layers * config.hidden_size_per_layer_input,
+        )
+        self.embed_scale_per_layer = math.sqrt(config.hidden_size_per_layer_input)
+
+        self.per_layer_model_projection = nn.Linear(
+            config.hidden_size,
+            config.num_hidden_layers * config.hidden_size_per_layer_input,
+            bias=False,
+        )
+        self.per_layer_projection_norm = NORM_REGISTRY["rmsnorm"](
+            config.hidden_size_per_layer_input, eps=config.rms_norm_eps
+        )
+
+        self._ple_input_scale = 1.0 / math.sqrt(2.0)
+        self._ple_proj_scale = 1.0 / math.sqrt(config.hidden_size)
+
+        self.layers = nn.ModuleList(
+            [
+                Gemma4DecoderLayer(
+                    config, i, output_new_cache_only, enable_masked_softmax
+                )
+                for i in range(self.num_layers)
+            ]
+        )
+
+    def compute_per_layer_inputs(
+        self,
+        tokens: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        bsz, seq_len, _ = hidden_states.shape
+        n_layers = self.config.num_hidden_layers
+        hpl = self.config.hidden_size_per_layer_input
+
+        per_layer_proj = (
+            self.per_layer_model_projection(hidden_states) * self._ple_proj_scale
+        )
+        per_layer_proj = per_layer_proj.view(bsz, seq_len, n_layers, hpl)
+        per_layer_proj = self.per_layer_projection_norm(per_layer_proj)
+        per_layer_embed = (
+            self.embed_tokens_per_layer(tokens) * self.embed_scale_per_layer
+        )
+        per_layer_embed = per_layer_embed.view(bsz, seq_len, n_layers, hpl)
+
+        combined = (per_layer_proj + per_layer_embed) * self._ple_input_scale
+        return combined.permute(2, 0, 1, 3)
diff --git a/examples/qualcomm/oss_scripts/gemma4/text_decoder/text_model.py b/examples/qualcomm/oss_scripts/gemma4/text_decoder/text_model.py
new file mode 100644
index 00000000000..7d3635fd59c
--- /dev/null
+++ b/examples/qualcomm/oss_scripts/gemma4/text_decoder/text_model.py
@@ -0,0 +1,129 @@
+# Copyright (c) Qualcomm Innovation Center, Inc.
+# All rights reserved
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Dict, List, Optional, Tuple
+
+import torch
+
+from executorch.examples.models.gemma4.text_decoder.gemma4_config import Gemma4Config
+from executorch.examples.qualcomm.oss_scripts.gemma4.text_decoder.cross_decoder import (
+    Gemma4CrossDecoder,
+)
+from executorch.examples.qualcomm.oss_scripts.gemma4.text_decoder.self_decoder import (
+    Gemma4SelfDecoder,
+)
+from executorch.examples.qualcomm.oss_scripts.llama.model import NORM_REGISTRY
+from torch import nn
+
+
+class Gemma4TextModel(nn.Module):
+
+    def __init__(
+        self,
+        config: Gemma4Config,
+        output_new_cache_only: bool = True,
+        enable_masked_softmax: bool = False,
+    ):
+        super().__init__()
+        self.config = config
+
+        self.self_decoder = Gemma4SelfDecoder(
+            config, output_new_cache_only, enable_masked_softmax
+        )
+        self.cross_decoder = Gemma4CrossDecoder(
+            config, output_new_cache_only, enable_masked_softmax
+        )
+        self.norm = NORM_REGISTRY["rmsnorm"](
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.final_logit_softcapping = config.final_logit_softcapping
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        atten_mask: torch.Tensor,
+        window_atten_mask: torch.Tensor,
+        input_pos: Optional[torch.Tensor],
+        global_freqs_cos: torch.Tensor,
+        global_freqs_sin: torch.Tensor,
+        local_freqs_cos: torch.Tensor,
+        local_freqs_sin: torch.Tensor,
+        k_caches: List[torch.Tensor],
+        v_caches: List[torch.Tensor],
+    ) -> Tuple[
+        torch.Tensor, List[Optional[torch.Tensor]], List[Optional[torch.Tensor]]
+    ]:
+        layer_types = self.config.layer_types
+        n_self = self.config.num_self_decoder_layers
+
+        hidden_states = (
+            self.self_decoder.embed_tokens(tokens) * self.self_decoder.embed_scale
+        )
+        per_layer_inputs = self.self_decoder.compute_per_layer_inputs(
+            tokens, hidden_states
+        )
+
+        # Self-decoder layers (KV I/O)
+        output_k_cache: List[Optional[torch.Tensor]] = []
+        output_v_cache: List[Optional[torch.Tensor]] = []
+        shared_kv: Dict[int, Tuple[torch.Tensor, torch.Tensor]] = {}
+
+        for i, layer in enumerate(self.self_decoder.layers):
+            k_cache = k_caches[i] if k_caches else None
+            v_cache = v_caches[i] if v_caches else None
+
+            if layer_types[i] == "sliding_attention":
+                fc, fs, mask = local_freqs_cos, local_freqs_sin, window_atten_mask
+            else:
+                fc, fs, mask = global_freqs_cos, global_freqs_sin, atten_mask
+
+            hidden_states, new_k, new_v = layer(
+                hidden_states,
+                per_layer_inputs[i],
+                fc,
+                fs,
+                mask,
+                k_cache,
+                v_cache,
+            )
+            output_k_cache.append(new_k)
+            output_v_cache.append(new_v)
+
+            # Donor layers expose their full updated K/V for YOCO sharing
+            if layer.attention.is_kv_donor_layer:
+                if layer.attention.output_new_cache_only and k_cache is not None:
+                    full_k = torch.cat([k_cache, new_k], dim=-1)
+                    full_v = torch.cat([v_cache, new_v], dim=2)
+                else:
+                    full_k = new_k
+                    full_v = new_v
+                shared_kv[i] = (full_k, full_v)
+
+        # Cross-decoder layers (no KV I/O, consume donor K/V)
+        hidden_states = self.cross_decoder(
+            hidden_states,
+            per_layer_inputs[n_self:],
+            global_freqs_cos,
+            global_freqs_sin,
+            local_freqs_cos,
+            local_freqs_sin,
+            atten_mask,
+            window_atten_mask,
+            layer_types,
+            shared_kv,
+        )
+
+        hidden_states = self.norm(hidden_states)
+        # Compute logits only for the last token(s)
+        logits = self.lm_head(hidden_states)
+
+        if self.final_logit_softcapping > 0:
+            logits = self.final_logit_softcapping * torch.tanh(
+                logits / self.final_logit_softcapping
+            )
+
+        return logits, output_k_cache, output_v_cache
diff --git a/examples/qualcomm/oss_scripts/llama/README.md b/examples/qualcomm/oss_scripts/llama/README.md
index 4522582f2f6..6892c489df8 100644
--- a/examples/qualcomm/oss_scripts/llama/README.md
+++ b/examples/qualcomm/oss_scripts/llama/README.md
@@ -14,6 +14,7 @@ This file provides you the instructions to run LLM Decoder model, VLM model, and
  1. Gemma 2B
  1. Gemma2 2B
  1. Gemma3 1B
+ 1. Gemma4 E2B
  1. GLM 1.5B
  1. Granite3.3 2B
  1. Phi4-mini-instruct
@@ -157,6 +158,12 @@ Default example using hybrid mode
 python examples/qualcomm/oss_scripts/llama/llama.py --build_folder build-android --device ${SERIAL_NUM} --soc_model ${SOC_MODEL} --temperature 0 --model_mode hybrid --max_seq_len 1024 --prefill_ar_len 128 --decoder_model gemma3-1b --prompt "I would like to learn python, could you teach me with a simple example?" --calib_tasks wikitext --calib_limit 1
 ```
 
+#### Gemma4 E2B
+Default example using hybrid mode
+```bash
+python examples/qualcomm/oss_scripts/llama/llama.py --build_folder build-android --device ${SERIAL_NUM} --soc_model ${SOC_MODEL} --temperature 0 --model_mode hybrid --max_seq_len 1024 --prefill_ar_len 128 --decoder_model gemma4-e2b --prompt "I would like to learn python, could you teach me with a simple example?" --calib_tasks wikitext --calib_limit 1 --embedding-quantize '4,32'
+```
+
 #### GLM 1.5B
 Default example using hybrid mode
 ```bash
diff --git a/examples/qualcomm/oss_scripts/llama/__init__.py b/examples/qualcomm/oss_scripts/llama/__init__.py
index 2e0b2278909..c39fac2aa54 100644
--- a/examples/qualcomm/oss_scripts/llama/__init__.py
+++ b/examples/qualcomm/oss_scripts/llama/__init__.py
@@ -17,6 +17,9 @@
 from executorch.examples.models.gemma import convert_weights as convert_gemma_weights
 from executorch.examples.models.gemma2 import convert_weights as convert_gemma2_weights
 from executorch.examples.models.gemma3 import convert_weights as convert_gemma3_weights
+from executorch.examples.models.gemma4 import (
+    convert_hf_to_custom as convert_gemma4_weights,
+)
 
 from executorch.examples.models.glm import convert_weights as convert_glm_weights
 from executorch.examples.models.granite import (
@@ -69,6 +72,7 @@
     CodegenQuantRecipe,
     Gemma2QuantRecipe,
     Gemma3QuantRecipe,
+    Gemma4QuantRecipe,
     Gemma_2BQuantRecipe,
     GLM_1_5B_InstructQuantRecipe,
     Granite_3_3_2B_InstructQuantRecipe,
@@ -602,3 +606,22 @@ class SmolVLM_500M(LLMModelConfig):
     r2 = False
     r3 = False
     quant_recipe = SmolVLMQuantRecipe
+
+
+@register_llm_model("gemma4-e2b")
+@dataclass(init=False, frozen=True)
+class Gemma4_E2B(LLMModelConfig):
+    repo_id: str = "google/gemma-4-e2b-it"
+    params_path: str = os.path.join(
+        BASE_DIR, "../../../models/gemma4/config/e2b_config.json"
+    )
+    convert_weights = convert_gemma4_weights
+    transform_weight = False
+    instruct_model = True
+    num_sharding = 4
+    masked_softmax = False
+    seq_mse_candidates = 0
+    r1 = False
+    r2 = False
+    r3 = False
+    quant_recipe = Gemma4QuantRecipe
diff --git a/examples/qualcomm/oss_scripts/llama/decoder_constants.py b/examples/qualcomm/oss_scripts/llama/decoder_constants.py
index cccf368459b..26e5aee1ad7 100644
--- a/examples/qualcomm/oss_scripts/llama/decoder_constants.py
+++ b/examples/qualcomm/oss_scripts/llama/decoder_constants.py
@@ -63,4 +63,5 @@
     "glm-1_5b": "glm",
     "smolvlm_500m_instruct": "smolvlm",
     "internvl3_1b": "internvl3",
+    "gemma4-e2b": "gemma4",
 }
diff --git a/examples/qualcomm/oss_scripts/llama/llama.py b/examples/qualcomm/oss_scripts/llama/llama.py
index edf719a6fe5..ad4e5257fe5 100755
--- a/examples/qualcomm/oss_scripts/llama/llama.py
+++ b/examples/qualcomm/oss_scripts/llama/llama.py
@@ -666,6 +666,11 @@ def export_llama(args) -> None:
                 "Pass --calib_tasks to match the previous --tasks behavior."
             )
         raise RuntimeError("Please provide --eval_tasks to eval perplexity")
+    if args.decoder_model == "gemma4-e2b" and not args.embedding_quantize:
+        raise RuntimeError(
+            "gemma4-e2b requires --embedding-quantize: "
+            "per-layer embedding is too large for HTP quantization."
+        )
     assert (
         args.decoder_model in SUPPORTED_LLM_MODELS
     ), f"Unknown decoder_model: {args.decoder_model}."
diff --git a/examples/qualcomm/oss_scripts/llama/model/__init__.py b/examples/qualcomm/oss_scripts/llama/model/__init__.py
index 8300afd4df3..966a224211b 100644
--- a/examples/qualcomm/oss_scripts/llama/model/__init__.py
+++ b/examples/qualcomm/oss_scripts/llama/model/__init__.py
@@ -7,10 +7,12 @@
 from .apply_rope import ROTARY_EMB_REGISTRY
 from .feed_forward import FeedForward_REGISTRY
 from .layernorm import NORM_REGISTRY
+from .static_llama import repeat_kv
 
 
 __all__ = [
     "FeedForward_REGISTRY",
+    "repeat_kv",
     "ROTARY_EMB_REGISTRY",
     "NORM_REGISTRY",
 ]
diff --git a/examples/qualcomm/oss_scripts/llama/model/feed_forward.py b/examples/qualcomm/oss_scripts/llama/model/feed_forward.py
index 2f36779cc71..fc4b43aa806 100644
--- a/examples/qualcomm/oss_scripts/llama/model/feed_forward.py
+++ b/examples/qualcomm/oss_scripts/llama/model/feed_forward.py
@@ -54,13 +54,13 @@ def __init__(self, args: ModelArgs):  # in MLP: intermediate_size= 4 * embed_dim
         # HF uses NewGelu, however, Gelu is a fused op in QNN and can run faster
         self.act = GELUActivation(use_gelu_python=False)
 
-    def prepare_feedfoward_conv(self):
+    def prepare_feedforward_conv(self):
         intermediate_size = 4 * self.dim
         self.fc_in_conv = torch.nn.Conv2d(self.dim, intermediate_size, 1, bias=True)
         self.fc_out_conv = torch.nn.Conv2d(self.hidden_dim, self.dim, 1, bias=True)
 
         self.forward_no_conv = self.forward
-        self.forward = self.forward_feedfoward_conv
+        self.forward = self.forward_feedforward_conv
 
         self.fc_in_conv.weight.data.copy_(self.fc_in.weight[:, :, None, None])
         self.fc_out_conv.weight.data.copy_(self.fc_out.weight[:, :, None, None])
@@ -71,7 +71,7 @@ def prepare_feedfoward_conv(self):
         del self.fc_in
         del self.fc_out
 
-    def forward_feedfoward_conv(self, x):
+    def forward_feedforward_conv(self, x):
         bsz, _, _ = x.size()
 
         x = torch.reshape(x, (bsz, -1, 1, self.dim))
@@ -105,14 +105,14 @@ def __init__(self, args: ModelArgs):  # in MLP: intermediate_size= 4 * embed_dim
         self.down_proj = torch.nn.Linear(args.hidden_dim, args.dim, bias=False)
         self.activation_fn = args.act_fn.get_function()
 
-    def prepare_feedfoward_conv(self):
+    def prepare_feedforward_conv(self):
         self.gate_up_proj_conv = torch.nn.Conv2d(
             self.dim, 2 * self.hidden_dim, 1, bias=False
         )
         self.down_proj_conv = torch.nn.Conv2d(self.hidden_dim, self.dim, 1, bias=False)
 
         self.forward_no_conv = self.forward
-        self.forward = self.forward_feedfoward_conv
+        self.forward = self.forward_feedforward_conv
 
         self.gate_up_proj_conv.weight.data.copy_(
             self.gate_up_proj.weight[:, :, None, None]
@@ -122,7 +122,7 @@ def prepare_feedfoward_conv(self):
         del self.gate_up_proj
         del self.down_proj
 
-    def forward_feedfoward_conv(self, x):
+    def forward_feedforward_conv(self, x):
         bsz, _, _ = x.size()
         x = torch.reshape(x, (bsz, -1, 1, self.dim))
         x = x.transpose(1, 3)  # Transpose right before and after Conv
diff --git a/examples/qualcomm/oss_scripts/llama/model/static_llama.py b/examples/qualcomm/oss_scripts/llama/model/static_llama.py
index d57ada7fef6..fa9af5c0761 100755
--- a/examples/qualcomm/oss_scripts/llama/model/static_llama.py
+++ b/examples/qualcomm/oss_scripts/llama/model/static_llama.py
@@ -536,13 +536,13 @@ def __init__(self, args: ModelArgs):
         self.w3 = nn.Linear(self.dim, self.hidden_dim, bias=False)
         self.act_fn = args.act_fn.get_function()
 
-    def prepare_feedfoward_conv(self):
+    def prepare_feedforward_conv(self):
         self.w1_conv = nn.Conv2d(self.dim, self.hidden_dim, 1, bias=False)
         self.w2_conv = nn.Conv2d(self.hidden_dim, self.dim, 1, bias=False)
         self.w3_conv = nn.Conv2d(self.dim, self.hidden_dim, 1, bias=False)
 
         self.forward_no_conv = self.forward
-        self.forward = self.forward_feedfoward_conv
+        self.forward = self.forward_feedforward_conv
         self.w1_conv.weight.data.copy_(self.w1.weight[:, :, None, None])
         self.w2_conv.weight.data.copy_(self.w2.weight[:, :, None, None])
         self.w3_conv.weight.data.copy_(self.w3.weight[:, :, None, None])
@@ -551,7 +551,7 @@ def prepare_feedfoward_conv(self):
         del self.w2
         del self.w3
 
-    def forward_feedfoward_conv(self, x):
+    def forward_feedforward_conv(self, x):
         bsz, _, _ = x.size()
         x = torch.reshape(x, (bsz, -1, 1, self.dim))
         x = x.transpose(1, 3)  # Transpose right before and after Conv
diff --git a/examples/qualcomm/oss_scripts/llama/qnn_llama_runner.cpp b/examples/qualcomm/oss_scripts/llama/qnn_llama_runner.cpp
index 9b8cdd7999e..3556bf40da6 100644
--- a/examples/qualcomm/oss_scripts/llama/qnn_llama_runner.cpp
+++ b/examples/qualcomm/oss_scripts/llama/qnn_llama_runner.cpp
@@ -140,6 +140,18 @@ std::string get_formatted_prompt(
       formatted_prompt.append("<end_of_turn>\n");
       formatted_prompt.append("<start_of_turn>model\n");
       break;
+    case example::DecoderModelVersion::kGemma4:
+      formatted_prompt.append("<bos>");
+      if (!system_prompt.empty()) {
+        formatted_prompt.append("<|turn>system\n");
+        formatted_prompt.append(system_prompt);
+        formatted_prompt.append("<turn|>\n");
+      }
+      formatted_prompt.append("<|turn>user\n");
+      formatted_prompt.append(prompt);
+      formatted_prompt.append("<turn|>\n");
+      formatted_prompt.append("<|turn>model\n");
+      break;
     case example::DecoderModelVersion::kGranite:
       if (!system_prompt.empty()) {
         formatted_prompt.append("<|start_of_role|>system<|end_of_role|>");
diff --git a/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp b/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp
index 6fb2255aa94..99f014fbabf 100644
--- a/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp
+++ b/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp
@@ -69,39 +69,66 @@ KVManager::KVManager(Metadata metadata, std::unique_ptr<MethodMeta> method_meta)
   Result<TensorInfo> attention_mask = method_meta->input_tensor_meta(1);
   attention_mask_dtype_ = attention_mask->scalar_type();
 
-  // inputs are [input_tokens, attention_mask, (sliding window attention_mask),
-  // (input_pos), kv_caches] search kv_cache in inputs
-  for (int i = 2; i < method_meta->num_inputs(); i++) {
-    Result<TensorInfo> tensor_meta = method_meta->input_tensor_meta(i);
-    // k_cache: [1, n_heads, head_dim, seq_len]
-    size_t tensor_nbytes = tensor_meta->nbytes();
-    size_t expected_tensor_nbytes = metadata_.head_dim * metadata_.num_heads *
-        metadata_.max_cache_len * getDtypeSize(tensor_meta->scalar_type());
-    if (tensor_nbytes != expected_tensor_nbytes) {
-      // Not a kv_cache tensor (e.g. input_pos, sliding window attention mask).
-      continue;
-    }
+  // inputs: [tokens, atten_mask, (window_atten_mask), (input_pos), k_caches...,
+  // v_caches...].
+  // outputs: [logits, k_cache_0..n-1, v_cache_0..n-1]
+  // k_cache shape: [bsz, n_kv_head, head_dim, ar_len]
+  // v_cache v_cache shape: [bsz, n_kv_head, ar_len, head_dim]
+  k_cache_.resize(metadata_.num_layers);
+  v_cache_.resize(metadata_.num_layers);
+  k_cache_in_nbytes_.resize(metadata_.num_layers);
+  k_cache_out_nbytes_.resize(metadata_.num_layers);
+  v_cache_in_nbytes_.resize(metadata_.num_layers);
+  v_cache_out_nbytes_.resize(metadata_.num_layers);
+
+  for (int layer = 0; layer < metadata_.num_layers; ++layer) {
+    // k output: index 1 + layer
+    Result<TensorInfo> k_out = method_meta->output_tensor_meta(1 + layer);
+    // v output: index 1 + num_layers + layer
+    Result<TensorInfo> v_out =
+        method_meta->output_tensor_meta(1 + metadata_.num_layers + layer);
+
     if (kv_cache_dtype_ == executorch::aten::ScalarType::Undefined) {
-      kv_cache_dtype_ = tensor_meta->scalar_type();
-    } else {
-      ET_CHECK_MSG(
-          tensor_meta->scalar_type() == kv_cache_dtype_,
-          "Currently mixed scalar type of kv_cache is not allowed");
+      kv_cache_dtype_ = k_out->scalar_type();
     }
+    ET_CHECK_MSG(
+        k_out->scalar_type() == kv_cache_dtype_,
+        "Mixed scalar type of kv_cache is not allowed (k layer %d)",
+        layer);
+    ET_CHECK_MSG(
+        v_out->scalar_type() == kv_cache_dtype_,
+        "Mixed scalar type of kv_cache is not allowed (v layer %d)",
+        layer);
+
+    // k_cache shape: [1, n_kv, head_dim, ar_len] -> head_dim at dim 2
+    // v_cache shape: [1, n_kv, ar_len, head_dim] -> head_dim at dim 3
+    const int64_t k_head_dim = k_out->sizes()[2];
+    const int64_t v_head_dim = v_out->sizes()[3];
+    k_head_dim_.push_back(k_head_dim);
+    v_head_dim_.push_back(v_head_dim);
+
+    const size_t dtype_size = getDtypeSize(kv_cache_dtype_);
+    k_cache_out_nbytes_[layer] =
+        metadata_.num_heads * k_head_dim * metadata_.max_ar_len * dtype_size;
+    v_cache_out_nbytes_[layer] =
+        metadata_.num_heads * v_head_dim * metadata_.max_ar_len * dtype_size;
+    k_cache_in_nbytes_[layer] =
+        metadata_.num_heads * k_head_dim * metadata_.max_cache_len * dtype_size;
+    v_cache_in_nbytes_[layer] =
+        metadata_.num_heads * v_head_dim * metadata_.max_cache_len * dtype_size;
   }
+
   ET_CHECK_MSG(
       kv_cache_dtype_ != executorch::aten::ScalarType::Undefined,
-      "kv_cache_dtype was not detected from method inputs");
-  k_cache_.resize(metadata_.num_layers);
-  v_cache_.resize(metadata_.num_layers);
+      "kv_cache_dtype was not detected from method outputs");
 
-  // Calculate cache size
-  size_t cache_in_bytes = metadata_.num_layers * metadata_.num_heads *
-      metadata_.head_dim * metadata_.max_cache_len *
-      getDtypeSize(kv_cache_dtype_);
-  size_t cache_out_bytes = metadata_.num_layers * metadata_.num_heads *
-      metadata_.head_dim * metadata_.max_ar_len * getDtypeSize(kv_cache_dtype_);
-  total_cache_size_ = 2 * (cache_in_bytes + cache_out_bytes);
+  // Total cache size across all layers (K + V, input + output)
+  total_cache_size_ = 0;
+  for (int layer = 0; layer < metadata_.num_layers; ++layer) {
+    total_cache_size_ += k_cache_in_nbytes_[layer] +
+        k_cache_out_nbytes_[layer] + v_cache_in_nbytes_[layer] +
+        v_cache_out_nbytes_[layer];
+  }
 };
 
 void KVManager::init_attention_mask(
@@ -269,15 +296,15 @@ void KVManager::update_attention_mask(
 
 void KVManager::init_cache(IMemAlloc* buffer_manager, int32_t ar_len) {
   cur_ar_len_ = ar_len;
-  const size_t cache_in_bytes = metadata_.num_heads * metadata_.head_dim *
-      metadata_.max_cache_len * getDtypeSize(kv_cache_dtype_);
-  const size_t cache_out_bytes = metadata_.num_heads * metadata_.head_dim *
-      metadata_.max_ar_len * getDtypeSize(kv_cache_dtype_);
   for (int layer = 0; layer < metadata_.num_layers; ++layer) {
-    k_cache_[layer].buffer = buffer_manager->allocate(cache_in_bytes);
-    k_cache_[layer].output_buffer = buffer_manager->allocate(cache_out_bytes);
-    v_cache_[layer].buffer = buffer_manager->allocate(cache_in_bytes);
-    v_cache_[layer].output_buffer = buffer_manager->allocate(cache_out_bytes);
+    k_cache_[layer].buffer =
+        buffer_manager->allocate(k_cache_in_nbytes_[layer]);
+    k_cache_[layer].output_buffer =
+        buffer_manager->allocate(k_cache_out_nbytes_[layer]);
+    v_cache_[layer].buffer =
+        buffer_manager->allocate(v_cache_in_nbytes_[layer]);
+    v_cache_[layer].output_buffer =
+        buffer_manager->allocate(v_cache_out_nbytes_[layer]);
   }
 }
 
@@ -286,14 +313,17 @@ void KVManager::rearrange_cache(int32_t ar_len_dst) {
   if (cur_ar_len_ == ar_len_dst)
     return;
   for (int layer = 0; layer < metadata_.num_layers; ++layer) {
-    rearrange_key(k_cache_[layer], ar_len_dst);
-    rearrange_value(v_cache_[layer], ar_len_dst);
+    rearrange_key(k_cache_[layer], ar_len_dst, k_head_dim_[layer]);
+    rearrange_value(v_cache_[layer], ar_len_dst, v_head_dim_[layer]);
   }
   // rearrange done.
   cur_ar_len_ = ar_len_dst;
 }
 
-void KVManager::rearrange_key(KVCache& k_cache, int32_t ar_len_dst) {
+void KVManager::rearrange_key(
+    KVCache& k_cache,
+    int32_t ar_len_dst,
+    int64_t head_dim) {
   const int32_t src_cache_num = (cur_ar_len_ == metadata_.context_len)
       ? metadata_.context_len
       : metadata_.context_len - cur_ar_len_;
@@ -304,18 +334,18 @@ void KVManager::rearrange_key(KVCache& k_cache, int32_t ar_len_dst) {
   size_t dst_cache_nbytes = dst_cache_num * getDtypeSize(kv_cache_dtype_);
   if (src_cache_num > dst_cache_num) {
     // copy from first dimension
-    for (int i = 0; i < metadata_.head_dim * metadata_.num_heads; i++) {
+    for (int i = 0; i < head_dim * metadata_.num_heads; i++) {
       std::memmove(k_cache_in_write_ptr, k_cache_in_read_ptr, dst_cache_nbytes);
       k_cache_in_read_ptr += src_cache_nbytes;
       k_cache_in_write_ptr += dst_cache_nbytes;
     }
   } else {
     k_cache_in_read_ptr +=
-        (metadata_.head_dim * metadata_.num_heads - 1) * src_cache_nbytes;
+        (head_dim * metadata_.num_heads - 1) * src_cache_nbytes;
     k_cache_in_write_ptr +=
-        (metadata_.head_dim * metadata_.num_heads - 1) * dst_cache_nbytes;
+        (head_dim * metadata_.num_heads - 1) * dst_cache_nbytes;
     // copy from last dimension
-    for (int i = 0; i < metadata_.head_dim * metadata_.num_heads; i++) {
+    for (int i = 0; i < head_dim * metadata_.num_heads; i++) {
       std::memmove(k_cache_in_write_ptr, k_cache_in_read_ptr, src_cache_nbytes);
       k_cache_in_read_ptr -= src_cache_nbytes;
       k_cache_in_write_ptr -= dst_cache_nbytes;
@@ -323,7 +353,10 @@ void KVManager::rearrange_key(KVCache& k_cache, int32_t ar_len_dst) {
   }
 }
 
-void KVManager::rearrange_value(KVCache& v_cache, int32_t ar_len_dst) {
+void KVManager::rearrange_value(
+    KVCache& v_cache,
+    int32_t ar_len_dst,
+    int64_t head_dim) {
   const int32_t src_cache_num = (cur_ar_len_ == metadata_.context_len)
       ? metadata_.context_len
       : metadata_.context_len - cur_ar_len_;
@@ -338,23 +371,23 @@ void KVManager::rearrange_value(KVCache& v_cache, int32_t ar_len_dst) {
       std::memmove(
           v_cache_in_write_ptr,
           v_cache_in_read_ptr,
-          dst_cache_nbytes * metadata_.head_dim);
-      v_cache_in_read_ptr += src_cache_nbytes * metadata_.head_dim;
-      v_cache_in_write_ptr += dst_cache_nbytes * metadata_.head_dim;
+          dst_cache_nbytes * head_dim);
+      v_cache_in_read_ptr += src_cache_nbytes * head_dim;
+      v_cache_in_write_ptr += dst_cache_nbytes * head_dim;
     }
   } else {
     v_cache_in_read_ptr +=
-        metadata_.head_dim * (metadata_.num_heads - 1) * src_cache_nbytes;
+        head_dim * (metadata_.num_heads - 1) * src_cache_nbytes;
     v_cache_in_write_ptr +=
-        metadata_.head_dim * (metadata_.num_heads - 1) * dst_cache_nbytes;
+        head_dim * (metadata_.num_heads - 1) * dst_cache_nbytes;
     // copy from last dimension
     for (int i = 0; i < metadata_.num_heads; i++) {
       std::memmove(
           v_cache_in_write_ptr,
           v_cache_in_read_ptr,
-          src_cache_nbytes * metadata_.head_dim);
-      v_cache_in_read_ptr -= src_cache_nbytes * metadata_.head_dim;
-      v_cache_in_write_ptr -= dst_cache_nbytes * metadata_.head_dim;
+          src_cache_nbytes * head_dim);
+      v_cache_in_read_ptr -= src_cache_nbytes * head_dim;
+      v_cache_in_write_ptr -= dst_cache_nbytes * head_dim;
     }
   }
 }
@@ -370,8 +403,9 @@ void KVManager::update_cache(
       cur_ar_len_,
       ar_len);
   for (int layer = 0; layer < metadata_.num_layers; ++layer) {
-    update_key(k_cache_[layer], n_past, n_update, selected);
-    update_value(v_cache_[layer], n_past, n_update, selected);
+    update_key(k_cache_[layer], n_past, n_update, selected, k_head_dim_[layer]);
+    update_value(
+        v_cache_[layer], n_past, n_update, selected, v_head_dim_[layer]);
   }
 }
 
@@ -379,7 +413,8 @@ void KVManager::update_key(
     KVCache& k_cache,
     int32_t n_past,
     int32_t n_update,
-    const std::vector<bool>& selected) {
+    const std::vector<bool>& selected,
+    int64_t head_dim) {
   std::byte* write_ptr = k_cache.buffer;
   std::byte* read_ptr = k_cache.output_buffer;
   const int32_t copy_size = n_update * getDtypeSize(kv_cache_dtype_);
@@ -388,7 +423,7 @@ void KVManager::update_key(
       : (metadata_.context_len - cur_ar_len_) * getDtypeSize(kv_cache_dtype_);
   const int32_t out_size = cur_ar_len_ * getDtypeSize(kv_cache_dtype_);
   const int32_t past_size = n_past * getDtypeSize(kv_cache_dtype_);
-  const int32_t n_iter = metadata_.head_dim * metadata_.num_heads;
+  const int32_t n_iter = head_dim * metadata_.num_heads;
 
   write_ptr += past_size;
   if (selected.empty()) {
@@ -423,21 +458,19 @@ void KVManager::update_value(
     KVCache& v_cache,
     int32_t n_past,
     int32_t n_update,
-    const std::vector<bool>& selected) {
+    const std::vector<bool>& selected,
+    int64_t head_dim) {
   std::byte* write_ptr = v_cache.buffer;
   std::byte* read_ptr = v_cache.output_buffer;
-  const int32_t copy_size =
-      n_update * metadata_.head_dim * getDtypeSize(kv_cache_dtype_);
-  const int32_t past_size =
-      n_past * metadata_.head_dim * getDtypeSize(kv_cache_dtype_);
+  const int32_t copy_size = n_update * head_dim * getDtypeSize(kv_cache_dtype_);
+  const int32_t past_size = n_past * head_dim * getDtypeSize(kv_cache_dtype_);
   const int32_t n_iter = metadata_.num_heads;
   const int32_t iter_size = (cur_ar_len_ == metadata_.context_len)
-      ? metadata_.context_len * metadata_.head_dim *
-          getDtypeSize(kv_cache_dtype_)
-      : (metadata_.context_len - cur_ar_len_) * metadata_.head_dim *
+      ? metadata_.context_len * head_dim * getDtypeSize(kv_cache_dtype_)
+      : (metadata_.context_len - cur_ar_len_) * head_dim *
           getDtypeSize(kv_cache_dtype_);
   const int32_t out_size =
-      cur_ar_len_ * metadata_.head_dim * getDtypeSize(kv_cache_dtype_);
+      cur_ar_len_ * head_dim * getDtypeSize(kv_cache_dtype_);
 
   write_ptr += past_size;
 
@@ -453,14 +486,13 @@ void KVManager::update_value(
       auto wp = write_ptr, rp = read_ptr;
       for (auto sel : selected) {
         if (sel) {
-          std::memcpy(
-              wp, rp, metadata_.head_dim * getDtypeSize(kv_cache_dtype_));
-          wp += metadata_.head_dim * getDtypeSize(kv_cache_dtype_);
+          std::memcpy(wp, rp, head_dim * getDtypeSize(kv_cache_dtype_));
+          wp += head_dim * getDtypeSize(kv_cache_dtype_);
           update_times--;
           if (update_times == 0)
             break;
         }
-        rp += metadata_.head_dim * getDtypeSize(kv_cache_dtype_);
+        rp += head_dim * getDtypeSize(kv_cache_dtype_);
       }
       write_ptr += iter_size;
       read_ptr += out_size;
diff --git a/examples/qualcomm/oss_scripts/llama/runner/kv_manager.h b/examples/qualcomm/oss_scripts/llama/runner/kv_manager.h
index 3b8e67dd38d..cbebaa5d8cb 100644
--- a/examples/qualcomm/oss_scripts/llama/runner/kv_manager.h
+++ b/examples/qualcomm/oss_scripts/llama/runner/kv_manager.h
@@ -171,21 +171,23 @@ class KVManager {
  private:
   // Helper functions to rearrange and update key and value caches
 
-  void rearrange_key(KVCache& k_cache, int32_t ar_len_dst);
+  void rearrange_key(KVCache& k_cache, int32_t ar_len_dst, int64_t head_dim);
 
-  void rearrange_value(KVCache& v_cache, int32_t ar_len_dst);
+  void rearrange_value(KVCache& v_cache, int32_t ar_len_dst, int64_t head_dim);
 
   void update_key(
       KVCache& k_cache,
       int32_t n_past,
       int32_t n_update,
-      const std::vector<bool>& selected);
+      const std::vector<bool>& selected,
+      int64_t head_dim);
 
   void update_value(
       KVCache& v_cache,
       int32_t n_past,
       int32_t n_update,
-      const std::vector<bool>& selected);
+      const std::vector<bool>& selected,
+      int64_t head_dim);
 
   // metadata
   Metadata metadata_;
@@ -200,5 +202,13 @@ class KVManager {
   // output: layer -> head * head_dim * max_ar_len
   std::vector<KVCache> k_cache_;
   std::vector<KVCache> v_cache_;
+  // Per-layer buffer sizes and head_dim (may differ when head_dim varies across
+  // layers)
+  std::vector<size_t> k_cache_in_nbytes_;
+  std::vector<size_t> k_cache_out_nbytes_;
+  std::vector<size_t> v_cache_in_nbytes_;
+  std::vector<size_t> v_cache_out_nbytes_;
+  std::vector<int64_t> k_head_dim_;
+  std::vector<int64_t> v_head_dim_;
 };
 } // namespace example
diff --git a/examples/qualcomm/oss_scripts/llama/runner/runner.cpp b/examples/qualcomm/oss_scripts/llama/runner/runner.cpp
index 0e1a27a751a..409f7c96f34 100644
--- a/examples/qualcomm/oss_scripts/llama/runner/runner.cpp
+++ b/examples/qualcomm/oss_scripts/llama/runner/runner.cpp
@@ -126,6 +126,9 @@ Runner::Runner(
   } else if (decoder_model_version == "gemma3") {
     decoder_model_version_ = DecoderModelVersion::kGemma3;
     cache_mode_ = CacheMode::HybridCache;
+  } else if (decoder_model_version == "gemma4") {
+    decoder_model_version_ = DecoderModelVersion::kGemma4;
+    cache_mode_ = CacheMode::HybridCache;
   } else if (decoder_model_version == "granite") {
     decoder_model_version_ = DecoderModelVersion::kGranite;
   } else if (decoder_model_version == "phi_4_mini") {
@@ -208,6 +211,8 @@ Error Runner::load() {
       decoder_model_version_ == DecoderModelVersion::kGemma2 ||
       decoder_model_version_ == DecoderModelVersion::kGemma3) {
     eos_ids->insert(tokenizer_->encode("<end_of_turn>", 0, 0).get()[0]);
+  } else if (decoder_model_version_ == DecoderModelVersion::kGemma4) {
+    eos_ids->insert(tokenizer_->encode("<turn|>", 0, 0).get()[0]);
   } else if (decoder_model_version_ == DecoderModelVersion::kCodegen) {
     eos_ids->insert(tokenizer_->encode("<|endoftext|>", 0, 0).get()[0]);
   } else if (decoder_model_version_ == DecoderModelVersion::kGlm) {
@@ -274,6 +279,14 @@ Error Runner::load() {
         sliding_window_evalue__, module_->get("get_sliding_window"));
     sliding_window = sliding_window_evalue__.toInt();
   }
+  // Gemma4 uses YOCO: only self-decoder layers have KV I/O; use
+  // get_n_self_layers when available so KVManager allocates the correct number
+  // of caches.
+  int64_t num_kv_layers = num_layers;
+  if (module_->method_names()->count("get_n_self_layers") > 0) {
+    ET_ASSIGN_OR_RETURN(num_layers_evalue__, module_->get("get_n_self_layers"));
+    num_kv_layers = num_layers_evalue__.toScalar().to<int64_t>();
+  }
   kv_manager_ = std::make_unique<KVManager>(
       KVManager::Metadata{
           context_len_,
@@ -281,7 +294,7 @@ Error Runner::load() {
           max_ar_len,
           max_cache_len,
           num_heads,
-          num_layers},
+          num_kv_layers},
       std::make_unique<MethodMeta>(
           std::move(module_->method_meta(token_generator_method_name).get())));
 
@@ -299,7 +312,7 @@ Error Runner::load() {
       PromptProcessor::Metadata{
           context_len_,
           num_heads,
-          num_layers,
+          num_kv_layers,
           prompt_processor_ar_len,
           vocab_size,
           use_int64_token,
@@ -317,7 +330,7 @@ Error Runner::load() {
         LhdTokenGenerator::Metadata{
             context_len_,
             num_heads,
-            num_layers,
+            num_kv_layers,
             token_generator_ar_len,
             vocab_size,
             use_int64_token,
@@ -339,7 +352,7 @@ Error Runner::load() {
         TokenGenerator::Metadata{
             context_len_,
             num_heads,
-            num_layers,
+            num_kv_layers,
             token_generator_ar_len,
             vocab_size,
             use_int64_token,
diff --git a/examples/qualcomm/oss_scripts/llama/runner/runner.h b/examples/qualcomm/oss_scripts/llama/runner/runner.h
index 86934558656..a996c92997d 100644
--- a/examples/qualcomm/oss_scripts/llama/runner/runner.h
+++ b/examples/qualcomm/oss_scripts/llama/runner/runner.h
@@ -44,6 +44,7 @@ enum DecoderModelVersion {
   kCodegen,
   kGlm,
   kGemma2,
+  kGemma4,
 };
 
 class Runner : public executorch::extension::llm::IRunner {
diff --git a/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py b/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py
index 4eecf308f88..9f7188de807 100644
--- a/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py
+++ b/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py
@@ -335,6 +335,30 @@ def __init__(self, verbose: bool = False):
         self.recipe.custom_quant_annotations.append(annotate_kv_8bit)
 
 
+class Gemma4QuantRecipe(StaticLLMQuantRecipe):
+    default_quant_dtype = QuantDtype.use_16a8w
+
+    def __init__(self, verbose: bool = False):
+        super().__init__()
+
+        self.recipe = QuantRecipe(
+            self.default_quant_dtype,
+            False,
+            act_observer=MinMaxObserver,
+            granularity=QuantGranularity.PER_TENSOR,
+            verbose=verbose,
+        ).add_node_target(
+            {
+                torch.ops.aten.conv2d.default,
+            },
+            QuantDtype.use_16a4w_block,
+            False,
+            act_observer=MinMaxObserver,
+            granularity=QuantGranularity.PER_BLOCK,
+            extra_kwargs={"block_size": (1, 64, 1, 1)},
+        )
+
+
 class GLM_1_5B_InstructQuantRecipe(StaticLLMQuantRecipe):
     default_quant_dtype = QuantDtype.use_16a4w
 
diff --git a/examples/qualcomm/oss_scripts/llama/tokenizer.py b/examples/qualcomm/oss_scripts/llama/tokenizer.py
index 954b73384fa..653b2702548 100644
--- a/examples/qualcomm/oss_scripts/llama/tokenizer.py
+++ b/examples/qualcomm/oss_scripts/llama/tokenizer.py
@@ -15,7 +15,6 @@
     AUDIO_ENCODER,
     VISION_ENCODER,
 )
-from executorch.examples.qualcomm.oss_scripts.llama.model.static_llama import ModelArgs
 from pytorch_tokenizers import get_tokenizer, TiktokenTokenizer
 from pytorch_tokenizers.llama2c import Llama2cTokenizer as SentencePieceTokenizer
 
@@ -77,8 +76,8 @@ def __init__(self, control_args: argparse.Namespace, config: LLMModelConfig):
             config.params_path if control_args.params is None else control_args.params
         )
         with open(params_path) as f:
-            model_args = ModelArgs(**json.load(f))
-        self.vocab_size = model_args.vocab_size
+            model_args = json.load(f)
+        self.vocab_size = model_args.get("text_config", model_args)["vocab_size"]
 
         self.runtime_tokenizer_path = self._init_tokenizer(
             control_args.tokenizer_model, control_args.tokenizer_bin
@@ -136,6 +135,10 @@ def _from_hf(self):
             # Override the default BOS and EOS token IDs for codegen2_1b
             tokenizer.bos_id = 1
             tokenizer.eos_id = 2
+        if self.decoder_model == "gemma4-e2b":
+            # Override the default BOS and EOS token IDs for gemma4_e2b
+            tokenizer.bos_id = 2
+            tokenizer.eos_id = 1
 
         return runtime_tokenizer_path, tokenizer, chat_template
 
diff --git a/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py b/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py
index 9a1ad50f854..c15124e6385 100644
--- a/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py
+++ b/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py
@@ -184,21 +184,49 @@ def __init__(
         )
 
         # load static llama model args
-        params_path = (
-            config.params_path if control_args.params is None else control_args.params
-        )
-        with open(params_path) as f:
-            self.model_args = process_model_args(
-                control_args, ModelArgs(**json.load(f)), self.quant_recipe, config, mode
+        if control_args.decoder_model == "gemma4-e2b":
+            from executorch.examples.models.gemma4.text_decoder.gemma4_config import (
+                Gemma4Config,
+            )
+
+            params_path = config.params_path
+            self.gemma4_config = Gemma4Config.from_json(params_path)
+            self.gemma4_config.use_kv_cache = True
+            self.gemma4_config.max_batch_size = 1
+            self.gemma4_config.max_seq_len = control_args.max_seq_len
+            self.gemma4_config.max_context_len = control_args.max_context_len
+            self.model_args = None
+        else:
+            params_path = (
+                config.params_path
+                if control_args.params is None
+                else control_args.params
             )
+            with open(params_path) as f:
+                self.model_args = process_model_args(
+                    control_args,
+                    ModelArgs(**json.load(f)),
+                    self.quant_recipe,
+                    config,
+                    mode,
+                )
         # prepare instance
         self.tok_embedding, self.decoder = self._prepare_model()
 
         # check if sharding required
         if self.decoder and self.config.num_sharding > 1:
+            layer_prefix_offsets = None
+            if self.control_args.decoder_model == "gemma4-e2b":
+                n_self = self.gemma4_config.num_self_decoder_layers
+                layer_prefix_offsets = {
+                    "model.self_decoder.layers": 0,
+                    "model.cross_decoder.layers": n_self,
+                }
             SplitGraph, setting = model_sharding.get_split_graph_pass(
                 self.meta["get_n_layers"],
                 shares=self.config.num_sharding,
+                pattern=self._get_sharding_get_pattern(),
+                layer_prefix_offsets=layer_prefix_offsets,
             )
             self.passes_job[SplitGraph] = setting
             self.dep_table[SplitGraph] = [FoldQDQ]
@@ -217,30 +245,58 @@ def __init__(
             else None
         )
 
+    def _get_sharding_get_pattern(self):
+        if self.control_args.decoder_model == "gemma4-e2b":
+            prefixes = [
+                "model.cross_decoder.layers",
+                "model.self_decoder.layers",
+            ]
+        else:
+            prefixes = [
+                "layers",
+            ]
+        prefix_alt = "|".join(re.escape(p) for p in prefixes)
+        return rf"^(?:{prefix_alt})\.(\d+)"
+
     def _prepare_model(self):  # noqa: C901
         if (instance := self._get_model_instance()) is None:
             return None, None
         tok_embedding, decoder = instance
         # load parameters for HF models
         if self.control_args.checkpoint is None:
-            checkpoint = download_and_convert_hf_checkpoint(
-                self.config.repo_id,
-                self.config.convert_weights.__func__,
-            )
-            state_dict = torch.load(
-                checkpoint, weights_only=True, map_location="cpu", mmap=True
-            )
-            if self.control_args.decoder_model in {
-                "gemma-2b",
-                "gemma2-2b",
-                "gemma3-1b",
-            }:
-                for k, v in state_dict.items():
-                    if "norm" not in k:
-                        continue
-                    # Llama does x.to(float16) * w whilst Gemma3 is (x * w).to(float16)
-                    # See https://github.com/huggingface/transformers/pull/29402
-                    state_dict[k] = v.float() + torch.ones(v.shape, dtype=torch.float32)
+            if self.control_args.decoder_model == "gemma4-e2b":
+                from executorch.examples.qualcomm.oss_scripts.gemma4.text_decoder.convert_weights import (
+                    remap_keys,
+                )
+                from huggingface_hub import snapshot_download
+
+                state_dict = self.config.convert_weights.__func__(
+                    snapshot_download(repo_id=self.config.repo_id),
+                    self.gemma4_config,
+                    torch.float32,
+                )
+                state_dict = remap_keys(state_dict)
+            else:
+                checkpoint = download_and_convert_hf_checkpoint(
+                    self.config.repo_id,
+                    self.config.convert_weights.__func__,
+                )
+                state_dict = torch.load(
+                    checkpoint, weights_only=True, map_location="cpu", mmap=True
+                )
+                if self.control_args.decoder_model in {
+                    "gemma-2b",
+                    "gemma2-2b",
+                    "gemma3-1b",
+                }:
+                    for k, v in state_dict.items():
+                        if "norm" not in k:
+                            continue
+                        # Llama does x.to(float16) * w whilst Gemma3 is (x * w).to(float16)
+                        # See https://github.com/huggingface/transformers/pull/29402
+                        state_dict[k] = v.float() + torch.ones(
+                            v.shape, dtype=torch.float32
+                        )
         else:
             state_dict = torch.load(
                 self.control_args.checkpoint,
@@ -308,8 +364,8 @@ def permute(w, heads, partial_rotary_dim):
         for layer in decoder.layers:
             if getattr(layer.attention, "prepare_attention_conv", None):
                 layer.attention.prepare_attention_conv()
-            if getattr(layer.feed_forward, "prepare_feedfoward_conv", None):
-                layer.feed_forward.prepare_feedfoward_conv()
+            if getattr(layer.feed_forward, "prepare_feedforward_conv", None):
+                layer.feed_forward.prepare_feedforward_conv()
 
         decoder = convert_linear_to_conv2d(decoder)
 
@@ -320,16 +376,10 @@ def permute(w, heads, partial_rotary_dim):
 
         # check embedding fallback
         if self.control_args.embedding_quantize:
-            decoder = get_quant_embedding_transform(
-                embedding_quantize=self.control_args.embedding_quantize
-            )(decoder)
             self.passes_job[I64toI32][QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY][
                 "skip_node"
             ] = {"tokens"}
             if self.apply_embedding:
-                tok_embedding = get_quant_embedding_transform(
-                    embedding_quantize=self.control_args.embedding_quantize
-                )(tok_embedding)
                 self.tok_embedding_passes_job[I64toI32][
                     QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY
                 ]["skip_node"] = {"tokens"}
@@ -368,16 +418,47 @@ def _get_model_instance(self) -> LlamaModel:
             # For gemma, we have preprocessed the weight of rmsnorm
             self.model_args.norm_type = "rmsnorm"
 
-        decoder: LlamaModel = LLM_VARIANT_ARCHS.get(
-            self.control_args.decoder_model, LlamaModel
-        )(
-            self.model_args,
-            ar_len=self.model_args.ar_len,
-            output_new_cache_only=True,
-            output_cache=True,
-            use_i64_token=use_i64_token,
-            **get_model_specific_kwargs(self.control_args, self.config),
-        )
+        if self.control_args.decoder_model == "gemma4-e2b":
+            from executorch.examples.qualcomm.oss_scripts.gemma4.model_wrapper import (
+                Gemma4TextModelWrapper,
+            )
+
+            ar_len = (
+                self.control_args.prefill_ar_len
+                if self.mode == Mode.PREFILL
+                else (
+                    self.control_args.max_context_len
+                    if self.mode == Mode.CALIBRATE
+                    else 1
+                )
+            )
+            extra_kwargs = {
+                "kv_io_bit_width": (
+                    self.quant_recipe.get_kv_io_bit_width()
+                    if self.quant_recipe
+                    else None
+                ),
+            }
+            decoder: Gemma4TextModelWrapper = Gemma4TextModelWrapper(
+                self.gemma4_config,
+                ar_len=ar_len,
+                output_new_cache_only=True,
+                output_cache=True,
+                use_i64_token=use_i64_token,
+                enable_masked_softmax=False,
+                **extra_kwargs,
+            )
+        else:
+            decoder: LlamaModel = LLM_VARIANT_ARCHS.get(
+                self.control_args.decoder_model, LlamaModel
+            )(
+                self.model_args,
+                ar_len=self.model_args.ar_len,
+                output_new_cache_only=True,
+                output_cache=True,
+                use_i64_token=use_i64_token,
+                **get_model_specific_kwargs(self.control_args, self.config),
+            )
 
         self.meta = decoder.get_metadata()
         # get example input
@@ -399,14 +480,27 @@ def _get_model_instance(self) -> LlamaModel:
             ),
         }
         # shape of k caches and v caches
-        self.kv_cache_shape = {
-            # single head, kv input
-            (self.meta["get_head_dim"], self.meta["get_max_context_len"]),
-            (self.meta["get_max_context_len"], self.meta["get_head_dim"]),
-            # single head, kv output
-            (self.meta["get_head_dim"], self.meta["get_ar_len"]),
-            (self.meta["get_ar_len"], self.meta["get_head_dim"]),
-        }
+        if self.control_args.decoder_model == "gemma4-e2b":
+            # Gemma 4 has per-layer head_dim: sliding=256, full=512
+            kv_head_dims = {
+                self.meta["get_head_dim"],
+                self.meta["get_global_head_dim"],
+            }
+            self.kv_cache_shape = set()
+            for head_dim in kv_head_dims:
+                self.kv_cache_shape.add((head_dim, self.meta["get_max_context_len"]))
+                self.kv_cache_shape.add((self.meta["get_max_context_len"], head_dim))
+                self.kv_cache_shape.add((head_dim, self.meta["get_ar_len"]))
+                self.kv_cache_shape.add((self.meta["get_ar_len"], head_dim))
+        else:
+            self.kv_cache_shape = {
+                # single head, kv input
+                (self.meta["get_head_dim"], self.meta["get_max_context_len"]),
+                (self.meta["get_max_context_len"], self.meta["get_head_dim"]),
+                # single head, kv output
+                (self.meta["get_head_dim"], self.meta["get_ar_len"]),
+                (self.meta["get_ar_len"], self.meta["get_head_dim"]),
+            }
 
         if self.apply_embedding:
             self.tok_embedding_export_input = (
@@ -452,7 +546,7 @@ def _save_output_kv_cache_quant_attrs(self):
                 ]
                 kv_idx += 1
 
-    def _tag_ios(self, node, fixed_point_type):
+    def _tag_ios(self, node, fixed_point_type):  # noqa: C901
         atten_mask_shape = {
             (
                 self.meta["get_max_batch_size"],
@@ -465,6 +559,10 @@ def _tag_ios(self, node, fixed_point_type):
         freq_shape = {
             (self.meta["get_ar_len"], self.meta["get_head_dim"] // 2),
         }
+        if self.control_args.decoder_model == "gemma4-e2b":
+            freq_shape.add(
+                (self.meta["get_ar_len"], self.meta["get_global_head_dim"] // 2)
+            )
 
         freq_op = {
             exir_ops.edge.aten.select.int,
@@ -473,7 +571,8 @@ def _tag_ios(self, node, fixed_point_type):
 
         if node.op == "placeholder":
             if (
-                len(users := list(node.users)) == 1
+                len(users := list(node.users)) > 0
+                and "args_" in node.name
                 and users[0].meta["val"].size()[-2:] in self.kv_cache_shape
             ):
                 quant_io_type = fixed_point_type["kv_type"]
@@ -498,6 +597,22 @@ def _tag_ios(self, node, fixed_point_type):
         if node.target in freq_op and node.meta["val"].size() in freq_shape:
             quant_io_type = fixed_point_type["io_type"]
 
+        if (
+            self.control_args.decoder_model == "gemma4-e2b"
+            and "stack_trace" in node.meta
+            and (
+                "per_layer_inputs[n_self:]" in node.meta["stack_trace"]
+                or "per_layer_inputs[i]" in node.meta["stack_trace"]
+            )
+        ):
+            quant_io_type = fixed_point_type["io_type"]
+
+        # tag donor kv
+        if self.control_args.decoder_model == "gemma4-e2b" and (
+            is_node_src_start_with_name(node, "full_k")
+            or is_node_src_start_with_name(node, "full_v")
+        ):
+            quant_io_type = fixed_point_type["io_type"]
         return quant_io_type
 
     def _quant_recipe_suggestion(
@@ -566,6 +681,16 @@ def quantize(self, request: Request):  # noqa: C901
                     f"unknown logits io bit width {self.quant_recipe.get_logits_output_bit_width()}"
                 )
 
+        # embedding fallback and quantization
+        if self.control_args.embedding_quantize:
+            self.decoder = get_quant_embedding_transform(
+                embedding_quantize=self.control_args.embedding_quantize
+            )(self.decoder)
+            if self.apply_embedding:
+                self.tok_embedding = get_quant_embedding_transform(
+                    embedding_quantize=self.control_args.embedding_quantize
+                )(self.tok_embedding)
+
         data = request.method_data[TEXT_DECODER]
 
         quantizer = make_quantizer(backend=data.backend, soc_model=data.soc_model)
@@ -852,7 +977,12 @@ def _set_attr(
                 if "args_" in node.name:
                     args_idx = int(node.name.split("_")[-1])
 
-                    if args_idx >= self.decode.meta["get_n_layers"]:
+                    n_cache_layers = (
+                        self.decode.meta["get_n_self_layers"]
+                        if self.control_args.decoder_model == "gemma4-e2b"
+                        else self.decode.meta["get_n_layers"]
+                    )
+                    if args_idx >= n_cache_layers:
                         v_input_cache_nodes.append(node)
                     else:
                         k_input_cache_nodes.append(node)
diff --git a/exir/lowered_backend_module.py b/exir/lowered_backend_module.py
index 75097718032..a4c2f2cfe79 100644
--- a/exir/lowered_backend_module.py
+++ b/exir/lowered_backend_module.py
@@ -884,10 +884,11 @@ def create_submodule_from_nodes(
         # all uses with a getitem call to the 0th index of the result
         with gm.graph.inserting_after(submodule_node):
             proxy_out = torch.fx.Proxy(submodule_node)[0].node  # type: ignore[index]
-            submodule_node.replace_all_uses_with(proxy_out)
-            proxy_out.meta["val"] = submodule_node.meta["val"]
+            submodule_node.replace_all_uses_with(proxy_out, propagate_meta=True)
             # Reset the args since it was overwritten in the previous line
             proxy_out.args = (submodule_node, 0)
+            proxy_out.meta.pop("nn_module_stack", None)
+            proxy_out.meta.pop("source_fn_stack", None)
     else:
         # fuse_as_graphmodule will automatically propagate the metadata of the
         # partition's last node to the getitem nodes that appear after the
diff --git a/extension/llm/custom_ops/model_sharding.py b/extension/llm/custom_ops/model_sharding.py
index 916b13a90b8..b744a170446 100644
--- a/extension/llm/custom_ops/model_sharding.py
+++ b/extension/llm/custom_ops/model_sharding.py
@@ -27,10 +27,28 @@ class SplitGraph(ExportPass):
     not load all llama model in one pte.
     """
 
-    def __init__(self, shard_layers: List[int], pattern=r"layers.(\d+)"):
+    def __init__(
+        self,
+        shard_layers: List[int],
+        pattern=r"layers\.(\d+)",
+        layer_prefix_offsets: dict = None,
+    ):
         super().__init__()
         self.shard_layers = shard_layers
         self.pattern = pattern
+        # Maps nn_module_stack prefix -> offset to convert local layer index to
+        # global layer index. Required when a model has multiple sub-decoders
+        # whose layers are independently numbered (e.g. Gemma4 self_decoder
+        # layers 0-14 and cross_decoder layers 0-19 that map to global 15-34).
+        self.layer_prefix_offsets = layer_prefix_offsets or {}
+
+    def _get_global_layer(self, full_qualified_name: str):
+        for prefix, offset in self.layer_prefix_offsets.items():
+            m = re.match(rf"^{re.escape(prefix)}\.(\d+)", full_qualified_name)
+            if m:
+                return int(m.group(1)) + offset
+        m = re.search(self.pattern, full_qualified_name)
+        return int(m.group(1)) if m else None
 
     def _insert_fallback_op(
         self, graph_module: torch.fx.GraphModule
@@ -52,11 +70,10 @@ def _insert_fallback_op(
             module_values_list = list(node.meta["nn_module_stack"].values())
             full_qualified_name = module_values_list[-1][0]
             # Search which layer this node belongs to
-            match = re.search(self.pattern, full_qualified_name)
-            if match is None:
+            cur_layer = self._get_global_layer(full_qualified_name)
+            if cur_layer is None:
                 continue
 
-            cur_layer = int(match.group(1))
             # Check the current node which is the last node of the layer
             if cur_layer in self.shard_layers and prev_layer == cur_layer - 1:
                 with graph_module.graph.inserting_after(prev_node):
@@ -84,19 +101,30 @@ def call(self, graph_module: torch.fx.GraphModule):
 
 
 def split_graph(
-    edge_program: ExportedProgram, num_layers: int, shares: int, pattern=r"layers.(\d+)"
+    edge_program: ExportedProgram,
+    num_layers: int,
+    shares: int,
+    pattern=r"layers\.(\d+)",
+    layer_prefix_offsets: dict = None,
 ):
     graph_module = edge_program.graph_module
     shard_layers = list(range(0, num_layers, int(num_layers / shares)))
-    return SplitGraph(shard_layers, pattern=pattern)(graph_module)
+    return SplitGraph(
+        shard_layers, pattern=pattern, layer_prefix_offsets=layer_prefix_offsets
+    )(graph_module)
 
 
-def get_split_graph_pass(num_layers: int, shares: int, pattern=r"layers.(\d+)"):
+def get_split_graph_pass(
+    num_layers: int,
+    shares: int,
+    pattern=r"layers\.(\d+)",
+    layer_prefix_offsets: dict = None,
+):
     shard_layers = list(range(0, num_layers, int(num_layers / shares)))
+    kwargs = {"shard_layers": shard_layers, "pattern": pattern}
+    if layer_prefix_offsets:
+        kwargs["layer_prefix_offsets"] = layer_prefix_offsets
     return SplitGraph, {
         QCOM_PASS_ACTIVATE_KEY: True,
-        QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY: {
-            "shard_layers": shard_layers,
-            "pattern": pattern,
-        },
+        QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY: kwargs,
     }