diff --git a/include/psa/crypto_builtin_key_derivation.h b/include/psa/crypto_builtin_key_derivation.h
index cd6d51df02..c598fa438e 100644
--- a/include/psa/crypto_builtin_key_derivation.h
+++ b/include/psa/crypto_builtin_key_derivation.h
@@ -105,7 +105,8 @@ typedef struct psa_tls12_prf_key_derivation_s {
 } psa_tls12_prf_key_derivation_t;
 #endif /* MBEDTLS_PSA_BUILTIN_ALG_TLS12_PRF) ||
         * MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS */
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_PBKDF2_HMAC)
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_PBKDF2_HMAC) || \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_PBKDF2_AES_CMAC_PRF_128)
 typedef enum {
     PSA_PBKDF2_STATE_INIT,             /* no input provided */
     PSA_PBKDF2_STATE_INPUT_COST_SET,   /* input cost has been set */
@@ -125,6 +126,7 @@ typedef struct {
     uint8_t MBEDTLS_PRIVATE(bytes_used);
     uint32_t MBEDTLS_PRIVATE(block_number);
 } psa_pbkdf2_key_derivation_t;
-#endif /* MBEDTLS_PSA_BUILTIN_ALG_PBKDF2_HMAC */
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_PBKDF2_HMAC ||
+        * MBEDTLS_PSA_BUILTIN_ALG_PBKDF2_AES_CMAC_PRF_128 */
 
 #endif /* PSA_CRYPTO_BUILTIN_KEY_DERIVATION_H */
diff --git a/include/psa/crypto_driver_contexts_key_derivation.h b/include/psa/crypto_driver_contexts_key_derivation.h
index 5b4e4745d6..32de4f7654 100644
--- a/include/psa/crypto_driver_contexts_key_derivation.h
+++ b/include/psa/crypto_driver_contexts_key_derivation.h
@@ -55,7 +55,8 @@ typedef union {
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS)
     psa_tls12_ecjpake_to_pms_t MBEDTLS_PRIVATE(tls12_ecjpake_to_pms);
 #endif
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_PBKDF2_HMAC)
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_PBKDF2_HMAC) || \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_PBKDF2_AES_CMAC_PRF_128)
     psa_pbkdf2_key_derivation_t MBEDTLS_PRIVATE(pbkdf2);
 #endif
 } psa_driver_key_derivation_context_t;