diff --git a/keccak-air/Cargo.toml b/keccak-air/Cargo.toml
index 2804df4c2..bfcf786ce 100644
--- a/keccak-air/Cargo.toml
+++ b/keccak-air/Cargo.toml
@@ -9,6 +9,7 @@ p3-air = { path = "../air" }
 p3-field = { path = "../field" }
 p3-matrix = { path = "../matrix" }
 p3-maybe-rayon = { path = "../maybe-rayon" }
+p3-uni-stark = { path = "../uni-stark" }
 p3-util = { path = "../util" }
 tracing = "0.1.37"
 
@@ -29,7 +30,6 @@ p3-poseidon = { path = "../poseidon" }
 p3-poseidon2 = { path = "../poseidon2" }
 p3-sha256 = { path = "../sha256", features = ["asm"] }
 p3-symmetric = { path = "../symmetric" }
-p3-uni-stark = { path = "../uni-stark" }
 rand = "0.8.5"
 tracing-subscriber = { version = "0.3.17", features = ["std", "env-filter"] }
 tracing-forest = { version = "0.1.6", features = ["ansi", "smallvec"] }
diff --git a/keccak-air/src/air.rs b/keccak-air/src/air.rs
index 2425207cc..a3905f617 100644
--- a/keccak-air/src/air.rs
+++ b/keccak-air/src/air.rs
@@ -3,6 +3,7 @@ use core::borrow::Borrow;
 use p3_air::{Air, AirBuilder, BaseAir};
 use p3_field::AbstractField;
 use p3_matrix::Matrix;
+use p3_uni_stark::MultiStageAir;
 
 use crate::columns::{KeccakCols, NUM_KECCAK_COLS};
 use crate::constants::rc_value_bit;
@@ -20,6 +21,8 @@ impl<F> BaseAir<F> for KeccakAir {
     }
 }
 
+impl<AB: AirBuilder> MultiStageAir<AB> for KeccakAir {}
+
 impl<AB: AirBuilder> Air<AB> for KeccakAir {
     #[inline]
     fn eval(&self, builder: &mut AB) {
diff --git a/uni-stark/src/traits.rs b/uni-stark/src/traits.rs
index 53334afb0..724dc182d 100644
--- a/uni-stark/src/traits.rs
+++ b/uni-stark/src/traits.rs
@@ -35,5 +35,3 @@ pub trait MultiStageAir<AB: AirBuilder>: Air<AB> {
         0
     }
 }
-
-impl<AB: AirBuilder, A: Air<AB>> MultiStageAir<AB> for A {}
diff --git a/uni-stark/tests/fib_air.rs b/uni-stark/tests/fib_air.rs
index b6d456afa..e22a23525 100644
--- a/uni-stark/tests/fib_air.rs
+++ b/uni-stark/tests/fib_air.rs
@@ -13,7 +13,7 @@ use p3_matrix::Matrix;
 use p3_merkle_tree::FieldMerkleTreeMmcs;
 use p3_poseidon2::{Poseidon2, Poseidon2ExternalMatrixGeneral};
 use p3_symmetric::{PaddingFreeSponge, TruncatedPermutation};
-use p3_uni_stark::{prove, verify, StarkConfig};
+use p3_uni_stark::{prove, verify, MultiStageAir, StarkConfig};
 use rand::thread_rng;
 
 /// For testing the public values feature
@@ -56,6 +56,8 @@ impl<AB: AirBuilderWithPublicValues> Air<AB> for FibonacciAir {
     }
 }
 
+impl<AB: AirBuilderWithPublicValues> MultiStageAir<AB> for FibonacciAir {}
+
 pub fn generate_trace_rows<F: PrimeField64>(a: u64, b: u64, n: usize) -> RowMajorMatrix<F> {
     assert!(n.is_power_of_two());
 
diff --git a/uni-stark/tests/mul_air.rs b/uni-stark/tests/mul_air.rs
index c5fe81ede..ae03b1178 100644
--- a/uni-stark/tests/mul_air.rs
+++ b/uni-stark/tests/mul_air.rs
@@ -2,7 +2,7 @@ use std::fmt::Debug;
 use std::marker::PhantomData;
 
 use itertools::Itertools;
-use p3_air::{Air, AirBuilder, BaseAir};
+use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir};
 use p3_baby_bear::{BabyBear, DiffusionMatrixBabyBear};
 use p3_challenger::{DuplexChallenger, HashChallenger, SerializingChallenger32};
 use p3_circle::CirclePcs;
@@ -21,7 +21,7 @@ use p3_poseidon2::{Poseidon2, Poseidon2ExternalMatrixGeneral};
 use p3_symmetric::{
     CompressionFunctionFromHasher, PaddingFreeSponge, SerializingHasher32, TruncatedPermutation,
 };
-use p3_uni_stark::{prove, verify, StarkConfig, StarkGenericConfig, Val};
+use p3_uni_stark::{prove, verify, MultiStageAir, StarkConfig, StarkGenericConfig, Val};
 use rand::distributions::{Distribution, Standard};
 use rand::{thread_rng, Rng};
 
@@ -116,6 +116,8 @@ impl<AB: AirBuilder> Air<AB> for MulAir {
     }
 }
 
+impl<AB: AirBuilderWithPublicValues> MultiStageAir<AB> for MulAir {}
+
 fn do_test<SC: StarkGenericConfig>(
     config: SC,
     air: MulAir,