Jacchang/pa ragged experimental

[Jacob0226]

Motivation

The decode kernel _paged_attention_kernel in aiter/csrc/cpp_itfs/pa/pa_kernels.cuh shows an increasing runtime as the batch size scales from 8 to 512, rising from 4.8% to 11% of total runtime.
This PR currently supports only head_size = 128 and kv_dtype = bf16.

Technical Details

This work introduces the following changes:

1. Added double buffering for K-cache loading.
2. Used 64 threads to load the continuous K-cache into LDS and then distributed the data to thread registers to match the MFMA layout.
3. Turn on non-temporal loads for KV cache.

Test Plan

• Unit test: aiter/op_tests/test_pa_ragged_experimental.py
  This script differs from the original test_pa_ragged.py. The new test reproduces the same KV-cache layout used in SGLang, making it more realistic. It also performs numerical verification against the original kernel (referred to as GOLDEN) to ensure correctness. The new implementation is labeled EXPERIMENTAL.

# Docker: lmsysorg/sglang:v0.5.5.post3-rocm700-mi35x
export BS=512
export CL=2048 # context length (prefill length)
export PageSize=1
export OUT=UT
rocprof-compute profile -n $OUT -- python aiter/op_tests/test_pa_ragged_experimental.py -n 512 -c 2048 --warmup 3 --page-size 1
rocprof-compute analyze -p workloads/$OUT/MI355/ --list-stats  > tmp.log 
# Find the kernel paged_attention_ll4mi_QKV_mfma16_kernel dispatch ID for GOLDEN and EXPERIMENTAL
# The second-to-last matched dispatch  ID is GOLDEN
# The last one is for matched dispatch ID EXPERIMENTAL

• Accuracy test:
  sglang/benchmark/gsm8k/bench_sglang.py --num-questions 200 --parallel 200 --num-shots 8

Test Result

Node: MI355X (smci355-ccs-aus-n10-09)

• Unit test:
  GOLDEN: 100 us
  EXPERIMENTAL: 88 us.
  Speedup: 1.13x.
• Accuracy test: GSM8K doesn't drop by using this PR.

Submission Checklist

• [V] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[fsx950223]

Format code.

aiter/.githooks/pre-commit

Line 75 in 1f5b378

echo "black $file"

[Jacob0226]

@fsx950223 Ready for re-review