[Feature]: Optional topk_weights arg in Fused MoE B2B GEMM

[sijiac]

Suggestion Description

We may not need topk_weights for the B2B GEMM fused MoE. For now, we can pass a dummy torch.ones() tensor to the kernel, which loads the argument and performs the calculation. Ideally, we should make this an optional argument and skip the associated logic when it's set to None.

cc @carlushuang

Operating System

No response

GPU

No response

ROCm Component

No response

[coconutruben]

@carlushuang to simultaneously address this and #195 it would be ideal if we can control when the topk_weight scaling happens: before GEMMs or after (current implementation). An interface that allows us to say

• topk_weights before GEMMs
• topk_weights after GEMMs (current)
• no topk_weights