*/
        CHEST_COLLECTOR_FILTER_CONFIG_CMD = 0x14,
 
+       /**
+        * @RX_BAID_ALLOCATION_CONFIG_CMD: Allocate/deallocate a BAID for an RX
+        *      blockack session, uses &struct iwl_rx_baid_cfg_cmd for the
+        *      command, and &struct iwl_rx_baid_cfg_resp as a response.
+        */
+       RX_BAID_ALLOCATION_CONFIG_CMD = 0x16,
+
        /**
         * @MONITOR_NOTIF: Datapath monitoring notification, using
         *      &struct iwl_datapath_monitor_notif
        u8 reserved[3];
 } __packed; /* RLC_CONFIG_CMD_API_S_VER_2 */
 
+/**
+ * enum iwl_rx_baid_action - BAID allocation/config action
+ * @IWL_RX_BAID_ACTION_ADD: add a new BAID session
+ * @IWL_RX_BAID_ACTION_MODIFY: modify the BAID session
+ * @IWL_RX_BAID_ACTION_REMOVE: remove the BAID session
+ */
+enum iwl_rx_baid_action {
+       IWL_RX_BAID_ACTION_ADD,
+       IWL_RX_BAID_ACTION_MODIFY,
+       IWL_RX_BAID_ACTION_REMOVE,
+}; /*  RX_BAID_ALLOCATION_ACTION_E_VER_1 */
+
+/**
+ * struct iwl_rx_baid_cfg_cmd_alloc - BAID allocation data
+ * @sta_id_mask: station ID mask
+ * @tid: the TID for this session
+ * @reserved: reserved
+ * @ssn: the starting sequence number
+ * @win_size: RX BA session window size
+ */
+struct iwl_rx_baid_cfg_cmd_alloc {
+       __le32 sta_id_mask;
+       u8 tid;
+       u8 reserved[3];
+       __le16 ssn;
+       __le16 win_size;
+} __packed; /* RX_BAID_ALLOCATION_ADD_CMD_API_S_VER_1 */
+
+/**
+ * struct iwl_rx_baid_cfg_cmd_modify - BAID modification data
+ * @sta_id_mask: station ID mask
+ * @baid: the BAID to modify
+ */
+struct iwl_rx_baid_cfg_cmd_modify {
+       __le32 sta_id_mask;
+       __le32 baid;
+} __packed; /* RX_BAID_ALLOCATION_MODIFY_CMD_API_S_VER_1 */
+
+/**
+ * struct iwl_rx_baid_cfg_cmd_remove - BAID removal data
+ * @baid: the BAID to remove
+ */
+struct iwl_rx_baid_cfg_cmd_remove {
+       __le32 baid;
+} __packed; /* RX_BAID_ALLOCATION_REMOVE_CMD_API_S_VER_1 */
+
+/**
+ * struct iwl_rx_baid_cfg_cmd - BAID allocation/config command
+ * @action: the action, from &enum iwl_rx_baid_action
+ */
+struct iwl_rx_baid_cfg_cmd {
+       __le32 action;
+       union {
+               struct iwl_rx_baid_cfg_cmd_alloc alloc;
+               struct iwl_rx_baid_cfg_cmd_modify modify;
+               struct iwl_rx_baid_cfg_cmd_remove remove;
+       }; /* RX_BAID_ALLOCATION_OPERATION_API_U_VER_1 */
+} __packed; /* RX_BAID_ALLOCATION_CONFIG_CMD_API_S_VER_1 */
+
+/**
+ * struct iwl_rx_baid_cfg_resp - BAID allocation response
+ * @baid: the allocated BAID
+ */
+struct iwl_rx_baid_cfg_resp {
+       __le32 baid;
+}; /* RX_BAID_ALLOCATION_RESPONSE_API_S_VER_1 */
+
 #endif /* __iwl_fw_api_datapath_h__ */
 
        }
 }
 
-static int iwl_mvm_fw_baid_op(struct iwl_mvm *mvm, struct iwl_mvm_sta *mvm_sta,
-                             bool start, int tid, u16 ssn, u16 buf_size)
+static int iwl_mvm_fw_baid_op_sta(struct iwl_mvm *mvm,
+                                 struct iwl_mvm_sta *mvm_sta,
+                                 bool start, int tid, u16 ssn,
+                                 u16 buf_size)
 {
        struct iwl_mvm_add_sta_cmd cmd = {
                .mac_id_n_color = cpu_to_le32(mvm_sta->mac_id_n_color),
        }
 }
 
+static int iwl_mvm_fw_baid_op_cmd(struct iwl_mvm *mvm,
+                                 struct iwl_mvm_sta *mvm_sta,
+                                 bool start, int tid, u16 ssn,
+                                 u16 buf_size, int baid)
+{
+       struct iwl_rx_baid_cfg_cmd cmd = {
+               .action = start ? cpu_to_le32(IWL_RX_BAID_ACTION_ADD) :
+                                 cpu_to_le32(IWL_RX_BAID_ACTION_REMOVE),
+       };
+       u32 cmd_id = WIDE_ID(DATA_PATH_GROUP, RX_BAID_ALLOCATION_CONFIG_CMD);
+       int ret;
+
+       BUILD_BUG_ON(sizeof(struct iwl_rx_baid_cfg_resp) != sizeof(baid));
+
+       if (start) {
+               cmd.alloc.sta_id_mask = cpu_to_le32(BIT(mvm_sta->sta_id));
+               cmd.alloc.tid = tid;
+               cmd.alloc.ssn = cpu_to_le16(ssn);
+               cmd.alloc.win_size = cpu_to_le16(buf_size);
+               baid = -EIO;
+       } else {
+               cmd.remove.baid = cpu_to_le32(baid);
+       }
+
+       ret = iwl_mvm_send_cmd_pdu_status(mvm, cmd_id, sizeof(cmd),
+                                         &cmd, &baid);
+       if (ret)
+               return ret;
+
+       if (!start) {
+               /* ignore firmware baid on remove */
+               baid = 0;
+       }
+
+       IWL_DEBUG_HT(mvm, "RX BA Session %sed in fw\n",
+                    start ? "start" : "stopp");
+
+       if (baid < 0 || baid >= ARRAY_SIZE(mvm->baid_map))
+               return -EINVAL;
+
+       return baid;
+}
+
+static int iwl_mvm_fw_baid_op(struct iwl_mvm *mvm, struct iwl_mvm_sta *mvm_sta,
+                             bool start, int tid, u16 ssn, u16 buf_size,
+                             int baid)
+{
+       if (fw_has_capa(&mvm->fw->ucode_capa,
+                       IWL_UCODE_TLV_CAPA_BAID_ML_SUPPORT))
+               return iwl_mvm_fw_baid_op_cmd(mvm, mvm_sta, start,
+                                             tid, ssn, buf_size, baid);
+
+       return iwl_mvm_fw_baid_op_sta(mvm, mvm_sta, start,
+                                     tid, ssn, buf_size);
+}
+
 int iwl_mvm_sta_rx_agg(struct iwl_mvm *mvm, struct ieee80211_sta *sta,
                       int tid, u16 ssn, bool start, u16 buf_size, u16 timeout)
 {
                        reorder_buf_size / sizeof(baid_data->entries[0]);
        }
 
-       baid = iwl_mvm_fw_baid_op(mvm, mvm_sta, start, tid, ssn, buf_size);
+       if (iwl_mvm_has_new_rx_api(mvm) && !start) {
+               baid = mvm_sta->tid_to_baid[tid];
+       } else {
+               /* we don't really need it in this case */
+               baid = -1;
+       }
+
+       baid = iwl_mvm_fw_baid_op(mvm, mvm_sta, start, tid, ssn, buf_size,
+                                 baid);
        if (baid < 0) {
                ret = baid;
                goto out_free;