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

add env

Browse Source
This commit is contained in:
chenhongmin.will
2025-03-01 14:44:25 +08:00
parent 6199b0b4b5
commit 7fafcd217d
2 changed files with 9 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
csrc/flash_api.cpp
View File

@@ -218,9 +218,12 @@ mha_fwd_kvcache_mla(
run_mha_fwd_splitkv_mla<cutlass::half_t, cutlass::half_t, 576>(params, stream);
}
#endif
#ifndef FLASH_MLA_DISABLE_FP8
else if (q_dtype == torch::kFloat8_e4m3fn) {
run_mha_fwd_splitkv_mla<cutlass::float_e4m3_t, cutlass::bfloat16_t, 576>(params, stream);
} else {
}
#endif
else {
TORCH_CHECK(false, "Unsupported tensor dtype for query");
}