diff --git a/core/wma/src/wma_mgmt.c b/core/wma/src/wma_mgmt.c
index 208d18ac0c9d..cd0270a58bf8 100644
--- a/core/wma/src/wma_mgmt.c
+++ b/core/wma/src/wma_mgmt.c
@@ -4377,19 +4377,19 @@ QDF_STATUS wma_mgmt_unified_cmd_send(struct wlan_objmgr_vdev *vdev,
 void wma_mgmt_nbuf_unmap_cb(struct wlan_objmgr_pdev *pdev,
 			    qdf_nbuf_t buf)
 {
-	tp_wma_handle wma_handle = cds_get_context(QDF_MODULE_ID_WMA);
-
-	if (!wma_handle) {
-		WMA_LOGE("%s: wma handle is NULL", __func__);
-		return;
-	}
+	struct wlan_objmgr_psoc *psoc;
+	qdf_device_t dev;
 
 	if (!buf)
 		return;
 
-	if (wmi_service_enabled(wma_handle->wmi_handle,
-				wmi_service_mgmt_tx_wmi)) {
-		qdf_nbuf_unmap_single(wma_handle->qdf_dev, buf,
-				      QDF_DMA_TO_DEVICE);
+	psoc = wlan_pdev_get_psoc(pdev);
+	if (!psoc) {
+		WMA_LOGE("%s: Psoc handle NULL", __func__);
+		return;
 	}
+
+	dev = wlan_psoc_get_qdf_dev(psoc);
+	if (wlan_psoc_nif_fw_ext_cap_get(psoc, WLAN_SOC_CEXT_WMI_MGMT_REF))
+		qdf_nbuf_unmap_single(dev, buf, QDF_DMA_TO_DEVICE);
 }