[Performance, Triton Kernel Args] extend_attention, optimize kern args to _fwd_kernel

[HaiShaw]

Motivation

Speedup:
Llama-3.1-8B, TP=8, FP8, b32/i1024: prefill throughput +2.82%
Grok-1, TP=8, FP8, b32/i1024: prefill throughput +2.23%

Modifications

Setting optimal kernel arguments to _fwd_kernel of extend_attention on ROCm.

Checklist

• Format your code according to the Contributor Guide.
• Add unit tests as outlined in the Contributor Guide.
• Update documentation as needed, including docstrings or example tutorials.