Qualcomm AI Engine Direct - Add smart mask kv updator for llama3.2

[chunit-quic]

• Add flag to use smart mask or shift pointer
• Add llama3_2 python with smart mask updator
• Change Memory class to IoMgrBase
• Change HybridMemory class to ShiftPointerIoMgr

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure

As of commit 593b866 with merge base e00eaea ():

NEW FAILURE - The following job has failed:

• pull / test-phi-3-mini-runner-linux / linux-job (gh)
  RuntimeError: Command docker exec -t 8ad39757bbd4d4124889cdc1610d79ee386091257aca88e348d65555e3cf6b4d /exec failed with exit code 2

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

[chunit-quic]

Hi @cccclai,

We provide a new KV cache updating mechanism that benefits long-length settings.
For example, with the SM8650 SoC and a KV length of 4096, it generates 47 tokens/sec (compared to 40 tokens/sec previously).

python examples/qualcomm/oss_scripts/llama3_2/llama.py -b build-android -H ${HOST} -s ${DEVICE} -m "SM8650" --checkpoint ${Llama3.2-1B-Instruct}/consolidated.00.pth --params ${Llama3.2-1B-Instruct}/params.json --tokenizer_model ${Llama3.2-1B-Instruct/tokenizer.model} --prompt ${PROMPT} --ptq 16a4w --temperature 0 --model_size 1B --model_mode hybrid --prefill_seq_len 128 --kv_seq_len 4096 --kv_updator smart_mask

Following this, a document PR for the updators will be submitted recently.
Feel free to ask us if you have any questions. Thank you! :D

[facebook-github-bot]

@cccclai has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[cccclai]

I guess this PR needs to be rebased as #7618 is merged?

[chunit-quic]

I guess this PR needs to be rebased as #7618 is merged?

Thanks for pointing out! Just rebased.

[facebook-github-bot]

@cccclai has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[cccclai]

Import fail...I feel like it's still not rebased successfully

[chunit-quic]

Import fail...I feel like it's still not rebased successfully

I have rebased to the latest main branch again. Because I cannot check the import error, Feel free to let me know if it still fails.

[facebook-github-bot]

@cccclai has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[cccclai]

Thanks!

[jds250]

Hi @cccclai,

We provide a new KV cache updating mechanism that benefits long-length settings. For example, with the SM8650 SoC and a KV length of 4096, it generates 47 tokens/sec (compared to 40 tokens/sec previously).

python examples/qualcomm/oss_scripts/llama3_2/llama.py -b build-android -H ${HOST} -s ${DEVICE} -m "SM8650" --checkpoint ${Llama3.2-1B-Instruct}/consolidated.00.pth --params ${Llama3.2-1B-Instruct}/params.json --tokenizer_model ${Llama3.2-1B-Instruct/tokenizer.model} --prompt ${PROMPT} --ptq 16a4w --temperature 0 --model_size 1B --model_mode hybrid --prefill_seq_len 128 --kv_seq_len 4096 --kv_updator smart_mask

Following this, a document PR for the updators will be submitted recently. Feel free to ask us if you have any questions. Thank you! :D

Hi, I am a little confused about why smart_mask updator is more beneficial for long-length settings? Can you help me? Thanks!

[chunit-quic]

Hi, I am a little confused about why smart_mask updator is more beneficial for long-length settings? Can you help me? Thanks!

Sure! The smart_mask updator is better for long-length settings because it uses shared memory, reducing data transfer between CPU and HTP.

[jds250]

Hi, I am a little confused about why smart_mask updator is more beneficial for long-length settings? Can you help me? Thanks!

Sure! The smart_mask updator is better for long-length settings because it uses shared memory, reducing data transfer between CPU and HTP.

Thanks for your reply! Except using shared memory, smark_mask also reduce the memory allocation’s overhead? It seems that shift pointers needs to malloc new space for new kv cache every inference iteration. Aslo I wanna know why we can’t just keep all kv cache in the HTP side, then we don’t have the data transfer cost.

[chunit-quic]

It seems that shift pointers needs to malloc new space for new kv cache every inference iteration.

No both of them allocate only once. You can check this init function implementations of class.

why we can’t just keep all kv cache in the HTP side, then we don’t have the data transfer cost.

That's what the shared buffer help to achieve. :)

[jds250]

It seems that shift pointers needs to malloc new space for new kv cache every inference iteration.

No both of them allocate only once. You can check this init function implementations of class.

why we can’t just keep all kv cache in the HTP side, then we don’t have the data transfer cost.

That's what the shared buffer help to achieve. :)

Oh I see, thanks for your help. BTW I find in the static_llama.py, here implement prepare_feedfoward_conv, it seems like we convert original Linear to a Conv layer, why we should do that? Is convolution more effiecient that Linear with HTP?

[chunit-quic]

why we should do that? Is convolution more effiecient that Linear with HTP?

The reason is that we have specific optimization passes for convolution operations. Since in some case linear is mathematically equivalent to convolution, this conversion results in performance improvements.