Skip to content

[KMCompiler][Ascend] Add custom ops for vllm-ascend moe_gating_top_k replication - #1159

Open
liuxiao0909c wants to merge 2 commits into
flagos-ai:triton_v3.5.xfrom
liuxiao0909c:triton_v3.5.x_grouped_topk
Open

liuxiao0909c wants to merge 2 commits into
flagos-ai:triton_v3.5.xfrom
liuxiao0909c:triton_v3.5.x_grouped_topk

Conversation

@liuxiao0909c

@liuxiao0909c liuxiao0909c commented Sep 12, 2026

Copy link
Copy Markdown

Custom operator interfaces are defined as specified in https://www.hiascend.com/document/detail/zh/canncommercial/latest/API/ascendcopapi/atlasascendc_api_07_0003.html, and all custom operator functions are implemented with reference to the implementation under /usr/local/Ascend/cann-9.1.0/arm64-linux/asc/impl/basic_api/dav_c220.

custom算子 对应的Ascend C API 必要性
sort32 Sort32 可以用{score, index}的形式输出排序结果,比triton中pack为int64排序再unpack更加高效
duplicate_bitwise_mask Duplicate 用tle.dsa.extract_slice/insert_slice这种拼接形式填充得到的tensor,用于sort32时,有正确性问题,这种在原tensor上修改存储内容的形式没有正确性问题
gather_mask_custom_pattern GatherMask 配合duplicate_bitwise_mask,输入输出使用同一个tensor
pair_reduce_sum_continuous_mask PairReduceSum 配合duplicate_bitwise_mask,输入输出使用同一个tensor
gather_mask_builtin_pattern GatherMask 比trion中构造索引再gather高效
mrgsort MrgSort 四路归并后再取topk比vllm C实现中循环topk次,每次调用max规约操作更加高效

…replication

Custom operator interfaces are defined as specified in
https://www.hiascend.com/document/detail/zh/canncommercial/latest/API/ascendcopapi/atlasascendc_api_07_0003.html,
and all custom operator functions are implemented with reference to
the implementation under /usr/local/Ascend/cann-9.1.0/arm64-linux/asc/impl/basic_api/dav_c220.

sort32: corresponding to the Sort32<float> in Ascend C API
mrgsort: corresponding to the MrgSort<float> in Ascend C API
gather_mask_builtin_pattern/gather_mask_custom_pattern: corresponding to
the GatherMask in Ascend C API (builtin fixed mode and user-definded
mode)
pair_reduce_sum_continuous_mask: corresponding to the PairReduceSum in
Ascend C API (only continuous mode)
duplicate_bitwise_mask: corresponding to the Duplicate in Ascend C API
(only bitwise mode for tensor high-dimensional slicing computation)
@CLAassistant

CLAassistant commented Sep 12, 2026

Copy link
Copy Markdown

CLA assistant check
All committers have signed the CLA.

@liuxiao0909c
liuxiao0909c force-pushed the triton_v3.5.x_grouped_topk branch from 31c7197 to b9e4be4 Compare September 14, 2026 11:01
@kateyijian

Copy link
Copy Markdown
Collaborator

LGTM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CORE DOC Improvements or additions to documentation tle triton_v3.5.x

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants