Integrate chunked prefill into t3k Llama3-70B

[cglagovichTT]

Ticket

#15873

Problem description

Llama3-7-B on T3K supports up to 128k context length with batch=1. However, it is limited such that the maximum prompt length is 32k. That means that the full 128k context can be used by prefill len + decode len only if prefill len <= 32k.

This constraint is caused by the fact that prefill activations for long sequences become large, leading to OOM on device DRAM. For example, the activations are 128k * 8k * 2B = 2 GB for input len 128k.

What's changed

The solution to this problem is to implement chunked prefill. Given some prompt_len and a chunk_size (<= 32k), the prompt is split into prompt_len / chunk_size chunks and iteratively prefilled. Chunked prefill will reduce the size of the activations which should solve the OOM error.

This PR uses the new chunked_scaled_dot_product_attention kernel in Llama to enable chunked prefill. It modifies the attention and model tests, and creates a new test in the T3K demo pipeline which tests chunked prefill as used by llama_generation, which is the entrypoint that vLLM will use.

Note that this functionality was added to the "old" Llama codebase since it was an urgent request. I expect we will be adding chunked prefill to the llama family folder sometime soon.

Checklist

• All post commit https://github.com/tenstorrent/tt-metal/actions/runs/12303936776
• T3K pipelines https://github.com/tenstorrent/tt-metal/actions/runs/12303156746
  • ttnn test fails but this PR should not have affected it
  • Kicking T3K pipelines again after some changes https://github.com/tenstorrent/tt-metal/actions/runs/12354567605
  • Many T3K pipelines are failing but these changes only affect old-codebase llama3-70b, which are green.