diff --git a/drivers/gpu/nvgpu/gk20a/buddy_allocator_priv.h b/drivers/gpu/nvgpu/gk20a/buddy_allocator_priv.h
index 55de9bc92..1a6aed483 100644
--- a/drivers/gpu/nvgpu/gk20a/buddy_allocator_priv.h
+++ b/drivers/gpu/nvgpu/gk20a/buddy_allocator_priv.h
@@ -46,11 +46,10 @@ struct gk20a_buddy {
 
 	/*
 	 * Size of the PDE this buddy is using. This allows for grouping like
-	 * sized allocations into the same PDE.
+	 * sized allocations into the same PDE. This uses the gmmu_pgsz_gk20a
+	 * enum except for the BALLOC_PTE_SIZE_ANY specifier.
 	 */
-#define BALLOC_PTE_SIZE_ANY	0x0
-#define BALLOC_PTE_SIZE_SMALL	0x1
-#define BALLOC_PTE_SIZE_BIG	0x2
+#define BALLOC_PTE_SIZE_ANY	-1
 	int pte_size;
 };
 
diff --git a/drivers/gpu/nvgpu/gk20a/gk20a_allocator_buddy.c b/drivers/gpu/nvgpu/gk20a/gk20a_allocator_buddy.c
index 976058006..4ed559716 100644
--- a/drivers/gpu/nvgpu/gk20a/gk20a_allocator_buddy.c
+++ b/drivers/gpu/nvgpu/gk20a/gk20a_allocator_buddy.c
@@ -111,6 +111,7 @@ static struct gk20a_buddy *balloc_new_buddy(struct gk20a_buddy_allocator *a,
 	new_buddy->start = start;
 	new_buddy->order = order;
 	new_buddy->end = start + (1 << order) * a->blk_size;
+	new_buddy->pte_size = BALLOC_PTE_SIZE_ANY;
 
 	return new_buddy;
 }
@@ -132,7 +133,7 @@ static void __balloc_buddy_list_add(struct gk20a_buddy_allocator *a,
 	 * without cycling through the entire list.
 	 */
 	if (a->flags & GPU_ALLOC_GVA_SPACE &&
-	    b->pte_size == BALLOC_PTE_SIZE_BIG)
+	    b->pte_size == gmmu_page_size_big)
 		list_add_tail(&b->buddy_entry, list);
 	else
 		list_add(&b->buddy_entry, list);
@@ -474,7 +475,7 @@ static struct gk20a_buddy *__balloc_find_buddy(struct gk20a_buddy_allocator *a,
 		return NULL;
 
 	if (a->flags & GPU_ALLOC_GVA_SPACE &&
-	    pte_size == BALLOC_PTE_SIZE_BIG)
+	    pte_size == gmmu_page_size_big)
 		bud = list_last_entry(balloc_get_order_list(a, order),
 				      struct gk20a_buddy, buddy_entry);
 	else