[feature] support mctlassEx bf16 moe

[metax-yi1zhang]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

Test Plan

Test Result

(Optional) Documentation Update

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[gemini-code-assist]

Code Review

This pull request introduces support for bf16 MoE (Mixture of Experts) utilizing the mctlass library, including custom operator registration and routing logic in the fused MoE execution path. The review feedback highlights critical issues where the code assumes mctlass is successfully imported based solely on environment variables, which could lead to runtime crashes if the library is missing; it is recommended to explicitly verify that the mctlass operators are not None. Additionally, several type annotations for scale and bias parameters should be updated to allow None values to prevent static type-checking errors.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The bf16 moe block checks mx_envs.MACA_VLLM_ENABLE_MCTLASS_FUSED_MOE and mx_envs.MACA_VLLM_ENABLE_MCTLASS_PYTHON_API, but it does not verify if the optional mctlassEx library was successfully imported (i.e., mctlass_ops.mctlass_moe_gemm is not None). If these environment variables are enabled but mctlassEx is not installed, this will result in an AssertionError and crash the execution instead of falling back to the default Triton implementation.

Please add a check to ensure mctlass_ops.mctlass_moe_gemm is not None before entering this block.

            if (
                hidden_states.dtype == torch.bfloat16
                and not self.quant_config.use_int4_w4a8
                and not self.quant_config.use_int4_w4a16
                and not self.quant_config.use_int8_w8a8
                and not self.quant_config.use_int8_w8a16
                and mx_envs.MACA_VLLM_ENABLE_MCTLASS_FUSED_MOE
                and mx_envs.MACA_VLLM_ENABLE_MCTLASS_PYTHON_API
                and getattr(mctlass_ops, "mctlass_moe_gemm", None) is not None
            ):

[gemini-code-assist]

Similar to the issue in triton_moe.py, this elif block will call mctlass_ops.cutlass_moe_mm_bf16 which asserts that mctlass_moe_gemm is not None. If mctlassEx is not installed, this will crash instead of falling back to the else block (Triton kernel).

Please add a check to ensure mctlass_ops.mctlass_moe_gemm is not None in the condition.

    elif (
        A.dtype == torch.bfloat16 and B.dtype == torch.bfloat16 and C.dtype == torch.bfloat16
        and A_scale is None and B_scale is None
        and mx_envs.MACA_VLLM_ENABLE_MCTLASS_FUSED_MOE
        and mx_envs.MACA_VLLM_ENABLE_MCTLASS_PYTHON_API
        and getattr(mctlass_ops, "mctlass_moe_gemm", None) is not None
    ):