Qualcomm AI Engine Direct - GA Static Gemma3-1B

[DannyYuyang-quic]

Summary:

• e2e script for GA Static Gemma3-1B
  • perf: 16a4w block quant token rate in kv mode: ~= 110 tokens/sec(SM8750), max_seq_len=1024
  • acc: PPL ~= (fp:21.375 -> htp:23.086) in wikitext dataset
  • add model params config
  • add End-to-End example in README
• add new architecture:
  • add new class to support global/local ROPE static llama architecture required by Gemma3
  • enable global/local static llama architecture support in runner
• refactoring:
  • refactor attention mask to improve integration with global/local ROPE static llama model
  • refactor kv_inference and prefill_inference for better readability
• Unitest:
  • add unit test for Gemma3-1B
  • improve readability of memory size constant in unit test
• LLM model config visualization
  • support tabular LLMmodelConfig visulization

Test plan

python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${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?" --tasks wikitext --limit 1

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures, 1 Cancelled Job, 5 Unrelated Failures

As of commit 9c80e2f with merge base 8496f27 ():

NEW FAILURES - The following jobs have failed:

• pull / test-moshi-linux / linux-job (gh)
  RuntimeError: Command docker exec -t e2f831ba3c2d50b58e1a0026a0147ea18de93897b4ccd6010e06390aa6d5a756 /exec failed with exit code 1
• pull / unittest-editable / macos / macos-job (gh)
  backends/xnnpack/test/ops/test_conv2d.py::TestConv2d::test_fp16_conv2d

CANCELLED JOB - The following job was cancelled. Please retry:

• trunk / test-models-macos-coreml (mv3) / macos-job (gh)
  ##[error]The operation was canceled.

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-openvino-linux / linux-job (gh) (similar failure)
  ##[error]The operation was canceled.

BROKEN TRUNK - The following jobs failed but was present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Build Presets / zephyr (zephyr) / build (gh) (trunk failure)
• pull / test-binary-size-linux-gcc / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / test-samsung-models-linux / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / test-setup-linux-gcc / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.

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

[DannyYuyang-quic]

@pytorchbot label "release notes: qualcomm"

[DannyYuyang-quic]

Hi @cccclai,
This PR enables support for Gemma3-1B in the static version.

Both accuracy and performance in Hybrid/KV mode are promising.
However, since Gemma3 uses global/local attention mechanism, the implementation of lookahead decoding is a bit tricky, so we temporarily block lookahead decoding only for Gemma3 now, we plan to enable lookahead decoding support for Gemma3 in a future update.

cc: @haowhsu-quic

[facebook-github-bot]

@cccclai has imported this pull request. If you are a Meta employee, you can view this in D82034374.

[jackzhxng]

Are the changes in examples/models/gemma3 and examples/models/llama relevant?

[cccclai]

Are the changes in examples/models/gemma3 and examples/models/llama relevant?

Yes because we reuse the config from the etllm for qualcomm llm models as well

[cccclai]

There are some merge conflicts, can you resolve it?

[facebook-github-bot]

@cccclai has imported this pull request. If you are a Meta employee, you can view this in D82034374.

[jackzhxng]

Looks good. Trunk errors look like flakes

[mergennachin]

@DannyYuyang-quic

There's a bug in this code that was uncovered in our internal testing:

executorch/examples/qualcomm/oss_scripts/llama/model/static_llama.py", line 266, in forward_sha
    attn = attn + atten_mask
           ~~~~~^~~~~~~~~~~~
TypeError: unsupported operand type(s) for +: 'Tensor' and 'AttentionMask'

[mergennachin]

atten_mask is sometimes Tensor and sometimes AttenionMask object.

cc @cccclai

[mergennachin]

Can you fix this asap, otherwise, I'll revert this PR

[haowhsu-quic]

May I know the details of internal test scenario? I tested the latest mainline without masked_softmax and everything works fine.

[haowhsu-quic]

The atten_mask is now a wrapper class for causal / sliding attention. We do unwrap it when perform lowering process in llama.py.
Would it be possible you use the inputs from LlamaModel.get_example_inputs() to do the inference? If so, please add asterisk to the attention_mask:

tokens, atten_mask, pos_ids, k_cache, v_cache = model.get_example_inputs()
logits, new_k_caches, new_v_caches = module(
  tokens,
  *atten_mask,
  pos_ids,
  *k_caches,
  *v_caches,
)

[cccclai]

I'll try to fix it, looks like just an internal reference need to update