diff --git a/src/gallium/drivers/asahi/agx_state.c b/src/gallium/drivers/asahi/agx_state.c
index 9236a35ef3d..bdbf3fea28f 100644
--- a/src/gallium/drivers/asahi/agx_state.c
+++ b/src/gallium/drivers/asahi/agx_state.c
@@ -938,9 +938,9 @@ agx_update_vs(struct agx_context *ctx)
    memcpy(key.attributes, ctx->attributes,
           sizeof(key.attributes[0]) * AGX_MAX_ATTRIBS);
 
-   for (unsigned i = 0; i < key.num_vbufs; ++i) {
+   u_foreach_bit(i, ctx->vb_mask) {
       assert((ctx->vertex_buffers[i].stride & 0x3) == 0);
-      key.vbuf_strides[i] = ctx->vertex_buffers[i].stride / 4; // TODO: alignment
+      key.vbuf_strides[i] = ctx->vertex_buffers[i].stride / 4;
    }
 
    struct asahi_shader_key akey = {
diff --git a/src/gallium/drivers/asahi/agx_uniforms.c b/src/gallium/drivers/asahi/agx_uniforms.c
index 7ef9ddfd413..13650373732 100644
--- a/src/gallium/drivers/asahi/agx_uniforms.c
+++ b/src/gallium/drivers/asahi/agx_uniforms.c
@@ -71,7 +71,7 @@ agx_push_location_direct(struct agx_context *ctx, struct agx_push push,
       struct agx_ptr ptr = agx_pool_alloc_aligned(&batch->pool, count * sizeof(uint64_t), 8);
       uint64_t *addresses = ptr.cpu;
 
-      for (unsigned i = 0; i < count; ++i) {
+      u_foreach_bit(i, ctx->vb_mask) {
          struct pipe_vertex_buffer vb = ctx->vertex_buffers[i];
          assert(!vb.is_user_buffer);