diff --git a/drivers/infiniband/hw/hns/hns_roce_hw_v2.c b/drivers/infiniband/hw/hns/hns_roce_hw_v2.c
index 90f56c1372f3c645f8ad96437fab91d7f8190d06..1cb7ce19b00070d6c813fb6639e35250d3663f1e 100644
--- a/drivers/infiniband/hw/hns/hns_roce_hw_v2.c
+++ b/drivers/infiniband/hw/hns/hns_roce_hw_v2.c
@@ -2440,7 +2440,7 @@ static int set_mtpt_pbl(struct hns_roce_v2_mpt_entry *mpt_entry,
 
 	i = 0;
 	for_each_sg(mr->umem->sg_head.sgl, sg, mr->umem->nmap, entry) {
-		len = sg_dma_len(sg) >> PAGE_SHIFT;
+		len = sg_dma_len(sg) >> mr->umem->page_shift;
 		for (j = 0; j < len; ++j) {
 			page_addr = sg_dma_address(sg) +
 				(j << mr->umem->page_shift);
diff --git a/drivers/infiniband/hw/hns/hns_roce_mr.c b/drivers/infiniband/hw/hns/hns_roce_mr.c
index b420c480f0f46dc113b249e7817d02154b4e76d2..328065e9af5ac6beab42ffa40f7a104d9eb22820 100644
--- a/drivers/infiniband/hw/hns/hns_roce_mr.c
+++ b/drivers/infiniband/hw/hns/hns_roce_mr.c
@@ -1128,7 +1128,7 @@ int hns_roce_ib_umem_write_mr(struct hns_roce_dev *hr_dev,
 
 	pbl_bt_sz = 1 << (hr_dev->caps.pbl_ba_pg_sz + PAGE_SHIFT);
 	for_each_sg(umem->sg_head.sgl, sg, umem->nmap, entry) {
-		len = sg_dma_len(sg) >> PAGE_SHIFT;
+		len = sg_dma_len(sg) >> umem->page_shift;
 		for (k = 0; k < len; ++k) {
 			page_addr = sg_dma_address(sg) +
 				    (k << umem->page_shift);