diff --git a/evm/src/all_stark.rs b/evm/src/all_stark.rs
index c3c9653c..eb1d7087 100644
--- a/evm/src/all_stark.rs
+++ b/evm/src/all_stark.rs
@@ -87,11 +87,11 @@ mod tests {
         let keccak_stark = KeccakStark::<F, D> {
             f: Default::default(),
         };
-        let keccak_rows = (NUM_ROUNDS + 1).next_power_of_two();
+        let keccak_rows = (2 * NUM_ROUNDS + 1).next_power_of_two();
         let keccak_looked_col = 3;
 
         let mut rng = ChaCha8Rng::seed_from_u64(0x6feb51b7ec230f25);
-        let num_inputs = 1;
+        let num_inputs = 2;
         let keccak_inputs = (0..num_inputs)
             .map(|_| [0u64; INPUT_LIMBS].map(|_| rng.gen()))
             .collect_vec();
diff --git a/evm/src/keccak/keccak_stark.rs b/evm/src/keccak/keccak_stark.rs
index cbc2408b..771513a6 100644
--- a/evm/src/keccak/keccak_stark.rs
+++ b/evm/src/keccak/keccak_stark.rs
@@ -345,8 +345,13 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
                 let input_bits_combined_hi = (32..64)
                     .rev()
                     .fold(P::ZEROS, |acc, z| acc.doubles() + input_bits[z]);
-                yield_constr.constraint_transition(output_lo - input_bits_combined_lo);
-                yield_constr.constraint_transition(output_hi - input_bits_combined_hi);
+                let is_last_round = vars.local_values[reg_step(NUM_ROUNDS - 1)];
+                yield_constr.constraint_transition(
+                    (P::ONES - is_last_round) * (output_lo - input_bits_combined_lo),
+                );
+                yield_constr.constraint_transition(
+                    (P::ONES - is_last_round) * (output_hi - input_bits_combined_hi),
+                );
             }
         }
     }
@@ -492,10 +497,13 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
                     reduce_with_powers_ext_circuit(builder, &input_bits[0..32], two);
                 let input_bits_combined_hi =
                     reduce_with_powers_ext_circuit(builder, &input_bits[32..64], two);
-                let diff = builder.sub_extension(output_lo, input_bits_combined_lo);
-                yield_constr.constraint_transition(builder, diff);
-                let diff = builder.sub_extension(output_hi, input_bits_combined_hi);
-                yield_constr.constraint_transition(builder, diff);
+                let is_last_round = vars.local_values[reg_step(NUM_ROUNDS - 1)];
+                let diff = builder.sub_extension(input_bits_combined_lo, output_lo);
+                let filtered_diff = builder.mul_sub_extension(is_last_round, diff, diff);
+                yield_constr.constraint_transition(builder, filtered_diff);
+                let diff = builder.sub_extension(input_bits_combined_hi, output_hi);
+                let filtered_diff = builder.mul_sub_extension(is_last_round, diff, diff);
+                yield_constr.constraint_transition(builder, filtered_diff);
             }
         }
     }