diff --git a/drivers/net/wireless/intel/iwlwifi/mld/key.c b/drivers/net/wireless/intel/iwlwifi/mld/key.c
index 04192c5f07ff..2cb312ebd7a7 100644
--- a/drivers/net/wireless/intel/iwlwifi/mld/key.c
+++ b/drivers/net/wireless/intel/iwlwifi/mld/key.c
@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: GPL-2.0 OR BSD-3-Clause
 /*
- * Copyright (C) 2024 Intel Corporation
+ * Copyright (C) 2024, 2026 Intel Corporation
  */
 #include "key.h"
 #include "iface.h"
@@ -12,7 +12,6 @@ static u32 iwl_mld_get_key_flags(struct iwl_mld *mld,
 				 struct ieee80211_sta *sta,
 				 struct ieee80211_key_conf *key)
 {
-	struct iwl_mld_vif *mld_vif = iwl_mld_vif_from_mac80211(vif);
 	bool pairwise = key->flags & IEEE80211_KEY_FLAG_PAIRWISE;
 	bool igtk = key->keyidx == 4 || key->keyidx == 5;
 	u32 flags = 0;
@@ -38,9 +37,6 @@ static u32 iwl_mld_get_key_flags(struct iwl_mld *mld,
 		break;
 	}
 
-	if (!sta && vif->type == NL80211_IFTYPE_STATION)
-		sta = mld_vif->ap_sta;
-
 	/* If we are installing an iGTK (in AP or STA mode), we need to tell
 	 * the firmware this key will en/decrypt MGMT frames.
 	 * Same goes if we are installing a pairwise key for an MFP station.
@@ -92,11 +88,7 @@ static u32 iwl_mld_get_key_sta_mask(struct iwl_mld *mld,
 		return BIT(link->mcast_sta.sta_id);
 	}
 
-	/* for client mode use the AP STA also for group keys */
-	if (!sta && vif->type == NL80211_IFTYPE_STATION)
-		sta = mld_vif->ap_sta;
-
-	/* STA should be non-NULL now */
+	/* STA should be non-NULL */
 	if (WARN_ON(!sta))
 		return 0;
 
diff --git a/drivers/net/wireless/intel/iwlwifi/mld/mac80211.c b/drivers/net/wireless/intel/iwlwifi/mld/mac80211.c
index c02994e63b14..4b5f50ff2b05 100644
--- a/drivers/net/wireless/intel/iwlwifi/mld/mac80211.c
+++ b/drivers/net/wireless/intel/iwlwifi/mld/mac80211.c
@@ -2282,9 +2282,20 @@ static int iwl_mld_mac80211_set_key(struct ieee80211_hw *hw,
 				    struct ieee80211_sta *sta,
 				    struct ieee80211_key_conf *key)
 {
+	struct iwl_mld_vif *mld_vif = iwl_mld_vif_from_mac80211(vif);
 	struct iwl_mld *mld = IWL_MAC80211_GET_MLD(hw);
 	int ret;
 
+	/*
+	 * FW always needs the AP STA for client mode.
+	 * Note that during removal this could already
+	 * be NULL (mac80211 removes keys after STAs)
+	 * but then we'll already have removed the key
+	 * and set hw_key_idx = STA_KEY_IDX_INVALID.
+	 */
+	if (!sta && vif->type == NL80211_IFTYPE_STATION)
+		sta = mld_vif->ap_sta;
+
 	switch (cmd) {
 	case SET_KEY:
 		ret = iwl_mld_set_key_add(mld, vif, sta, key);