feat: add EP memory dispatch and combine - #86
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0aa50c0656
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // printf("aivId_ %d, time %d, %d, %d\n", aivId_, | ||
| // (usedTime_[1] - usedTime_[0])/ 1000 , (usedTime_[2] - usedTime_[1])/ 1000, (usedTime_[3] - usedTime_[2])/ 1000); | ||
| // printf("aivId_ %d, time %d\n", aivId_, (usedTime_[3] - usedTime_[0])/ 1000); | ||
| return; |
There was a problem hiding this comment.
Every valid memory-dispatch launch reaches this unconditional return on each AIV, so AllToAllDispatch, CalCumSum, and LocalWindowCopy are all unreachable. Consequently, TileXRMoeEpDispatchMemory* reports a successful launch without producing any dispatch outputs, and a subsequent memory combine can wait indefinitely for data that was never sent.
Useful? React with 👍 / 👎.
| Copy(outTensorInt32, xInTensorInt32, uint64_t(64), uint8_t(blockCntPerToken_), {1, 1, 16, 15}); | ||
|
|
||
| Copy(outTensorInt32[64], xInTensorInt32[64], uint64_t(56), uint8_t(blockCntPerToken_), {1, 1, 16, 15}); |
There was a problem hiding this comment.
Once the early return is removed, these two Copy operations consume 480 source bytes for every blockCntPerToken_ repeat, but the non-quantized source tensor is allocated with only hAlignSize_ bytes while the metadata starts at that exact boundary, and the MXFP8 source has only quantTensorBytes_ rather than the rounded compact-payload size. Ordinary FP16/BF16 dispatch therefore reads and writes outside its queue tensor for every route, while MXFP8 does so whenever the final 480-byte block is partial, potentially corrupting adjacent UB buffers and the emitted assist metadata.
Useful? React with 👍 / 👎.
| const uint32_t compareCount = static_cast<uint32_t>(AlignUp(flagFloatCount, 64U)); | ||
| pipe_->InitBuffer(packedCheckFlagBuf_, compareCount * sizeof(float)); | ||
| pipe_->InitBuffer(packedCheckCompareBuf_, AlignUp(compareCount * sizeof(uint8_t), 256U)); | ||
| pipe_->InitBuffer(packedInputQueue_, 1, compactPayloadBytes_); |
There was a problem hiding this comment.
When quantMode is 3 or 4 and sharedExpertX is non-null, this queue is sized for the compact FP8 communication payload, but AddSharedExpertX copies a full FP16/BF16 row (rowBytes_) into it. For a typical h=1024, for example, the queue is 1,440 bytes while the copy is 2,048 bytes, so this supported API combination overwrites neighboring UB buffers before accumulation.
Useful? React with 👍 / 👎.
| Copy(packedFloat, sourceFloat, 64U, static_cast<uint8_t>(blockCntPerToken_), {1, 1, 16, 15}); | ||
| Copy(packedFloat[64], sourceFloat[64], 56U, static_cast<uint8_t>(blockCntPerToken_), {1, 1, 16, 15}); |
There was a problem hiding this comment.
For non-quantized combine rows whose byte length is not divisible by 480, these copies read blockCntPerToken_ * 480 bytes from inputLocal, although that queue is allocated with only align32(rowBytes_) bytes. For example, h=256 creates a 512-byte input tensor but the two-block pack reads 960 bytes, causing an out-of-bounds UB read on every sent row; the source buffer must be sized and padded to the compact payload length.
Useful? React with 👍 / 👎.
| DataCopyPadExtParams<int32_t> expertIdsCntCopyPadParams{true, 0U, uint8_t(rightPadding), -1}; | ||
| DataCopyExtParams expertIdsCntParams{1U, static_cast<uint32_t>(expertIdsMask * sizeof(uint32_t)), 0U, 0U, 0U}; | ||
| SyncFunc<AscendC::HardEvent::V_MTE2>(); | ||
| DataCopyPad(validExpertIdsTensor_, expertIdsGMTensor_, expertIdsCntParams, expertIdsCntCopyPadParams); |
There was a problem hiding this comment.
For a token mask with an inactive token anywhere except the suffix, TokenActiveMaskCal reduces the mask to only its number of true entries and this copy then takes that many expert-ID rows from the start of the tensor. For example, [true, false, true] dispatches tokens 0 and 1 rather than tokens 0 and 2; the active token's output is lost, and memory combine can wait indefinitely for its route while data was instead written for the masked token.
Useful? React with 👍 / 👎.
2 participants
Summary
Tests
test_tilexr_ep_api_sources(direct C++14 build/run): passedtest_tilexr_ep_kernel_sources(direct C++14 build/run): passedtest_tilexr_ep_layout(direct C++14 build/run): passedtest_tilexr_ep_host_validationwith the project ACL stub contract (direct C++14 build/run): passedgit diff origin/main...HEAD --check: passedValidation gaps