Repository navigation
[Feature][Model] Enable MiniMax-M3 FP8 MSA index score on A5 - #15918
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request integrates native FP8 support for MiniMax-M3 inference on Ascend 950 (A5) hardware. By enabling the bundled AscendC MsaIndexScore operator, the implementation replaces the previous Triton-based scoring path with a more optimized AscendC version, ensuring compatibility with FP8 tensors and paged cache structures. The changes include kernel registration, build system updates, and runtime logic to handle dtype casting and windowed score flushing, significantly improving performance and stability for A5 inference. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
There was a problem hiding this comment.
Code Review
Suggested PR Title:
[Attention][Feature] Add Ascend 950 and FP8 support for MsaIndexScoreSuggested PR Summary:
### What this PR does / why we need it?
This pull request adds support for the Ascend 950 (A5) platform and FP8 data types (`float8_e4m3fn`, `float8_e5m2`, and `hifloat8`) to the `MsaIndexScore` operator. It introduces the `arch35` kernel implementation, vendors the `catlass` headers locally under `op_kernel/catlass` to remove the repository-level submodule dependency, and updates the host-side tiling, shape inference, and PyTorch adapter to support the new platform and data types. Additionally, it adds support for non-contiguous PageAttention keys along the physical page axis, handles wide block tables (>256 columns) via sliding window flushing, and updates the minimax_m3 model integration to use the AscendC path by default on A5.
Feedback on the code changes:
- In `csrc/attention/msa_index_score/op_kernel/arch35/msa_seg_row_max_epilogue.h` at line 225, `RoundMode::CAST_RINT` is used during float-to-half conversion, which rounds scores to the nearest integer and destroys fractional precision. It should be changed to `RoundMode::CAST_NONE`.
- In the same file at line 641, the loop condition `j < M_PAD / VA` evaluates to `j < 0` due to integer division (8 / 16), causing the loop and the critical `TransDataTo5HD` transpose operation to be skipped entirely. The loop should be removed as the entire block fits within a single transpose tile.
### Does this PR introduce _any_ user-facing change?
Yes, it enables the AscendC-based `MsaIndexScore` path by default on Ascend 950 (A5) and introduces support for FP8 data types in the operator interface.
### How was this patch tested?
The changes were tested using the end-to-end accuracy self-check in `examples/test_aclnn_msa_index_score.cpp` with 40 test cases covering FP8, mixed-batch padding, zero-length sequences, non-contiguous keys, and wide block tables. Unit tests were also added in `tests/ut/models/minimax_m3/test_msa_m3.py`.| const uint32_t rowOff = mOff + p0; | ||
| IssueS16(gS, rowOff, rows, stage); | ||
| if (hasPrev) { | ||
| WaitS16(prevStage); |
There was a problem hiding this comment.
Using RoundMode::CAST_RINT rounds the float scores to the nearest integer during the float-to-half conversion. Since these are attention/block scores, rounding them to integers will destroy the fractional precision and severely degrade the accuracy of the subsequent TopK selection. Use RoundMode::CAST_NONE instead to perform a standard float-to-half cast.
AscendC::Cast(ubRed16_[col * MSA_A5_S_PING_ROWS], dst0, AscendC::RoundMode::CAST_NONE, mSub_);| params.srcStride = 0; | ||
| params.dstStride = 0; | ||
| AscendC::DataCopyPadExtParams<float> pad; | ||
| pad.isPad = false; | ||
| pad.leftPadding = 0; | ||
| pad.rightPadding = 0; | ||
| pad.paddingValue = 0; | ||
| AscendC::DataCopyPad(ubDeqScale_, gScale_[scaleOff], params, pad); | ||
| } |
There was a problem hiding this comment.
The loop condition j < M_PAD / VA evaluates to j < 8 / 16 which is j < 0 due to integer division. As a result, this loop never executes, and the critical TransDataTo5HD transpose operation is completely skipped, leaving ubOutH uninitialized. Since M_PAD (8) and VA (16) are constants and the entire block fits within a single 16x16 transpose tile, the outer loop over j is unnecessary and should be removed.
for (uint32_t i = 0; i < VA; ++i) {
srcList[i] = ubPad[i * M_PAD];
}
for (uint32_t k = 0; k < VA; ++k) {
dstList[k] = ubOutH[k * N_PAD];
}
AscendC::TransDataTo5HD(dstList, srcList, params);|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. Tip 💡 Consider Linking a Related Issue or RFCYour PR title contains the [Feature] tag, indicating a bug fix or new feature. Linking a related issue or RFC in the PR description is strongly encouraged — it gives reviewers helpful context and speeds up the review. You can use any of these keywords:
🙏 Thanks for helping us keep the project well-organized! |
|
please provide performance data (triton vs ascendc) |
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
6609bf9 to
01f7b56
Compare
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
01f7b56 to
248dfe9
Compare
…7.1rc (#15926) ### What this PR does / why we need it? Backports #15918 to `releases/v0.27.1rc` so MiniMax-M3 FP8 on A5 can use the bundled AscendC `npu_msa_index_score` implementation. - Syncs the MsaIndexScore operator from [ops-transformer PR #10672](https://gitcode.com/cann/ops-transformer/pull/10672) at `36bf16e9cd774740505019b61d25aa87410cb2b3`, including the Ascend 950 arch35 FP8 kernel. - Adds `msa_index_score` to the Ascend 950 build list. - Routes A5 index scoring through AscendC while retaining the A5 Triton top-k/decode path. - Includes the wide block-table stage flush fix and its regression assertions. - Adds the FP8 E4M3 clamp/cast helper and test import needed by this release baseline. ### Does this PR introduce _any_ user-facing change? Yes. MiniMax-M3 FP8 inference on A5 uses the AscendC MSA index score operator instead of the previous Triton score implementation. ### How was this patch tested? - Verified all 54 imported operator files match ops-transformer PR #10672 at `36bf16e9cd774740505019b61d25aa87410cb2b3` byte-for-byte, excluding its standalone `torch_extension` directory. - Ran Ruff format and lint checks on the changed Python integration and unit-test files. - Ran Python `compileall`, `bash -n csrc/build_aclnn.sh`, and `git diff --check`. - Added unit assertions for A5 FP8 registration, Ascend 950 build inclusion, the AscendC route, and wide block-table flushing. - vLLM main: vllm-project/vllm@ba07e4a --------- Signed-off-by: zonghaoxin <z00946994@china.huawei.com> Co-authored-by: zonghaoxin <z00946994@china.huawei.com>
AuroraEmiya
left a comment
There was a problem hiding this comment.
Add test uts for the imported op
|
This PR including profound name-only changes ,check if it is necessary :> |
|
/rerun Rerun (failed jobs only):
|
|
This pull request has conflicts, please resolve those before we can evaluate the pull request. |
Signed-off-by: zonghaoxin <z00946994@china.huawei.com>
Signed-off-by: zonghaoxin <z00946994@china.huawei.com>
Signed-off-by: zonghaoxin <z00946994@china.huawei.com>
Signed-off-by: zonghaoxin <z00946994@china.huawei.com>
Signed-off-by: zonghaoxin <z00946994@china.huawei.com>
Signed-off-by: Haoxin Zong <534687988@qq.com>
Signed-off-by: Haoxin Zong <534687988@qq.com>
Signed-off-by: Haoxin Zong <534687988@qq.com>
Signed-off-by: Haoxin Zong <534687988@qq.com>
Signed-off-by: Haoxin Zong <534687988@qq.com>
Signed-off-by: Haoxin Zong <534687988@qq.com>
Signed-off-by: Haoxin Zong <534687988@qq.com>
e02efe5 to
0c25b07
Compare
|
/rerun Rerun (failed jobs only):
|
…hado/vllm-ascend into main_fix_mrv2_eagle3_mamba * 'main_fix_mrv2_eagle3_mamba' of https://github.com/windshado/vllm-ascend: (42 commits) Update vllm_ascend/worker/v2/model_states/mamba_hybrid.py [Feature][Kimi K3 DSPark] Enable TP for context_proj (vllm-project#16344) [BugFix][SpecDecode] Refresh replicated PCP draft graph cache mappings (vllm-project#16300) [Feature][Model] Integrate Triton KeyPool indexing for GLM-5.3-Flash (vllm-project#16253) [BugFix][Offloader] Re-bind params to NZ static buffers after npu_format_cast (vllm-project#15415) [Feature][Model] Integrate AscendC KDA and causal convolution for GLM-5.3-Flash (vllm-project#16251) [Performance][Communicator] Replace per-layer F.pad with cat of a persistent zero block in MoE prepare (vllm-project#16343) [Feature][Operator] Add DeepSeek V4.1 sparse attention operators (vllm-project#16422) [Doc][Misc] Document batch invariance scheduling limitations (vllm-project#16232) [CI][MRV2] Enable mrv2 dspark e2e test (vllm-project#16319) [BugFix] Precast MoE gate weight_fp32 to avoid aclop Cast (vllm-project#16189) [Feature][MRV2][310P] MRv2 adapting MTP on the 310P for Qwen3.5 (vllm-project#16043) [Revert] Revert "[Feature][MRV1][MRV2] Refactor Host-Side Parameter Updates for ACL Graph Replay." (vllm-project#15908) (vllm-project#16409) [Feature][Ops] Add Triton KeyPool compression and pooled indexing (vllm-project#16243) [Feature][Attention] Support NoPE in the shared SFA backend (vllm-project#16252) [Performance][Model] Reuse fused mHC operators for GLM-5.3-Flash (vllm-project#16321) [Feature][Model] Enable MiniMax-M3 FP8 MSA index score on A5 (vllm-project#15918) [Performance][KDA] Reduce preprocessing copies and redundant output masks (vllm-project#16067) [Feature][Model][MTP] Support speculative decoding for GLM-5.3-Flash (vllm-project#16214) [BugFix][Model] Skip unused hash-router bias when loading DeepSeek-V4 weights (vllm-project#16259) ...
…oject#15918) ### What this PR does / why we need it? This PR enables the bundled AscendC `MsaIndexScore` path for MiniMax-M3 on Ascend 950 (A5) with native FP8 index query/key tensors and an FP8 index-key paged cache. - Sync the operator implementation from [cann/ops-transformer#10672](https://gitcode.com/cann/ops-transformer/pull/10672) at `2e685d24ed2e54ea1f1547dc6e5ad29862f152ef`. - Add the Ascend 950 `arch35` FP8 kernels, tiling, Catlass dependencies, dtype registration, examples, and documentation. - Add `msa_index_score` to the Ascend 950 custom-op build list. - Route A5 prefill index scoring through `torch.ops._C_ascend.npu_msa_index_score`, while retaining the existing A5 Triton TopK and decode paths. - Cast Index-Q to E4M3 when the index-key cache is E4M3 and the query dtype does not already match. - Include the A5 wide-`block_table` fix: score columns beyond the 256-column UB window are flushed in windows, with width-257 FP8/BF16 regression coverage. The MiniMax-M3 A5 FP8 contract used by this integration is: query/key use `torch.float8_e4m3fn`, `scale=None`, and the operator emits FP32 scores. ### Does this PR introduce _any_ user-facing change? Yes. MiniMax-M3 FP8 inference on Ascend 950 uses the bundled AscendC MSA index-score implementation for prefill, while decode continues to use the existing A5 Triton path. The public Python API is unchanged. ### How was this patch tested? Static validation completed: - Verified all 54 imported operator files match PR vllm-project#10672 head `2e685d24` byte-for-byte (excluding its standalone `torch_extension`; vLLM Ascend keeps its in-tree Torch adapter). - `ruff check` passed for the modified MiniMax-M3 Python files and unit tests. - `ruff format --check` passed. - Python AST/compile checks passed. - `bash -n csrc/build_aclnn.sh` passed. - `git diff --check` passed. - Added unit assertions for A5 FP8 registration/build wiring and the width-257 windowed-flush regression. #### A5 prefill IndexScore operator performance Compared the AscendC `MsaIndexScore` kernel with the Triton `_index_block_score_kernel` on MiniMax-M3-MXFP8, using Ascend 950, TP4 × DP2, and the same P+D mixed-load request orchestration. Only prefill IndexScore device time is reported. Each value is the mean device duration per IndexScore kernel call across 8 ranks. MiniMax-M3 invokes IndexScore in 57 sparse-attention layers, so 57 calls form one prefill chunk. For 128K, only chunks captured by both implementations are compared to keep the context lengths matched. | Input length | Matched prefill scope | AscendC | Triton | Duration reduction | Speedup | | --- | --- | ---: | ---: | ---: | ---: | | 16K | Chunk 0 | 165.646 µs/call | 214.812 µs/call | 22.89% | 1.297× | | 128K | Chunk 0 | 166.359 µs/call | 213.549 µs/call | 22.10% | 1.284× | | 128K | Chunk 1 | 363.359 µs/call | 529.388 µs/call | 31.36% | 1.457× | | 128K | Chunk 2 | 560.355 µs/call | 837.372 µs/call | 33.08% | 1.494× | | 128K | First 3 matched chunks combined | 363.358 µs/call | 526.770 µs/call | 31.02% | 1.450× | These are matched operator-level profiling results, not end-to-end TTFT, throughput, or serving-latency improvements. - vLLM main: vllm-project/vllm@b2f6858 --------- Signed-off-by: zonghaoxin <z00946994@china.huawei.com> Signed-off-by: Haoxin Zong <534687988@qq.com> Co-authored-by: zonghaoxin <z00946994@china.huawei.com> Signed-off-by: tianming2009 <13246728590@163.com>
…oject#15918) ### What this PR does / why we need it? This PR enables the bundled AscendC `MsaIndexScore` path for MiniMax-M3 on Ascend 950 (A5) with native FP8 index query/key tensors and an FP8 index-key paged cache. - Sync the operator implementation from [cann/ops-transformer#10672](https://gitcode.com/cann/ops-transformer/pull/10672) at `2e685d24ed2e54ea1f1547dc6e5ad29862f152ef`. - Add the Ascend 950 `arch35` FP8 kernels, tiling, Catlass dependencies, dtype registration, examples, and documentation. - Add `msa_index_score` to the Ascend 950 custom-op build list. - Route A5 prefill index scoring through `torch.ops._C_ascend.npu_msa_index_score`, while retaining the existing A5 Triton TopK and decode paths. - Cast Index-Q to E4M3 when the index-key cache is E4M3 and the query dtype does not already match. - Include the A5 wide-`block_table` fix: score columns beyond the 256-column UB window are flushed in windows, with width-257 FP8/BF16 regression coverage. The MiniMax-M3 A5 FP8 contract used by this integration is: query/key use `torch.float8_e4m3fn`, `scale=None`, and the operator emits FP32 scores. ### Does this PR introduce _any_ user-facing change? Yes. MiniMax-M3 FP8 inference on Ascend 950 uses the bundled AscendC MSA index-score implementation for prefill, while decode continues to use the existing A5 Triton path. The public Python API is unchanged. ### How was this patch tested? Static validation completed: - Verified all 54 imported operator files match PR vllm-project#10672 head `2e685d24` byte-for-byte (excluding its standalone `torch_extension`; vLLM Ascend keeps its in-tree Torch adapter). - `ruff check` passed for the modified MiniMax-M3 Python files and unit tests. - `ruff format --check` passed. - Python AST/compile checks passed. - `bash -n csrc/build_aclnn.sh` passed. - `git diff --check` passed. - Added unit assertions for A5 FP8 registration/build wiring and the width-257 windowed-flush regression. #### A5 prefill IndexScore operator performance Compared the AscendC `MsaIndexScore` kernel with the Triton `_index_block_score_kernel` on MiniMax-M3-MXFP8, using Ascend 950, TP4 × DP2, and the same P+D mixed-load request orchestration. Only prefill IndexScore device time is reported. Each value is the mean device duration per IndexScore kernel call across 8 ranks. MiniMax-M3 invokes IndexScore in 57 sparse-attention layers, so 57 calls form one prefill chunk. For 128K, only chunks captured by both implementations are compared to keep the context lengths matched. | Input length | Matched prefill scope | AscendC | Triton | Duration reduction | Speedup | | --- | --- | ---: | ---: | ---: | ---: | | 16K | Chunk 0 | 165.646 µs/call | 214.812 µs/call | 22.89% | 1.297× | | 128K | Chunk 0 | 166.359 µs/call | 213.549 µs/call | 22.10% | 1.284× | | 128K | Chunk 1 | 363.359 µs/call | 529.388 µs/call | 31.36% | 1.457× | | 128K | Chunk 2 | 560.355 µs/call | 837.372 µs/call | 33.08% | 1.494× | | 128K | First 3 matched chunks combined | 363.358 µs/call | 526.770 µs/call | 31.02% | 1.450× | These are matched operator-level profiling results, not end-to-end TTFT, throughput, or serving-latency improvements. - vLLM main: vllm-project/vllm@b2f6858 --------- Signed-off-by: zonghaoxin <z00946994@china.huawei.com> Signed-off-by: Haoxin Zong <534687988@qq.com> Co-authored-by: zonghaoxin <z00946994@china.huawei.com> Signed-off-by: like-0517 <ithwlike@126.com>
What this PR does / why we need it?
This PR enables the bundled AscendC
MsaIndexScorepath for MiniMax-M3 on Ascend 950 (A5) with native FP8 index query/key tensors and an FP8 index-key paged cache.2e685d24ed2e54ea1f1547dc6e5ad29862f152ef.arch35FP8 kernels, tiling, Catlass dependencies, dtype registration, examples, and documentation.msa_index_scoreto the Ascend 950 custom-op build list.torch.ops._C_ascend.npu_msa_index_score, while retaining the existing A5 Triton TopK and decode paths.block_tablefix: score columns beyond the 256-column UB window are flushed in windows, with width-257 FP8/BF16 regression coverage.The MiniMax-M3 A5 FP8 contract used by this integration is: query/key use
torch.float8_e4m3fn,scale=None, and the operator emits FP32 scores.Does this PR introduce any user-facing change?
Yes. MiniMax-M3 FP8 inference on Ascend 950 uses the bundled AscendC MSA index-score implementation for prefill, while decode continues to use the existing A5 Triton path. The public Python API is unchanged.
How was this patch tested?
Static validation completed:
2e685d24byte-for-byte (excluding its standalonetorch_extension; vLLM Ascend keeps its in-tree Torch adapter).ruff checkpassed for the modified MiniMax-M3 Python files and unit tests.ruff format --checkpassed.bash -n csrc/build_aclnn.shpassed.git diff --checkpassed.A5 prefill IndexScore operator performance
Compared the AscendC
MsaIndexScorekernel with the Triton_index_block_score_kernelon MiniMax-M3-MXFP8, using Ascend 950, TP4 × DP2, and the same P+D mixed-load request orchestration. Only prefill IndexScore device time is reported.Each value is the mean device duration per IndexScore kernel call across 8 ranks. MiniMax-M3 invokes IndexScore in 57 sparse-attention layers, so 57 calls form one prefill chunk. For 128K, only chunks captured by both implementations are compared to keep the context lengths matched.
These are matched operator-level profiling results, not end-to-end TTFT, throughput, or serving-latency improvements.