diff --git a/dumux/freeflow/compositional/navierstokesncmodel.hh b/dumux/freeflow/compositional/navierstokesncmodel.hh
index c9093182576d9f23161720cb71b2b6224959a4c1..b7958d01555de0bc15257a648123c231f97c919f 100644
--- a/dumux/freeflow/compositional/navierstokesncmodel.hh
+++ b/dumux/freeflow/compositional/navierstokesncmodel.hh
@@ -26,7 +26,7 @@
  * \f[
  *    \frac{\partial \left(\varrho X^\kappa\right)}{\partial t}
  *    + \nabla \cdot \left( \varrho {\boldsymbol{v}} X^\kappa
- *    - (D^\kappa + D_\text{t}) \varrho \frac{M^\kappa}{M} \textbf{grad}\, x^\kappa \right)
+ *    - (D^\kappa + D_\text{t}) \varrho \textbf{grad}\, X^\kappa \right)
  *    - q^\kappa = 0
  * \f]
  *
@@ -35,7 +35,7 @@
  *    \frac{\partial \varrho_g}{\partial t}
  *    + \nabla \cdot \left(
  *        \varrho {\boldsymbol{v}}
- *        - \sum_\kappa (D^\kappa + D_\text{t}) \varrho \frac{M^\kappa}{M} \textbf{grad}\, x^\kappa
+ *        - \sum_\kappa (D^\kappa + D_\text{t}) \varrho \textbf{grad}\, X^\kappa
  *      \right)
  *    - q = 0
  * \f]
diff --git a/dumux/freeflow/compositional/staggered/fluxvariables.hh b/dumux/freeflow/compositional/staggered/fluxvariables.hh
index 32ae6f46100e3d356db2722479f726afadd2a4d4..de283bab6fdc6f9fff9e24100201bdcedd8138e4 100644
--- a/dumux/freeflow/compositional/staggered/fluxvariables.hh
+++ b/dumux/freeflow/compositional/staggered/fluxvariables.hh
@@ -86,7 +86,13 @@ public:
 
             flux[compIdx] = ParentType::advectiveFluxForCellCenter(problem, elemVolVars, elemFaceVars, scvf, upwindTerm);
 
-            flux[compIdx] += useMoles ? diffusiveFluxes[compIdx] : diffusiveFluxes[compIdx]*FluidSystem::molarMass(compIdx);
+            //check for the reference system and adapt units of the diffusive flux accordingly.
+            if (MolecularDiffusionType::referenceSystemFormulation() == ReferenceSystemFormulation::massAveraged)
+            {
+                flux[compIdx] += useMoles ? diffusiveFluxes[compIdx]/FluidSystem::molarMass(compIdx) : diffusiveFluxes[compIdx];
+            }
+            else
+                flux[compIdx] += useMoles ? diffusiveFluxes[compIdx] : diffusiveFluxes[compIdx]*FluidSystem::molarMass(compIdx);
         }
 
         // in case one balance is substituted by the total mass balance
diff --git a/dumux/freeflow/compositional/volumevariables.hh b/dumux/freeflow/compositional/volumevariables.hh
index a9ccb5ae877151ca33e4c94363e4f19e4d5aadae..8be88d881ca21b2bcac362d00a942fc7fbb7d571 100644
--- a/dumux/freeflow/compositional/volumevariables.hh
+++ b/dumux/freeflow/compositional/volumevariables.hh
@@ -245,6 +245,14 @@ public:
         return FluidSystem::molarMass(compIdx);
     }
 
+    /*!
+     * \brief Returns the average molar mass \f$\mathrm{[kg/mol]}\f$$ the of the fluid phase.
+     *
+     * \param phaseIdx The phase index
+     */
+    Scalar averageMolarMass(const int phaseIdx = 0) const
+    { return fluidState_.averageMolarMass(phaseIdx); }
+
    /*!
     * \brief Returns the diffusion coefficient \f$\mathrm{[m^2/s]}\f$
     *
diff --git a/dumux/freeflow/nonisothermal/localresidual.hh b/dumux/freeflow/nonisothermal/localresidual.hh
index 773157340b5f33b062713de600c739b1c6cce59c..0a5475a3a758fa19e6ac90e199958dd8a9eb7513 100644
--- a/dumux/freeflow/nonisothermal/localresidual.hh
+++ b/dumux/freeflow/nonisothermal/localresidual.hh
@@ -25,6 +25,7 @@
 #define DUMUX_FREE_FLOW_ENERGY_LOCAL_RESIDUAL_HH
 
 #include <dumux/discretization/method.hh>
+#include <dumux/flux/referencesystem.hh>
 
 namespace Dumux {
 
@@ -158,11 +159,11 @@ public:
             const bool insideIsUpstream = scvf.directionSign() == sign(diffusiveFlux[compIdx]);
             const auto& upstreamVolVars = insideIsUpstream ? elemVolVars[scvf.insideScvIdx()] : elemVolVars[scvf.outsideScvIdx()];
 
-            // always use a mass-based calculation for the energy balance
-            if (FluxVariables::useMoles)
-                diffusiveFlux[compIdx] *= elemVolVars[scvf.insideScvIdx()].molarMass(compIdx);
+            if (FluxVariables::MolecularDiffusionType::referenceSystemFormulation() == ReferenceSystemFormulation::massAveraged)
+                flux[localEnergyBalanceIdx] += diffusiveFlux[compIdx] * upstreamVolVars.componentEnthalpy(compIdx);
+            else
+                flux[localEnergyBalanceIdx] += diffusiveFlux[compIdx] * upstreamVolVars.componentEnthalpy(compIdx)* elemVolVars[scvf.insideScvIdx()].molarMass(compIdx);
 
-            flux[localEnergyBalanceIdx] += diffusiveFlux[compIdx] * upstreamVolVars.componentEnthalpy(compIdx);
         }
     }
 };