Question of O_accum and LSE_accum shape

https://github.com/deepseek-ai/FlashMLA/blob/b31bfe72a83ea205467b3271a5845440a03ed7cb/csrc/flash_api.cpp#L184

Hi all,

Just wondering why the shape of O_accum and LSE_accum change from [numsplit, batch, head, seqlen_q, (dim)] to [**batch + sm_part**, ... ] I can get the idea that sm_parts should be the first dimension of output cz it replaced the batchsize, and I also get that sm_parts is not the exact number of loops cz some batch are merge into one sm_parts, but why is the addition work here?