return NULL;
 }
 
-static int ath11k_mac_setup_bcn_tmpl_ema(struct ath11k_vif *arvif)
+static int ath11k_mac_setup_bcn_tmpl_ema(struct ath11k_vif *arvif,
+                                        struct ath11k_vif *tx_arvif)
 {
-       struct ath11k_vif *tx_arvif;
        struct ieee80211_ema_beacons *beacons;
        int ret = 0;
        bool nontx_vif_params_set = false;
        u32 params = 0;
        u8 i = 0;
 
-       tx_arvif = ath11k_mac_get_tx_arvif(arvif);
-
        beacons = ieee80211_beacon_get_template_ema_list(tx_arvif->ar->hw,
                                                         tx_arvif->vif, 0);
        if (!beacons || !beacons->cnt) {
        return ret;
 }
 
-static int ath11k_mac_setup_bcn_tmpl_mbssid(struct ath11k_vif *arvif)
+static int ath11k_mac_setup_bcn_tmpl_mbssid(struct ath11k_vif *arvif,
+                                           struct ath11k_vif *tx_arvif)
 {
        struct ath11k *ar = arvif->ar;
        struct ath11k_base *ab = ar->ab;
-       struct ath11k_vif *tx_arvif;
        struct ieee80211_hw *hw = ar->hw;
        struct ieee80211_vif *vif = arvif->vif;
        struct ieee80211_mutable_offsets offs = {};
        struct sk_buff *bcn;
        int ret;
 
-       tx_arvif = ath11k_mac_get_tx_arvif(arvif);
-       if (tx_arvif && tx_arvif != arvif) {
+       if (tx_arvif != arvif) {
                ar = tx_arvif->ar;
                ab = ar->ab;
                hw = ar->hw;
                vif = tx_arvif->vif;
-       } else {
-               tx_arvif = arvif;
        }
 
        bcn = ieee80211_beacon_get_template(hw, vif, &offs, 0);
 static int ath11k_mac_setup_bcn_tmpl(struct ath11k_vif *arvif)
 {
        struct ieee80211_vif *vif = arvif->vif;
+       struct ath11k_vif *tx_arvif;
 
        if (arvif->vdev_type != WMI_VDEV_TYPE_AP)
                return 0;
        /* Target does not expect beacon templates for the already up
         * non-transmitting interfaces, and results in a crash if sent.
         */
-       if (vif->mbssid_tx_vif &&
-           arvif != ath11k_vif_to_arvif(vif->mbssid_tx_vif) && arvif->is_up)
-               return 0;
+       tx_arvif = ath11k_mac_get_tx_arvif(arvif);
+       if (tx_arvif) {
+               if (arvif != tx_arvif && arvif->is_up)
+                       return 0;
 
-       if (vif->bss_conf.ema_ap && vif->mbssid_tx_vif)
-               return ath11k_mac_setup_bcn_tmpl_ema(arvif);
+               if (vif->bss_conf.ema_ap)
+                       return ath11k_mac_setup_bcn_tmpl_ema(arvif, tx_arvif);
+       } else {
+               tx_arvif = arvif;
+       }
 
-       return ath11k_mac_setup_bcn_tmpl_mbssid(arvif);
+       return ath11k_mac_setup_bcn_tmpl_mbssid(arvif, tx_arvif);
 }
 
 void ath11k_mac_bcn_tx_event(struct ath11k_vif *arvif)