[QUESTION] dsv32 mqa_logits kernel not considering causal masking?

[carlushuang]

Thanks for dsv32 great work!

By analysis the fp8_mqa_logits and fp8_paged_mqa_logits function, looks like after the q@k, we don't consider causal masking before topk?

I know the q/k feed to mqa_logits kernel is different from the q/k feed to MLA attention kernel, but we use the output of the mqa_logits kernel (and topk 2048) as indexer into the real MLA's kvcache, hence during the real attention computation we need consider causal, in prefill or decode(MTP) case.

vLLM prefill dispatch using torch.ops._C.top_k_per_row seems not considering causal, decode dispatch looks like considered causal in MTP case after the logits kernel, before topk.

not sure if it is suppose to let the framework side to do causal before topk, or actually causal is not important during the indexer kernel?

[zheanxu]

To apply causal masking, you can set clean_logits=True. In this case, the subsequent clean_logits kernel will perform the masking. See:

DeepGEMM/csrc/apis/attention.hpp

Line 120 in c9f8b34

smxx_clean_logits(logits, cu_seq_len_k_start, cu_seq_len_k_end, 1, seq_len, seq_len_kv, aligned_seq_len_kv);

A more efficient approach is to set clean_logits=False and incorporate causal masking directly into the topk operation — i.e., ignore masked positions during topk selection. However, this requires implementing your own topk kernel.