diff --git a/android/TerminalApp/java/com/android/virtualization/terminal/BaseActivity.kt b/android/TerminalApp/java/com/android/virtualization/terminal/BaseActivity.kt
index 70bc5e4..229cdbb 100644
--- a/android/TerminalApp/java/com/android/virtualization/terminal/BaseActivity.kt
+++ b/android/TerminalApp/java/com/android/virtualization/terminal/BaseActivity.kt
@@ -30,10 +30,6 @@
                     TerminalExceptionHandler(applicationContext)
             }
         }
-    }
-
-    public override fun onResume() {
-        super.onResume()
 
         if (
             applicationContext.checkSelfPermission(Manifest.permission.POST_NOTIFICATIONS) !=
diff --git a/android/composd/aidl/android/system/composd/IIsolatedCompilationService.aidl b/android/composd/aidl/android/system/composd/IIsolatedCompilationService.aidl
index 3748899..d7757db 100644
--- a/android/composd/aidl/android/system/composd/IIsolatedCompilationService.aidl
+++ b/android/composd/aidl/android/system/composd/IIsolatedCompilationService.aidl
@@ -35,7 +35,7 @@
      * callback, unless the returned ICompilationTask is cancelled. The caller should maintain
      * a reference to the ICompilationTask until compilation completes or is cancelled.
      */
-    ICompilationTask startStagedApexCompile(ICompilationTaskCallback callback);
+    ICompilationTask startStagedApexCompile(ICompilationTaskCallback callback, String os);
 
     /**
      * Run odrefresh in a test instance of CompOS until completed or failed.
diff --git a/android/composd/src/instance_manager.rs b/android/composd/src/instance_manager.rs
index d1b0b99..a7154ec 100644
--- a/android/composd/src/instance_manager.rs
+++ b/android/composd/src/instance_manager.rs
@@ -39,10 +39,11 @@
         Self { service, state: Default::default() }
     }
 
-    pub fn start_current_instance(&self) -> Result<CompOsInstance> {
+    pub fn start_current_instance(&self, os: &str) -> Result<CompOsInstance> {
         let mut vm_parameters = new_vm_parameters()?;
         vm_parameters.name = String::from("Composd");
         vm_parameters.prefer_staged = true;
+        vm_parameters.os = os.to_owned();
         self.start_instance(CURRENT_INSTANCE_DIR, vm_parameters)
     }
 
diff --git a/android/composd/src/service.rs b/android/composd/src/service.rs
index 3cc40af..1e38eee 100644
--- a/android/composd/src/service.rs
+++ b/android/composd/src/service.rs
@@ -51,9 +51,10 @@
     fn startStagedApexCompile(
         &self,
         callback: &Strong<dyn ICompilationTaskCallback>,
+        os: &str,
     ) -> binder::Result<Strong<dyn ICompilationTask>> {
         check_permissions()?;
-        to_binder_result(self.do_start_staged_apex_compile(callback))
+        to_binder_result(self.do_start_staged_apex_compile(callback, os))
     }
 
     fn startTestCompile(
@@ -76,8 +77,10 @@
     fn do_start_staged_apex_compile(
         &self,
         callback: &Strong<dyn ICompilationTaskCallback>,
+        os: &str,
     ) -> Result<Strong<dyn ICompilationTask>> {
-        let comp_os = self.instance_manager.start_current_instance().context("Starting CompOS")?;
+        let comp_os =
+            self.instance_manager.start_current_instance(os).context("Starting CompOS")?;
 
         let target_dir_name = PENDING_ARTIFACTS_SUBDIR.to_owned();
         let task = OdrefreshTask::start(
diff --git a/android/composd_cmd/composd_cmd.rs b/android/composd_cmd/composd_cmd.rs
index 6281bd0..c944c17 100644
--- a/android/composd_cmd/composd_cmd.rs
+++ b/android/composd_cmd/composd_cmd.rs
@@ -39,7 +39,11 @@
 #[derive(Parser)]
 enum Actions {
     /// Compile classpath for real. Output can be used after a reboot.
-    StagedApexCompile {},
+    StagedApexCompile {
+        /// OS for the VM.
+        #[clap(long, default_value = "microdroid")]
+        os: String,
+    },
 
     /// Compile classpath in a debugging VM. Output is ignored.
     TestCompile {
@@ -59,7 +63,7 @@
     ProcessState::start_thread_pool();
 
     match action {
-        Actions::StagedApexCompile {} => run_staged_apex_compile()?,
+        Actions::StagedApexCompile { os } => run_staged_apex_compile(&os)?,
         Actions::TestCompile { prefer_staged, os } => run_test_compile(prefer_staged, &os)?,
     }
 
@@ -116,8 +120,8 @@
     }
 }
 
-fn run_staged_apex_compile() -> Result<()> {
-    run_async_compilation(|service, callback| service.startStagedApexCompile(callback))
+fn run_staged_apex_compile(os: &str) -> Result<()> {
+    run_async_compilation(|service, callback| service.startStagedApexCompile(callback, os))
 }
 
 fn run_test_compile(prefer_staged: bool, os: &str) -> Result<()> {
diff --git a/android/virtmgr/src/aidl.rs b/android/virtmgr/src/aidl.rs
index 6aecc75..b5cf643 100644
--- a/android/virtmgr/src/aidl.rs
+++ b/android/virtmgr/src/aidl.rs
@@ -420,47 +420,83 @@
     }
 }
 
-fn find_partition(path: Option<&Path>) -> binder::Result<String> {
-    let path = match path {
-        Some(path) => path,
-        None => return Ok("system".to_owned()),
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+enum CallingPartition {
+    Odm,
+    Product,
+    System,
+    SystemExt,
+    Vendor,
+    Unknown,
+}
+
+impl CallingPartition {
+    fn as_str(&self) -> &'static str {
+        match self {
+            CallingPartition::Odm => "odm",
+            CallingPartition::Product => "product",
+            CallingPartition::System => "system",
+            CallingPartition::SystemExt => "system_ext",
+            CallingPartition::Vendor => "vendor",
+            CallingPartition::Unknown => "[unknown]",
+        }
+    }
+}
+
+impl std::fmt::Display for CallingPartition {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.as_str())
+    }
+}
+
+fn find_partition(path: Option<&Path>) -> binder::Result<CallingPartition> {
+    let Some(path) = path else {
+        return Ok(CallingPartition::System);
     };
     if path.starts_with("/system/system_ext/") {
-        return Ok("system_ext".to_owned());
-    } else if path.starts_with("/system/product/") {
-        return Ok("product".to_owned());
+        return Ok(CallingPartition::SystemExt);
     }
-    let mut components = path.components();
-    match components.nth(1) {
-        Some(std::path::Component::Normal(partition)) => {
-            if partition != "apex" {
-                return Ok(partition.to_string_lossy().into_owned());
-            }
+    if path.starts_with("/system/product/") {
+        return Ok(CallingPartition::Product);
+    }
+    let partition = {
+        let mut components = path.components();
+        let Some(std::path::Component::Normal(partition)) = components.nth(1) else {
+            return Err(anyhow!("Can't find partition in '{}'", path.display()))
+                .or_service_specific_exception(-1);
+        };
 
-            // If path is under /apex, find a partition of the preinstalled .apex path
-            let apex_name = match components.next() {
-                Some(std::path::Component::Normal(name)) => name.to_string_lossy(),
-                _ => {
-                    return Err(anyhow!("Can't find apex name for '{}'", path.display()))
-                        .or_service_specific_exception(-1)
-                }
+        // If path is under /apex, find a partition of the preinstalled .apex path
+        if partition == "apex" {
+            let Some(std::path::Component::Normal(apex_name)) = components.next() else {
+                return Err(anyhow!("Can't find apex name for '{}'", path.display()))
+                    .or_service_specific_exception(-1);
             };
-
             let apex_info_list = ApexInfoList::load()
                 .context("Failed to get apex info list")
                 .or_service_specific_exception(-1)?;
-
-            for apex_info in apex_info_list.list.iter() {
-                if apex_info.name == apex_name {
-                    return Ok(apex_info.partition.to_lowercase());
-                }
-            }
-
-            Err(anyhow!("Can't find apex info for '{apex_name}'")).or_service_specific_exception(-1)
+            apex_info_list
+                .list
+                .iter()
+                .find(|apex_info| apex_info.name.as_str() == apex_name)
+                .map(|apex_info| apex_info.partition.to_lowercase())
+                .ok_or(anyhow!("Can't find apex info for {apex_name:?}"))
+                .or_service_specific_exception(-1)?
+        } else {
+            partition.to_string_lossy().into_owned()
         }
-        _ => Err(anyhow!("Can't find partition in '{}'", path.display()))
-            .or_service_specific_exception(-1),
-    }
+    };
+    Ok(match partition.as_str() {
+        "odm" => CallingPartition::Odm,
+        "product" => CallingPartition::Product,
+        "system" => CallingPartition::System,
+        "system_ext" => CallingPartition::SystemExt,
+        "vendor" => CallingPartition::Vendor,
+        _ => {
+            warn!("unknown partition for '{}'", path.display());
+            CallingPartition::Unknown
+        }
+    })
 }
 
 impl VirtualizationService {
@@ -472,13 +508,13 @@
         &self,
         config: &VirtualMachineConfig,
         calling_exe_path: Option<&Path>,
+        calling_partition: CallingPartition,
     ) -> binder::Result<(VmContext, Cid, PathBuf)> {
         let name = match config {
             VirtualMachineConfig::RawConfig(config) => &config.name,
             VirtualMachineConfig::AppConfig(config) => &config.name,
         };
-        let calling_partition = find_partition(calling_exe_path)?;
-        let early_vm = find_early_vm_for_partition(&calling_partition, name)
+        let early_vm = find_early_vm_for_partition(calling_partition, name)
             .or_service_specific_exception(-1)?;
         let calling_exe_path = match calling_exe_path {
             Some(path) => path,
@@ -487,17 +523,7 @@
                     .or_service_specific_exception(-1)
             }
         };
-        let expected_exe_path = Path::new(&early_vm.path);
-        if expected_exe_path != calling_exe_path
-            && Path::new("/system").join(expected_exe_path) != calling_exe_path
-        {
-            return Err(anyhow!(
-                "VM '{name}' in partition '{calling_partition}' must be created with '{}', not '{}'",
-                &early_vm.path,
-                calling_exe_path.display()
-            ))
-            .or_service_specific_exception(-1);
-        }
+        early_vm.check_exe_paths_match(calling_exe_path)?;
 
         let cid = early_vm.cid as Cid;
         let temp_dir = PathBuf::from(format!("/mnt/vm/early/{cid}"));
@@ -567,6 +593,8 @@
         let requester_uid = get_calling_uid();
         let requester_debug_pid = get_calling_pid();
 
+        let calling_partition = find_partition(CALLING_EXE_PATH.as_deref())?;
+
         check_config_features(config)?;
 
         if cfg!(early) {
@@ -578,7 +606,7 @@
 
         // Allocating VM context checks the MANAGE_VIRTUAL_MACHINE permission.
         let (vm_context, cid, temporary_directory) = if cfg!(early) {
-            self.create_early_vm_context(config, CALLING_EXE_PATH.as_deref())?
+            self.create_early_vm_context(config, CALLING_EXE_PATH.as_deref(), calling_partition)?
         } else {
             self.create_vm_context(requester_debug_pid, config)?
         };
@@ -668,7 +696,8 @@
 
             // In a protected VM, we require custom kernels to come from a trusted source
             // (b/237054515).
-            check_label_for_kernel_files(&kernel, &initrd).or_service_specific_exception(-1)?;
+            check_label_for_kernel_files(&kernel, &initrd, calling_partition)
+                .or_service_specific_exception(-1)?;
 
             // Check if partition images are labeled incorrectly. This is to prevent random images
             // which are not protected by the Android Verified Boot (e.g. bits downloaded by apps)
@@ -686,15 +715,14 @@
                         !is_safe_raw_partition(&partition.label)
                     }
                 })
-                .try_for_each(check_label_for_partition)
+                .try_for_each(|partition| check_label_for_partition(partition, calling_partition))
                 .or_service_specific_exception(-1)?;
         }
 
         // Check if files for payloads and bases are NOT coming from /vendor and /odm, as they may
         // have unstable interfaces.
         // TODO(b/316431494): remove once Treble interfaces are stabilized.
-        check_partitions_for_files(config, &find_partition(CALLING_EXE_PATH.as_deref())?)
-            .or_service_specific_exception(-1)?;
+        check_partitions_for_files(config, calling_partition).or_service_specific_exception(-1)?;
 
         let zero_filler_path = temporary_directory.join("zero.img");
         write_zero_filler(&zero_filler_path)
@@ -1335,7 +1363,10 @@
     Ok(vm_config)
 }
 
-fn check_partition_for_file(fd: &ParcelFileDescriptor, calling_partition: &str) -> Result<()> {
+fn check_partition_for_file(
+    fd: &ParcelFileDescriptor,
+    calling_partition: CallingPartition,
+) -> Result<()> {
     let path = format!("/proc/self/fd/{}", fd.as_raw_fd());
     let link = fs::read_link(&path).context(format!("can't read_link {path}"))?;
 
@@ -1346,7 +1377,8 @@
     }
 
     let is_fd_vendor = link.starts_with("/vendor") || link.starts_with("/odm");
-    let is_caller_vendor = calling_partition == "vendor" || calling_partition == "odm";
+    let is_caller_vendor =
+        calling_partition == CallingPartition::Vendor || calling_partition == CallingPartition::Odm;
 
     if is_fd_vendor != is_caller_vendor {
         bail!("{} can't be used for VM client in {calling_partition}", link.display());
@@ -1357,7 +1389,7 @@
 
 fn check_partitions_for_files(
     config: &VirtualMachineRawConfig,
-    calling_partition: &str,
+    calling_partition: CallingPartition,
 ) -> Result<()> {
     config
         .disks
@@ -1511,7 +1543,7 @@
 ///
 /// App private data files are deliberately excluded, to avoid arbitrary payloads being run on
 /// user devices (W^X).
-fn check_label_is_allowed(context: &SeContext) -> Result<()> {
+fn check_label_is_allowed(context: &SeContext, calling_partition: CallingPartition) -> Result<()> {
     match context.selinux_type()? {
         | "apk_data_file" // APKs of an installed app
         | "shell_data_file" // test files created via adb shell
@@ -1520,27 +1552,42 @@
         | "virtualizationservice_data_file" // files created by VS / VirtMgr
         | "vendor_microdroid_file" // immutable dm-verity protected partition (/vendor/etc/avf/microdroid/.*)
          => Ok(()),
+        // It is difficult to require specific label types for vendor initiated VM's files, so we
+        // allow anything with a vendor prefix.
+        t if calling_partition == CallingPartition::Vendor && t.starts_with("vendor_")  => Ok(()),
         _ => bail!("Label {} is not allowed", context),
     }
 }
 
-fn check_label_for_partition(partition: &Partition) -> Result<()> {
+fn check_label_for_partition(
+    partition: &Partition,
+    calling_partition: CallingPartition,
+) -> Result<()> {
     let file = partition.image.as_ref().unwrap().as_ref();
-    check_label_is_allowed(&getfilecon(file)?)
+    check_label_is_allowed(&getfilecon(file)?, calling_partition)
         .with_context(|| format!("Partition {} invalid", &partition.label))
 }
 
-fn check_label_for_kernel_files(kernel: &Option<File>, initrd: &Option<File>) -> Result<()> {
+fn check_label_for_kernel_files(
+    kernel: &Option<File>,
+    initrd: &Option<File>,
+    calling_partition: CallingPartition,
+) -> Result<()> {
     if let Some(f) = kernel {
-        check_label_for_file(f, "kernel")?;
+        check_label_for_file(f, "kernel", calling_partition)?;
     }
     if let Some(f) = initrd {
-        check_label_for_file(f, "initrd")?;
+        check_label_for_file(f, "initrd", calling_partition)?;
     }
     Ok(())
 }
-fn check_label_for_file(file: &File, name: &str) -> Result<()> {
-    check_label_is_allowed(&getfilecon(file)?).with_context(|| format!("{} file invalid", name))
+fn check_label_for_file(
+    file: &File,
+    name: &str,
+    calling_partition: CallingPartition,
+) -> Result<()> {
+    check_label_is_allowed(&getfilecon(file)?, calling_partition)
+        .with_context(|| format!("{} file invalid", name))
 }
 
 /// Implementation of the AIDL `IVirtualMachine` interface. Used as a handle to a VM.
@@ -2258,15 +2305,38 @@
     early_vm: Vec<EarlyVm>,
 }
 
-static EARLY_VMS_CACHE: LazyLock<Mutex<HashMap<String, Vec<EarlyVm>>>> =
+impl EarlyVm {
+    /// Verifies that the provided executable path matches the expected path stored in the XML
+    /// configuration.
+    /// If the provided path starts with `/system`, it will be stripped before comparison.
+    fn check_exe_paths_match<P: AsRef<Path>>(&self, calling_exe_path: P) -> binder::Result<()> {
+        let actual_path = calling_exe_path.as_ref();
+        if Path::new(&self.path)
+            == Path::new("/").join(actual_path.strip_prefix("/system").unwrap_or(actual_path))
+        {
+            return Ok(());
+        }
+        Err(Status::new_service_specific_error_str(
+            -1,
+            Some(format!(
+                "Early VM '{}' executable paths do not match. Expected: {}. Found: {:?}.",
+                self.name,
+                self.path,
+                actual_path.display()
+            )),
+        ))
+    }
+}
+
+static EARLY_VMS_CACHE: LazyLock<Mutex<HashMap<CallingPartition, Vec<EarlyVm>>>> =
     LazyLock::new(|| Mutex::new(HashMap::new()));
 
-fn range_for_partition(partition: &str) -> Result<Range<Cid>> {
+fn range_for_partition(partition: CallingPartition) -> Range<Cid> {
     match partition {
-        "system" => Ok(100..200),
-        "system_ext" | "product" => Ok(200..300),
-        "vendor" | "odm" => Ok(300..400),
-        _ => Err(anyhow!("Early VMs are not supported for {partition}")),
+        CallingPartition::System => 100..200,
+        CallingPartition::SystemExt | CallingPartition::Product => 200..300,
+        CallingPartition::Vendor | CallingPartition::Odm => 300..400,
+        CallingPartition::Unknown => 0..0,
     }
 }
 
@@ -2301,10 +2371,10 @@
     Ok(())
 }
 
-fn get_early_vms_in_partition(partition: &str) -> Result<Vec<EarlyVm>> {
+fn get_early_vms_in_partition(partition: CallingPartition) -> Result<Vec<EarlyVm>> {
     let mut cache = EARLY_VMS_CACHE.lock().unwrap();
 
-    if let Some(result) = cache.get(partition) {
+    if let Some(result) = cache.get(&partition) {
         return Ok(result.clone());
     }
 
@@ -2317,10 +2387,10 @@
         }
     }
 
-    validate_cid_range(&early_vms, &range_for_partition(partition)?)
+    validate_cid_range(&early_vms, &range_for_partition(partition))
         .with_context(|| format!("CID validation for {partition} failed"))?;
 
-    cache.insert(partition.to_owned(), early_vms.clone());
+    cache.insert(partition, early_vms.clone());
 
     Ok(early_vms)
 }
@@ -2343,7 +2413,7 @@
     found_vm.ok_or_else(|| anyhow!("Can't find a VM named '{name}'"))
 }
 
-fn find_early_vm_for_partition(partition: &str, name: &str) -> Result<EarlyVm> {
+fn find_early_vm_for_partition(partition: CallingPartition, name: &str) -> Result<EarlyVm> {
     let early_vms = get_early_vms_in_partition(partition)
         .with_context(|| format!("Failed to get early VMs from {partition}"))?;
 
@@ -2359,23 +2429,30 @@
     #[test]
     fn test_is_allowed_label_for_partition() -> Result<()> {
         let expected_results = vec![
-            ("u:object_r:system_file:s0", true),
-            ("u:object_r:apk_data_file:s0", true),
-            ("u:object_r:app_data_file:s0", false),
-            ("u:object_r:app_data_file:s0:c512,c768", false),
-            ("u:object_r:privapp_data_file:s0:c512,c768", false),
-            ("invalid", false),
-            ("user:role:apk_data_file:severity:categories", true),
-            ("user:role:apk_data_file:severity:categories:extraneous", false),
+            (CallingPartition::System, "u:object_r:system_file:s0", true),
+            (CallingPartition::System, "u:object_r:apk_data_file:s0", true),
+            (CallingPartition::System, "u:object_r:app_data_file:s0", false),
+            (CallingPartition::System, "u:object_r:app_data_file:s0:c512,c768", false),
+            (CallingPartition::System, "u:object_r:privapp_data_file:s0:c512,c768", false),
+            (CallingPartition::System, "invalid", false),
+            (CallingPartition::System, "user:role:apk_data_file:severity:categories", true),
+            (
+                CallingPartition::System,
+                "user:role:apk_data_file:severity:categories:extraneous",
+                false,
+            ),
+            (CallingPartition::System, "u:object_r:vendor_unknowable:s0", false),
+            (CallingPartition::Vendor, "u:object_r:vendor_unknowable:s0", true),
         ];
 
-        for (label, expected_valid) in expected_results {
+        for (calling_partition, label, expected_valid) in expected_results {
             let context = SeContext::new(label)?;
-            let result = check_label_is_allowed(&context);
-            if expected_valid {
-                assert!(result.is_ok(), "Expected label {} to be allowed, got {:?}", label, result);
-            } else if result.is_ok() {
-                bail!("Expected label {} to be disallowed", label);
+            let result = check_label_is_allowed(&context, calling_partition);
+            if expected_valid != result.is_ok() {
+                bail!(
+                    "Expected label {label} to be {} for {calling_partition} partition",
+                    if expected_valid { "allowed" } else { "disallowed" }
+                );
             }
         }
         Ok(())
@@ -2666,7 +2743,7 @@
     fn test_symlink_to_system_ext_supported() -> Result<()> {
         let link_path = Path::new("/system/system_ext/file");
         let partition = find_partition(Some(link_path)).unwrap();
-        assert_eq!("system_ext", partition);
+        assert_eq!(CallingPartition::SystemExt, partition);
         Ok(())
     }
 
@@ -2674,11 +2751,44 @@
     fn test_symlink_to_product_supported() -> Result<()> {
         let link_path = Path::new("/system/product/file");
         let partition = find_partition(Some(link_path)).unwrap();
-        assert_eq!("product", partition);
+        assert_eq!(CallingPartition::Product, partition);
         Ok(())
     }
 
     #[test]
+    fn early_vm_exe_paths_match_succeeds_with_same_paths() {
+        let early_vm = EarlyVm {
+            name: "vm_demo_native_early".to_owned(),
+            cid: 123,
+            path: "/system_ext/bin/vm_demo_native_early".to_owned(),
+        };
+        let calling_exe_path = "/system_ext/bin/vm_demo_native_early";
+        assert!(early_vm.check_exe_paths_match(calling_exe_path).is_ok())
+    }
+
+    #[test]
+    fn early_vm_exe_paths_match_succeeds_with_calling_exe_path_from_system() {
+        let early_vm = EarlyVm {
+            name: "vm_demo_native_early".to_owned(),
+            cid: 123,
+            path: "/system_ext/bin/vm_demo_native_early".to_owned(),
+        };
+        let calling_exe_path = "/system/system_ext/bin/vm_demo_native_early";
+        assert!(early_vm.check_exe_paths_match(calling_exe_path).is_ok())
+    }
+
+    #[test]
+    fn early_vm_exe_paths_match_fails_with_unmatched_paths() {
+        let early_vm = EarlyVm {
+            name: "vm_demo_native_early".to_owned(),
+            cid: 123,
+            path: "/system_ext/bin/vm_demo_native_early".to_owned(),
+        };
+        let calling_exe_path = "/system/etc/system_ext/bin/vm_demo_native_early";
+        assert!(early_vm.check_exe_paths_match(calling_exe_path).is_err())
+    }
+
+    #[test]
     fn test_duplicated_early_vms() -> Result<()> {
         let tmp_dir = tempfile::TempDir::new()?;
         let tmp_dir_path = tmp_dir.path().to_owned();
diff --git a/build/debian/build.sh b/build/debian/build.sh
index 6facfcf..616dd00 100755
--- a/build/debian/build.sh
+++ b/build/debian/build.sh
@@ -269,7 +269,7 @@
 	            --extract "${dsc_file}"
 	pushd "linux-${debian_kver%-*}" > /dev/null
 
-	local kpatches_src="$SCRIPT_DIR/kernel_patches"
+	local kpatches_src="$SCRIPT_DIR/kernel/patches"
 	cp -r "${kpatches_src}/avf" debian/patches/
 	cat "${kpatches_src}/series" >> debian/patches/series
 	./debian/rules orig
@@ -282,9 +282,8 @@
 
 	# 2. Define our custom flavour and regenerate control file
 	# NOTE: Our flavour extends Debian's `cloud` config on the `none` featureset.
-	cat > debian/config/${debian_arch}/config.${debarch_flavour} <<EOF
-# TODO: Add our custom kernel config to this file
-EOF
+	cp "$SCRIPT_DIR/kernel/config" \
+	   debian/config/${debian_arch}/config.${debarch_flavour}
 
 	sed -z "s;\[base\]\nflavours:;[base]\nflavours:\n ${debarch_flavour};" \
 	    -i debian/config/${debian_arch}/none/defines
diff --git a/build/debian/fai_config/files/usr/local/bin/enable_display/AVF b/build/debian/fai_config/files/usr/local/bin/enable_display/AVF
index 69dce6a..76f9f97 100644
--- a/build/debian/fai_config/files/usr/local/bin/enable_display/AVF
+++ b/build/debian/fai_config/files/usr/local/bin/enable_display/AVF
@@ -1,4 +1,6 @@
 #!/bin/bash
 sudo systemd-run --collect -E XDG_SESSION_TYPE=wayland --uid=1000 -p PAMName=login -p TTYPath=/dev/tty7 sleep 1d
 systemctl --user start weston
-export DISPLAY=:0
\ No newline at end of file
+export DISPLAY=:0
+export MESA_LOADER_DRIVER_OVERRIDE=zink
+export LIBGL_ALWAYS_SOFTWARE=1
\ No newline at end of file
diff --git a/build/debian/fai_config/scripts/AVF/20-useradd b/build/debian/fai_config/scripts/AVF/20-useradd
index 1c93772..b92648a 100755
--- a/build/debian/fai_config/scripts/AVF/20-useradd
+++ b/build/debian/fai_config/scripts/AVF/20-useradd
@@ -1,4 +1,4 @@
 #!/bin/bash
 
-$ROOTCMD useradd -m -u 1000 -N -G sudo -s /usr/bin/bash droid
+$ROOTCMD useradd -m -u 1000 -N -G sudo,video,render -s /usr/bin/bash droid
 $ROOTCMD echo 'droid ALL=(ALL) NOPASSWD:ALL' >> $target/etc/sudoers
diff --git a/build/debian/kernel/config b/build/debian/kernel/config
new file mode 100644
index 0000000..1ba603c
--- /dev/null
+++ b/build/debian/kernel/config
@@ -0,0 +1 @@
+CONFIG_DRM=m
diff --git a/build/debian/kernel_patches/avf/arm64-balloon.patch b/build/debian/kernel/patches/avf/arm64-balloon.patch
similarity index 100%
rename from build/debian/kernel_patches/avf/arm64-balloon.patch
rename to build/debian/kernel/patches/avf/arm64-balloon.patch
diff --git a/build/debian/kernel_patches/series b/build/debian/kernel/patches/series
similarity index 100%
rename from build/debian/kernel_patches/series
rename to build/debian/kernel/patches/series
diff --git a/build/debian/vm_config.json.aarch64 b/build/debian/vm_config.json.aarch64
index 96254f8..463583f 100644
--- a/build/debian/vm_config.json.aarch64
+++ b/build/debian/vm_config.json.aarch64
@@ -35,5 +35,8 @@
     "console_out": true,
     "console_input_device": "ttyS0",
     "network": true,
-    "auto_memory_balloon": true
+    "auto_memory_balloon": true,
+    "gpu": {
+        "backend": "2d"
+    }
 }
diff --git a/build/debian/vm_config.json.x86_64 b/build/debian/vm_config.json.x86_64
index c34a0f2..bc4e00a 100644
--- a/build/debian/vm_config.json.x86_64
+++ b/build/debian/vm_config.json.x86_64
@@ -44,5 +44,8 @@
     "console_out": true,
     "console_input_device": "ttyS0",
     "network": true,
-    "auto_memory_balloon": true
+    "auto_memory_balloon": true,
+    "gpu": {
+        "backend": "2d"
+    }
 }
diff --git a/guest/pvmfw/README.md b/guest/pvmfw/README.md
index 766a923..652ca90 100644
--- a/guest/pvmfw/README.md
+++ b/guest/pvmfw/README.md
@@ -461,7 +461,12 @@
   - `secretkeeper_protection`: pvmfw defers rollback protection to the guest
   - `supports_uefi_boot`: pvmfw boots the VM as a EFI payload (experimental)
   - `trusty_security_vm`: pvmfw skips rollback protection
-- `"com.android.virt.page_size"`: the guest page size in KiB (optional, defaults to 4)
+- `"com.android.virt.page_size"`: (optional) the guest page size in KiB, defaults to 4
+- `"com.android.virt.name"`: (optional) VM name, used as the
+  [`component_name`][dice-comp-name] (defaults to `"vm_entry"`) in the guest
+  DICE certificate and to identify special VMs
+
+[dice-comp-name]: https://cs.android.com/android/platform/superproject/main/+/main:external/open-dice/docs/android.md;l=81;drc=6d511e9533eac05d64d47fcd78ac5d881e72c3de
 
 ## Development
 
diff --git a/guest/pvmfw/avb/Android.bp b/guest/pvmfw/avb/Android.bp
index 141c1d2..0d55d7c 100644
--- a/guest/pvmfw/avb/Android.bp
+++ b/guest/pvmfw/avb/Android.bp
@@ -37,6 +37,7 @@
         ":test_image_with_one_hashdesc",
         ":test_image_with_non_initrd_hashdesc",
         ":test_image_with_initrd_and_non_initrd_desc",
+        ":test_image_with_name",
         ":test_image_with_invalid_page_size",
         ":test_image_with_negative_page_size",
         ":test_image_with_overflow_page_size",
@@ -123,6 +124,20 @@
 }
 
 avb_add_hash_footer {
+    name: "test_image_with_name",
+    src: ":unsigned_test_image",
+    partition_name: "boot",
+    private_key: ":pvmfw_sign_key",
+    salt: "2134",
+    props: [
+        {
+            name: "com.android.virt.name",
+            value: "test_vm_name",
+        },
+    ],
+}
+
+avb_add_hash_footer {
     name: "test_image_with_invalid_page_size",
     src: ":unsigned_test_image",
     partition_name: "boot",
diff --git a/guest/pvmfw/avb/src/error.rs b/guest/pvmfw/avb/src/error.rs
index 1307e15..eb82837 100644
--- a/guest/pvmfw/avb/src/error.rs
+++ b/guest/pvmfw/avb/src/error.rs
@@ -30,6 +30,8 @@
     UnknownVbmetaProperty,
     /// VBMeta has invalid page_size property.
     InvalidPageSize,
+    /// VBMeta has invalid VM name property.
+    InvalidVmName,
 }
 
 impl From<SlotVerifyError<'_>> for PvmfwVerifyError {
@@ -54,6 +56,7 @@
             }
             Self::UnknownVbmetaProperty => write!(f, "Unknown vbmeta property"),
             Self::InvalidPageSize => write!(f, "Invalid page_size property"),
+            Self::InvalidVmName => write!(f, "Invalid name property"),
         }
     }
 }
diff --git a/guest/pvmfw/avb/src/verify.rs b/guest/pvmfw/avb/src/verify.rs
index 8810696..6a3d7de 100644
--- a/guest/pvmfw/avb/src/verify.rs
+++ b/guest/pvmfw/avb/src/verify.rs
@@ -17,7 +17,7 @@
 use crate::ops::{Ops, Payload};
 use crate::partition::PartitionName;
 use crate::PvmfwVerifyError;
-use alloc::vec::Vec;
+use alloc::{string::String, vec::Vec};
 use avb::{
     Descriptor, DescriptorError, DescriptorResult, HashDescriptor, PartitionData, SlotVerifyError,
     SlotVerifyNoDataResult, VbmetaData,
@@ -47,6 +47,8 @@
     pub rollback_index: u64,
     /// Page size of kernel, if present.
     pub page_size: Option<usize>,
+    /// Name of the guest payload, if present.
+    pub name: Option<String>,
 }
 
 impl VerifiedBootData<'_> {
@@ -238,6 +240,18 @@
     Ok(Some(size))
 }
 
+/// Returns the indicated payload name, if present.
+fn read_name(vbmeta_data: &VbmetaData) -> Result<Option<String>, PvmfwVerifyError> {
+    let Some(property) = vbmeta_data.get_property_value("com.android.virt.name") else {
+        return Ok(None);
+    };
+    let name = str::from_utf8(property).map_err(|_| PvmfwVerifyError::InvalidVmName)?;
+    if name.is_empty() {
+        return Err(PvmfwVerifyError::InvalidVmName);
+    }
+    Ok(Some(name.into()))
+}
+
 /// Verifies the given initrd partition, and checks that the resulting contents looks like expected.
 fn verify_initrd(
     ops: &mut Ops,
@@ -275,6 +289,7 @@
     let hash_descriptors = HashDescriptors::get(&descriptors)?;
     let capabilities = Capability::get_capabilities(vbmeta_image)?;
     let page_size = read_page_size(vbmeta_image)?;
+    let name = read_name(vbmeta_image)?;
 
     if initrd.is_none() {
         hash_descriptors.verify_no_initrd()?;
@@ -286,6 +301,7 @@
             capabilities,
             rollback_index,
             page_size,
+            name,
         });
     }
 
@@ -307,5 +323,6 @@
         capabilities,
         rollback_index,
         page_size,
+        name,
     })
 }
diff --git a/guest/pvmfw/avb/tests/api_test.rs b/guest/pvmfw/avb/tests/api_test.rs
index 3027c47..b3899d9 100644
--- a/guest/pvmfw/avb/tests/api_test.rs
+++ b/guest/pvmfw/avb/tests/api_test.rs
@@ -27,6 +27,7 @@
 use utils::*;
 
 const TEST_IMG_WITH_ONE_HASHDESC_PATH: &str = "test_image_with_one_hashdesc.img";
+const TEST_IMG_WITH_NAME_PATH: &str = "test_image_with_name.img";
 const TEST_IMG_WITH_INVALID_PAGE_SIZE_PATH: &str = "test_image_with_invalid_page_size.img";
 const TEST_IMG_WITH_NEGATIVE_PAGE_SIZE_PATH: &str = "test_image_with_negative_page_size.img";
 const TEST_IMG_WITH_OVERFLOW_PAGE_SIZE_PATH: &str = "test_image_with_overflow_page_size.img";
@@ -102,6 +103,7 @@
         capabilities: vec![],
         rollback_index: 0,
         page_size: None,
+        name: None,
     };
     assert_eq!(expected_boot_data, verified_boot_data);
 
@@ -147,6 +149,7 @@
         capabilities: vec![Capability::RemoteAttest],
         rollback_index: 0,
         page_size: None,
+        name: None,
     };
     assert_eq!(expected_boot_data, verified_boot_data);
 
@@ -247,6 +250,18 @@
 }
 
 #[test]
+fn kernel_has_expected_valid_name() {
+    let kernel = fs::read(TEST_IMG_WITH_NAME_PATH).unwrap();
+    assert_eq!(read_name(&kernel), Ok(Some("test_vm_name".to_owned())));
+}
+
+#[test]
+fn kernel_has_expected_missing_name() {
+    let kernel = fs::read(TEST_IMG_WITH_ONE_HASHDESC_PATH).unwrap();
+    assert_eq!(read_name(&kernel), Ok(None));
+}
+
+#[test]
 fn kernel_has_expected_page_size_invalid() {
     let kernel = fs::read(TEST_IMG_WITH_INVALID_PAGE_SIZE_PATH).unwrap();
     assert_eq!(read_page_size(&kernel), Err(PvmfwVerifyError::InvalidPageSize));
@@ -483,6 +498,7 @@
         capabilities: vec![],
         rollback_index: 5,
         page_size: None,
+        name: None,
     };
     assert_eq!(expected_boot_data, verified_boot_data);
     Ok(())
diff --git a/guest/pvmfw/avb/tests/utils.rs b/guest/pvmfw/avb/tests/utils.rs
index 7282f3e..227daa2 100644
--- a/guest/pvmfw/avb/tests/utils.rs
+++ b/guest/pvmfw/avb/tests/utils.rs
@@ -28,6 +28,7 @@
 use std::{
     fs,
     mem::{size_of, transmute, MaybeUninit},
+    string::String,
 };
 
 const MICRODROID_KERNEL_IMG_PATH: &str = "microdroid_kernel";
@@ -134,6 +135,7 @@
         capabilities,
         rollback_index: 1,
         page_size,
+        name: None,
     };
     assert_eq!(expected_boot_data, verified_boot_data);
 
@@ -166,12 +168,23 @@
         capabilities,
         rollback_index: expected_rollback_index,
         page_size,
+        name: None,
     };
     assert_eq!(expected_boot_data, verified_boot_data);
 
     Ok(())
 }
 
+pub fn read_name(kernel: &[u8]) -> Result<Option<String>, PvmfwVerifyError> {
+    let public_key = load_trusted_public_key().unwrap();
+    let verified_boot_data = verify_payload(
+        kernel,
+        None, // initrd
+        &public_key,
+    )?;
+    Ok(verified_boot_data.name)
+}
+
 pub fn read_page_size(kernel: &[u8]) -> Result<Option<usize>, PvmfwVerifyError> {
     let public_key = load_trusted_public_key().unwrap();
     let verified_boot_data = verify_payload(
diff --git a/guest/pvmfw/src/dice.rs b/guest/pvmfw/src/dice.rs
index f49fedb..49a3807 100644
--- a/guest/pvmfw/src/dice.rs
+++ b/guest/pvmfw/src/dice.rs
@@ -16,6 +16,7 @@
 extern crate alloc;
 
 use alloc::format;
+use alloc::string::String;
 use alloc::vec::Vec;
 use ciborium::cbor;
 use ciborium::Value;
@@ -83,6 +84,7 @@
     pub mode: DiceMode,
     pub security_version: u64,
     pub rkp_vm_marker: bool,
+    component_name: String,
 }
 
 impl PartialInputs {
@@ -90,12 +92,13 @@
         let code_hash = to_dice_hash(data)?;
         let auth_hash = hash(data.public_key)?;
         let mode = to_dice_mode(data.debug_level);
+        let component_name = data.name.clone().unwrap_or(String::from("vm_entry"));
         // We use rollback_index from vbmeta as the security_version field in dice certificate.
         let security_version = data.rollback_index;
         let rkp_vm_marker = data.has_capability(Capability::RemoteAttest)
             || data.has_capability(Capability::TrustySecurityVm);
 
-        Ok(Self { code_hash, auth_hash, mode, security_version, rkp_vm_marker })
+        Ok(Self { code_hash, auth_hash, mode, security_version, rkp_vm_marker, component_name })
     }
 
     pub fn write_next_bcc(
@@ -156,7 +159,7 @@
 
     fn generate_config_descriptor(&self, instance_hash: Option<Hash>) -> Result<Vec<u8>> {
         let mut config = Vec::with_capacity(4);
-        config.push((cbor!(COMPONENT_NAME_KEY)?, cbor!("vm_entry")?));
+        config.push((cbor!(COMPONENT_NAME_KEY)?, cbor!(self.component_name.as_str())?));
         config.push((cbor!(SECURITY_VERSION_KEY)?, cbor!(self.security_version)?));
         if self.rkp_vm_marker {
             config.push((cbor!(RKP_VM_MARKER_KEY)?, Value::Null))
@@ -200,6 +203,7 @@
         kernel_digest: [1u8; size_of::<Digest>()],
         initrd_digest: Some([2u8; size_of::<Digest>()]),
         public_key: b"public key",
+        name: None,
         capabilities: vec![],
         rollback_index: 42,
         page_size: None,
@@ -249,12 +253,13 @@
     }
 
     #[test]
-    fn rkp_vm_config_descriptor_has_rkp_vm_marker() {
+    fn rkp_vm_config_descriptor_has_rkp_vm_marker_and_component_name() {
         let vb_data =
             VerifiedBootData { capabilities: vec![Capability::RemoteAttest], ..BASE_VB_DATA };
         let inputs = PartialInputs::new(&vb_data).unwrap();
         let config_map = decode_config_descriptor(&inputs, Some(HASH));
 
+        assert_eq!(config_map.get(&COMPONENT_NAME_KEY).unwrap().as_text().unwrap(), "vm_entry");
         assert!(config_map.get(&RKP_VM_MARKER_KEY).unwrap().is_null());
     }
 
diff --git a/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachine.java b/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachine.java
index 5f634ef..af313a1 100644
--- a/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachine.java
+++ b/libs/framework-virtualization/src/android/system/virtualmachine/VirtualMachine.java
@@ -276,28 +276,16 @@
 
         @Override
         public void onTrimMemory(int level) {
-            int percent;
+            /* Treat level < TRIM_MEMORY_UI_HIDDEN as generic low-memory warnings */
+            int percent = 10;
 
-            switch (level) {
-                case ComponentCallbacks2.TRIM_MEMORY_RUNNING_CRITICAL:
-                    percent = 50;
-                    break;
-                case ComponentCallbacks2.TRIM_MEMORY_RUNNING_LOW:
-                    percent = 30;
-                    break;
-                case ComponentCallbacks2.TRIM_MEMORY_RUNNING_MODERATE:
-                    percent = 10;
-                    break;
-                case ComponentCallbacks2.TRIM_MEMORY_BACKGROUND:
-                case ComponentCallbacks2.TRIM_MEMORY_MODERATE:
-                case ComponentCallbacks2.TRIM_MEMORY_COMPLETE:
-                    /* Release as much memory as we can. The app is on the LMKD LRU kill list. */
-                    percent = 50;
-                    break;
-                default:
-                    /* Treat unrecognised messages as generic low-memory warnings. */
-                    percent = 30;
-                    break;
+            if (level >= ComponentCallbacks2.TRIM_MEMORY_UI_HIDDEN) {
+                percent = 30;
+            }
+
+            if (level >= ComponentCallbacks2.TRIM_MEMORY_BACKGROUND) {
+                /* Release as much memory as we can. The app is on the LMKD LRU kill list. */
+                percent = 50;
             }
 
             synchronized (mLock) {
diff --git a/libs/service-compos/java/com/android/server/compos/IsolatedCompilationJobService.java b/libs/service-compos/java/com/android/server/compos/IsolatedCompilationJobService.java
index adc0300..3033991 100644
--- a/libs/service-compos/java/com/android/server/compos/IsolatedCompilationJobService.java
+++ b/libs/service-compos/java/com/android/server/compos/IsolatedCompilationJobService.java
@@ -174,7 +174,7 @@
             }
 
             try {
-                ICompilationTask composTask = composd.startStagedApexCompile(this);
+                ICompilationTask composTask = composd.startStagedApexCompile(this, "microdroid");
                 mMetrics.onCompilationStarted();
                 mTask.set(composTask);
                 composTask.asBinder().linkToDeath(this, 0);
diff --git a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
index e2956f2..003c3f0 100644
--- a/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
+++ b/tests/benchmark_hostside/java/android/avf/test/AVFHostTestCase.java
@@ -19,6 +19,7 @@
 import static com.android.tradefed.device.TestDevice.MicrodroidBuilder;
 import static com.android.tradefed.testtype.DeviceJUnit4ClassRunner.TestMetrics;
 
+import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 import static com.google.common.truth.TruthJUnit.assume;
 
@@ -113,16 +114,78 @@
 
     @Test
     public void testBootWithCompOS() throws Exception {
-        composTestHelper(true);
+        composTestHelper(true, "microdroid");
+    }
+
+    @Test
+    public void testBootWithCompOS_os_android15_66() throws Exception {
+        composTestHelper(true, "android15_66");
+    }
+
+    @Test
+    public void testBootWithCompOS_os_microdroid_16k() throws Exception {
+        composTestHelper(true, "microdroid_16k");
     }
 
     @Test
     public void testBootWithoutCompOS() throws Exception {
-        composTestHelper(false);
+        composTestHelper(false, null);
     }
 
     @Test
     public void testNoLongHypSections() throws Exception {
+        noLongHypSectionsHelper("microdroid");
+    }
+
+    @Test
+    public void testNoLongHypSections_os_android15_66() throws Exception {
+        noLongHypSectionsHelper("android15_66");
+    }
+
+    @Test
+    public void testNoLongHypSections_os_microdroid_16k() throws Exception {
+        noLongHypSectionsHelper("microdroid_16k");
+    }
+
+    @Test
+    public void testPsciMemProtect() throws Exception {
+        psciMemProtectHelper("microdroid");
+    }
+
+    @Test
+    public void testPsciMemProtect_os_android15_66() throws Exception {
+        psciMemProtectHelper("android15_66");
+    }
+
+    @Test
+    public void testPsciMemProtect_os_microdroid_16k() throws Exception {
+        psciMemProtectHelper("microdroid_16k");
+    }
+
+    @Test
+    public void testCameraAppStartupTime() throws Exception {
+        String[] launchIntentPackages = {
+            "com.android.camera2",
+            "com.google.android.GoogleCamera/com.android.camera.CameraLauncher"
+        };
+        String launchIntentPackage = findSupportedPackage(launchIntentPackages);
+        assume().withMessage("No supported camera package").that(launchIntentPackage).isNotNull();
+        appStartupHelper(launchIntentPackage);
+    }
+
+    @Test
+    public void testSettingsAppStartupTime() throws Exception {
+        String[] launchIntentPackages = {"com.android.settings"};
+        String launchIntentPackage = findSupportedPackage(launchIntentPackages);
+        assume().withMessage("No supported settings package").that(launchIntentPackage).isNotNull();
+        appStartupHelper(launchIntentPackage);
+    }
+
+    private void noLongHypSectionsHelper(String osKey) throws Exception {
+        assumeKernelSupported(osKey);
+        assumeVmTypeSupported(osKey, true);
+        String os = SUPPORTED_OSES.get(osKey);
+
         String[] hypEvents = {"hyp_enter", "hyp_exit"};
 
         assumeTrue(
@@ -130,7 +193,7 @@
                 KvmHypTracer.isSupported(getDevice(), hypEvents));
 
         KvmHypTracer tracer = new KvmHypTracer(getDevice(), hypEvents);
-        String result = tracer.run(COMPOSD_CMD_BIN + " test-compile");
+        String result = tracer.run(COMPOSD_CMD_BIN + " test-compile --os " + os);
         assertWithMessage("Failed to test compilation VM.")
                 .that(result)
                 .ignoringCase()
@@ -141,8 +204,11 @@
         CLog.i("Hypervisor traces parsed successfully.");
     }
 
-    @Test
-    public void testPsciMemProtect() throws Exception {
+    public void psciMemProtectHelper(String osKey) throws Exception {
+        assumeKernelSupported(osKey);
+        assumeVmTypeSupported(osKey, true);
+        String os = SUPPORTED_OSES.get(osKey);
+
         String[] hypEvents = {"psci_mem_protect"};
 
         assumeTrue(
@@ -151,7 +217,12 @@
         KvmHypTracer tracer = new KvmHypTracer(getDevice(), hypEvents);
 
         /* We need to wait for crosvm to die so all the VM pages are reclaimed */
-        String result = tracer.run(COMPOSD_CMD_BIN + " test-compile && killall -w crosvm || true");
+        String result =
+                tracer.run(
+                        COMPOSD_CMD_BIN
+                                + " test-compile --os "
+                                + os
+                                + " && killall -w crosvm || true");
         assertWithMessage("Failed to test compilation VM.")
                 .that(result)
                 .ignoringCase()
@@ -176,25 +247,6 @@
                 .isGreaterThan(0);
     }
 
-    @Test
-    public void testCameraAppStartupTime() throws Exception {
-        String[] launchIntentPackages = {
-            "com.android.camera2",
-            "com.google.android.GoogleCamera/com.android.camera.CameraLauncher"
-        };
-        String launchIntentPackage = findSupportedPackage(launchIntentPackages);
-        assume().withMessage("No supported camera package").that(launchIntentPackage).isNotNull();
-        appStartupHelper(launchIntentPackage);
-    }
-
-    @Test
-    public void testSettingsAppStartupTime() throws Exception {
-        String[] launchIntentPackages = {"com.android.settings"};
-        String launchIntentPackage = findSupportedPackage(launchIntentPackages);
-        assume().withMessage("No supported settings package").that(launchIntentPackage).isNotNull();
-        appStartupHelper(launchIntentPackage);
-    }
-
     private void appStartupHelper(String launchIntentPackage) throws Exception {
         assumeTrue(
                 "Skip on non-protected VMs",
@@ -471,8 +523,14 @@
         throw new IllegalArgumentException("Failed to get boot time info.");
     }
 
-    private void composTestHelper(boolean isWithCompos) throws Exception {
+    private void composTestHelper(boolean isWithCompos, String osKey) throws Exception {
         assumeFalse("Skip on CF; too slow", isCuttlefish());
+        if (isWithCompos) {
+            assumeKernelSupported(osKey);
+            assumeVmTypeSupported(osKey, true);
+        } else {
+            assertThat(osKey).isNull();
+        }
 
         List<Double> bootDmesgTime = new ArrayList<>(ROUND_COUNT);
 
@@ -480,7 +538,8 @@
             reInstallApex(REINSTALL_APEX_TIMEOUT_SEC);
             try {
                 if (isWithCompos) {
-                    compileStagedApex(COMPILE_STAGED_APEX_TIMEOUT_SEC);
+                    String os = SUPPORTED_OSES.get(osKey);
+                    compileStagedApex(COMPILE_STAGED_APEX_TIMEOUT_SEC, os);
                 }
             } finally {
                 // If compilation fails, we still have a staged APEX, and we need to reboot to
@@ -518,7 +577,7 @@
         getDevice().enableAdbRoot();
     }
 
-    private void compileStagedApex(int timeoutSec) throws Exception {
+    private void compileStagedApex(int timeoutSec, String os) throws Exception {
 
         long timeStart = System.currentTimeMillis();
         long timeEnd = timeStart + timeoutSec * 1000L;
@@ -530,7 +589,7 @@
 
                 String result =
                         android.runWithTimeout(
-                                3 * 60 * 1000, COMPOSD_CMD_BIN + " staged-apex-compile");
+                                3 * 60 * 1000, COMPOSD_CMD_BIN + " staged-apex-compile --os " + os);
                 assertWithMessage("Failed to compile staged APEX. Reason: " + result)
                         .that(result)
                         .ignoringCase()
diff --git a/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java b/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
index ad37dda..fcef19a 100644
--- a/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
+++ b/tests/hostside/helper/java/com/android/microdroid/test/host/MicrodroidHostTestCaseBase.java
@@ -18,6 +18,7 @@
 
 import static com.android.tradefed.testtype.DeviceJUnit4ClassRunner.TestLogData;
 
+import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 
 import static org.junit.Assume.assumeFalse;
@@ -267,4 +268,27 @@
     protected boolean isPkvmHypervisor() throws DeviceNotAvailableException {
         return "kvm.arm-protected".equals(getDevice().getProperty("ro.boot.hypervisor.version"));
     }
+
+    protected TestDevice getAndroidDevice() {
+        TestDevice androidDevice = (TestDevice) getDevice();
+        assertThat(androidDevice).isNotNull();
+        return androidDevice;
+    }
+
+    protected void assumeKernelSupported(String osKey) throws Exception {
+        String os = SUPPORTED_OSES.get(osKey);
+        assumeTrue(
+                "Skipping test as OS \"" + os + "\" is not supported",
+                getSupportedOSList().contains(os));
+    }
+
+    protected void assumeVmTypeSupported(String os, boolean protectedVm) throws Exception {
+        // TODO(b/376870129): remove this check
+        if (protectedVm) {
+            assumeFalse("pVMs with 16k kernel are not supported yet :(", os.endsWith("_16k"));
+        }
+        assumeTrue(
+                "Microdroid is not supported for specific VM protection type",
+                getAndroidDevice().supportsMicrodroid(protectedVm));
+    }
 }
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
index 7864f3f..59a57f1 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
@@ -1474,12 +1474,6 @@
         }
     }
 
-    private TestDevice getAndroidDevice() {
-        TestDevice androidDevice = (TestDevice) getDevice();
-        assertThat(androidDevice).isNotNull();
-        return androidDevice;
-    }
-
     // The TradeFed Dockerfile sets LD_LIBRARY_PATH to a directory with an older libc++.so, which
     // breaks binaries that are linked against a newer libc++.so. Binaries commonly use DT_RUNPATH
     // to find an adjacent libc++.so (e.g. `$ORIGIN/../lib64`), but LD_LIBRARY_PATH overrides
@@ -1490,23 +1484,6 @@
         return runUtil;
     }
 
-    private void assumeKernelSupported(String osKey) throws Exception {
-        String os = SUPPORTED_OSES.get(osKey);
-        assumeTrue(
-                "Skipping test as OS \"" + os + "\" is not supported",
-                getSupportedOSList().contains(os));
-    }
-
-    private void assumeVmTypeSupported(String os, boolean protectedVm) throws Exception {
-        // TODO(b/376870129): remove this check
-        if (protectedVm) {
-            assumeFalse("pVMs with 16k kernel are not supported yet :(", os.endsWith("_16k"));
-        }
-        assumeTrue(
-                "Microdroid is not supported for specific VM protection type",
-                getAndroidDevice().supportsMicrodroid(protectedVm));
-    }
-
     private void assumeArm64Supported() throws Exception {
         CommandRunner android = new CommandRunner(getDevice());
         String abi = android.run("getprop", "ro.product.cpu.abi");