[Bugfix] Fix wrong num_experts in invoke_fused_moe_wna16 kernels

[weiguangli-io]

Summary

Both invoke_fused_moe_wna16_cuda_kernel and invoke_fused_moe_wna16_triton_kernel pass B.size(1) (output dimension N) as num_experts to get_moe_wna16_block_config. The expert dimension is actually B.size(0).

This causes incorrect block configuration and can lead to a CUDA runtime crash:

RuntimeError: BLOCK_SIZE_K // group_size must be one of [1, 2, 4, 8]

Fix: Change num_experts=B.size(1) to num_experts=B.size(0) in both kernel invocation functions.

Fixes #36008

Test plan

• The fix is a straightforward parameter correction (dimension 1 → dimension 0) with no behavioral change for correctly-shaped inputs
• Verified that tensor B has shape (num_experts, N, K // pack_factor), confirming B.size(0) is the expert count

[gemini-code-assist]

Code Review

This pull request correctly fixes a bug in the invoke_fused_moe_wna16_cuda_kernel and invoke_fused_moe_wna16_triton_kernel functions. The num_experts parameter passed to get_moe_wna16_block_config was incorrectly using B.size(1) instead of B.size(0). This change aligns with the tensor shape documentation, where the first dimension of the weight tensor B represents the number of experts, and resolves a potential CUDA runtime crash. The fix is straightforward and correct.

[weiguangli-io]

Friendly ping — this PR has been open for a few weeks. Would appreciate a review when you get a chance. Happy to address any feedback. Thanks!

[weiguangli-io]

Closing this PR as it has been open for 8+ weeks with no maintainer review despite follow-up. The fix is straightforward, so feel free to reopen if this becomes relevant. Happy to rebase if needed. Thanks!