[Performance] Fix K/V prefetch offsets and register spilling in flash attention tests

[charithaintc]

Please review these guidelines to help with the review process:

• Have you provided a meaningful PR description?
• Have you added a test, a reproducer, or a reference to an issue with a reproducer?
• Have you tested your changes locally for CPU and GPU devices?
• Have you made sure that new changes do not introduce compiler warnings?
• If this PR is a work in progress, are you filing the PR as a draft?
• Have you organized your commits logically and ensured each can be built by itself?