[DannyYuyang-quic]

Summary

Gemma4 Text model enablement

Test plan

python examples/qualcomm/oss_scripts/llama/llama.py --build_folder build-android --device ${SERIAL_NUM} --soc_model ${SOC_MODEL} --temperature 0 --model_mode kv--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'

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20979

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 21633d2 with merge base fb5eedc ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[DannyYuyang-quic]

@psiddh Hi,
This PR adds support for the Gemma4 text model.

We chose a standalone implementation because Gemma4's architecture differs substantially from the current static_llama architecture.

Once more Gemma 4-like architectures (e.g., YOCO or per-layer embedding variants) are supported, perhaps around 4-5 models, it may make sense to revisit the design and refactor them into a unified architecture (static llama path).

cc: @shewu-quic @winskuo-quic @haowhsu-quic @chenweng-quic @abhinaykukkadapu

[DannyYuyang-quic]

@pytorchbot label "release notes: qualcomm"

[psiddh]

cc @digantdesai @Gasoonjia @JacobSzwejbka

[psiddh]

There are merge conflicts, can you pls resolve.

[psiddh]

Ca yo check why pull / test-sqnr-static-llm-qnn-linux (smollm2_135m) / linux-job (pull_request) and Test QNN Windows MSVC build / build-qnn-windows-msvc / windows-job (pull_request) are failing

[DannyYuyang-quic]

@psiddh Hi, it looks like those failures are not related to this PR. The failures seem to be caused by the branch being based on an older commit. like: ./.ci/scripts/build-qnn-windows-msvc.ps1 hadn't been added to the main branch at that point, which caused build-qnn-windows-msvc to fail. I just rebased the branch. I think we can rerun the CI.

[DannyYuyang-quic]

@psiddh Looks like the smollm2_135m SQNR CI regression was caused by changing prefix_token_id from eot_token_id to bos_id, which changes the calibration inputs.

This change is actually correct. According to lm_eval, the prefix token should be the tokenizer's BOS token when available, falling back to EOS/EOT otherwise.

I think we need to change the SQNR threshold.

Reference:

• https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/models/huggingface.py#L169-L171
• https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/models/huggingface.py#L608

[digantdesai]

@claude review this PR.

[claude]

Claude finished @digantdesai's task in 4m 20s —— View job

---

Claude PR Review — Gemma4 Text model enablement

Reviewed the full diff against origin/main. Overall this is a well-structured enablement: the standalone Gemma4 decoder is cleanly separated, the static_llama-style flat-KV convention is respected, and the shared-infra changes (model_sharding, kv_manager) are generalized rather than special-cased where it matters. Nice work on the per-layer heterogeneous head_dim handling.

Below are findings ordered by importance.

---

🔴 Blocking / high-priority

1. Branch is stale — this PR deletes unrelated operator tests.
The branch is 1 commit behind main, and the backends/qualcomm/tests/test_qnn_delegate.py diff removes 7 operator tests that exist on main (test_qnn_backend_reflection_pad3d, test_qnn_backend_scatter_value, test_qnn_backend_sort, test_qnn_backend_sort_and_index, plus quantized variants). I confirmed these exist on origin/main but not on this branch:

origin/main: 7 matches   HEAD: 0 matches

These were added to main after this branch diverged, so the diff clobbers them. This is almost certainly the cause of the unittest / linux CI failures. Please rebase onto latest main so this PR only touches the intended Gemma4 additions (gemma4-e2b LlmSpecs + the --embedding-quantize cmd). The smollm2_135m SQNR threshold change (19.5 → 19) is legitimate per your prefix_token_id analysis; the operator-test deletions are not.

2. kv_manager.cpp now hard-assumes output tensor ordering for ALL models.
The KV-cache discovery was rewritten from scanning inputs by matching byte-size to indexing outputs by position:

Result<TensorInfo> k_out = method_meta->output_tensor_meta(1 + layer);
Result<TensorInfo> v_out = method_meta->output_tensor_meta(1 + metadata_.num_layers + layer);

This assumes every decoder emits outputs in exactly [logits, k_0..k_{n-1}, v_0..v_{n-1}] order. That's a behavioral change affecting every existing QNN LLM routed through KVManager, not just Gemma4. Please confirm (or note in the PR) that all currently-supported decoders produce this exact output layout — otherwise a model with interleaved (k,v) outputs or a different logits position would silently read the wrong tensor metadata. A defensive ET_CHECK on num_outputs() == 1 + 2*num_layers would make a mismatch fail loudly instead of corrupting cache sizing. examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp:80-120

---

🟡 Should address

3. _tag_ios placeholder condition change affects all models. examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py:576

len(users := list(node.users)) > 0
and "args_" in node.name
and users[0].meta["val"].size()[-2:] in self.kv_cache_shape

This relaxed the previous == 1 to > 0. The new "args_" in node.name guard narrows scope to cache inputs, but a KV placeholder that previously had 2+ users was not tagged before and is now. Since this isn't gated on gemma4-e2b, please confirm it doesn't shift IO tagging for existing models (the passing SQNR CI for other models would be the evidence — worth a sentence in the PR).

4. Embedding-quantize transform relocated — also affects all models. The get_quant_embedding_transform call moved out of _prepare_model (llm_wrappers.py:~381, now removed) into quantize() (:687-695). This is a shared-path change; the reordering relative to convert_linear_to_conv2d / pass registration matters. Assuming existing-model CI passes this is fine, but it deserves a callout in the PR description since it's not Gemma4-specific.

5. convert_hf_to_custom new param has a type/default mismatch. examples/models/gemma4/text_decoder/convert_weights.py:230

get_weight_mapping: Optional[Callable] = False,

Default should be None (it's typed Optional[Callable], False is a bool), and the docstring wasn't updated for the new arg. The if not get_weight_mapping: branch works but reads oddly against a callable. Suggest = None + if get_weight_mapping is None:. Fix this →

---

🟢 Minor / nits

6. kv_io_bit_width can silently become None. In llm_wrappers.py, extra_kwargs["kv_io_bit_width"] is set to None when there's no quant_recipe. In model_wrapper.py:80, kwargs.get("kv_io_bit_width", 32) returns None (not 32) when the key is present-but-None, so get_metadata()["get_kv_io_bit_width"] propagates None. If a default of 32 is intended, use kwargs.get("kv_io_bit_width") or 32.

7. Misleading comment in text_model.py:121. # Compute logits only for the last token(s) precedes logits = self.lm_head(hidden_states), which computes logits for all ar_len positions (no slicing). Either slice or reword.

8. _tag_ios string-matching on stack_trace/source names is fragile. Matching "per_layer_inputs[n_self:]" / "per_layer_inputs[i]" in node.meta["stack_trace"] (:603-611) couples IO tagging to source-line text; it'll break silently if the wrapper is refactored or line variables renamed. Works for now, but worth a comment flagging the coupling so a future edit doesn't quietly drop the tag.

9. Good catches worth keeping: the layers.(\d+) → layers\.(\d+) regex escape fix in model_sharding.py and the prepare_feedfoward_conv → prepare_feedforward_conv typo fix are correct and consistently applied across all call sites.

---

Summary

The Gemma4 model code itself is solid and well-isolated. The two things that most need attention before merge: (1) re-rebase to stop deleting unrelated operator tests (this is failing CI), and (2) confirm the kv_manager.cpp output-ordering assumption holds for all existing decoders. Items 3–5 are shared-path changes that should be explicitly called out as validated against existing-model CI.

Review generated with Claude Code (Opus 4.8).
· branch dev1/danny/gemma4

[psiddh]

@DannyYuyang-quic Can you look at the review comments ?

[DannyYuyang-quic]

@DannyYuyang-quic Can you look at the review comments ?

Response to Review Comments:

1. Branch is stale — this PR deletes unrelated operator tests.

Rebasing onto latest main so the test_qnn_delegate.py diff only contains the gemma4-e2b LlmSpecs.

2. kv_manager.cpp now hard-assumes output tensor ordering for ALL models.

To address concerns about incorrect KV-cache ordering, I added a defensive check that validates the head_dim obtained from both the K-cache and V-cache is the same. This helps ensure that the correct K/V cache tensors are being identified rather than tensors from an unexpected output ordering.

3. _tag_ios placeholder condition change affects all models.

Moving the embedding quantization transform into the quantize flow is more consistent with the overall quantization flow.

4. Embedding-quantize transform relocated:

Move Embedding-quantize transform into quantize flow will make more sense then before.

5~8:

Addressed the review comments by updating documentation/comments and removing unused or unnecessary changes.

[DannyYuyang-quic]

defensive check that validates the head_dim obtained from both the K-cache and V-cache is the same.

Since the cache shapes are:

• output k cache: [1, n_kv, head_dim, ar_len]
• output v cache: [1, n_kv, ar_len, head_dim]

the head_dim inferred from each tensor should match. This provides a sanity check that the selected K/V cache tensors are consistent with the expected shape.