diff --git a/attention/sparse_flash_mla/op_kernel/arch22/sparse_flash_mla_csa_block_vector.h b/attention/sparse_flash_mla/op_kernel/arch22/sparse_flash_mla_csa_block_vector.h index ec94e71574..b2eb1002d2 100644 --- a/attention/sparse_flash_mla/op_kernel/arch22/sparse_flash_mla_csa_block_vector.h +++ b/attention/sparse_flash_mla/op_kernel/arch22/sparse_flash_mla_csa_block_vector.h @@ -333,7 +333,7 @@ __aicore__ inline void SMLAVectorBlock::CopySinksIn() repeatParams.blockLen = constInfo.qHeadNum; repeatParams.srcStride = 0U; repeatParams.dstStride = 0U; - for (uint32_t i = 1U; i <= 256U / constInfo.qHeadNum; i++) { + for (uint32_t i = 1U; i < 256U / constInfo.qHeadNum; i++) { DataCopy(sinksBrcbUb[constInfo.qHeadNum * BLOCK_ELEMENT_NUM * i], sinksBrcbUb, repeatParams); } PipeBarrier(); diff --git a/attention/sparse_flash_mla/op_kernel/arch22/sparse_flash_mla_swa_block_vector.h b/attention/sparse_flash_mla/op_kernel/arch22/sparse_flash_mla_swa_block_vector.h index 0a1a842b45..0918bcb9d9 100644 --- a/attention/sparse_flash_mla/op_kernel/arch22/sparse_flash_mla_swa_block_vector.h +++ b/attention/sparse_flash_mla/op_kernel/arch22/sparse_flash_mla_swa_block_vector.h @@ -289,7 +289,7 @@ __aicore__ inline void SWAVectorBlock::CopySinksIn() repeatParams.blockLen = constInfo.qHeadNum; repeatParams.srcStride = 0U; repeatParams.dstStride = 0U; - for (uint32_t i = 1U; i <= 256U / constInfo.qHeadNum; i++) { + for (uint32_t i = 1U; i < 256U / constInfo.qHeadNum; i++) { DataCopy(sinksBrcbUb[constInfo.qHeadNum * BLOCK_ELEMENT_NUM * i], sinksBrcbUb, repeatParams); } PipeBarrier(); diff --git a/attention/sparse_flash_mla/op_kernel/arch35/sparse_flash_mla_csa_block_vector.h b/attention/sparse_flash_mla/op_kernel/arch35/sparse_flash_mla_csa_block_vector.h index 47447b1549..9bf57e0b06 100644 --- a/attention/sparse_flash_mla/op_kernel/arch35/sparse_flash_mla_csa_block_vector.h +++ b/attention/sparse_flash_mla/op_kernel/arch35/sparse_flash_mla_csa_block_vector.h @@ -487,6 +487,11 @@ __aicore__ inline void CSABlockVec::ProcessVec1( ConstInfo &constInfo) { bmm1ResBuf.WaitCrossCore(); + if (unlikely(runInfo.halfMRealSize == 0)) { + bmm1ResBuf.SetCrossCore(); + outputBuf.SetCrossCore(); + return; + } LocalTensor sumUb = this->softmaxSumBuf[runInfo.multiCoreIdxMod2].template Get(); LocalTensor maxUb = this->softmaxMaxBuf[runInfo.multiCoreIdxMod2].template Get(); diff --git a/attention/sparse_flash_mla/op_kernel/arch35/sparse_flash_mla_kvcache.h b/attention/sparse_flash_mla/op_kernel/arch35/sparse_flash_mla_kvcache.h index 3b87b151a4..e381986e2d 100644 --- a/attention/sparse_flash_mla/op_kernel/arch35/sparse_flash_mla_kvcache.h +++ b/attention/sparse_flash_mla/op_kernel/arch35/sparse_flash_mla_kvcache.h @@ -162,7 +162,11 @@ __aicore__ inline void ComputeSouterParam(RunParamStr& runParam, const ConstInfo } runParam.cubeMOuterOffset = cubeSOuterOffset * constInfo.gSize; - runParam.halfMRealSize = (runParam.mRealSize + 1) >> 1; + if (!IS_SPLIT_G && runParam.mRealSize <= 32U) { + runParam.halfMRealSize = runParam.mRealSize; + } else { + runParam.halfMRealSize = (runParam.mRealSize + 1) >> 1; + } runParam.firstHalfMRealSize = runParam.halfMRealSize; if (constInfo.subBlockIdx == 1) { runParam.halfMRealSize = runParam.mRealSize - runParam.halfMRealSize; diff --git a/attention/sparse_flash_mla/tests/pytest/sparse_flash_mla_paramset.py b/attention/sparse_flash_mla/tests/pytest/sparse_flash_mla_paramset.py index bc7ca74beb..c1670b8e41 100644 --- a/attention/sparse_flash_mla/tests/pytest/sparse_flash_mla_paramset.py +++ b/attention/sparse_flash_mla/tests/pytest/sparse_flash_mla_paramset.py @@ -275,6 +275,94 @@ "ori_win_left": [127], "ori_win_right": [0] }, + "csa_g32_test_prefill":{ + "testcase_name": ["csa_g32_test_prefill"], + "layout_q": ["TND"], + "layout_kv": ["PA_BBND"], + "q_type": [torch.bfloat16], + "ori_kv_type": [torch.bfloat16], + "cmp_kv_type": [torch.bfloat16], + "B": [1], + "S1": [2], + "T1": [2], + "T2": [256], + "N1": [32], + "N2": [1], + "D": [512], + "K": [512], + "block_num1": [2], + "block_num2": [1], + "block_size1": [128], + "block_size2": [128], + "cu_seqlens_q": [[0, 2]], + "cu_seqlens_ori_kv": [[0, 256]], + "cu_seqlens_cmp_kv": [[0, 64]], + "seqused_ori_kv": [[256]], + "seqused_cmp_kv": [[64]], + "cmp_residual_kv": [[0]], + "softmax_scale": [0.04419417], + "cmp_ratio": [4], + "return_softmax_lse": [False], + "ori_mask_mode": [4], + "cmp_mask_mode": [3], + "ori_win_left": [127], + "ori_win_right": [0] + }, + "csa_g16_test_prefill":{ + "testcase_name": ["csa_g16_test_prefill"], + "layout_q": ["TND"], + "layout_kv": ["PA_BBND"], + "q_type": [torch.bfloat16], + "ori_kv_type": [torch.bfloat16], + "cmp_kv_type": [torch.bfloat16], + "B": [1], + "S1": [2], + "T1": [2], + "T2": [256], + "N1": [16], + "N2": [1], + "D": [512], + "K": [512], + "block_num1": [2], + "block_num2": [1], + "block_size1": [128], + "block_size2": [128], + "cu_seqlens_q": [[0, 2]], + "cu_seqlens_ori_kv": [[0, 256]], + "cu_seqlens_cmp_kv": [[0, 64]], + "seqused_ori_kv": [[256]], + "seqused_cmp_kv": [[64]], + "cmp_residual_kv": [[0]], + "softmax_scale": [0.04419417], + "cmp_ratio": [4], + "return_softmax_lse": [False], + "ori_mask_mode": [4], + "cmp_mask_mode": [3], + "ori_win_left": [127], + "ori_win_right": [0] + }, + "swa_g16_test_prefill":{ + "testcase_name": ["swa_g16_test_prefill"], + "layout_q": ["TND"], + "layout_kv": ["PA_BBND"], + "q_type": [torch.bfloat16], + "ori_kv_type": [torch.bfloat16], + "B": [1], + "S1": [2], + "T1": [2], + "N1": [16], + "N2": [1], + "D": [512], + "block_num1": [2], + "block_size1": [128], + "cu_seqlens_q": [[0, 2]], + "seqused_ori_kv": [[256]], + "softmax_scale": [0.04419417], + "return_softmax_lse": [False], + "ori_mask_mode": [4], + "ori_win_left": [127], + "ori_win_right": [0] + }, "csa_g128_test_prefill":{ "testcase_name": ["csa_g128_test_prefill"], "layout_q": ["TND"],