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

fix mma0

Browse Source
This commit is contained in:
chenhongmin.will
2025-02-28 16:52:30 +08:00
parent 061af5fc56
commit fd1e662deb
2 changed files with 22 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
csrc/flash_fwd_mla_kernel.h
View File

@@ -328,8 +328,13 @@ __forceinline__ __device__ void compute_attn_1rowblock_splitkv_mla(const Flash_f
if (n_block % 2 == 1) {
// Double buffer for sK
constexpr int sK_offset = size(sK);
tSrK.data() = tSrK.data() + sK_offset / 8;
if constexpr (!Kernel_traits::Is_FP8) tOrVt.data() = tOrVt.data() + sK_offset / 8;
if constexpr (Kernel_traits::Is_FP8) {
tSrK.data() = tSrK.data() + sK_offset / 16;
} else {
tSrK.data() = tSrK.data() + sK_offset / 8;
tOrVt.data() = tOrVt.data() + sK_offset / 8;
}
}
// We need masking on S for the very last block when K and V has length not multiple of kBlockN.
@@ -392,8 +397,12 @@ __forceinline__ __device__ void compute_attn_1rowblock_splitkv_mla(const Flash_f
// Double buffer for sK
const int sK_offset = n_block % 2 == 0 ? size(sK) : -size(sK);
tSrK.data() = tSrK.data() + sK_offset / 8;
if constexpr (!Kernel_traits::Is_FP8) tOrVt.data() = tOrVt.data() + sK_offset / 8;
if constexpr (Kernel_traits::Is_FP8) {
tSrK.data() = tSrK.data() + sK_offset / 16;
} else {
tSrK.data() = tSrK.data() + sK_offset / 8;
tOrVt.data() = tOrVt.data() + sK_offset / 8;
}
}
cute::copy(softmax.row_max, tRow_maxsRow_max);
@@ -513,9 +522,9 @@ __forceinline__ __device__ void compute_attn_1rowblock_splitkv_mla(const Flash_f
if constexpr (Kernel_traits::Is_FP8) __syncthreads();
flash::gemm</*zero_init=*/false, /*wg_wait=*/0>(tiled_mma_o, tOrP, tOrVt, tOrO);
// Double buffer for sK
const int sK_offset = n_block % 2 == 0 ? size(sK) : -size(sK);
if constexpr (!Kernel_traits::Is_FP8) {
// Double buffer for sK
const int sK_offset = n_block % 2 == 0 ? size(sK) : -size(sK);
tOrVt.data() = tOrVt.data() + sK_offset / 8;
}
}