Ck moe bs splitk pr

[huaiguxu]

Proposed changes

Add split-k support for a8w8_moe_blk gemm1 only.

Checklist

Please put an x into the boxes that apply. You can also fill these out after creating the PR. If you're not sure, please don't hesitate to ask.

• [x ] I have added tests relevant to the introduced functionality, and the unit tests are passing locally
• I have added the test to REGRESSION_TESTS list defined at the top of CMakeLists.txt in tests/CMakeLists.txt, IF the test takes more than 30 seconds to run.
• I have added inline documentation which enables the maintainers with understanding the motivation
• I have removed the stale documentation which is no longer relevant after this pull request
• (If this change is user-facing) I have added release notes which provide the end users with a brief summary of the improvement from this pull request
• I have run clang-format on all changed files
• Any dependent changes have been merged

Discussion

The change is due to low CU utilization in gemm1 decode stage with TP mode.

[coderfeli]

LGTM

[illsilin]

Hi @huaiguxu, please resolve the conflicts and make sure all CI tests are passing.

[illsilin]

Hi @coderfeli and @huaiguxu, please resolve conflicts and make sure CI passes or we'll close this PR.

[ammallya]

Error importing due to merge conflicts – please reopen the PR on ROCm/rocm-libraries