flash_fwd_mla_kernel.h offset calculation 偏移量计算问题

[ChengruiZhang]

In the file flash_fwd_mla_kernel.h, there are several double-buffering processes. In each of these processes, the target offsets are either sK_offset / 8 or sK_offset, where sK_offset is equal to 576 * 64. However, I am unclear about the purpose of the /8 operation.

If the iteration is intended to proceed along the N-dimension, the offset in matrix K should naturally be 576 * 64. Why is the additional division by 8 necessary?

[beginlner]

For instance, tSrK.data() is a wgmma shared memory descriptor, as described in https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-shared-memory-layout-matrix-descriptor. Shifting the matrix start address by sK_offset bytes corresponds to incrementing the descriptor by (sK_offset >> 4) * sizeof(bf16) = sK_offset / 8.

[ChengruiZhang]

For instance, tSrK.data() is a wgmma shared memory descriptor, as described in https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-shared-memory-layout-matrix-descriptor. Shifting the matrix start address by sK_offset bytes corresponds to incrementing the descriptor by (sK_offset >> 4) * sizeof(bf16) = sK_offset / 8.

Thanks for your response!
But I am still confused about the offset value. After we compute one of the blocks "BK", we need to load the next "BK", then the offset should be 576 * 64 * sizeof(fp16) Bytes == sK_offset, so the descriptor offset will be sK_offset / 8. But why some of the tKsK.data() offsets sK_offset？

Is it caused by warp groups? 1 warp group load 576 * 64 * sizeof(bf16) * 2(double_buffer) * 4 = sK_offset * 8?

[lileilai]

For instance, tSrK.data() is a wgmma shared memory descriptor, as described in https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-shared-memory-layout-matrix-descriptor. Shifting the matrix start address by sK_offset bytes corresponds to incrementing the descriptor by (sK_offset >> 4) * sizeof(bf16) = sK_offset / 8.

Thanks for your response! But I am still confused about the offset value. After we compute one of the blocks "BK", we need to load the next "BK", then the offset should be 576 * 64 * sizeof(fp16) Bytes == sK_offset, so the descriptor offset will be sK_offset / 8. But why some of the tKsK.data() offsets sK_offset？

Is it caused by warp groups? 1 warp group load 576 * 64 * sizeof(bf16) * 2(double_buffer) * 4 = sK_offset * 8?

Also confused the offset calculation, Do you have more clear explanation

[yangpc615]

The stride of GMMA::DescriptorIterator is 16bytes, so (sK_offset >> 4) * sizeof(bfloat16)

[RookieArt]

For instance, tSrK.data() is a wgmma shared memory descriptor, as described in https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-shared-memory-layout-matrix-descriptor. Shifting the matrix start address by sK_offset bytes corresponds to incrementing the descriptor by (sK_offset >> 4) * sizeof(bf16) = sK_offset / 8.

Thanks for your response! But I am still confused about the offset value. After we compute one of the blocks "BK", we need to load the next "BK", then the offset should be 576 * 64 * sizeof(fp16) Bytes == sK_offset, so the descriptor offset will be sK_offset / 8. But why some of the tKsK.data() offsets sK_offset？
Is it caused by warp groups? 1 warp group load 576 * 64 * sizeof(bf16) * 2(double_buffer) * 4 = sK_offset * 8?

Also confused the offset calculation, Do you have more clear explanation

As to tKsK, which is smem, the address is offset by number of elements, which is sK_offset. While tOrVt is register memory, corresponding to descriptor, the address is offset by (bytes >> 4), which is (sK_offset >> 4) * sizeof(bf16) = sK_offset / 8