diff --git a/include/psa/crypto_sizes.h b/include/psa/crypto_sizes.h
index c9de062..10151c6 100644
--- a/include/psa/crypto_sizes.h
+++ b/include/psa/crypto_sizes.h
@@ -129,9 +129,9 @@
  *                            the ciphertext, return 0.
  *                            If the AEAD algorithm is not recognized, return 0.
  */
-#define PSA_AEAD_TAG_LENGTH(alg)                                        \
-    (PSA_ALG_IS_AEAD(alg) ?                                             \
-     (((alg) & PSA_ALG_AEAD_TAG_LENGTH_MASK) >> PSA_AEAD_TAG_LENGTH_OFFSET) : \
+#define PSA_AEAD_TAG_LENGTH(key_type, key_bits, alg)                        \
+    (PSA_AEAD_NONCE_LENGTH(key_type, alg) ?                                 \
+     ((alg) & PSA_ALG_AEAD_TAG_LENGTH_MASK) >> PSA_AEAD_TAG_LENGTH_OFFSET : \
      0)
 
 /** The maximum tag size for all supported AEAD algorithms, in bytes.
@@ -254,9 +254,9 @@
  *                            algorithm.
  *                            If the AEAD algorithm is not recognized, return 0.
  */
-#define PSA_AEAD_ENCRYPT_OUTPUT_SIZE(alg, plaintext_length)       \
-    (PSA_AEAD_TAG_LENGTH(alg) != 0 ?                              \
-     (plaintext_length) + PSA_AEAD_TAG_LENGTH(alg) :              \
+#define PSA_AEAD_ENCRYPT_OUTPUT_SIZE(key_type, alg, plaintext_length) \
+    (PSA_AEAD_NONCE_LENGTH(key_type, alg) ?                           \
+     (plaintext_length) + PSA_AEAD_TAG_LENGTH(key_type, 0, alg) :     \
      0)
 
 /** A sufficient output buffer size for psa_aead_encrypt(), for any of the
@@ -300,9 +300,9 @@
  *                            algorithm.
  *                            If the AEAD algorithm is not recognized, return 0.
  */
-#define PSA_AEAD_DECRYPT_OUTPUT_SIZE(alg, ciphertext_length)      \
-    (PSA_AEAD_TAG_LENGTH(alg) != 0 ?                              \
-     (ciphertext_length) - PSA_AEAD_TAG_LENGTH(alg) :             \
+#define PSA_AEAD_DECRYPT_OUTPUT_SIZE(key_type, alg, ciphertext_length) \
+    (PSA_AEAD_NONCE_LENGTH(key_type, alg) ?                            \
+     (ciphertext_length) - PSA_AEAD_TAG_LENGTH(key_type, 0, alg) :     \
      0)
 
 /** A sufficient output buffer size for psa_aead_decrypt(), for any of the
@@ -396,10 +396,12 @@
  * to emit output without delay. However, hardware may not always be
  * capable of this. So for modes based on a block cipher, allow the
  * implementation to delay the output until it has a full block. */
-#define PSA_AEAD_UPDATE_OUTPUT_SIZE(alg, input_length)                              \
-    (PSA_ALG_IS_AEAD_ON_BLOCK_CIPHER(alg) ?                                         \
-     PSA_ROUND_UP_TO_MULTIPLE(PSA_BLOCK_CIPHER_BLOCK_MAX_SIZE, (input_length)) :    \
-     (input_length))
+#define PSA_AEAD_UPDATE_OUTPUT_SIZE(key_type, alg, input_length)                             \
+    (PSA_AEAD_NONCE_LENGTH(key_type, alg) ?                                                  \
+         PSA_ALG_IS_AEAD_ON_BLOCK_CIPHER(alg) ?                                              \
+         PSA_ROUND_UP_TO_MULTIPLE(PSA_BLOCK_CIPHER_BLOCK_LENGTH(key_type), (input_length)) : \
+         (input_length) : \
+     0)
 
 /** A sufficient output buffer size for psa_aead_update(), for any of the
  *  supported key types and AEAD algorithms.
@@ -429,9 +431,9 @@
  *                            specified algorithm.
  *                            If the AEAD algorithm is not recognized, return 0.
  */
-#define PSA_AEAD_FINISH_OUTPUT_SIZE(alg)                                \
-    (PSA_ALG_IS_AEAD_ON_BLOCK_CIPHER(alg) ?                             \
-     PSA_BLOCK_CIPHER_BLOCK_MAX_SIZE :                                  \
+#define PSA_AEAD_FINISH_OUTPUT_SIZE(key_type, alg)                                  \
+    (PSA_AEAD_NONCE_LENGTH(key_type, alg) && PSA_ALG_IS_AEAD_ON_BLOCK_CIPHER(alg) ? \
+     PSA_BLOCK_CIPHER_BLOCK_LENGTH(key_type) :                                      \
      0)
 
 /** A sufficient ciphertext buffer size for psa_aead_finish(), for any of the
@@ -456,9 +458,9 @@
  *                            specified algorithm.
  *                            If the AEAD algorithm is not recognized, return 0.
  */
-#define PSA_AEAD_VERIFY_OUTPUT_SIZE(alg)                                \
-    (PSA_ALG_IS_AEAD_ON_BLOCK_CIPHER(alg) ?                             \
-     PSA_BLOCK_CIPHER_BLOCK_MAX_SIZE :                                  \
+#define PSA_AEAD_VERIFY_OUTPUT_SIZE(key_type, alg)                                  \
+    (PSA_AEAD_NONCE_LENGTH(key_type, alg) && PSA_ALG_IS_AEAD_ON_BLOCK_CIPHER(alg) ? \
+     PSA_BLOCK_CIPHER_BLOCK_LENGTH(key_type) :                                      \
      0)
 
 /** A sufficient plaintext buffer size for psa_aead_verify(), for any of the