diff --git a/aosp_beryllium.mk b/aosp_beryllium.mk
index b56bc69..cddde76 100644
--- a/aosp_beryllium.mk
+++ b/aosp_beryllium.mk
@@ -1,22 +1,34 @@
-DB845C_KERNEL_DIR := device/linaro/dragonboard-kernel/android-5.4
-DB845C_MODS := $(wildcard $(DB845C_KERNEL_DIR)/*.ko)
+ifndef TARGET_KERNEL_USE
+
+KERNEL_MODS := $(wildcard device/linaro/dragonboard-kernel/android-5.4/*.ko)
 
 # Local modules are built from following tree:
 # https://github.com/pundiramit/linux/tree/android-5.4-modules
-LOCAL_KERNEL_DIR := device/xiaomi/beryllium/prebuilt-kernel/android-5.4
-LOCAL_MODS := $(wildcard $(LOCAL_KERNEL_DIR)/*.ko)
+LOCAL_MODS := $(wildcard device/xiaomi/beryllium/prebuilt-kernel/android-5.4/*.ko)
 
 # Skip copying modules broken on android-5.4
 # Use local module copy instead
-DB845C_SKIP_MODS := %/msm.ko
+SKIP_MODS := %/msm.ko
 
 # BT modules go to vendor partition
-DB845C_ONLY_VENDOR := %/btqca.ko %/hci_uart.ko
-BOARD_VENDOR_KERNEL_MODULES := $(filter $(DB845C_ONLY_VENDOR),$(DB845C_MODS))
+ONLY_VENDOR := %/btqca.ko %/hci_uart.ko
+BOARD_VENDOR_KERNEL_MODULES := $(filter $(ONLY_VENDOR),$(KERNEL_MODS))
 
-BOARD_VENDOR_RAMDISK_KERNEL_MODULES := $(filter-out $(DB845C_SKIP_MODS) $(DB845C_ONLY_VENDOR),$(DB845C_MODS))
+BOARD_VENDOR_RAMDISK_KERNEL_MODULES := $(filter-out $(SKIP_MODS) $(ONLY_VENDOR),$(KERNEL_MODS))
 BOARD_VENDOR_RAMDISK_KERNEL_MODULES += $(LOCAL_MODS)
 
+else # ifdef TARGET_KERNEL_USE
+
+KERNEL_MODS := $(wildcard device/xiaomi/beryllium/prebuilt-kernel/android-$(TARGET_KERNEL_USE)/*.ko)
+
+# Following modules go to vendor partition
+ONLY_VENDOR := %/btqca.ko %/hci_uart.ko
+
+BOARD_VENDOR_KERNEL_MODULES := $(filter $(ONLY_VENDOR),$(KERNEL_MODS))
+BOARD_VENDOR_RAMDISK_KERNEL_MODULES := $(filter-out $(ONLY_VENDOR),$(KERNEL_MODS))
+
+endif
+
 # Inherit the full_base and device configurations
 $(call inherit-product, $(SRC_TARGET_DIR)/product/core_64_bit.mk)
 $(call inherit-product, device/xiaomi/beryllium/aosp_beryllium/device.mk)