diff --git a/backends/vulkan/runtime/graph/ops/glsl/softmax.glsl b/backends/vulkan/runtime/graph/ops/glsl/softmax.glsl
index 86a2229c416..9b44d5c5a94 100644
--- a/backends/vulkan/runtime/graph/ops/glsl/softmax.glsl
+++ b/backends/vulkan/runtime/graph/ops/glsl/softmax.glsl
@@ -23,8 +23,10 @@ layout(std430) buffer;
 ${layout_declare_tensor(B, "w", "tout", DTYPE, STORAGE)}
 ${layout_declare_tensor(B, "r", "tin", DTYPE, STORAGE)}
 
-${layout_declare_ubo(B, "ivec3", "tout_limits")}
-${layout_declare_ubo(B, "ivec4", "tin_sizes")}
+layout(push_constant) uniform restrict Block {
+  ivec4 tin_sizes;
+  ivec3 tout_limits;
+};
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
diff --git a/backends/vulkan/runtime/graph/ops/impl/Softmax.cpp b/backends/vulkan/runtime/graph/ops/impl/Softmax.cpp
index 5e645e29e3d..2d683719ba2 100644
--- a/backends/vulkan/runtime/graph/ops/impl/Softmax.cpp
+++ b/backends/vulkan/runtime/graph/ops/impl/Softmax.cpp
@@ -139,9 +139,9 @@ void add_softmax_node(
       // Inputs and Outputs
       {{out, vkapi::kWrite}, {in, vkapi::kRead}},
       // Shader params buffers
-      {graph.logical_limits_ubo(out), graph.sizes_ubo(in)},
-      // Push Constants
       {},
+      // Push Constants
+      {graph.sizes_pc_of(in), graph.logical_limits_pc_of(out)},
       // Specialization Constants
       {graph.packed_dim_of(out), reduce_dim_xyz, group_dim},
       // Resize Args