diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml
index e1af2b19d4..e96df1aaef 100644
--- a/.github/workflows/rust.yml
+++ b/.github/workflows/rust.yml
@@ -313,6 +313,7 @@ jobs:
         path: |
           cairo_programs/**/*.memory
           cairo_programs/**/*.trace
+          cairo_programs/**/*.air_public_input
         key: ${{ matrix.program-target }}-reference-trace-cache-${{ hashFiles('cairo_programs/**/*.cairo', 'examples/wasm-demo/src/array_sum.cairo') }}
         restore-keys: ${{ matrix.program-target }}-reference-trace-cache-
 
@@ -352,7 +353,7 @@ jobs:
         include:
         - program-target: cairo_proof_programs
           programs-dir: cairo_programs/proof_programs
-          extra-args: '--proof_mode'
+          extra-args: '--proof_mode --air_public_input {program}.rs.air_public_input '
         - program-target: cairo_test_programs
           programs-dir: cairo_programs
           extra-args: ''
@@ -390,6 +391,7 @@ jobs:
         path: |
           cairo_programs/**/*.memory
           cairo_programs/**/*.trace
+          cairo_programs/**/*.air_public_input
         key: ${{ matrix.program-target }}-release-trace-cache-${{ github.sha }}
 
 
@@ -464,6 +466,7 @@ jobs:
         path: |
           cairo_programs/**/*.memory
           cairo_programs/**/*.trace
+          cairo_programs/**/*.air_public_input
         key: ${{ matrix.program-target }}-reference-trace-cache-${{ hashFiles('cairo_programs/**/*.cairo', 'examples/wasm-demo/src/array_sum.cairo') }}
         fail-on-cache-miss: true
 
@@ -473,12 +476,14 @@ jobs:
         path: |
           cairo_programs/**/*.memory
           cairo_programs/**/*.trace
+          cairo_programs/**/*.air_public_input
         key: ${{ matrix.program-target }}-release-trace-cache-${{ github.sha }}
         fail-on-cache-miss: true
 
     - name: Run comparison script
       run: |
         if [ ${{ matrix.program-target }} = cairo_proof_programs ]; then
-          PROOF=proof_mode
+          PROOF=proof_mode 
+          AIR_PUBLIC_INPUT=air_public_input
         fi
-        ./vm/src/tests/compare_vm_state.sh trace memory $PROOF
+        ./vm/src/tests/compare_vm_state.sh trace memory $PROOF $AIR_PUBLIC_INPUT
diff --git a/Makefile b/Makefile
index 2f4258753a..926a75cf2d 100644
--- a/Makefile
+++ b/Makefile
@@ -225,7 +225,7 @@ cairo_bench_programs: $(COMPILED_BENCHES)
 cairo_1_test_contracts: $(CAIRO_1_COMPILED_CASM_CONTRACTS)
 cairo_2_test_contracts: $(CAIRO_2_COMPILED_CASM_CONTRACTS)
 
-cairo_proof_trace: $(CAIRO_TRACE_PROOF) $(CAIRO_MEM_PROOF)
+cairo_proof_trace: $(CAIRO_TRACE_PROOF) $(CAIRO_MEM_PROOF) $(CAIRO_AIR_PUBLIC_INPUT)
 cairo-vm_proof_trace: $(CAIRO_RS_TRACE_PROOF) $(CAIRO_RS_MEM_PROOF) $(CAIRO_RS_AIR_PUBLIC_INPUT)
 
 cairo_trace: $(CAIRO_TRACE) $(CAIRO_MEM)