Complete Implementation of ML-KEM (FIPS 203)

Every code listing from this chapter of Applied Post-Quantum Cryptography — 7 in total, 7 runnable here. Edit any cell and press Run.

The book's snippets build on each other down the chapter, but a Sage Cell kernel runs one cell and keeps no state afterwards, so each cell replays the earlier listings with apqc_book. That call is the only thing added to the book's own code.

← the playground

Listing 1 — SageMath experiment

For a teaching implementation, we can first simulate the expansion by using a simple deterministic pseudo-random generator.

Listing 2 — Experiment 1: verify ExpandA

Use the same seed twice:

Listing 3 — Experiment 2: inspect CBD statistics

Generate many samples from the CBD distribution and check that the average is close to zero.

Listing 4 — Experiment 4: one MultiplyNTTs base product

The NTT-domain multiplication in FIPS 203 is not 256 independent scalar products. Each coordinate pair is multiplied modulo a quadratic relation. The following check implements one such pair; it tests the algebraic base case, not the full standardized transform or its serialization.

Listing 5 — A complete SageMath implementation

The experiments above illustrate individual pieces. The companion file sage/fips203_mlkem.sage assembles all of them into a complete, byte-exact implementation of FIPS 203: every one of the standard's twenty-one numbered algorithms appears as its own function, named after the standard and annotated with its algorithm number, for all three parameter sets.

The algebraic objects are genuine Sage objects rather than opaque byte arrays. The coefficient domain is the quotient ring itself,

and the NTT domain T_q is carried as a Sage vector over 𝔽_q:

Listing 6 — A complete SageMath implementation

With that in place, the transform is the standard's butterfly network written directly over 𝔽_q:

Listing 7 — A complete SageMath implementation

Key generation is then the seven phases of this chapter, line for line: