Potential Bug in **_attn_fwd_tma** Function

[qinxiangyujiayou]

Describe the bug

Hi, I believe I found a potential bug in the _attn_fwd_tma function. Specifically, it concerns how the offset_y variable is computed.

Here is the relevant part of the code:

start_m = tl.program_id(0)
off_hz = tl.program_id(1)
off_z = off_hz // H
off_h = off_hz % H
offset_y = off_z + off_h * N_CTX

The calculation of offset_y is a bit confusing. Based on the current implementation, it is computed as:

offset_y = off_z + off_h * N_CTX

However, I believe it should be:

offset_y = off_hz * N_CTX

Could you clarify if this is a bug or explain the reasoning behind the current implementation? Thanks for your help!

Environment details

The latest version of Triton