SIGN IN SIGN UP

Qualcomm AI Engine Direct - Optimize the performance for AR-N model (#9079)

Summary:
  - Fix the bug of rms norm builder
- Use HuggingFace version RoPE to improve the performance due to stride
= 1 in StrideSlice Op
  - Modificate the axis order of the conv in qkv, feedforward and output
- Original (AR:128, CL:2048): QNN_RmsNorm (1,1,128,2048) -> QNN_Reshape
(1,128,2048,1)->QNN_Transpose (1,128,1,2048)->self.output->
QNN_Transpose(1,128,2048,1) -> QNN_Reshape (1,1,128,2048)
- New: QNN_RmsNorm (1,1,128,2048) -> QNN_Reshape
(1,128,1,2048)->QNN_Transpose (1,1,128,2048)->self.output->
QNN_Transpose(1,128,1,2048) -> QNN_Reshape (1,1,128,2048)


## Test Result:
- Verify the output for story llama with smart mask, CL=128,
prefill_ar_n=16, prompt="Once"
Note that using Hugging Face RoPE will slightly affect accuracy
  - Original (mainline)
```
INFO:root:Results[0]:
Once upon a time, there was a little girl named Lily. She loved to play with her toys and her favorite toy was a big, red ball. One day, Lily's mom asked her to help her with the laundry. Lily was happy to help and she put all the clothes in the washing machine. 
After the clothes were washed, Lily's mom asked her to help her hang them up to dry. Lily saw a big, black rake and asked her mom what it was. Her mom told her it was a rake and that it helps to
```
  - Optimized (this PR)
```
INFO:root:Results[0]:
Once upon a time, there was a little girl named Lily. She loved to play with her toys and her favorite toy was a big, red ball. One day, Lily's mom asked her to help her with the laundry. Lily was happy to help and she put all the clothes in the washing machine. 
After the clothes were washed, Lily's mom asked her to help her hang them up to dry. Lily saw a big, black iron on the counter and asked her mom what it was for. Her mom explained that it was used to make clothes smooth
```

- Verify the performance for llama 3.2 1B with shift pointer, CL=2048,
prefill_ar_n=256
  - Original (mainline)
```
I 00:00:02.048851 executorch:runner.cpp:354] Prompt Processor: total 256 tokens (AR-256 * 1 iters)
I 00:00:36.606984 executorch:runner.cpp:456] 	Prompt Tokens: 256    Generated Tokens: 1791
I 00:00:36.607049 executorch:runner.cpp:462] 	Model Load Time:		2.012000 (seconds)
I 00:00:36.607062 executorch:runner.cpp:472] 	Total inference time:		34.592000 (seconds)		 Rate: 	51.774977 (tokens/second)
I 00:00:36.607072 executorch:runner.cpp:480] 		Prompt evaluation:	0.293000 (seconds)		 Rate: 	873.720137 (tokens/second)
I 00:00:36.607080 executorch:runner.cpp:491] 		Generated 1791 tokens:	34.299000 (seconds)		 Rate: 	52.217266 (tokens/second)
I 00:00:36.607089 executorch:runner.cpp:499] 	Time to first generated token:	0.293000 (seconds)
I 00:00:36.607099 executorch:runner.cpp:506] 	Sampling time over 1791 tokens:	1.473000 (seconds)
```
  - Optimized (this PR)
```
I 00:00:01.827440 executorch:runner.cpp:354] Prompt Processor: total 256 tokens (AR-256 * 1 iters)
I 00:00:03.143673 executorch:runner.cpp:456] 	Prompt Tokens: 256    Generated Tokens: 64
I 00:00:03.143686 executorch:runner.cpp:462] 	Model Load Time:		1.791000 (seconds)
I 00:00:03.143698 executorch:runner.cpp:472] 	Total inference time:		1.350000 (seconds)		 Rate: 	47.407407 (tokens/second)
I 00:00:03.143706 executorch:runner.cpp:480] 		Prompt evaluation:	0.126000 (seconds)		 Rate: 	2031.746032 (tokens/second)
I 00:00:03.143715 executorch:runner.cpp:491] 		Generated 64 tokens:	1.224000 (seconds)		 Rate: 	52.287582 (tokens/second)
I 00:00:03.143723 executorch:runner.cpp:499] 	Time to first generated token:	0.126000 (seconds)
I 00:00:03.143733 executorch:runner.cpp:506] 	Sampling time over 64 tokens:	0.058000 (seconds)
```
S
shewu-quic committed
baf35d2c3ea48d5d7b43ea2552bb3b7bb69fe298
Parent: ebea003
Committed by GitHub <noreply@github.com> on 3/13/2025, 4:04:47 AM