已合并
QSFA oom bugfix #7676
郑文惠创建于 6月27日
QSFA oom bugfix #7676
已合并
共 1 个文件变更+9-8
| @@ -115,7 +115,8 @@ private: | |||
| 115 | ConstInfo &constInfo); | 115 | ConstInfo &constInfo); |
| 116 | __aicore__ inline void CopyOutMrgeResult(Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &outputL1, | 116 | __aicore__ inline void CopyOutMrgeResult(Buffer<BufferType::L1, SyncType::CROSS_CORE_SYNC_FORWARD> &outputL1, |
| 117 | int64_t mte2Size, int64_t mte3Size, int64_t s2keyOffset, int64_t mergeMte3Idx, const RunInfo &runInfo); | 117 | int64_t mte2Size, int64_t mte3Size, int64_t s2keyOffset, int64_t mergeMte3Idx, const RunInfo &runInfo); |
| 118 | - __aicore__ inline void CopyInSingleKv(LocalTensor<KV_T> kvInUb, int64_t startRow, int64_t keyOffset); | 118 | + __aicore__ inline void CopyInSingleKv(LocalTensor<KV_T> kvInUb, int64_t startRow, |
| 119 | + int64_t keyOffset, uint32_t combineBytes); | ||
| 119 | /* VEC2_RES_T 表示bmm2ResUb当前的类型，VEC2_RES_T = Q_T那么不需要做Cast。另外，无效行场景当前默认需要做Cast */ | 120 | /* VEC2_RES_T 表示bmm2ResUb当前的类型，VEC2_RES_T = Q_T那么不需要做Cast。另外，无效行场景当前默认需要做Cast */ |
| 120 | using VEC2_RES_T = T; | 121 | using VEC2_RES_T = T; |
| 121 | template <typename VEC2_RES_T> | 122 | template <typename VEC2_RES_T> |
| @@ -212,7 +213,8 @@ TEMPLATES_DEF_NO_DEFAULT __aicore__ inline int64_t QSFAVectorService<TEMPLATE_AR | |||
| 212 | } | 213 | } |
| 213 | 214 | ||
| 214 | TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void | 215 | TEMPLATES_DEF_NO_DEFAULT __aicore__ inline void |
| 215 | -QSFAVectorService<TEMPLATE_ARGS>::CopyInSingleKv(LocalTensor<KV_T> kvInUb, int64_t startRow, int64_t keyOffset) | 216 | +QSFAVectorService<TEMPLATE_ARGS>::CopyInSingleKv(LocalTensor<KV_T> kvInUb, int64_t startRow, |
| 217 | + int64_t keyOffset, uint32_t combineBytes) | ||
| 216 | { | 218 | { |
| 217 | if (keyOffset < 0) { | 219 | if (keyOffset < 0) { |
| 218 | return; | 220 | return; |
| @@ -224,7 +226,6 @@ QSFAVectorService<TEMPLATE_ARGS>::CopyInSingleKv(LocalTensor<KV_T> kvInUb, int64 | |||
| 224 | intriParams.srcStride = 0; | 226 | intriParams.srcStride = 0; |
| 225 | DataCopyPadExtParams<KV_T> padParams; | 227 | DataCopyPadExtParams<KV_T> padParams; |
| 226 | // 当前仅支持COMBINE模式 | 228 | // 当前仅支持COMBINE模式 |
| 227 | - uint32_t combineBytes = 672; | ||
| 228 | intriParams.blockLen = combineBytes; | 229 | intriParams.blockLen = combineBytes; |
| 229 | uint32_t combineDim = combineBytes / sizeof(KV_T); | 230 | uint32_t combineDim = combineBytes / sizeof(KV_T); |
| 230 | uint32_t combineDimAlign = CeilAlign(combineBytes, BUFFER_SIZE_BYTE_32B) / sizeof(KV_T); | 231 | uint32_t combineDimAlign = CeilAlign(combineBytes, BUFFER_SIZE_BYTE_32B) / sizeof(KV_T); |
| @@ -243,13 +244,13 @@ TEMPLATES_DEF_NO_DEFAULT __aicore__ inline uint32_t QSFAVectorService<TEMPLATE_A | |||
| 243 | if (unlikely(keyOffset0 < 0 && keyOffset1 < 0)) { | 244 | if (unlikely(keyOffset0 < 0 && keyOffset1 < 0)) { |
| 244 | return 0; | 245 | return 0; |
| 245 | } | 246 | } |
| 246 | - uint32_t combineBytes = constInfo.dSizeVInput; | 247 | + uint32_t combineBytes = constInfo.dSizeVInput * sizeof(KV_T); |
 | |||
| 247 | - int64_t keySrcStride = (keyOffset0 > keyOffset1 ? (keyOffset0 - keyOffset1) : | 248 | + int64_t keySrcStride = (keyOffset0 > keyOffset1 ? (keyOffset0 - keyOffset1) * sizeof(KV_T): |
| 248 | - (keyOffset1 - keyOffset0)) - combineBytes; | 249 | + (keyOffset1 - keyOffset0)) * sizeof(KV_T) - combineBytes; |
| 249 | if (keySrcStride >= INT32_MAX || keySrcStride < 0 || constInfo.sparseBlockSize > 1) { | 250 | if (keySrcStride >= INT32_MAX || keySrcStride < 0 || constInfo.sparseBlockSize > 1) { |
| 250 | // stride溢出、stride为负数、s2超长等异常场景，还原成2条搬运指令 | 251 | // stride溢出、stride为负数、s2超长等异常场景，还原成2条搬运指令 |
| 251 | - CopyInSingleKv(kvInUb, startRow, keyOffset0); | 252 | + CopyInSingleKv(kvInUb, startRow, keyOffset0, combineBytes); |
| 252 | - CopyInSingleKv(kvInUb, startRow + 1, keyOffset1); | 253 | + CopyInSingleKv(kvInUb, startRow + 1, keyOffset1, combineBytes); |
| 253 | } else { | 254 | } else { |
| 254 | DataCopyExtParams intriParams; | 255 | DataCopyExtParams intriParams; |
| 255 | intriParams.blockCount = (keyOffset0 >= 0) + (keyOffset1 >= 0); | 256 | intriParams.blockCount = (keyOffset0 >= 0) + (keyOffset1 >= 0); |
🟡 Medium Priority
变更行 245 将
combineBytes从元素数改为字节数（constInfo.dSizeVInput * sizeof(KV_T)），但行 246-247 的keySrcStride计算公式未同步调整：uint32_t combineBytes = constInfo.dSizeVInput * sizeof(KV_T); // 字节 int64_t keySrcStride = abs(keyOffset0 - keyOffset1) - combineBytes; // 元素 - 字节 = 单位混乱
keyOffset0/keyOffset1由GetkeyOffset返回，是元素偏移量（见行 205：(boIdx * s2Size + s2Idx) * dSizeVInput），而combineBytes已是字节。两者直接相减导致单位不一致。对于
sizeof(KV_T) > 1（如 bf16/fp32）：对比
sparse_flash_mla/op_kernel/arch35/sparse_flash_mla_scfa_block_vector.h第 305 行的平行代码，已正确使用* sizeof(KV_T)： int64_t keySrcStride = (keyOffset0 > keyOffset1 ?(keyOffset0 - keyOffset1) : (keyOffset1 - keyOffset0)) * sizeof(KV_T) - combineBytes;
触发条件：
sizeof(KV_T) > 1且稀疏索引中相邻两个 token 在 KV cache 中的偏移差 ≥ 3 倍dSizeVInput。建议：在 keyOffset 差值上乘以 sizeof(KV_T) 使其单位转换为字节，与 combineBytes 保持一致。参照 sparse_flash_mla 中的平行实现修复。