diff --git a/bin/cros_make_image_bootable b/bin/cros_make_image_bootable
index e670ff9..ac11988 100755
--- a/bin/cros_make_image_bootable
+++ b/bin/cros_make_image_bootable
@@ -233,6 +233,7 @@
     --private="${private}" \
     --public="${public}" \
     --enable_serial="${FLAGS_enable_serial}" \
+    --image_type="${FLAGS_image_type}" \
     ${vblock} \
     ${enable_rootfs_verification_flag} \
     ${enable_bootcache_flag}
diff --git a/build_kernel_image.sh b/build_kernel_image.sh
index 0f517af..9f8a3c0 100755
--- a/build_kernel_image.sh
+++ b/build_kernel_image.sh
@@ -59,6 +59,8 @@
   "Enable serial port for printks. Example values: ttyS0"
 DEFINE_integer loglevel 7 \
   "The loglevel to add to the kernel command line."
+DEFINE_string image_type "base" \
+  "Type of image we're building for (base/factory_install)."
 
 # Parse flags
 FLAGS "$@" || exit 1
@@ -70,6 +72,7 @@
 # N.B.  Ordering matters for some of the libraries below, because
 # some of the files contain initialization used by later files.
 . "${BUILD_LIBRARY_DIR}/board_options.sh" || exit 1
+. "${BUILD_LIBRARY_DIR}/disk_layout_util.sh" || exit 1
 
 rootdigest() {
   local digest=${table#*root_hexdigest=}
@@ -99,8 +102,14 @@
   :
 }
 
+get_base_root() {
+  echo 'PARTUUID=%U/PARTNROFF=1'
+}
+
 load_board_specific_script "${BOARD}" "build_kernel_image.sh"
 
+base_root=$(get_base_root)
+
 device_mapper_args=
 # Even with a rootfs_image, root= is not changed unless specified.
 if [[ -n "${FLAGS_rootfs_image}" && -n "${FLAGS_rootfs_hash}" ]]; then
@@ -142,8 +151,6 @@
   if [[ ${FLAGS_enable_rootfs_verification} -eq ${FLAGS_TRUE} ]]; then
     if [[ ${FLAGS_enable_bootcache} -eq ${FLAGS_TRUE} ]]; then
       base_root='254:0'  # major:minor numbers for /dev/dm-0
-    else
-      base_root='PARTUUID=%U/PARTNROFF=1'  # kern_guid + 1
     fi
     table=${table//HASH_DEV/${base_root}}
     table=${table//ROOT_DEV/${base_root}}
@@ -173,7 +180,7 @@
 # production use.  If +%d can be added upstream, then we can use:
 #   root_dev=PARTUID=uuid+1
 dev_wait=0
-root_dev="PARTUUID=%U/PARTNROFF=1"
+root_dev=${base_root}
 if [[ ${FLAGS_enable_rootfs_verification} -eq ${FLAGS_TRUE} ]]; then
   dev_wait=1
   if [[ ${FLAGS_enable_bootcache} -eq ${FLAGS_TRUE} ]]; then
@@ -261,6 +268,33 @@
   error "Unknown arch: ${FLAGS_arch}"
 fi
 
+already_seen_rootfs=0
+for p in $(get_partitions "${FLAGS_image_type}"); do
+  format=$(get_format "${FLAGS_image_type}" "${p}")
+  if [[ "${format}" == "ubi" ]]; then
+    type=$(get_type "${FLAGS_image_type}" "${p}")
+    # cgpt.py ensures that the rootfs partitions are compatible, in that if
+    # one is ubi then both are, and they have the same number of reserved
+    # blocks. We only want to attach one of them in boot to save time, so
+    # attach %P and get the information for whichever rootfs comes first.
+    if [[ "${type}" == "rootfs" ]]; then
+      if [[ "${already_seen_rootfs}" -ne 0 ]]; then
+        continue
+      fi
+      already_seen_rootfs=1
+      partname='%P'
+    else
+      partname="${p}"
+    fi
+    echo "ubi.mtd=${partname},0,0,${partname}" \
+        >> "${FLAGS_working_dir}/config.txt"
+    fs_format=$(get_filesystem_format "${FLAGS_image_type}" "${p}")
+    if [[ "${fs_format}" != "ubifs" ]]; then
+      echo "ubi.block=${partname},0" >> "${FLAGS_working_dir}/config.txt"
+    fi
+  fi
+done
+
 config_file="${FLAGS_working_dir}/config.txt"
 modify_kernel_command_line "${config_file}"
 # Create and sign the kernel blob
diff --git a/build_library/cgpt.py b/build_library/cgpt.py
index 8bf8508..54acb47 100755
--- a/build_library/cgpt.py
+++ b/build_library/cgpt.py
@@ -33,6 +33,12 @@
 class ConflictingOptions(Exception):
   """Conflicting Options"""
 
+class MismatchedRootfsFormat(Exception):
+  """Rootfs partitions in different formats"""
+
+class MismatchedRootfsBlocks(Exception):
+  """Rootfs partitions have different numbers of reserved erase blocks"""
+
 COMMON_LAYOUT = 'common'
 BASE_LAYOUT = 'base'
 # Blocks of the partition entry array.
@@ -821,6 +827,25 @@
   return partition.get('fs_format')
 
 
+def GetFormat(options, image_type, layout_filename, num):
+  """Returns the format of a given partition for a given layout type.
+
+  Args:
+    options: Flags passed to the script
+    image_type: Type of image eg base/test/dev/factory_install
+    layout_filename: Path to partition configuration file
+    num: Number of the partition you want to read from
+
+  Returns:
+    Format of the selected partition's filesystem
+  """
+
+  partitions = GetPartitionTableFromConfig(options, layout_filename, image_type)
+  partition = GetPartitionByNumber(partitions, num)
+
+  return partition.get('format')
+
+
 def GetFilesystemOptions(options, image_type, layout_filename, num):
   """Returns the filesystem options of a given partition and layout type.
 
@@ -952,15 +977,45 @@
     ))
 
 
-def DoParseOnly(options, image_type, layout_filename):
-  """Parses a layout file only, used before reading sizes to check for errors.
+def CheckRootfsPartitionsMatch(partitions):
+  """Checks that rootfs partitions are substitutable with each other.
+
+  This function asserts that either all rootfs partitions are in the same format
+  or none have a format, and it asserts that have the same number of reserved
+  erase blocks.
+  """
+  partition_format = None
+  reserved_erase_blocks = -1
+  for partition in partitions:
+    if partition.get('type') == 'rootfs':
+      new_format = partition.get('format', '')
+      new_reserved_erase_blocks = partition.get('reserved_erase_blocks', 0)
+
+      if partition_format is None:
+        partition_format = new_format
+        reserved_erase_blocks = new_reserved_erase_blocks
+
+      if new_format != partition_format:
+        raise MismatchedRootfsFormat(
+            'mismatched rootfs formats: "%s" and "%s"' %
+            (partition_format, new_format))
+
+      if reserved_erase_blocks != new_reserved_erase_blocks:
+        raise MismatchedRootfsBlocks(
+            'mismatched rootfs reserved erase block counts: %s and %s' %
+            (reserved_erase_blocks, new_reserved_erase_blocks))
+
+
+def Validate(options, image_type, layout_filename):
+  """Validates a layout file, used before reading sizes to check for errors.
 
   Args:
     options: Flags passed to the script
     image_type: Type of image eg base/test/dev/factory_install
     layout_filename: Path to partition configuration file
   """
-  GetPartitionTableFromConfig(options, layout_filename, image_type)
+  partitions = GetPartitionTableFromConfig(options, layout_filename, image_type)
+  CheckRootfsPartitionsMatch(partitions)
 
 
 def main(argv):
@@ -981,6 +1036,10 @@
           'usage': ['<image_type>', '<disk_layout>', '<partition_num>'],
           'func': GetPartitionSize,
       },
+      'readformat': {
+          'usage': ['<image_type>', '<disk_layout>', '<partition_num>'],
+          'func': GetFormat,
+      },
       'readfsformat': {
           'usage': ['<image_type>', '<disk_layout>', '<partition_num>'],
           'func': GetFilesystemFormat,
@@ -1013,9 +1072,9 @@
           'usage': ['<image_type>', '<disk_layout>'],
           'func': DoDebugOutput,
       },
-      'parseonly': {
+      'validate': {
           'usage': ['<image_type>', '<disk_layout>'],
-          'func': DoParseOnly,
+          'func': Validate,
       },
   }
 
diff --git a/build_library/disk_layout_util.sh b/build_library/disk_layout_util.sh
index d2c5771..6789edd 100644
--- a/build_library/disk_layout_util.sh
+++ b/build_library/disk_layout_util.sh
@@ -93,6 +93,14 @@
   cgpt_py readfsformat "${image_type}" "${DISK_LAYOUT_PATH}" ${part_id}
 }
 
+get_format() {
+  local image_type=$1
+  local part_id=$2
+  get_disk_layout_path
+
+  cgpt_py readformat "${image_type}" "${DISK_LAYOUT_PATH}" "${part_id}"
+}
+
 get_partitions() {
   local image_type=$1
   get_disk_layout_path
@@ -136,7 +144,7 @@
   local image_type=$1
   get_disk_layout_path
 
-  cgpt_py parseonly "${image_type}" "${DISK_LAYOUT_PATH}" > /dev/null
+  cgpt_py validate "${image_type}" "${DISK_LAYOUT_PATH}" > /dev/null
 }
 
 get_disk_layout_type() {