With partial channel scanning set, wifi_scan_networks_with_ssid should already be returning after finding the AP, since it is sending a directed probe packet only on a single channel and listening for the response.
Alternatively, you may want to try setting your own scan event callback, like in wifi_scan(), then shutting off WiFi when the SSID is found, but I am not sure if this will be faster.
wifi_reg_event_handler(WIFI_EVENT_SCAN_RESULT_REPORT, wifi_scan_each_report_hdl, NULL);
Ultimately, this forum is more catered to the general/open source side of the code, and I am unable to see exact details of how lower level code operates. I would suggest you consider an NDA, which would get you in contact with a specialized FAE from whom you can get more detailed information on WiFi / BLE operation.