[Hardware][TPU][V1] Multi-LoRA Optimisations for the V1 TPU backend

[Akshat-Tripathi]

Summary

This PR optimises the Multi-LoRA implementation from #14238. This one should be merged in after it.

This includes several kernel optimisations:

• Block size tuning 2bb8868 d7338f8
• Faster mask creation 2aacb34
• Allowing for some blocks to be skipped 6ee0b57
• Adding LoRA Laning eb804a0
• Splitting the Pallas kernel into shrink/expand variants de6746a
• Removing masking when only 1 LoRA adapter is used aad109b

And a few general ones:

• Pre-transposing the LoRA adapters used in the expand op a82f3fe
• Reducing recompilations 5638e7d

Things left/RFC

• There are still a few recompilations at the start of a run that I need to track down
• LogitsProcessorWithLoRA introduces a long (~1.5 second) stall when it's enabled, but not much activity seems to happen on the CPU or TPU during this time. I've disabled this for now.
• It seems LogitsProcessorWithLoRA is always created even if there's no LoRA adapter that needs it, is there a reason for this?
• I have microbenchmarks for the kernels, but I'm not sure what the right place to put them is.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Akshat-Tripathi.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[Akshat-Tripathi]

I've got some performance numbers using the MLPerf Llama2-70B inference benchmark. I retokenised the dataset for Llama3.1.

Model	Parameters	Without LoRA (tok/s)	With LoRA (tok/s)
Llama3.1	8B	1621.62	1426.01
Llama3.1	70B	432.964	326.709

[psyhtest]

Benchmarking LoRA against baseline (no LoRA) throughput

We use NVIDIA's GenAI-Perf tool to force fixed-length inputs and outputs to produce "heatmap" plots as below. On TPU-v6e and H100 instances, we vary the inputs from 128 to 8k. On L4 instances, we vary the inputs from 128 to 2k.

We calculate the LoRA slowdown as ((LoRA throughput / baseline throughput) - 1) * 100%.

Llama3.1-8B

1x TPU-v6e

The LoRA slowdown varies from -8.4% to -23.9%.

1x GPU-L4

The LoRA slowdown varies from -17.3% to -32.8%.

1x GPU-H100

The LoRA slowdown varies from -10.0% to -51.8%.

Llama3.1-70B

8x TPU-v6e

The LoRA slowdown varies from -20.7% to -46.3%.

8x GPU-L4

The LoRA slowdown varies from -13.8% (second best: -25.1%) to -49.7%.

4x GPU-H100

Unable to launch VMs due to persistent unavailability across multiple zones and regions.