Question of O_accum and LSE_accum shape

[mingyangHao]

FlashMLA/csrc/flash_api.cpp

Line 184 in b31bfe7

at::Tensor softmax_lse_accum = torch::empty({batch_size + params.num_sm_parts, num_heads, seqlen_q}, opts.dtype(at::kFloat));

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?

[beginlner]

Because there is at most one split job between two consecutive SMs, num_sm_parts additional jobs in total.