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 fp8 ut

Browse Source
This commit is contained in:
chenhongmin.will
2025-02-25 23:29:18 +08:00
parent dfe8ffc75a
commit 870418802a
3 changed files with 28 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
csrc/flash_api.cpp
View File

@@ -70,8 +70,7 @@ mha_fwd_kvcache_mla(
const at::Tensor &tile_scheduler_metadata, // num_sm_parts x TileSchedulerMetaDataSize
const at::Tensor &num_splits, // batch_size + 1
c10::optional<const at::Tensor> &descale_q, // batch_size
c10::optional<const at::Tensor> &descale_k, // batch_size
c10::optional<const at::Tensor> &descale_v // batch_size
c10::optional<const at::Tensor> &descale_k // batch_size
) {
auto dprops = at::cuda::getCurrentDeviceProperties();
bool is_sm90 = dprops->major == 9 && dprops->minor == 0;