diff --git a/ml-kem/src/kem.rs b/ml-kem/src/kem.rs
index 229bcfb..2e18806 100644
--- a/ml-kem/src/kem.rs
+++ b/ml-kem/src/kem.rs
@@ -14,7 +14,7 @@ use crate::{Encoded, EncodedSizeUser, Seed};
 use zeroize::{Zeroize, ZeroizeOnDrop};
 
 // Re-export traits from the `kem` crate
-pub use ::kem::{Decapsulate, Encapsulate};
+pub use ::kem::{Decapsulate, Encapsulate, KeyInit, KeySizeUser};
 
 /// A shared key resulting from an ML-KEM transaction
 pub(crate) type SharedKey = B32;
diff --git a/ml-kem/src/lib.rs b/ml-kem/src/lib.rs
index 7100681..3e4d59f 100644
--- a/ml-kem/src/lib.rs
+++ b/ml-kem/src/lib.rs
@@ -21,12 +21,18 @@
 //! computers.
 //!
 //! ```
-//! # use ml_kem::*;
-//! # use ::kem::{Decapsulate, Encapsulate};
-//! let mut rng = rand::rng();
+//! use ml_kem::{
+//!     MlKem768Params,
+//!     kem::{Decapsulate, Encapsulate, KeyInit}
+//! };
+//!
+//! type DecapsulationKey = ml_kem::kem::DecapsulationKey<MlKem768Params>;
 //!
 //! // Generate a (decapsulation key, encapsulation key) pair
-//! let (dk, ek) = MlKem768::generate(&mut rng);
+//! let mut rng = rand::rng();
+//! let seed = DecapsulationKey::generate_key_with_rng(&mut rng);
+//! let dk = DecapsulationKey::new(&seed);
+//! let ek = dk.encapsulator();
 //!
 //! // Encapsulate a shared key to the holder of the decapsulation key, receive the shared
 //! // secret `k_send` and the encapsulated form `ct`.