diff --git a/src/asahi/lib/shaders/geometry.cl b/src/asahi/lib/shaders/geometry.cl
index 96e8618e177..ec4e0f970e0 100644
--- a/src/asahi/lib/shaders/geometry.cl
+++ b/src/asahi/lib/shaders/geometry.cl
@@ -476,7 +476,7 @@ libagx_build_gs_draw(global struct agx_geometry_params *p, uint vertices,
    descriptor[3] = 0;                         /* index bias */
    descriptor[4] = 0;                         /* start instance */
 
-   if (state->heap_bottom > 1024 * 1024 * 128) {
+   if (state->heap_bottom > state->heap_size) {
       global uint *foo = (global uint *)(uintptr_t)0xdeadbeef;
       *foo = 0x1234;
    }
@@ -535,7 +535,7 @@ libagx_gs_setup_indirect(global struct agx_gs_setup_indirect_params *gsi,
    *(gsi->vertex_buffer) = (uintptr_t)(state->heap + state->heap_bottom);
    state->heap_bottom += align(vertex_buffer_size, 4);
 
-   if (state->heap_bottom > 1024 * 1024 * 128) {
+   if (state->heap_bottom > state->heap_size) {
       global uint *foo = (global uint *)(uintptr_t)0x1deadbeef;
       *foo = 0x1234;
    }
diff --git a/src/gallium/drivers/asahi/agx_state.c b/src/gallium/drivers/asahi/agx_state.c
index e8d493b9e38..9b655713e7f 100644
--- a/src/gallium/drivers/asahi/agx_state.c
+++ b/src/gallium/drivers/asahi/agx_state.c
@@ -3929,13 +3929,16 @@ agx_batch_geometry_state(struct agx_batch *batch)
    struct agx_context *ctx = batch->ctx;
 
    if (!batch->geometry_state) {
+      uint32_t size = 128 * 1024 * 1024;
+
       if (!ctx->heap) {
          ctx->heap = pipe_buffer_create(ctx->base.screen, PIPE_BIND_GLOBAL,
-                                        PIPE_USAGE_DEFAULT, 1024 * 1024 * 128);
+                                        PIPE_USAGE_DEFAULT, size);
       }
 
       struct agx_geometry_state state = {
          .heap = agx_resource(ctx->heap)->bo->ptr.gpu,
+         .heap_size = size,
       };
 
       agx_batch_writes(batch, agx_resource(ctx->heap), 0);