__le32 beacon_template;
 } __packed; /* IBSS_MAC_DATA_API_S_VER_1 */
 
+/**
+ * enum iwl_mac_data_policy - policy of the data path for this MAC
+ * @TWT_SUPPORTED: twt is supported
+ */
+enum iwl_mac_data_policy {
+       TWT_SUPPORTED   = BIT(0),
+};
+
 /**
  * struct iwl_mac_data_sta - configuration data for station MAC context
  * @is_assoc: 1 for associated state, 0 otherwise
  * @bi: beacon interval in TU, applicable only when associated
  * @reserved1: reserved
  * @dtim_interval: DTIM interval in TU, applicable only when associated
- * @reserved2: reserved
+ * @data_policy: see &enum iwl_mac_data_policy
  * @listen_interval: in beacon intervals, applicable only when associated
  * @assoc_id: unique ID assigned by the AP during association
  * @assoc_beacon_arrive_time: TSF of first beacon after association
        __le32 bi;
        __le32 reserved1;
        __le32 dtim_interval;
-       __le32 reserved2;
+       __le32 data_policy;
        __le32 listen_interval;
        __le32 assoc_id;
        __le32 assoc_beacon_arrive_time;
-} __packed; /* STA_MAC_DATA_API_S_VER_1 */
+} __packed; /* STA_MAC_DATA_API_S_VER_2 */
 
 /**
  * struct iwl_mac_data_go - configuration data for P2P GO MAC context
 struct iwl_mac_data_p2p_sta {
        struct iwl_mac_data_sta sta;
        __le32 ctwin;
-} __packed; /* P2P_STA_MAC_DATA_API_S_VER_1 */
+} __packed; /* P2P_STA_MAC_DATA_API_S_VER_2 */
 
 /**
  * struct iwl_mac_data_pibss - Pseudo IBSS config data
 
                cmd.filter_flags |= cpu_to_le32(MAC_FILTER_IN_PROBE_REQUEST);
 
        if (vif->bss_conf.assoc && vif->bss_conf.he_support &&
-           !iwlwifi_mod_params.disable_11ax)
+           !iwlwifi_mod_params.disable_11ax) {
+               struct iwl_mvm_vif *mvmvif = iwl_mvm_vif_from_mac80211(vif);
+               u8 sta_id = mvmvif->ap_sta_id;
+
                cmd.filter_flags |= cpu_to_le32(MAC_FILTER_IN_11AX);
+               if (sta_id != IWL_MVM_INVALID_STA) {
+                       struct ieee80211_sta *sta;
+
+                       sta = rcu_dereference_protected(mvm->fw_id_to_mac_id[sta_id],
+                               lockdep_is_held(&mvm->mutex));
+
+                       /*
+                        * TODO: we should check the ext cap IE but it is
+                        * unclear why the spec requires two bits (one in HE
+                        * cap IE, and one in the ext cap IE). In the meantime
+                        * rely on the HE cap IE only.
+                        */
+                       if (sta && (sta->he_cap.he_cap_elem.mac_cap_info[0] &
+                                   IEEE80211_HE_MAC_CAP0_TWT_RES))
+                               ctxt_sta->data_policy |=
+                                       cpu_to_le32(TWT_SUPPORTED);
+               }
+       }
+
 
        return iwl_mvm_mac_ctxt_send_cmd(mvm, &cmd);
 }