Qualcomm AI Engine Direct - HF LLM Optimization Stage 1

[winskuo-quic]

Summary

As HF use case is increasing, we aim to align accuracy and performance of HF LLM with Static LLM(static_llama.py).
Currently, there are many gaps, and we have a series of PR to fill the gaps.
For now, we will first focus on Llama3.2 1B decode mode.

What is done

• Annotate KV 8bit IO. This largely reduces the KV IO size from 16bits to 8 bits, significantly reducing runtimes both memory consumption and performance. In order to reuse the existing, annotate_kv_8bit that is used by static llama, we have also made corresponding changes to HF-LLM so it has the same pattern as static llama.
• HF LLM Quant Recipe: In order to align performance and accuracy with static llama, a hf llm will also have its own quant_recipe under examples/qualcomm/oss_scripts/llm_utils/hf_llm_quant_recipe.py. The quantization technique is mostly migrated from executorch/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py, with some minor changes, especially the weight naming as HF-LLM and Static-LLM structure naming is different.
• Reuse qnn_llama_runner instead of using Llama Main. This is done by aligning the IO shapes of HF llama in AoT stage, meaning we don't need to modify qnn_llama_runner at all. This provides us more flexibility on model modification and reusing a lot of features already enabled in qnn_llama_runner.
• Custom Static Cache: In mainline, we can only dump all KV caches, so the output shape is directly affected by seq_len. This increases the memory consumption and is also bad for performance. To resolve this issue, we have implemented QnnCustomStaticCache so mode now only outputs the new pair of kv-cache.
• Precompute Rope: This can greatly improve the accuracy especially for model that has huge inv_freq Additionally, since we are precomputing rope, we can also improve performance.
• Convert_Linear_to_Conv added to the stage of annotation: This might be counter intuitive because a linear node is decomposed to a set of nodes {conv2d, reshape, squeeze, permute}. All these should be sharing the same quant_attrs with since nodes since conv2d is equivalent to linear here, and reshape, squeeze, permute is just changing the shape but not values, so they should be able to share the same quant_attrs. However, this becomes tricky when this pass is mixed with annotate_kv_8bit in the mainline. This issue is happens during v_proj linear node. When these 2 are used together, the permute and view_copy will be tagged as 16bit, so it becomes conv(16bit) ->permute(16bit) -> view_copy(16bit) -> cast(8bit). This is bad and different from static llama since if we insert cast(8bit) right after conv(16bit), rest of the ops can run in 8bits, which is faster. After this change, the behavior now aligns with static llama, which is conv(16bit) -> cast(8bit) ->permute(8bit) -> view_copy(8bit).
• R3 is removed on purpose for this PR. First, without R3, all models are still generating pretty accurate outputs. Second, to make R3 well align with static_llama, the matrix might be spreading around. It would be better to enable this feature in later PRs when supporting advanced features.
• FuseConsecutiveReshape: This fold's extra view_copy node and also allow the model shape to match mha_to_sha.

Optimization Result Compared with Llama3.2 1B_instruct (HF uses non-instruct verison)

For now, accuracy is not compared as HF flow has not yet support SQNR evaluation. Official accuracy features will be supported in Upcoming PRs. However, this PR did ensure HF output is reasonable and readable.

llama3.2-1b	Static (Baseline)	HF Before (Mainline)	HF Optimized (This PR)
performance (tok/sec)	64	19	55

Sample Script

python examples/qualcomm/oss_scripts/hf_causal_lm.py --prompt "Simply put, the theory of relativity states that" --soc_model SM8750 --device $DEVICE --build_folder build-android/ --decoder_model llama3_2-1b --max_seq_len 1024

Upcoming Features

Performance & Accuracy

• As we are now filling the gaps, HF seq=1024 improves from 16 to 55 tokens, and it is now getting close to static llm's performance. For stage 2, we will be focusing on using Qualcomm SDK tools such as qhas to do detailed profiling to compare HF and Static model, understanding the gap, analyze and see if we can fill rest of the gaps.

Feature

• Reusing Eval Tools under Static Llama
• Reorg model specific configs.
• Adding R3 and migrating more StaticLlama source model optimization.

Test plan

• Test for pass:

1. pytest backends/qualcomm/tests/rework/passes/test.py -k "test_fuse_consecutive_reshape"
2. pytest backends/qualcomm/tests/rework/passes/test.py -k "test_convert_linear_to_conv2d"

• E2E test updated: python backends/qualcomm/tests/test_qnn_delegate.py TestExampleLLMScript.test_hf_causal_lm --device $DEVICE--soc_model SM8750 --build_folder build-android --executorch_root . --artifact_dir ./llama3_2_1b --model_name llama3_2-1b

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 8fdf21c with merge base d750618 ():
💚 Looks good so far! There are no failures yet. 💚

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

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.