diff --git a/src/main/scala/Gateway.scala b/src/main/scala/Gateway.scala
index 92c8f51e5..7876662f8 100644
--- a/src/main/scala/Gateway.scala
+++ b/src/main/scala/Gateway.scala
@@ -259,19 +259,21 @@ object Preprocess {
 class Preprocess(bundles: Seq[DifftestBundle], config: GatewayConfig) extends Module {
   val in = IO(Input(MixedVec(bundles)))
   val out = IO(Output(MixedVec(bundles)))
-
+  val numCores = bundles.count(_.isUniqueIdentifier)
   out := in
 
   if (config.hasDutZone || config.isSquash || config.isBatch) {
     // Special fix for int writeback. Work for single-core only
     if (in.exists(_.desiredCppName == "wb_int")) {
-      require(in.count(_.isUniqueIdentifier) == 1, "only single-core is supported yet")
+      if (config.isSquash || config.isBatch) {
+        require(numCores == 1, "only single-core is supported yet")
+      }
       val writebacks = in.filter(_.desiredCppName == "wb_int").map(_.asInstanceOf[DiffIntWriteback])
       val numPhyRegs = writebacks.head.numElements
-      val wb_int = Reg(Vec(numPhyRegs, UInt(64.W)))
+      val wb_int = Reg(Vec(numCores, Vec(numPhyRegs, UInt(64.W))))
       for (wb <- writebacks) {
         when(wb.valid) {
-          wb_int(wb.address) := wb.data
+          wb_int(wb.coreid)(wb.address) := wb.data
         }
       }
 
@@ -283,7 +285,7 @@ class Preprocess(bundles: Seq[DifftestBundle], config: GatewayConfig) extends Mo
         when(c.valid && c.skip) {
           wb_for_skip.valid := true.B
           wb_for_skip.address := c.wpdest
-          wb_for_skip.data := wb_int(c.wpdest)
+          wb_for_skip.data := wb_int(c.coreid)(c.wpdest)
           for (wb <- writebacks) {
             when(wb.valid && wb.address === c.wpdest) {
               wb_for_skip.data := wb.data