Why Does FA3 Use Registers Instead of Directly Accessing SMEM with WGMMA on SM90?

[ziyuhuang123]

I learned that SM90's WGMMA can fetch operands directly from SMEM instead of requiring them to be loaded into registers. However, I noticed that FA3 still fetches operands from registers. What is the reason for this?

For example:

flash::gemm</*zero_init=*/true, /*wg_wait=*/-1>(tiled_mma0, tSrQ, tSrK(_, _, _, smem_pipe_read_k.index()), tSrS);

https://github.com/Dao-AILab/flash-attention/blob/0dfb28174333d9eefb7c1dd4292690a8458d1e89/hopper/mainloop_fwd_sm90_tma_gmma_ws.hpp#L724C9-L724C122