llama_3.2_1b example error

[Davidqian123]

When I try to run applications/llama_3.2_1b, it throws such error:

AIE Configuration (✔ = AIE NPU / ✘ = CPU):
[Model] KV Cache ✔
[Decode] GEMV ✔
[Attention] Rope ✔
[Attention] QKV GEMM ✔
[Attention] Fused MHA ✔
[FFN] GEMM ✔
[FFN] Elementwise Mul ✔
[FFN] SiLU ✔
[Transformer] Residual Addition ✔
[Transformer] Pre Norm ✔
[Transformer] Post Norm ✔
[Transformer] Final Norm ✔
[Transformer] Final GEMM ✘

Traceback (most recent call last):
  File "/home/nexaai/Desktop/code/IRON/applications/llama_3.2_1b/inference.py", line 319, in <module>
    inference(
  File "/home/nexaai/Desktop/code/IRON/applications/llama_3.2_1b/inference.py", line 143, in inference
    AIEOperatorBase.compile_all_operators()
  File "/home/nexaai/Desktop/code/IRON/applications/llama_3.2_1b/src/operator/aie_base.py", line 35, in compile_all_operators
    op.compile()
  File "/home/nexaai/Desktop/code/IRON/applications/llama_3.2_1b/src/operator/aie_base.py", line 204, in compile
    comp.compile(compilation_rules, work_list)
  File "/home/nexaai/Desktop/code/IRON/applications/llama_3.2_1b/src/compilation.py", line 496, in compile
    success, artifacts = apply_rules(rules, remaining)
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/nexaai/Desktop/code/IRON/applications/llama_3.2_1b/src/compilation.py", line 483, in apply_rules
    artifacts = rule.compile(artifacts)
                ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/nexaai/Desktop/code/IRON/applications/llama_3.2_1b/src/compilation.py", line 210, in compile
    mlir_code = callback_function(
                ^^^^^^^^^^^^^^^^^^
  File "/home/nexaai/Desktop/code/IRON/example/mha/mha.py", line 155, in fused_mha
    num_q_blocks % number_of_pipelines == 0
AssertionError: Number of Q blocks must be divisible by number of pipelines (for now)

How to solve it?

[andrej]

Hi David,

Good catch. Thanks for reporting this.

Our under-development fused MHA kernel currently places relatively tight restrictions on input sizes. We have been testing with --prompt_len 2048 -- please try if this works for you as well. Adding padding or making the fused MHA more general is on our to-do list.

Until then, you can also try the non-fused MHA: In the configuration (applications/llama_3.2_1b/configs/llama32_1b.json) choose use_regular_mha to be true, and use_fused_mha to be false. This will be slower, but should support different prompt sizes.

Once you get things running, be sure to also check out PR #11, which speeds things up significantly. We expect to merge this shortly.

This model mapping is still brand new and we're expecting to make quick strides in usability and performance in the upcoming weeks. Look at it primarily as an example of the operators that we already have and how you might string them together, and less as a final polished application.

Best,
André

[Davidqian123]

Thanks for your reply, I have already successfully run the model, hope more performance improvement in the future!

[andrej]

To do before closing this: Merge AARInternal PR #87 into public repo

[kurtis-b-1]

To do before closing this: Merge AARInternal PR #87 into public repo

Created #25