DeepSeek/FlashMLA
1
0
Fork 0
mirror of https://github.com/deepseek-ai/FlashMLA synced 2025-06-26 18:15:54 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

change to use per_tensor

Browse Source
This commit is contained in:
chenhongmin.will
2025-02-26 10:17:29 +08:00
parent 4b314cd655
commit f6fab1b915
3 changed files with 12 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
csrc/flash_api.cpp
View File

@@ -118,8 +118,8 @@ mha_fwd_kvcache_mla(
TORCH_CHECK(descale_k.stride(-1) == 1);
TORCH_CHECK(descale_q.dtype() == torch::kFloat);
TORCH_CHECK(descale_k.dtype() == torch::kFloat);
CHECK_SHAPE(descale_q, batch_size);
CHECK_SHAPE(descale_k, batch_size);
CHECK_SHAPE(descale_q, 1);
CHECK_SHAPE(descale_k, 1);
}
if (seqlen_q_ori == 1) { is_causal = false; }