diff --git a/ui/build/kati.go b/ui/build/kati.go
index b3e820e..56e9a88 100644
--- a/ui/build/kati.go
+++ b/ui/build/kati.go
@@ -129,9 +129,7 @@
 		"TARGET_DEVICE_DIR="+config.TargetDeviceDir(),
 		"KATI_PACKAGE_MK_DIR="+config.KatiPackageMkDir())
 
-	runKati(ctx, config, katiBuildSuffix, args, func(env *Environment) {
-		env.Unset("DIST_DIR")
-	})
+	runKati(ctx, config, katiBuildSuffix, args, func(env *Environment) {})
 }
 
 func runKatiPackage(ctx Context, config Config) {
@@ -170,6 +168,7 @@
 
 		if config.Dist() {
 			env.Set("DIST", "true")
+			env.Set("DIST_DIR", config.DistDir())
 		}
 	})
 }
@@ -184,7 +183,5 @@
 		"-f", "build/make/core/cleanbuild.mk",
 		"SOONG_MAKEVARS_MK=" + config.SoongMakeVarsMk(),
 		"TARGET_DEVICE_DIR=" + config.TargetDeviceDir(),
-	}, func(env *Environment) {
-		env.Unset("DIST_DIR")
-	})
+	}, func(env *Environment) {})
 }