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 memory consumption and improve 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 1a317f2 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.

[harshs-qti]

This class/file need not be named as "HF". This is how we want the Quant recipe to be applied in any non-static llama flow. This is our eventual target. linear to conv is a graph change that should happen during lowering or to_backend() stage

[harshs-qti]

Also, isn't this supposed to be an abstract class?

[winskuo-quic]

This class/file need not be named as "HF". This is how we want the Quant recipe to be applied in any non-static llama flow. This is our eventual target. linear to conv is a graph change that should happen during lowering or to_backend() stage

Sounds good. I have removed HF from the receipe class and file name.
For the linear_to_conv explanation at top, I forgot to remove that when updating the linear_to_conv behavior to annotation_pass, so we actually target conv node in recipe instead of linear. Thanks for the catch here.
I have removed the linear_to_conv expplanation since you have mentioned all non-static-llama flow uses this file, which means other flow can turn off linear_to_conv

[winskuo-quic]

Also, isn't this supposed to be an abstract class?

I believe it was abstract in runtime behavior but not in declaration.
I have changed to abstract class now.
I only changed this file and not static_llm_quant_recipe.py since it is out of scope of this PR.

[harshs-qti]

can you check how different the Quant config with GenAI refactor as well as the LLM Config integration is. It would be nice to align or be close before creating this interface

[winskuo-quic]

I have checked GenAI flow and believe this should be fine. GenAI flow as it is still reusing quant recipe logic.

[harshs-qti]

Adding such flags one at a time is not very scalable. Also, get is trying to return list of passes while passing flags seems to already have knowledge of what passes are present. This seems inverted. Specific passes can be disabled by the caller as needed after get()

[winskuo-quic]

I totally agree with you, and this is actually on my list too.
Worth noting how deep it goes: the parameter is threaded through five layers, and the GPU override accepts it without ever reading it, since annotation is skipped for GPU entirely.
Doing it properly means a design refactor of the annotation pass pipeline, and I'd like to keep this PR focused on the HF optimizations with minimal changes to pass structure. I think it makes more sense to target it in a follow-up PR.

[harshs-qti]

can we replicate the "note" from Llama3_2_1B_HFQuantRecipe in all the recipes

[winskuo-quic]

Added to all other recipes.