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

init fp8

Browse Source
This commit is contained in:
chenhongmin.will 2025-02-24 21:12:36 +08:00
parent bcb90f2afd
commit dae0690055
2 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
csrc/flash_fwd_mla_fp8_sm90.cu Normal file
View File

@ -0,0 +1,3 @@
#include "flash_fwd_mla_kernel.h"
template void run_mha_fwd_splitkv_mla<cutlass::float_e4m3_t, 576>(Flash_fwd_mla_params &params, cudaStream_t stream);

1
setup.py
View File

@ -37,6 +37,7 @@ ext_modules.append(
sources=[ sources=[
"csrc/flash_api.cpp", "csrc/flash_api.cpp",
"csrc/flash_fwd_mla_bf16_sm90.cu", "csrc/flash_fwd_mla_bf16_sm90.cu",
"csrc/flash_fwd_mla_fp8_sm90.cu",
], ],
extra_compile_args={ extra_compile_args={
"cxx": cxx_args, "cxx": cxx_args,