diff --git a/dumux/porousmediumflow/compositional/localresidual.hh b/dumux/porousmediumflow/compositional/localresidual.hh
index 6c6c7ce5904bba69832a12dbe4d10d89129d0d9b..6ce336f7115b54ef06938aa18ba44d2c6c543805 100644
--- a/dumux/porousmediumflow/compositional/localresidual.hh
+++ b/dumux/porousmediumflow/compositional/localresidual.hh
@@ -96,19 +96,19 @@ public:
                 for (int compIdx = 0; compIdx < numComponents; ++compIdx)
                 {
                     auto eqIdx = conti0EqIdx + compIdx;
-                    auto s = volVars.porosity()
-                             * volVars.saturation(phaseIdx)
-                             * volVars.molarDensity(phaseIdx)
-                             * volVars.moleFraction(phaseIdx, compIdx);
-
                     if (eqIdx != replaceCompEqIdx)
-                        storage[eqIdx] += s;
-
-                    // in case one balance is substituted by the total mass balance
-                    if (replaceCompEqIdx < numComponents)
-                        storage[replaceCompEqIdx] += s;
+                        storage[eqIdx] += volVars.porosity()
+                                          * volVars.saturation(phaseIdx)
+                                          * volVars.molarDensity(phaseIdx)
+                                          * volVars.moleFraction(phaseIdx, compIdx);
                 }
 
+                // in case one balance is substituted by the total mole balance
+                if (replaceCompEqIdx < numComponents)
+                    storage[replaceCompEqIdx] = volVars.molarDensity(phaseIdx)
+                                                * volVars.porosity()
+                                                * volVars.saturation(phaseIdx);
+
                 //! The energy storage in the fluid phase with index phaseIdx
                 EnergyLocalResidual::fluidPhaseStorage(storage, scv, volVars, phaseIdx);
             }
@@ -124,19 +124,19 @@ public:
                 for (int compIdx = 0; compIdx < numComponents; ++compIdx)
                 {
                     auto eqIdx = conti0EqIdx + compIdx;
-                    auto s = volVars.porosity()
-                             * volVars.saturation(phaseIdx)
-                             * volVars.density(phaseIdx)
-                             * volVars.massFraction(phaseIdx, compIdx);
-
                     if (eqIdx != replaceCompEqIdx)
-                        storage[eqIdx] += s;
-
-                    // in case one balance is substituted by the total mass balance
-                    if (replaceCompEqIdx < numComponents)
-                        storage[replaceCompEqIdx] += s;
+                        storage[eqIdx] += volVars.porosity()
+                                          * volVars.saturation(phaseIdx)
+                                          * volVars.density(phaseIdx)
+                                          * volVars.massFraction(phaseIdx, compIdx);
                 }
 
+                // in case one balance is substituted by the total mass balance
+                if (replaceCompEqIdx < numComponents)
+                    storage[replaceCompEqIdx] = volVars.density(phaseIdx)
+                                                * volVars.porosity()
+                                                * volVars.saturation(phaseIdx);
+
                 //! The energy storage in the fluid phase with index phaseIdx
                 EnergyLocalResidual::fluidPhaseStorage(storage, scv, volVars, phaseIdx);
             }
@@ -190,10 +190,6 @@ public:
                     if (eqIdx != replaceCompEqIdx)
                         flux[eqIdx] += advFlux;
 
-                    // in case one balance is substituted by the total mass balance
-                    if (replaceCompEqIdx < numComponents)
-                        flux[replaceCompEqIdx] += advFlux;
-
                     // diffusive fluxes (only for the component balances)
                     if (phaseIdx != compIdx)
                     {
@@ -205,6 +201,15 @@ public:
                     }
                 }
 
+                // in case one balance is substituted by the total mole balance
+                if (replaceCompEqIdx < numComponents)
+                {
+                    auto upwindTermTotalBalance = [phaseIdx](const VolumeVariables& volVars)
+                    { return volVars.molarDensity(phaseIdx)*volVars.mobility(phaseIdx); };
+
+                    flux[replaceCompEqIdx] = fluxVars.advectiveFlux(phaseIdx, upwindTermTotalBalance);
+                }
+
                 //! Add advective phase energy fluxes. For isothermal model the contribution is zero.
                 EnergyLocalResidual::heatConvectionFlux(flux, fluxVars, phaseIdx);
             }
@@ -232,10 +237,6 @@ public:
                     if (eqIdx != replaceCompEqIdx)
                         flux[eqIdx] += advFlux;
 
-                    // in case one balance is substituted by the total mass balance
-                    if (replaceCompEqIdx < numComponents)
-                        flux[replaceCompEqIdx] += advFlux;
-
                     // diffusive fluxes (only for the component balances)
                     if (phaseIdx != compIdx)
                     {
@@ -247,6 +248,15 @@ public:
                     }
                 }
 
+                // in case one balance is substituted by the total mass balance
+                if (replaceCompEqIdx < numComponents)
+                {
+                    auto upwindTermTotalBalance = [phaseIdx](const VolumeVariables& volVars)
+                    { return volVars.density(phaseIdx)*volVars.mobility(phaseIdx); };
+
+                    flux[replaceCompEqIdx] = fluxVars.advectiveFlux(phaseIdx, upwindTermTotalBalance);
+                }
+
                 //! Add advective phase energy fluxes. For isothermal model the contribution is zero.
                 EnergyLocalResidual::heatConvectionFlux(flux, fluxVars, phaseIdx);
             }