diff --git a/rtl/serv_bufreg2.v b/rtl/serv_bufreg2.v
index 4afffb7..fe6d4c2 100644
--- a/rtl/serv_bufreg2.v
+++ b/rtl/serv_bufreg2.v
@@ -6,7 +6,7 @@ module serv_bufreg2
    input wire 	      i_init,
    input wire 	      i_cnt_done,
    input wire [1:0]   i_lsb,
-   input wire 	      i_byte_valid,
+   input wire [1:0]   i_bytecnt,
    output wire 	      o_sh_done,
    //Control
    input wire 	      i_op_b_sel,
@@ -23,9 +23,24 @@ module serv_bufreg2
 
    reg [31:0] 	 dat;
 
+   /*
+    Before a store operation, the data to be written needs to be shifted into
+    place. Depending on the address alignment, we need to shift different
+    amounts. One formula for calculating this is to say that we shift when
+    i_lsb + i_bytecnt < 4. Unfortunately, the synthesis tools don't seem to be
+    clever enough so the hideous expression below is used to achieve the same
+    thing in a more optimal way.
+    */
+   wire byte_valid
+     = (!i_lsb[0] & !i_lsb[1])         |
+       (!i_bytecnt[0] & !i_bytecnt[1]) |
+       (!i_bytecnt[1] & !i_lsb[1])     |
+       (!i_bytecnt[1] & !i_lsb[0])     |
+       (!i_bytecnt[0] & !i_lsb[1]);
+
    assign o_op_b = i_op_b_sel ? i_rs2 : i_imm;
 
-   wire 	 dat_en = i_shift_op | (i_en & i_byte_valid);
+   wire 	 dat_en = i_shift_op | (i_en & byte_valid);
 
    /* The dat register has three different use cases for store, load and
     shift operations.
diff --git a/rtl/serv_mem_if.v b/rtl/serv_mem_if.v
index 7cdb39f..be68499 100644
--- a/rtl/serv_mem_if.v
+++ b/rtl/serv_mem_if.v
@@ -10,7 +10,6 @@ module serv_mem_if
    //State
    input wire [1:0]  i_bytecnt,
    input wire [1:0]  i_lsb,
-   output wire 	     o_byte_valid,
    output wire 	     o_misalign,
    //Control
    input wire 	     i_signed,
@@ -26,21 +25,6 @@ module serv_mem_if
 
    reg signbit;
 
-   /*
-    Before a store operation, the data to be written needs to be shifted into
-    place. Depending on the address alignment, we need to shift different
-    amounts. One formula for calculating this is to say that we shift when
-    i_lsb + i_bytecnt < 4. Unfortunately, the synthesis tools don't seem to be
-    clever enough so the hideous expression below is used to achieve the same
-    thing in a more optimal way.
-    */
-   assign o_byte_valid
-     = (!i_lsb[0] & !i_lsb[1])         |
-       (!i_bytecnt[0] & !i_bytecnt[1]) |
-       (!i_bytecnt[1] & !i_lsb[1])     |
-       (!i_bytecnt[1] & !i_lsb[0])     |
-       (!i_bytecnt[0] & !i_lsb[1]);
-
    wire dat_valid =
 	i_mdu_op |
 	i_word |
diff --git a/rtl/serv_top.v b/rtl/serv_top.v
index b6523ad..0b39662 100644
--- a/rtl/serv_top.v
+++ b/rtl/serv_top.v
@@ -153,7 +153,6 @@ module serv_top
    wire          mem_half;
    wire [1:0] 	 mem_bytecnt;
    wire 	 sh_done;
-   wire 	 byte_valid;
 
    wire 	 mem_misalign;
 
@@ -402,7 +401,7 @@ module serv_top
       .i_init       (init),
       .i_cnt_done   (cnt_done),
       .i_lsb        (lsb),
-      .i_byte_valid (byte_valid),
+      .i_bytecnt    (mem_bytecnt),
       .o_sh_done    (sh_done),
       //Control
       .i_op_b_sel   (op_b_sel),
@@ -525,7 +524,6 @@ module serv_top
       //State
       .i_bytecnt    (mem_bytecnt),
       .i_lsb        (lsb),
-      .o_byte_valid (byte_valid),
       .o_misalign   (mem_misalign),
       //Control
       .i_mdu_op     (mdu_op),