[Issue]: Index Mismatch Between Main PA Kernel _paged_attention_kernel and Reduce Kernel _paged_attention_ll4mi_reduce_kernel

Problem Description

The paged_attention_v1 kernel API accepts an mtp parameter suggesting multi-token prediction support,
but values > 1 appear to be incorrect due to multiple implementation bugs.

The main kernel and reduce kernel use incompatible indexing schemes for the intermediate buffers (exp_sums, max_logits, tmp_out).

Main kernel _paged_attention_kernel writes to:
https://github.com/ROCm/aiter/blob/d2f5f2791336a9839c65241b77a88f0a8407a054/csrc/cpp_itfs/pa/pa_kernels.cuh#L660-L668

Reduce kernel_paged_attention_ll4mi_reduce_kernel reads from:
https://github.com/ROCm/aiter/blob/d2f5f2791336a9839c65241b77a88f0a8407a054/csrc/cpp_itfs/pa/pa_kernels.cuh#L926-L928

The indices don't match, meaning the reduce kernel reads incorrect data, silently producing wrong results.

Operating System

Ubuntu 22.04.5 LTS (Jammy Jellyfish)

CPU

AMD EPYC 9575F 64-Core Processor

GPU

8 x AMD Instinct MI355X

ROCm Version

7.0.51831-a3e329ad8

ROCm Component

No response

Steps to Reproduce

Replicator script:
paged_attention_v1.issue1.mi355.py

Logs from running replicator script:
paged_attention_v1.issue1.log

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

No response

Additional Information

My goal is to enable CUDA graph support for speculative decoding in the ROCM_AITER_FA attention backend by adding spec_as_decode support and changing _cudagraph_support from UNIFORM_SINGLE_TOKEN_DECODE to UNIFORM_BATCH. This is expected to speed up speculative decoding: removes graph breaks and allows routing speculative tokens through the more efficient decode path instead of the extend path.
The AITER PA API appears to support MTP>1, but the implementation seems to be buggy.