Add KV Cache Layout Control and Attention Tests

[morgandu]

This PR is to enable KV Cache Layout Control to allow permutation experiments seeking performance improvement.

Checklist

• Add KV Cache Layout Control
• Add Unit Tests for Attention
• Update Docstring
• Microbenchmark - Performance
• E2E tests with JetStream v0.2.2 - Accuracy and Performance

Microbenchmark - Performance

Setup

• v5e4
• llama2-7b
• quantization=int8
• quantize_kvcache=true
• per_device_batch_size=24

Results

Summary

Throughput Before	Throughput After	Throughput Improvement
1469.555	2150.441	46.33%

Before

• default layout, i.e. 1,2,0,3 for both prefill/ar key/value

AutoRegressive results:
        AR step average time: 65.326 ms
        AR step average time per seq: 0.680 ms
        AR global batch size: 96
        AR throughput: 1469.555 tokens/second
        AR memory bandwidth per device: 227.221 GB/s

Xprof: https://xprof.corp.google.com/overview_page/morgandu-14884055542603397573

• HBM BW util of qk_product's dot_general is 26%
• HBM BW util of wv_product's dot_general fusions are around 27%

After

• layout control, i.e. 0,2,3,1 for both prefill/ar key/value

AutoRegressive results:
        AR step average time: 44.642 ms
        AR step average time per seq: 0.465 ms
        AR global batch size: 96
        AR throughput: 2150.441 tokens/second
        AR memory bandwidth per device: 332.499 GB/s

Xprof: https://xprof.corp.google.com/overview_page/morgandu-16660388030961792706

• HBM BW util of qk_product's dot_general is 89.12%
• HBM BW util of wv_product's dot_general fusions are around 85%

NOTE:

• There is an expensive dynamic-update-slice ops in the After run, this is likely due to layout not being optimal.

E2E tests with JetStream v0.2.2 - Accuracy and Performance

Setup

• Llama2-7b
• v5e-8
• JetStream release v0.2.2
• openorca chat dataset (base acc is not accountable, perf only)
• Same maybe sub optimal layout as mentioned in above Note

Results

Summary

• No Acc regression after introducing layout control
• Performance improvement is observed as below
  • For base model, quantized kv cache only, full warmup, throughput improved 33.22%
  • For chat model, quantized kv cache only, full warmup, throughput improved 26.80%
  • NOTE: The model mode chat and quant mode w-b16-kv-i8 combination is the main mode of interest, as weight quantization is still dev in progress
• Full warmup yield better result than existing partial warmup, especially for chat mode model, throughput improved about 20% purely from full warmup

Details

Model Name	Model Mode	Quant Mode	Batch Size	Exit Rouge1	Layout Rouge1	Exit Thrpt	Layout Thrpt	Layout Improv
llama2-7b	base	w-b16-kv-b16	10	-	-	2586.307341	3335.753476	28.98%
llama2-7b	base	w-b16-kv-i8	24	-	-	2348.965528	3129.199496	33.22%
llama2-7b	base	w-i8-kv-i8	24	-	-	2491.088673	3347.33998	34.37%
llama2-7b	chat	w-b16-kv-b16	10	45.3415	45.3326	2003.580613	2470.159128	23.29%
llama2-7b	chat	w-b16-kv-i8	24	45.4035	45.24	1470.583332	1864.755825	26.80%
llama2-7b	chat	w-i8-kv-i8	24	44.8522	44.7558	1470.031543	1930.533052	31.33%

[vipannalla]

lgtm, few comments. Also, the microbenchmark results looks great, were you able to do a full end-2-end on JetStream?

[vipannalla]

LGTM