if (vht_cap & (IEEE80211_VHT_CAP_SU_BEAMFORMEE_CAPABLE)) {
                nsts = vht_cap & IEEE80211_VHT_CAP_BEAMFORMEE_STS_MASK;
                nsts >>= IEEE80211_VHT_CAP_BEAMFORMEE_STS_SHIFT;
+               if (nsts > (ar->num_rx_chains - 1))
+                       nsts = ar->num_rx_chains - 1;
                value |= SM(nsts, WMI_TXBF_STS_CAP_OFFSET);
        }
 
 static void ath11k_set_vht_txbf_cap(struct ath11k *ar, u32 *vht_cap)
 {
        bool subfer, subfee;
-       int sound_dim = 0;
+       int sound_dim = 0, nsts = 0;
 
        subfer = !!(*vht_cap & (IEEE80211_VHT_CAP_SU_BEAMFORMER_CAPABLE));
        subfee = !!(*vht_cap & (IEEE80211_VHT_CAP_SU_BEAMFORMEE_CAPABLE));
                subfer = false;
        }
 
+       if (ar->num_rx_chains < 2) {
+               *vht_cap &= ~(IEEE80211_VHT_CAP_SU_BEAMFORMEE_CAPABLE);
+               subfee = false;
+       }
+
        /* If SU Beaformer is not set, then disable MU Beamformer Capability */
        if (!subfer)
                *vht_cap &= ~(IEEE80211_VHT_CAP_MU_BEAMFORMER_CAPABLE);
        sound_dim >>= IEEE80211_VHT_CAP_SOUNDING_DIMENSIONS_SHIFT;
        *vht_cap &= ~IEEE80211_VHT_CAP_SOUNDING_DIMENSIONS_MASK;
 
-       /* TODO: Need to check invalid STS and Sound_dim values set by FW? */
+       nsts = (*vht_cap & IEEE80211_VHT_CAP_BEAMFORMEE_STS_MASK);
+       nsts >>= IEEE80211_VHT_CAP_BEAMFORMEE_STS_SHIFT;
+       *vht_cap &= ~IEEE80211_VHT_CAP_BEAMFORMEE_STS_MASK;
 
        /* Enable Sounding Dimension Field only if SU BF is enabled */
        if (subfer) {
                *vht_cap |= sound_dim;
        }
 
-       /* Use the STS advertised by FW unless SU Beamformee is not supported*/
-       if (!subfee)
-               *vht_cap &= ~(IEEE80211_VHT_CAP_BEAMFORMEE_STS_MASK);
+       /* Enable Beamformee STS Field only if SU BF is enabled */
+       if (subfee) {
+               if (nsts > (ar->num_rx_chains - 1))
+                       nsts = ar->num_rx_chains - 1;
+
+               nsts <<= IEEE80211_VHT_CAP_BEAMFORMEE_STS_SHIFT;
+               nsts &=  IEEE80211_VHT_CAP_BEAMFORMEE_STS_MASK;
+               *vht_cap |= nsts;
+       }
 }
 
 static struct ieee80211_sta_vht_cap