diff --git a/src/main/scala/huancun/Common.scala b/src/main/scala/huancun/Common.scala
index db17f09..711226d 100644
--- a/src/main/scala/huancun/Common.scala
+++ b/src/main/scala/huancun/Common.scala
@@ -39,6 +39,7 @@ class SourceBReq(implicit p: Parameters) extends HuanCunBundle {
   val tag = UInt(tagBits.W)
   val param = UInt(3.W)
   val clients = UInt(clientBits.W)
+  val alias = aliasBitsOpt.map(w => UInt(w.W))
 }
 class SinkCReq(implicit p: Parameters) extends InnerTask {
   val size = UInt(msgSizeBits.W)
diff --git a/src/main/scala/huancun/HuanCun.scala b/src/main/scala/huancun/HuanCun.scala
index 5f7dfe7..4f3aa04 100644
--- a/src/main/scala/huancun/HuanCun.scala
+++ b/src/main/scala/huancun/HuanCun.scala
@@ -187,6 +187,7 @@ class HuanCun(implicit p: Parameters) extends LazyModule with HasHuanCunParamete
     val inclusion = if (cacheParams.inclusive) "Inclusive" else "Non-inclusive"
     val prefetch = "prefetch: " + cacheParams.prefetch.nonEmpty
     println(s"====== ${inclusion} ${cacheParams.name} ($sizeStr) $prefetch ======")
+    aliasBitsOpt.foreach(bits => println(s"aliasBits: $bits"))
     node.in.zip(node.out).foreach {
       case ((in, edgeIn), (out, edgeOut)) =>
         require(in.params.dataBits == out.params.dataBits)
diff --git a/src/main/scala/huancun/SourceB.scala b/src/main/scala/huancun/SourceB.scala
index 2a210cd..f5077c3 100644
--- a/src/main/scala/huancun/SourceB.scala
+++ b/src/main/scala/huancun/SourceB.scala
@@ -54,7 +54,7 @@ class SourceB(implicit p: Parameters) extends HuanCunModule {
   io.b.bits.size := offsetBits.U
   io.b.bits.source := getSourceId(chosenClient)
   io.b.bits.address := Cat(taskLatch.tag, taskLatch.set, 0.U(offsetBits.W))
-  io.b.bits.mask := ~0.U(beatBytes.W)
+  io.b.bits.mask := taskLatch.alias.getOrElse(0.U)
   io.b.bits.data := 0.U
   io.b.bits.corrupt := 0.U
 }
diff --git a/src/main/scala/huancun/noninclusive/Directory.scala b/src/main/scala/huancun/noninclusive/Directory.scala
index 81f6106..8fa0354 100644
--- a/src/main/scala/huancun/noninclusive/Directory.scala
+++ b/src/main/scala/huancun/noninclusive/Directory.scala
@@ -29,7 +29,7 @@ class SelfDirEntry(implicit p: Parameters) extends HuanCunBundle {
 
 class ClientDirEntry(implicit p: Parameters) extends HuanCunBundle {
   val state = UInt(stateBits.W)
-  val alias = aliasBitsOpt.map(_ => UInt(aliasBitsOpt.get.W))
+  val alias = aliasBitsOpt.map(bits => UInt(bits.W))
 }
 
 class SelfDirResult(implicit p: Parameters) extends SelfDirEntry {
diff --git a/src/main/scala/huancun/noninclusive/MSHR.scala b/src/main/scala/huancun/noninclusive/MSHR.scala
index f65b8f0..8ba3f49 100644
--- a/src/main/scala/huancun/noninclusive/MSHR.scala
+++ b/src/main/scala/huancun/noninclusive/MSHR.scala
@@ -662,15 +662,8 @@ class MSHR()(implicit p: Parameters) extends BaseMSHR[DirResult, SelfDirWrite, S
     ),
     io.dirResult.valid
   )
-  ob.set := {
-    if (!hasAliasBits) req.set
-    else
-      Cat(
-        req.set.head(setBits - log2Ceil(clientCacheParams.sets)),
-        probe_alias,
-        req.set(clientSetBits - 1, 0)
-      )
-  }
+  ob.set := req.set
+  ob.alias.foreach(_ := probe_alias)
   assert(ob.set.getWidth == req.set.getWidth)
   ob.param := Mux(
     req.fromB,