How to support topk>2 (topk=6 is needed in our experiment) in MoE model?

we find that

@staticmethod
def supports_config(config: DSMoEConfig) -> bool:
if config.input_dtype != config.output_dtype:
return False

    if config.input_dtype != torch.float16 and config.input_dtype != torch.bfloat16:
        return False

    if config.top_k != 1 and config.top_k != 2:
        return False

    return True

in "deepspeed.inference.v2.modules.implementations.moe.cutlass_multi_gemm.DSMultiGemmMoE"

It means DSMoE only supoort topk=1 or topk =2.

But our experiment need support for topk=6

Start at deepspeed.inference.v2.modules.implementations.moe.cutlass_multi_gemm.DSMultiGemmMoE.supports_config and trace how top_k is handled by the MoE implementation. Determine the existing validation and execution constraints for values above 2; done means top_k=6 is supported without breaking the current top_k=1 and top_k=2 cases.