[MoE Refactoring][Bugfix]Wrap WNA16 Triton kernel into mk and change compressed tensor kernel selection

[zyongye]

Refactor WNA16 Triton kernel into modular kernel format.
Previously, we factor out the WNA16 from TritonExperts and that caused WNA16 lora module can't select the correct kernel (see issue). This PR fixes the problem.

[gemini-code-assist]

Code Review

This pull request refactors the WNA16 Triton kernel into a modular kernel class TritonWNA16Experts and fixes a bug in the compressed tensor kernel selection to use this new class. The changes are logical, but there is significant code duplication in the new TritonWNA16Experts class, which inherits from TritonExperts. I've provided suggestions to refactor this to improve maintainability and adhere to the DRY principle. Specifically, the apply method is almost a complete copy of its parent's implementation, and the moe_sum method is also duplicated.

[gemini-code-assist]

The TritonWNA16Experts class, particularly the apply method, contains a large amount of duplicated code from its parent class TritonExperts. This significantly impacts maintainability, as any changes to the base logic would need to be manually synchronized in this class as well.

To adhere to the Don't Repeat Yourself (DRY) principle, I recommend refactoring this. A possible approach is to extract the common logic into a protected base method in TritonExperts and have TritonWNA16Experts override only the parts that differ, such as the kernel invocation logic.

For example, you could introduce helper methods in TritonExperts that TritonWNA16Experts can override:

class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
    # ...

    def _get_kernel_fn(self):
        return invoke_fused_moe_triton_kernel

    def _get_first_gemm_kwargs(self, ...):
        # returns kwargs for the first GEMM
        return { ... }

    def _get_second_gemm_kwargs(self, ...):
        # returns kwargs for the second GEMM
        return { ... }

    def apply(self, ...):
        # ... common setup ...

        kernel_fn = self._get_kernel_fn()
        first_gemm_kwargs = self._get_first_gemm_kwargs(...)
        kernel_fn(**first_gemm_kwargs)

        # ... common activation and quantization ...

        second_gemm_kwargs = self._get_second_gemm_kwargs(...)
        kernel_fn(**second_gemm_kwargs)

        # ... common finalization ...

class TritonWNA16Experts(TritonExperts):
    def _get_kernel_fn(self):
        return invoke_fused_moe_wna16_triton_kernel

    def _get_first_gemm_kwargs(self, ...):
        # returns kwargs for the first WNA16 GEMM
        return { ... }

    def _get_second_gemm_kwargs(self, ...):
        # returns kwargs for the second WNA16 GEMM
        return { ... }

    # No need to override apply() if all differences are in helper methods.

This refactoring would make the code much cleaner, more maintainable, and less prone to bugs from unsynchronized changes.

[gemini-code-assist]

This moe_sum method is identical to the one in the parent class TritonExperts. It can be removed to avoid code duplication and improve maintainability.

[mergify]

Hi @zyongye, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[mergify]

Hi @zyongye, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[JartX]

@zyongye works, problem solved for me :)

[bnellnm]

Looks ok to me. Can you update tests/kernels/moe/modular_kernel_tools/mk_objects.py so that the new experts class gets tested?

[bnellnm]

Doesn't this kernel only work for 16-bit float activation types?

[zyongye]

After discussion with @bnellnm , we decided to defer adding tests to the other PR sincethe current test structure lacks int4 test infra.