add xpu op grouped topk

[mayuyuace]

add xpu grouped topk kernel

[jikunshang]

can we do some benchmark among grouped_topk_native, grouped_topk_native with @torch.compile, grouped_topk

[mayuyuace]

grouped_topk_native, grouped_topk_native with @torch.compile has been added.

[dbyoung18]

Suggest follow upstream's structure, put moe related kernels under vllm-xpu-kernels/csrc/xpu/moe
reference: https://github.com/vllm-project/vllm/tree/main/csrc/moe

[mayuyuace]

Suggest follow upstream's structure, put moe related kernels under vllm-xpu-kernels/csrc/xpu/moe reference: https://github.com/vllm-project/vllm/tree/main/csrc/moe

Done.

[dbyoung18]

@jikunshang Considering later upstream, do u think we should follow vllm to separate moe kernels to another .so/module?
https://github.com/vllm-project/vllm/blob/main/setup.py#L324-L325
https://github.com/vllm-project/vllm/blob/main/vllm/_custom_ops.py#L16-L25

[jikunshang]

@jikunshang Considering later upstream, do u think we should follow vllm to separate moe kernels to another .so/module? https://github.com/vllm-project/vllm/blob/main/setup.py#L324-L325 https://github.com/vllm-project/vllm/blob/main/vllm/_custom_ops.py#L16-L25

agree, let's move this kernel into _moe.so since this is moe related kernel.

[mayuyuace]

@dbyoung18 @jikunshang
I have added _moe_C.abi3.so and dirs in benchmark & tests for moe.

[jikunshang]

@dbyoung18 @jikunshang I have added _moe_C.abi3.so and dirs in benchmark & tests for moe.

I have some discussion with @dbyoung18 about library name, we noticed that rocm add some non-cuda kernels in _rocm_C, see link maybe we should add this non-cuda kernel to _xpu_C extension. @dbyoung18 also have a similar kernel, he will try add _xpu_C extension first.

[mayuyuace]

@dbyoung18 @jikunshang I have added _moe_C.abi3.so and dirs in benchmark & tests for moe.

I have some discussion with @dbyoung18 about library name, we noticed that rocm add some non-cuda kernels in _rocm_C, see link maybe we should add this non-cuda kernel to _xpu_C extension. @dbyoung18 also have a similar kernel, he will try add _xpu_C extension first.

Since the name will affects all ops not only grouped_topk, maybe changing .so name should be done in another PR?

[dbyoung18]

@dbyoung18 @jikunshang I have added _moe_C.abi3.so and dirs in benchmark & tests for moe.

I have some discussion with @dbyoung18 about library name, we noticed that rocm add some non-cuda kernels in _rocm_C, see link maybe we should add this non-cuda kernel to _xpu_C extension. @dbyoung18 also have a similar kernel, he will try add _xpu_C extension first.

Since the name will affects all ops not only grouped_topk, maybe changing .so name should be done in another PR?

I just refactor cmake for _moe_C and submit a related kernel moe_sum, pls take as a reference.

[mayuyuace]

@dbyoung18 @jikunshang I have added _moe_C.abi3.so and dirs in benchmark & tests for moe.

I have some discussion with @dbyoung18 about library name, we noticed that rocm add some non-cuda kernels in _rocm_C, see link maybe we should add this non-cuda kernel to _xpu_C extension. @dbyoung18 also have a similar kernel, he will try add _xpu_C extension first.

Since the name will affects all ops not only grouped_topk, maybe changing .so name should be done in another PR?

I just refactor cmake for _moe_C and submit a related kernel moe_sum, pls take as a reference.

The PR for moe_sum looks like is same with what I did in this JIRA.

[dbyoung18]

@dbyoung18 @jikunshang I have added _moe_C.abi3.so and dirs in benchmark & tests for moe.

I have some discussion with @dbyoung18 about library name, we noticed that rocm add some non-cuda kernels in _rocm_C, see link maybe we should add this non-cuda kernel to _xpu_C extension. @dbyoung18 also have a similar kernel, he will try add _xpu_C extension first.

Since the name will affects all ops not only grouped_topk, maybe changing .so name should be done in another PR?

I just refactor cmake for _moe_C and submit a related kernel moe_sum, pls take as a reference.

The PR for moe_sum looks like is same with what I did in this JIRA.

1.Just noticed ur latest modifications. I made the change yesterday in parallel w/ u. The main parts of +_moe_C between ours are common.
2.@jikunshang Since grouped_topk is a part of moe, I would prefer put it under _moe_C not _xpu_C, what's ur options?
3.For UT, I think register all kernels to vllm-xpu-kernels/tests/register_ops.py is enough, as it's designed to replace _ipex_ops.py and for CUDA, moe related kernels also seen registered to _custom_ops.py . For single UT, it's ok to isolate according to catagory as vllm/tests/kernels.

I think our main philosophy in doing this is to align with the community as much as possible to reduce effort for later upstream.

[jikunshang]

vllm-project/vllm#23274 vllm add fused group_topk recently. Please take a look what we missed, thanks!

[mayuyuace]

vllm-project/vllm#23274 vllm add fused group_topk recently. Please take a look what we missed, thanks!

OK, I will rewrite a new kernel from vllm fused group_topk, and compare it with what we use now.
And I will create a new PR if the kernel performance from vllm is better.

[jikunshang]

can we close this?