diff --git a/backends/vulkan/runtime/graph/ops/glsl/addmm_naive_texture3d.glsl b/backends/vulkan/runtime/graph/ops/glsl/addmm_naive_texture3d.glsl
index bd210e210ce..3d5814eb6d0 100644
--- a/backends/vulkan/runtime/graph/ops/glsl/addmm_naive_texture3d.glsl
+++ b/backends/vulkan/runtime/graph/ops/glsl/addmm_naive_texture3d.glsl
@@ -21,13 +21,17 @@ ${layout_declare_tensor(B, "r", "mat1_tensor", DTYPE, "texture3d")}
 ${layout_declare_tensor(B, "r", "mat2_tensor", DTYPE, "texture3d")}
 $if HAS_BIAS:
   ${layout_declare_tensor(B, "r", "bias_tensor", DTYPE, "texture3d")}
-${layout_declare_ubo(B, "ivec4", "out_sizes")}
-${layout_declare_ubo(B, "ivec3", "out_limits")}
-${layout_declare_ubo(B, "ivec4", "mat1_sizes")}
-${layout_declare_ubo(B, "ivec4", "mat2_sizes")}
-$if HAS_BIAS:
-  ${layout_declare_ubo(B, "ivec4", "bias_sizes")}
-  ${layout_declare_ubo(B, "float", "alpha", "float", "beta")}
+
+layout(push_constant) uniform restrict Block {
+  ivec4 out_sizes;
+  ivec4 mat1_sizes;
+  ivec4 mat2_sizes;
+  ivec3 out_limits;
+  $if HAS_BIAS:
+    ivec4 bias_sizes;
+    float alpha;
+    float beta;
+};
 
 #include "indexing_utils.h"
 
diff --git a/backends/vulkan/runtime/graph/ops/impl/Linear.cpp b/backends/vulkan/runtime/graph/ops/impl/Linear.cpp
index 38d70271f4f..67c3f377f0c 100644
--- a/backends/vulkan/runtime/graph/ops/impl/Linear.cpp
+++ b/backends/vulkan/runtime/graph/ops/impl/Linear.cpp
@@ -182,16 +182,14 @@ void add_addmm_naive_texture_node(
       // Inputs and Outputs
       {{out, vkapi::kWrite}, {{mat1, mat2, self}, vkapi::kRead}},
       // Shader params buffers
-      {
-          graph.sizes_ubo(out),
-          graph.logical_limits_ubo(out),
-          graph.sizes_ubo(mat1),
-          graph.sizes_ubo(mat2),
-          graph.sizes_ubo(self),
-          graph.create_params_buffer(params),
-      },
-      // Push Constants
       {},
+      // Push Constants
+      {graph.sizes_pc_of(out),
+       graph.sizes_pc_of(mat1),
+       graph.sizes_pc_of(mat2),
+       graph.logical_limits_pc_of(out),
+       graph.sizes_pc_of(self),
+       PushConstantDataInfo(&params, sizeof(params))},
       // Specialization Constants
       {graph.hashed_layout_of(out),
        graph.hashed_layout_of(mat1),
diff --git a/backends/vulkan/runtime/graph/ops/impl/MatMul.cpp b/backends/vulkan/runtime/graph/ops/impl/MatMul.cpp
index 47ecf5f18d2..6c687ec67a8 100644
--- a/backends/vulkan/runtime/graph/ops/impl/MatMul.cpp
+++ b/backends/vulkan/runtime/graph/ops/impl/MatMul.cpp
@@ -162,14 +162,12 @@ void add_matmul_naive_texture3d_node(
       // Inputs and Outputs
       {{out, vkapi::kWrite}, {{mat1, mat2}, vkapi::kRead}},
       // Shader params buffers
-      {
-          graph.sizes_ubo(out),
-          graph.logical_limits_ubo(out),
-          graph.sizes_ubo(mat1),
-          graph.sizes_ubo(mat2),
-      },
-      // Push Constants
       {},
+      // Push Constants
+      {graph.sizes_pc_of(out),
+       graph.sizes_pc_of(mat1),
+       graph.sizes_pc_of(mat2),
+       graph.logical_limits_pc_of(out)},
       // Specialization Constants
       {graph.hashed_layout_of(out),
        graph.hashed_layout_of(mat1),