diff --git a/evm/src/generation/mod.rs b/evm/src/generation/mod.rs
index 5803d12c..5b0b3c8f 100644
--- a/evm/src/generation/mod.rs
+++ b/evm/src/generation/mod.rs
@@ -5,6 +5,7 @@ use plonky2::field::types::Field;
 use plonky2::hash::hash_types::RichField;
 
 use crate::all_stark::{AllStark, NUM_TABLES};
+use crate::config::StarkConfig;
 use crate::cpu::bootstrap_kernel::generate_bootstrap_kernel;
 use crate::cpu::columns::NUM_CPU_COLUMNS;
 use crate::cpu::kernel::global_metadata::GlobalMetadata;
@@ -45,6 +46,7 @@ pub struct GenerationInputs {
 pub(crate) fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
     all_stark: &AllStark<F, D>,
     inputs: GenerationInputs,
+    config: &StarkConfig,
 ) -> ([Vec<PolynomialValues<F>>; NUM_TABLES], PublicValues) {
     let mut state = GenerationState::<F>::default();
 
@@ -83,6 +85,7 @@ pub(crate) fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
         current_cpu_row,
         memory,
         keccak_inputs,
+        keccak_memory_inputs,
         logic_ops,
         ..
     } = state;
@@ -90,7 +93,9 @@ pub(crate) fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
 
     let cpu_trace = trace_rows_to_poly_values(cpu_rows);
     let keccak_trace = all_stark.keccak_stark.generate_trace(keccak_inputs);
-    let keccak_memory_trace = all_stark.keccak_memory_stark.generate_trace(vec![], 0); // TODO
+    let keccak_memory_trace = all_stark
+        .keccak_memory_stark
+        .generate_trace(keccak_memory_inputs, 1 << config.fri_config.cap_height);
     let logic_trace = all_stark.logic_stark.generate_trace(logic_ops);
     let memory_trace = all_stark.memory_stark.generate_trace(memory.log);
     let traces = [
diff --git a/evm/src/prover.rs b/evm/src/prover.rs
index 5a23297a..31e76a1c 100644
--- a/evm/src/prover.rs
+++ b/evm/src/prover.rs
@@ -53,7 +53,7 @@ where
     [(); LogicStark::<F, D>::COLUMNS]:,
     [(); MemoryStark::<F, D>::COLUMNS]:,
 {
-    let (traces, public_values) = generate_traces(all_stark, inputs);
+    let (traces, public_values) = generate_traces(all_stark, inputs, config);
     prove_with_traces(all_stark, config, traces, public_values, timing)
 }