diff --git a/kimchi/src/circuits/polynomials/keccak/gadget.rs b/kimchi/src/circuits/polynomials/keccak/gadget.rs
index 83b75e3294..a620645d81 100644
--- a/kimchi/src/circuits/polynomials/keccak/gadget.rs
+++ b/kimchi/src/circuits/polynomials/keccak/gadget.rs
@@ -56,7 +56,7 @@ impl<F: PrimeField + SquareRootField> CircuitGate<F> {
     }
 
     fn create_keccak_absorb(new_row: usize, root: bool, pad: bool, pad_bytes: usize) -> Self {
-        let mut coeffs = vec![F::zero(); 336];
+        let mut coeffs = vec![F::zero(); SPONGE_COEFFS];
         coeffs[0] = F::one(); // absorb
         if root {
             coeffs[2] = F::one(); // root
@@ -66,10 +66,10 @@ impl<F: PrimeField + SquareRootField> CircuitGate<F> {
             for i in 0..pad_bytes {
                 coeffs[140 - i] = F::one(); // flag for padding
                 if i == 0 {
-                    coeffs[335 - i] += F::from(0x80u8); // pad
+                    coeffs[SPONGE_COEFFS - 1 - i] += F::from(0x80u8); // pad
                 }
                 if i == pad_bytes - 1 {
-                    coeffs[335 - i] += F::one(); // pad
+                    coeffs[SPONGE_COEFFS - 1 - i] += F::one(); // pad
                 }
             }
         }
diff --git a/kimchi/src/linearization.rs b/kimchi/src/linearization.rs
index cf08b4da99..3ae105d8cd 100644
--- a/kimchi/src/linearization.rs
+++ b/kimchi/src/linearization.rs
@@ -9,6 +9,7 @@ use crate::circuits::lookup::{
     lookups::{LookupFeatures, LookupInfo, LookupPatterns},
 };
 use crate::circuits::polynomials::keccak;
+use crate::circuits::polynomials::keccak::circuitgates::KeccakRound;
 use crate::circuits::polynomials::{
     complete_add::CompleteAdd,
     endomul_scalar::EndomulScalar,
@@ -27,7 +28,6 @@ use crate::circuits::{
     constraints::FeatureFlags,
     expr::{Column, ConstantExpr, Expr, FeatureFlag, Linearization, PolishToken},
     gate::GateType,
-    wires::COLUMNS,
 };
 use ark_ff::{FftField, PrimeField, SquareRootField, Zero};
 
@@ -45,10 +45,13 @@ pub fn constraints_expr<const W: usize, F: PrimeField + SquareRootField>(
 
     // Set up powers of alpha. Only the max number of constraints matters.
     // The gate type argument can just be the zero gate.
-    powers_of_alpha.register(
-        ArgumentType::Gate(GateType::Zero),
-        VarbaseMul::<F>::CONSTRAINTS,
-    );
+    let mut max_exponents = VarbaseMul::<F>::CONSTRAINTS;
+    if let Some(feature_flags) = feature_flags {
+        if feature_flags.keccak_round {
+            max_exponents = KeccakRound::<F>::CONSTRAINTS;
+        }
+    }
+    powers_of_alpha.register(ArgumentType::Gate(GateType::Zero), max_exponents);
 
     let mut cache = expr::Cache::default();
 
@@ -271,7 +274,7 @@ pub fn constraints_expr<const W: usize, F: PrimeField + SquareRootField>(
 
 /// Adds the polynomials that are evaluated as part of the proof
 /// for the linearization to work.
-pub fn linearization_columns<F: FftField + SquareRootField>(
+pub fn linearization_columns<const W: usize, F: FftField + SquareRootField>(
     feature_flags: Option<&FeatureFlags>,
 ) -> std::collections::HashSet<Column> {
     let mut h = std::collections::HashSet::new();
@@ -308,12 +311,12 @@ pub fn linearization_columns<F: FftField + SquareRootField>(
     };
 
     // the witness polynomials
-    for i in 0..COLUMNS {
+    for i in 0..W {
         h.insert(Witness(i));
     }
 
     // the coefficient polynomials
-    for i in 0..COLUMNS {
+    for i in 0..W {
         h.insert(Coefficient(i));
     }
 
@@ -361,7 +364,7 @@ pub fn expr_linearization<const W: usize, F: PrimeField + SquareRootField>(
     feature_flags: Option<&FeatureFlags>,
     generic: bool,
 ) -> (Linearization<Vec<PolishToken<F>>>, Alphas<F>) {
-    let evaluated_cols = linearization_columns::<F>(feature_flags);
+    let evaluated_cols = linearization_columns::<W, F>(feature_flags);
 
     let (expr, powers_of_alpha) = constraints_expr::<W, F>(feature_flags, generic);