diff --git a/fpga/lib/pcie/rtl/dma_client_axis_sink.v b/fpga/lib/pcie/rtl/dma_client_axis_sink.v
index 3de0adcf1..85cefcf12 100644
--- a/fpga/lib/pcie/rtl/dma_client_axis_sink.v
+++ b/fpga/lib/pcie/rtl/dma_client_axis_sink.v
@@ -167,7 +167,7 @@ initial begin
         $finish;
     end
 
-    if (AXIS_DATA_WIDTH*2**PART_COUNT_WIDTH != SEG_COUNT*SEG_DATA_WIDTH) begin
+    if (AXIS_DATA_WIDTH*2**$clog2(PART_COUNT) != SEG_COUNT*SEG_DATA_WIDTH) begin
         $error("Error: AXI stream interface width must be a power of two fraction of RAM interface width (instance %m)");
         $finish;
     end
@@ -310,7 +310,11 @@ always @* begin
                     ram_mask_next = {SEG_COUNT{1'b1}};
                 end
 
-                ram_wr_cmd_be_int = (s_axis_write_data_tkeep & keep_mask_reg) << (addr_reg & ({PART_COUNT_WIDTH{1'b1}} << PART_OFFSET_WIDTH));
+                if (PART_COUNT > 1) begin
+                    ram_wr_cmd_be_int = (s_axis_write_data_tkeep & keep_mask_reg) << (addr_reg & ({PART_COUNT_WIDTH{1'b1}} << PART_OFFSET_WIDTH));
+                end else begin
+                    ram_wr_cmd_be_int = s_axis_write_data_tkeep & keep_mask_reg;
+                end
                 ram_wr_cmd_addr_int = {SEG_COUNT{addr_reg[RAM_ADDR_WIDTH-1:RAM_ADDR_WIDTH-SEG_ADDR_WIDTH]}};
                 ram_wr_cmd_data_int = {PART_COUNT{s_axis_write_data_tdata}};
                 for (i = 0; i < SEG_COUNT; i = i + 1) begin
diff --git a/fpga/lib/pcie/rtl/dma_client_axis_source.v b/fpga/lib/pcie/rtl/dma_client_axis_source.v
index ef8009bbf..b0493e228 100644
--- a/fpga/lib/pcie/rtl/dma_client_axis_source.v
+++ b/fpga/lib/pcie/rtl/dma_client_axis_source.v
@@ -166,7 +166,7 @@ initial begin
         $finish;
     end
 
-    if (AXIS_DATA_WIDTH*2**PART_COUNT_WIDTH != SEG_COUNT*SEG_DATA_WIDTH) begin
+    if (AXIS_DATA_WIDTH*2**$clog2(PART_COUNT) != SEG_COUNT*SEG_DATA_WIDTH) begin
         $error("Error: AXI stream interface width must be a power of two fraction of RAM interface width (instance %m)");
         $finish;
     end