feat(kernel): introduce unified moe kernel api

[borontion]

Summary

This PR moves moe backend concepts into tokenspeed kernel. MoE will be a platform-agnostic layer that invoke moe kernel api.

MoE entrypoint

Introduce a single moe_apply api. Under the hood it can launch a single or more multiple kernels depending on the implementation. This api is too coarse from kernel-level. Moving forward, we will replace it as a kernel execution plan.

See tokenspeed-kernel/python/tokenspeed_kernel/ops/moe/__init__.py.

Supported type and backend

Currently I keep the following weight dtype + implementation:

• Unquant (fp16): flashinfer cutlass
• fp8 (with scale): flashinfer cutlass
• mxfp4: triton / gluon / flashinfer trtllm
• nvfp4: flashinfer cutlass / flashinfer trtllm / flashinfer cutedsl

Other data types are not actively used in our current target models. Also drop the past triton backend for fp8/unquant, which has convoluted dependencies. A portable impl for fp8/unquant which can run amd gpu will be added later. Also renamed the backend name for 'triton_kernel' to 'triton'.

Kernel early selection

MoE is different from attention where the weights need to be preshuffled after model loaded. The preshuffle and actual moe impl are tightly coupled. So we need to determine the kernel implementation earlier, before engine accepts incoming requests.

This is done by adding a specific moe_plan utility function. It will select the corresbonding moe kernel based on various features and the current platform. The selection is encoded into a dict which includes the target kernel name, which will be passed to moe_apply.

AMD Gluon kernel

Decouple amd gluon kernel from other kernel utilit by copying used code into tokenspeed-kernel-amd module. This includes 3 parts

1. fp8 quant
2. fallback triton kernel routing
3. fallback triton kernel group gemm

See: tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/moe/utils.py for ported triton kernel code.

Currently I lower the priority for the gluon kernel due to known performance issue.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 682fdb6fb2

ℹ️ 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".

[antiagainst]

Awesome! Looks large change but mostly moving code around and fleshing out kernel interface and registration. Just two nits from me.