[GH-ISSUE #209] [QUESTION] dsv32 mqa_logits kernel not considering causal masking? #118

New issue

Open

opened 2026-10-26 21:13:00 +00:00 by bsakaguchi · 1 comment

bsakaguchi commented

2026-10-26 21:13:00 +00:00

Owner

Originally created by @carlushuang on GitHub (Oct 7, 2025).
Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/209

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?

Originally created by @carlushuang on GitHub (Oct 7, 2025). Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/209 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](https://github.com/vllm-project/vllm/blob/08d26a1b7edc200d8d117491eac3e28c0428e571/vllm/model_executor/models/deepseek_v2.py#L654C26-L654C39) using `torch.ops._C.top_k_per_row` seems not considering causal, [decode dispatch](https://github.com/vllm-project/vllm/blob/08d26a1b7edc200d8d117491eac3e28c0428e571/vllm/model_executor/models/deepseek_v2.py#L708) 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?

bsakaguchi commented

2026-10-26 21:13:01 +00:00

Author

Owner

@zheanxu commented on GitHub (Oct 22, 2025):

To apply causal masking, you can set clean_logits=True. In this case, the subsequent clean_logits kernel will perform the masking. See:
github.com/deepseek-ai/DeepGEMM@c9f8b34dcd/csrc/apis/attention.hpp (L120)

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.

@zheanxu commented on GitHub (Oct 22, 2025): To apply causal masking, you can set `clean_logits=True`. In this case, the subsequent `clean_logits` kernel will perform the masking. See: https://github.com/deepseek-ai/DeepGEMM/blob/c9f8b34dcdacc20aa746b786f983492c51072870/csrc/apis/attention.hpp#L120 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.

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

deepseek-ai/DeepGEMM#118

No description provided.

Rows
Columns