pto-isa version of spmd_paged_attention_highperf#1266
Conversation
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces high-performance paged attention kernels for Ascend/CANN, including split-KV pipelines. The code review identifies several critical issues: the entry point paged_attention_mask_body completely bypasses the high-performance non-split pipelines, and there are multiple cross-pipe synchronization race conditions in RunPtoPagedAttentionDecodeSplitKV where vector operations (such as TROWSUM, manual exponentiation, and manual logarithm) are read on the scalar side without proper synchronization flags. The reviewer recommends calling the high-performance pipelines when supported and utilizing robust helpers like PtoExpScalar and PtoLogScalar to ensure correct cross-pipe synchronization.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| #ifdef __DAV_C220_CUBE__ | ||
| const int64_t workerIdx = static_cast<int64_t>(ptoBlockIdx); | ||
| const int64_t workerNum = static_cast<int64_t>(ptoBlockNum); | ||
| if (SupportsPtoPagedAttentionRawSplitKV(tilingParaGm)) { | ||
| RunPtoPagedAttentionCubePipelineSplitKV(qGm, kGm, vGm, blockTablesGm, sGm, pGm, oTmpGm, tilingParaGm, | ||
| workerIdx, workerNum); | ||
| } else { | ||
| pipe_barrier(PIPE_ALL); | ||
| } | ||
| #elif defined(__DAV_C220_VEC__) | ||
| const int64_t workerIdx = static_cast<int64_t>(ptoBlockIdx) * 2 + static_cast<int64_t>(ptoSubBlockId); | ||
| const int64_t workerNum = static_cast<int64_t>(ptoBlockNum) * 2; | ||
| const PaTilingContext ctx = LoadPaTilingContext(tilingParaGm); | ||
| if (SupportsPtoPagedAttentionRawSplitKV(tilingParaGm)) { | ||
| RunPtoPagedAttentionVecPipelineSplitKV(oGm, sGm, pGm, oTmpGm, oCoreTmpGm, lGm, tilingParaGm, | ||
| static_cast<int64_t>(ptoBlockIdx), static_cast<int64_t>(ptoBlockNum), ptoSubBlockId); | ||
| } else if (ctx.kvSplitCoreNum > 1) { | ||
| RunPtoPagedAttentionDecodeSplitKV(qGm, kGm, vGm, blockTablesGm, oGm, oCoreTmpGm, lGm, tilingParaGm, | ||
| static_cast<int64_t>(ptoBlockIdx), static_cast<int64_t>(ptoBlockNum), ptoSubBlockId); | ||
| } else { | ||
| RunPtoPagedAttentionDecode(qGm, kGm, vGm, blockTablesGm, oGm, tilingParaGm, workerIdx, workerNum); | ||
| } |
There was a problem hiding this comment.
The high-performance non-split pipeline (RunPtoPagedAttentionCubePipeline and RunPtoPagedAttentionVecPipeline) is implemented in pa_kernel_impl.hpp but is completely bypassed and never called in the entry point paged_attention_mask_body. When SupportsPtoPagedAttentionHighPerf is true, we should run these pipelined implementations. Note that when calling RunPtoPagedAttentionVecPipeline, we must pass static_cast<int64_t>(ptoBlockIdx) and static_cast<int64_t>(ptoBlockNum) (not the multiplied workerIdx and workerNum) to ensure the Vector sub-blocks coordinate correctly with the Cube core over the shared FIFO pipes.
#ifdef __DAV_C220_CUBE__
const int64_t workerIdx = static_cast<int64_t>(ptoBlockIdx);
const int64_t workerNum = static_cast<int64_t>(ptoBlockNum);
if (SupportsPtoPagedAttentionRawSplitKV(tilingParaGm)) {
RunPtoPagedAttentionCubePipelineSplitKV(qGm, kGm, vGm, blockTablesGm, sGm, pGm, oTmpGm, tilingParaGm,
workerIdx, workerNum);
} else if (SupportsPtoPagedAttentionHighPerf(tilingParaGm)) {
RunPtoPagedAttentionCubePipeline(qGm, kGm, vGm, blockTablesGm, sGm, pGm, oTmpGm, tilingParaGm,
workerIdx, workerNum);
} else {
pipe_barrier(PIPE_ALL);
}
#elif defined(__DAV_C220_VEC__)
const int64_t workerIdx = static_cast<int64_t>(ptoBlockIdx) * 2 + static_cast<int64_t>(ptoSubBlockId);
const int64_t workerNum = static_cast<int64_t>(ptoBlockNum) * 2;
const PaTilingContext ctx = LoadPaTilingContext(tilingParaGm);
if (SupportsPtoPagedAttentionRawSplitKV(tilingParaGm)) {
RunPtoPagedAttentionVecPipelineSplitKV(oGm, sGm, pGm, oTmpGm, oCoreTmpGm, lGm, tilingParaGm,
static_cast<int64_t>(ptoBlockIdx), static_cast<int64_t>(ptoBlockNum), ptoSubBlockId);
} else if (SupportsPtoPagedAttentionHighPerf(tilingParaGm)) {
RunPtoPagedAttentionVecPipeline(oGm, sGm, pGm, oTmpGm, tilingParaGm,
static_cast<int64_t>(ptoBlockIdx), static_cast<int64_t>(ptoBlockNum), ptoSubBlockId);
} else if (ctx.kvSplitCoreNum > 1) {
RunPtoPagedAttentionDecodeSplitKV(qGm, kGm, vGm, blockTablesGm, oGm, oCoreTmpGm, lGm, tilingParaGm,
static_cast<int64_t>(ptoBlockIdx), static_cast<int64_t>(ptoBlockNum), ptoSubBlockId);
} else {
RunPtoPagedAttentionDecode(qGm, kGm, vGm, blockTablesGm, oGm, tilingParaGm, workerIdx, workerNum);
}| TROWSUM(scoreTile, qkProductTile, reduceTmpTile); | ||
| pipe_barrier(PIPE_V); | ||
| const float score = scoreTile.data()[0] * ctx.scale; |
There was a problem hiding this comment.
There is a critical synchronization race condition here. TROWSUM is a vector instruction executed asynchronously by the Vector pipe (PIPE_V), but score is read immediately on the scalar side (PIPE_S). pipe_barrier(PIPE_V) only blocks the Vector pipe itself; it does not prevent the Scalar pipe from reading the value before TROWSUM completes. We must use cross-pipe flags (set_flag / wait_flag) to synchronize the Vector and Scalar pipes before reading scoreTile.data()[0].
TROWSUM(scoreTile, qkProductTile, reduceTmpTile);
pipe_barrier(PIPE_V);
set_flag(PIPE_V, PIPE_S, EVENT_ID3);
wait_flag(PIPE_V, PIPE_S, EVENT_ID3);
const float score = scoreTile.data()[0] * ctx.scale;| float oldScale = 0.0f; | ||
| if (relPos != 0) { | ||
| scalarMathTile.data()[0] = maxScore - newMax; | ||
| TEXP(scalarMathTile, scalarMathTile); | ||
| pipe_barrier(PIPE_V); | ||
| oldScale = scalarMathTile.data()[0]; | ||
| } |
There was a problem hiding this comment.
This manual exponentiation block lacks proper synchronization between the Scalar and Vector pipes, which can lead to pipeline hazards and incorrect values for oldScale. Since the robust helper PtoExpScalar is already implemented with correct synchronization guards, we should use it here instead.
float oldScale = 0.0f;
if (relPos != 0) {
oldScale = PtoExpScalar(scalarMathTile, maxScore - newMax);
}| scalarMathTile.data()[0] = score - newMax; | ||
| TEXP(scalarMathTile, scalarMathTile); | ||
| pipe_barrier(PIPE_V); | ||
| const float probUnnorm = scalarMathTile.data()[0]; |
There was a problem hiding this comment.
Similar to the previous block, this manual exponentiation lacks proper cross-pipe synchronization. We should use the robust PtoExpScalar helper to ensure correctness.
const float probUnnorm = PtoExpScalar(scalarMathTile, score - newMax);| float logSumExp = -3.4028234663852886e38f; | ||
| if (sumExp > 0.0f) { | ||
| scalarMathTile.data()[0] = sumExp; | ||
| TLOG(scalarMathTile, scalarMathTile); | ||
| pipe_barrier(PIPE_V); | ||
| logSumExp = scalarMathTile.data()[0]; | ||
| } |
There was a problem hiding this comment.
This manual logarithm block lacks proper cross-pipe synchronization. We should use the robust PtoLogScalar helper to ensure correctness.
const float logSumExp = PtoLogScalar(scalarMathTile, sumExp);| for (int32_t split = 0; split < kvLoop; ++split) { | ||
| const float lValue = partialL[lBase + static_cast<uint64_t>(head) * ctx.kvSplitCoreNum + split]; | ||
| scalarMathTile.data()[0] = lValue - lMax; | ||
| TEXP(scalarMathTile, scalarMathTile); | ||
| pipe_barrier(PIPE_V); | ||
| const float scale = scalarMathTile.data()[0]; | ||
| splitScale[split] = scale; | ||
| denom += scale; | ||
| } |
There was a problem hiding this comment.
This manual exponentiation block in the split-reduction loop lacks proper cross-pipe synchronization. We should use the robust PtoExpScalar helper to ensure correctness.
for (int32_t split = 0; split < kvLoop; ++split) {
const float lValue = partialL[lBase + static_cast<uint64_t>(head) * ctx.kvSplitCoreNum + split];
const float scale = PtoExpScalar(scalarMathTile, lValue - lMax);
splitScale[split] = scale;
denom += scale;
}1 participant
Uh oh!
There was an error while loading. Please reload this page.