diff --git a/ceno_zkvm/examples/fibonacci_elf.rs b/ceno_zkvm/examples/fibonacci_elf.rs
index f03696eef..d8e20ee2c 100644
--- a/ceno_zkvm/examples/fibonacci_elf.rs
+++ b/ceno_zkvm/examples/fibonacci_elf.rs
@@ -96,8 +96,13 @@ fn main() {
     };
 
     let reg_init = initial_registers();
-    config.generate_fixed_traces(&zkvm_cs, &mut zkvm_fixed_traces, &reg_init);
-    mmu_config.generate_fixed_traces(&zkvm_cs, &mut zkvm_fixed_traces, &program_data_init);
+    config.generate_fixed_traces(&zkvm_cs, &mut zkvm_fixed_traces);
+    mmu_config.generate_fixed_traces(
+        &zkvm_cs,
+        &mut zkvm_fixed_traces,
+        &reg_init,
+        &program_data_init,
+    );
     dummy_config.generate_fixed_traces(&zkvm_cs, &mut zkvm_fixed_traces);
 
     let pk = zkvm_cs
@@ -218,10 +223,16 @@ fn main() {
 
     // assign table circuits
     config
-        .assign_table_circuit(&zkvm_cs, &mut zkvm_witness, &reg_final, &[])
+        .assign_table_circuit(&zkvm_cs, &mut zkvm_witness)
         .unwrap();
     mmu_config
-        .assign_table_circuit(&zkvm_cs, &mut zkvm_witness, &program_data_final)
+        .assign_table_circuit(
+            &zkvm_cs,
+            &mut zkvm_witness,
+            &reg_final,
+            &program_data_final,
+            &[],
+        )
         .unwrap();
     // assign program circuit
     zkvm_witness
diff --git a/ceno_zkvm/examples/riscv_opcodes.rs b/ceno_zkvm/examples/riscv_opcodes.rs
index 8e3703984..89cfaa8d0 100644
--- a/ceno_zkvm/examples/riscv_opcodes.rs
+++ b/ceno_zkvm/examples/riscv_opcodes.rs
@@ -2,7 +2,7 @@ use std::{panic, time::Instant};
 
 use ceno_zkvm::{
     declare_program,
-    instructions::riscv::{Rv32imConfig, constants::EXIT_PC},
+    instructions::riscv::{MmuConfig, Rv32imConfig, constants::EXIT_PC},
     scheme::{mock_prover::MockProver, prover::ZKVMProver},
     state::GlobalState,
     tables::{MemFinalRecord, ProgramTableCircuit, init_public_io, initial_registers},
@@ -106,6 +106,7 @@ fn main() {
     let mut zkvm_cs = ZKVMConstraintSystem::default();
 
     let config = Rv32imConfig::<E>::construct_circuits(&mut zkvm_cs);
+    let mmu_config = MmuConfig::<E>::construct_circuits(&mut zkvm_cs);
     let prog_config = zkvm_cs.register_table_circuit::<ExampleProgramTableCircuit<E>>();
     zkvm_cs.register_global_state::<GlobalState>();
 
@@ -119,7 +120,8 @@ fn main() {
 
     let reg_init = initial_registers();
 
-    config.generate_fixed_traces(&zkvm_cs, &mut zkvm_fixed_traces, &reg_init);
+    config.generate_fixed_traces(&zkvm_cs, &mut zkvm_fixed_traces);
+    mmu_config.generate_fixed_traces(&zkvm_cs, &mut zkvm_fixed_traces, &reg_init, &[]);
 
     let pk = zkvm_cs
         .clone()
@@ -219,7 +221,16 @@ fn main() {
 
         // assign table circuits
         config
-            .assign_table_circuit(&zkvm_cs, &mut zkvm_witness, &reg_final, &public_io_final)
+            .assign_table_circuit(&zkvm_cs, &mut zkvm_witness)
+            .unwrap();
+        mmu_config
+            .assign_table_circuit(
+                &zkvm_cs,
+                &mut zkvm_witness,
+                &reg_final,
+                &[],
+                &public_io_final,
+            )
             .unwrap();
 
         // assign program circuit
diff --git a/ceno_zkvm/src/instructions/riscv/rv32im.rs b/ceno_zkvm/src/instructions/riscv/rv32im.rs
index 24d9c4fa4..4ba008a09 100644
--- a/ceno_zkvm/src/instructions/riscv/rv32im.rs
+++ b/ceno_zkvm/src/instructions/riscv/rv32im.rs
@@ -20,9 +20,8 @@ use crate::{
     },
     structs::{ZKVMConstraintSystem, ZKVMFixedTraces, ZKVMWitnesses},
     tables::{
-        AndTableCircuit, LtuTableCircuit, MemFinalRecord, MemInitRecord, OrTableCircuit,
-        PowTableCircuit, PubIOCircuit, RegTableCircuit, TableCircuit, U5TableCircuit,
-        U8TableCircuit, U14TableCircuit, U16TableCircuit, XorTableCircuit,
+        AndTableCircuit, LtuTableCircuit, OrTableCircuit, PowTableCircuit, TableCircuit,
+        U5TableCircuit, U8TableCircuit, U14TableCircuit, U16TableCircuit, XorTableCircuit,
     },
 };
 use ceno_emul::{CENO_PLATFORM, InsnKind, InsnKind::*, StepRecord};
@@ -112,10 +111,6 @@ pub struct Rv32imConfig<E: ExtensionField> {
     pub xor_table_config: <XorTableCircuit<E> as TableCircuit<E>>::TableConfig,
     pub ltu_config: <LtuTableCircuit<E> as TableCircuit<E>>::TableConfig,
     pub pow_config: <PowTableCircuit<E> as TableCircuit<E>>::TableConfig,
-
-    // RW tables.
-    pub reg_config: <RegTableCircuit<E> as TableCircuit<E>>::TableConfig,
-    pub public_io_config: <PubIOCircuit<E> as TableCircuit<E>>::TableConfig,
 }
 
 impl<E: ExtensionField> Rv32imConfig<E> {
@@ -186,12 +181,6 @@ impl<E: ExtensionField> Rv32imConfig<E> {
         let ltu_config = cs.register_table_circuit::<LtuTableCircuit<E>>();
         let pow_config = cs.register_table_circuit::<PowTableCircuit<E>>();
 
-        // RW tables
-        let reg_config = cs.register_table_circuit::<RegTableCircuit<E>>();
-
-        // RO tables
-        let public_io_config = cs.register_table_circuit::<PubIOCircuit<E>>();
-
         Self {
             // alu opcodes
             add_config,
@@ -252,9 +241,6 @@ impl<E: ExtensionField> Rv32imConfig<E> {
             xor_table_config,
             ltu_config,
             pow_config,
-
-            reg_config,
-            public_io_config,
         }
     }
 
@@ -262,7 +248,6 @@ impl<E: ExtensionField> Rv32imConfig<E> {
         &self,
         cs: &ZKVMConstraintSystem<E>,
         fixed: &mut ZKVMFixedTraces<E>,
-        reg_init: &[MemInitRecord],
     ) {
         // alu
         fixed.register_opcode_circuit::<AddInstruction<E>>(cs);
@@ -323,10 +308,6 @@ impl<E: ExtensionField> Rv32imConfig<E> {
         fixed.register_table_circuit::<XorTableCircuit<E>>(cs, &self.xor_table_config, &());
         fixed.register_table_circuit::<LtuTableCircuit<E>>(cs, &self.ltu_config, &());
         fixed.register_table_circuit::<PowTableCircuit<E>>(cs, &self.pow_config, &());
-
-        fixed.register_table_circuit::<RegTableCircuit<E>>(cs, &self.reg_config, reg_init);
-
-        fixed.register_table_circuit::<PubIOCircuit<E>>(cs, &self.public_io_config, &());
     }
 
     pub fn assign_opcode_circuit(
@@ -440,8 +421,6 @@ impl<E: ExtensionField> Rv32imConfig<E> {
         &self,
         cs: &ZKVMConstraintSystem<E>,
         witness: &mut ZKVMWitnesses<E>,
-        reg_final: &[MemFinalRecord],
-        public_io_final: &[MemFinalRecord],
     ) -> Result<(), ZKVMError> {
         witness.assign_table_circuit::<U16TableCircuit<E>>(cs, &self.u16_range_config, &())?;
         witness.assign_table_circuit::<U14TableCircuit<E>>(cs, &self.u14_range_config, &())?;
@@ -453,15 +432,6 @@ impl<E: ExtensionField> Rv32imConfig<E> {
         witness.assign_table_circuit::<LtuTableCircuit<E>>(cs, &self.ltu_config, &())?;
         witness.assign_table_circuit::<PowTableCircuit<E>>(cs, &self.pow_config, &())?;
 
-        // assign register finalization.
-        witness
-            .assign_table_circuit::<RegTableCircuit<E>>(cs, &self.reg_config, reg_final)
-            .unwrap();
-
-        witness
-            .assign_table_circuit::<PubIOCircuit<E>>(cs, &self.public_io_config, public_io_final)
-            .unwrap();
-
         Ok(())
     }
 }
diff --git a/ceno_zkvm/src/instructions/riscv/rv32im/mmu.rs b/ceno_zkvm/src/instructions/riscv/rv32im/mmu.rs
index d98b424a9..2f5d8480d 100644
--- a/ceno_zkvm/src/instructions/riscv/rv32im/mmu.rs
+++ b/ceno_zkvm/src/instructions/riscv/rv32im/mmu.rs
@@ -3,18 +3,30 @@ use ff_ext::ExtensionField;
 use crate::{
     error::ZKVMError,
     structs::{ZKVMConstraintSystem, ZKVMFixedTraces, ZKVMWitnesses},
-    tables::{MemFinalRecord, MemInitRecord, ProgramDataCircuit, TableCircuit},
+    tables::{
+        MemFinalRecord, MemInitRecord, ProgramDataCircuit, PubIOCircuit, RegTableCircuit,
+        TableCircuit,
+    },
 };
 
 pub struct MmuConfig<E: ExtensionField> {
+    pub reg_config: <RegTableCircuit<E> as TableCircuit<E>>::TableConfig,
     pub program_data_config: <ProgramDataCircuit<E> as TableCircuit<E>>::TableConfig,
+    pub public_io_config: <PubIOCircuit<E> as TableCircuit<E>>::TableConfig,
 }
 
 impl<E: ExtensionField> MmuConfig<E> {
     pub fn construct_circuits(cs: &mut ZKVMConstraintSystem<E>) -> Self {
+        let reg_config = cs.register_table_circuit::<RegTableCircuit<E>>();
+
         let program_data_config = cs.register_table_circuit::<ProgramDataCircuit<E>>();
+
+        let public_io_config = cs.register_table_circuit::<PubIOCircuit<E>>();
+
         Self {
+            reg_config,
             program_data_config,
+            public_io_config,
         }
     }
 
@@ -22,25 +34,42 @@ impl<E: ExtensionField> MmuConfig<E> {
         &self,
         cs: &ZKVMConstraintSystem<E>,
         fixed: &mut ZKVMFixedTraces<E>,
+        reg_init: &[MemInitRecord],
         program_data_init: &[MemInitRecord],
     ) {
+        fixed.register_table_circuit::<RegTableCircuit<E>>(cs, &self.reg_config, reg_init);
+
         fixed.register_table_circuit::<ProgramDataCircuit<E>>(
             cs,
             &self.program_data_config,
             program_data_init,
         );
+
+        fixed.register_table_circuit::<PubIOCircuit<E>>(cs, &self.public_io_config, &());
     }
 
     pub fn assign_table_circuit(
         &self,
         cs: &ZKVMConstraintSystem<E>,
         witness: &mut ZKVMWitnesses<E>,
+        reg_final: &[MemFinalRecord],
         program_data_final: &[MemFinalRecord],
+        public_io_final: &[MemFinalRecord],
     ) -> Result<(), ZKVMError> {
+        witness.assign_table_circuit::<RegTableCircuit<E>>(cs, &self.reg_config, reg_final)?;
+
         witness.assign_table_circuit::<ProgramDataCircuit<E>>(
             cs,
             &self.program_data_config,
             program_data_final,
-        )
+        )?;
+
+        witness.assign_table_circuit::<PubIOCircuit<E>>(
+            cs,
+            &self.public_io_config,
+            public_io_final,
+        )?;
+
+        Ok(())
     }
 }