diff --git a/drivers/net/wireless/intel/iwlwifi/fw/api/datapath.h b/drivers/net/wireless/intel/iwlwifi/fw/api/datapath.h
index d268a9cba5ce..43619acc29fd 100644
--- a/drivers/net/wireless/intel/iwlwifi/fw/api/datapath.h
+++ b/drivers/net/wireless/intel/iwlwifi/fw/api/datapath.h
@@ -304,21 +304,33 @@ struct iwl_rx_baid_cfg_cmd_alloc {
 
 /**
  * struct iwl_rx_baid_cfg_cmd_modify - BAID modification data
- * @sta_id_mask: station ID mask
- * @baid: the BAID to modify
+ * @old_sta_id_mask: old station ID mask
+ * @new_sta_id_mask: new station ID mask
+ * @tid: TID of the BAID
  */
 struct iwl_rx_baid_cfg_cmd_modify {
-	__le32 sta_id_mask;
+	__le32 old_sta_id_mask;
+	__le32 new_sta_id_mask;
+	__le32 tid;
+} __packed; /* RX_BAID_ALLOCATION_MODIFY_CMD_API_S_VER_2 */
+
+/**
+ * struct iwl_rx_baid_cfg_cmd_remove_v1 - BAID removal data
+ * @baid: the BAID to remove
+ */
+struct iwl_rx_baid_cfg_cmd_remove_v1 {
 	__le32 baid;
-} __packed; /* RX_BAID_ALLOCATION_MODIFY_CMD_API_S_VER_1 */
+} __packed; /* RX_BAID_ALLOCATION_REMOVE_CMD_API_S_VER_1 */
 
 /**
  * struct iwl_rx_baid_cfg_cmd_remove - BAID removal data
- * @baid: the BAID to remove
+ * @sta_id_mask: the station mask of the BAID to remove
+ * @tid: the TID of the BAID to remove
  */
 struct iwl_rx_baid_cfg_cmd_remove {
-	__le32 baid;
-} __packed; /* RX_BAID_ALLOCATION_REMOVE_CMD_API_S_VER_1 */
+	__le32 sta_id_mask;
+	__le32 tid;
+} __packed; /* RX_BAID_ALLOCATION_REMOVE_CMD_API_S_VER_2 */
 
 /**
  * struct iwl_rx_baid_cfg_cmd - BAID allocation/config command
@@ -329,9 +341,10 @@ struct iwl_rx_baid_cfg_cmd {
 	union {
 		struct iwl_rx_baid_cfg_cmd_alloc alloc;
 		struct iwl_rx_baid_cfg_cmd_modify modify;
+		struct iwl_rx_baid_cfg_cmd_remove_v1 remove_v1;
 		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 */
+	}; /* RX_BAID_ALLOCATION_OPERATION_API_U_VER_2 */
+} __packed; /* RX_BAID_ALLOCATION_CONFIG_CMD_API_S_VER_2 */
 
 /**
  * struct iwl_rx_baid_cfg_resp - BAID allocation response
diff --git a/drivers/net/wireless/intel/iwlwifi/mvm/sta.c b/drivers/net/wireless/intel/iwlwifi/mvm/sta.c
index 9932ad342218..c7f9d3870f21 100644
--- a/drivers/net/wireless/intel/iwlwifi/mvm/sta.c
+++ b/drivers/net/wireless/intel/iwlwifi/mvm/sta.c
@@ -2660,8 +2660,12 @@ static int iwl_mvm_fw_baid_op_cmd(struct iwl_mvm *mvm,
 		cmd.alloc.ssn = cpu_to_le16(ssn);
 		cmd.alloc.win_size = cpu_to_le16(buf_size);
 		baid = -EIO;
+	} else if (iwl_fw_lookup_cmd_ver(mvm->fw, cmd_id, 1) == 1) {
+		cmd.remove_v1.baid = cpu_to_le32(baid);
+		BUILD_BUG_ON(sizeof(cmd.remove_v1) > sizeof(cmd.remove));
 	} else {
-		cmd.remove.baid = cpu_to_le32(baid);
+		cmd.remove.sta_id_mask = cpu_to_le32(BIT(mvm_sta->sta_id));
+		cmd.remove.tid = cpu_to_le32(tid);
 	}
 
 	ret = iwl_mvm_send_cmd_pdu_status(mvm, cmd_id, sizeof(cmd),