Repository: vllm-project/vllm
File: vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py (Python)

Code before:
            FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit")),
            (swiglu_alpha, swiglu_limit),
            2,
        )
        intermediate_cache1 = matmul_ogs(
            hidden_states,
            w1,
            quant_config.w1_bias,
            routing_data,
            gather_indx=gather_indx,
            precision_config=quant_config.w1_precision,
            gammas=gammas if apply_router_weight_on_input else None,
            fused_activation=act,
        )
        intermediate_cache3 = matmul_ogs(
            intermediate_cache1,
            w2,
            quant_config.w2_bias,
            routing_data,
            scatter_indx=scatter_indx,
            precision_config=quant_config.w2_precision,
            gammas=None if apply_router_weight_on_input else gammas,
            y=output_tensor,
        )

    elif quant_config.use_mxfp4_w4a4:
        from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4
        from aiter.ops.triton.quant_moe import downcast_to_static_fp8

        assert quant_config.w1_precision is not None, (
            "w1_precision in quant config can't be None"
        )
        assert quant_config.w2_precision is not None, (
            "w2_precision in quant config can't be None"
        )

        hidden_states = downcast_to_static_fp8(
            hidden_states, quant_config.w1_precision.flex_ctx.lhs_data.scale
        )

        intermediate_cache1 = moe_gemm_a8w4(
            hidden_states,
            w1.storage.data,
            None,
            quant_config.w1_precision.weight_scale.storage.data,
            quant_config.w1_precision.flex_ctx.lhs_data.scale,
            quant_config.w2_precision.flex_ctx.lhs_data.scale,
            quant_config.w1_bias,
            routing_data,
            gather_indx=gather_indx,
            gammas=gammas if apply_router_weight_on_input else None,

Review comment (suggestion):
why is it `use_mxfp4_w4a4` in config but `a8w4` kernel

Code after:
    # Shape check, only check non-mxfp4
    assert hidden_states.shape[-1] == w1.shape[-2]
    assert w2.shape[-1] == w1.shape[1]

    E, _, N = w1.shape

    if global_num_experts == -1:
        global_num_experts = E

    gammas = routing_data.gate_scal if routing_data else None

    from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4
    from aiter.ops.triton.quant_moe import downcast_to_static_fp8

    assert quant_config.w1_precision is not None, (
        "w1_precision in quant config can't be None"
    )
    assert quant_config.w2_precision is not None, (
        "w2_precision in quant config can't be None"
    )

    hidden_states = downcast_to_static_fp8(
        hidden_states, quant_config.w1_precision.flex_ctx.lhs_data.scale
    )

    intermediate_cache1 = moe_gemm_a8w4(
        hidden_states,
        w1.storage.data,
        None,
        quant_config.w1_precision.weight_scale.storage.data,
        quant_config.w1_precision.flex_ctx.lhs_data.scale,
        quant_config.w2_precision.flex_ctx.lhs_data.scale,
        quant_config.w1_bias,
        routing_data,
        gather_indx=gather_indx,
        gammas=gammas if apply_router_weight_on_input else None,
        swizzle_mx_scale="CDNA4_SCALE",
        out_dtype=torch.float8_e4m3fn,
        apply_swiglu=True,
        alpha=swiglu_alpha,
        limit=swiglu_limit,
        unpadded_N=unpadded_N_w1,
        unpadded_K=unpadded_K_w1,
    )

    intermediate_cache3 = moe_gemm_a8w4(
        intermediate_cache1,
        w2.storage.data,
        None,
        quant_config.w2_precision.weight_scale.storage.data,
        quant_config.w2_precision.flex_ctx.lhs_data.scale,