[QST]some question about blackwell fmha

[liuqi123123]

when ElementAccumulatorQK=float, acc_qk = q@k, will exist in tensor memory with float type, when next gemm acc_qk @ v, type of acc_qk is float, type of v is cutlass::half_t, acc_qk should trans to half_t firstly, but I don't see the transform code in sm100_fmha_fwd_kernel_tma_warpspecialized.hpp, where is it ? And question2, for the code: tOrP0.data() = tOrP0.data().get() + uint32_t(TmemAllocation::P0); why should tOrP0 offset 32 before matmul with v ? shouldn't be tmem_base_addr? I see the tOrP0 layout is tmem_16b o ((_128,_16),_1,(_4,_2)):((_131072,_1),_0,(_16,_64))
,why is the row coordinate stride _131072 , not the original _65536 ?
I am very confused and would appreciate it if you could answer my question. Thank you.

[khvsreddy-code]

I have removed all special formatting like bolding and italics so you can copy and paste the text easily.

Where is the Float to Half Transformation?

The transformation is not a separate piece of code because it happens during the Softmax write-back phase. In the SM100 architecture, the first GEMM results (Q times K) stay in TMEM as float. The kernel then runs a Softmax operation. This Softmax loads the float values into registers, calculates the exponential and normalization, and then stores the result back into TMEM into the tOrP0 buffer.

Since tOrP0 is defined as a 16-bit tensor (half), the store instruction itself performs the conversion from float to half. Look for the store or copy call inside the Softmax functor or the section between the two GEMMs.

Why the offset of 32 (P0)?

TMEM is organized into columns and rows. The base address points to the start of the TMEM allocation for that warp or thread block. However, different tensors need their own space.

The offset TmemAllocation::P0 (which is 32) shifts the starting column of the pointer. This means columns 0 to 31 might be used for the S matrix (the float accumulator), while column 32 starts the P matrix (the half-precision probabilities). If you used the base address for both, the second GEMM would overwrite the first one before it was ready.

Why is the stride 131072 instead of 65536?

In the TMEM addressing system, 65536 (2 to the power of 16) represents exactly one row. A stride of 131072 is exactly two rows.

This is because of how the Blackwell hardware (SM100) reads 16-bit data for the second GEMM (P times V). To maximize speed, the hardware often expects data to be interleaved or placed in specific banks. By using a stride of two rows, the layout is likely skipping every other row or organizing the data to match the specific requirements of the tcgen05.mma instruction. This ensures that when the math core requests a block of data, the TMEM banks can deliver it in a single cycle without bank conflicts.

[liuqi123123]

thanks for your reply! I relearned the code. I find where the transform occur and I think that the TMEM col offset of P0 tensor is 32 because softmax_step use tStS(shape 128*2) to storage the row_max and row_sum. So P0 tensor should get offset. In addition,
This is the first time I've heard that TMEM can also experience bank conflict, it's interesting .Is there any document to explain it ? I don't find it in https://docs.nvidia.com/cuda/parallel-thread-execution/