diff --git a/opencl/test/unit_test/os_interface/linux/drm_mock.cpp b/opencl/test/unit_test/os_interface/linux/drm_mock.cpp
index b6b3968ca4..9d65fba25f 100644
--- a/opencl/test/unit_test/os_interface/linux/drm_mock.cpp
+++ b/opencl/test/unit_test/os_interface/linux/drm_mock.cpp
@@ -119,7 +119,7 @@ int DrmMock::ioctl(unsigned long request, void *arg) {
         }
 
         if (receivedContextParamRequest.param == I915_CONTEXT_PARAM_VM) {
-            static_cast<drm_i915_gem_context_param *>(arg)->value = 1u;
+            static_cast<drm_i915_gem_context_param *>(arg)->value = this->StoredRetValForVmId;
             return 0u;
         }
     }
diff --git a/opencl/test/unit_test/os_interface/linux/drm_mock.h b/opencl/test/unit_test/os_interface/linux/drm_mock.h
index 88d2f56eb1..4e607d3671 100644
--- a/opencl/test/unit_test/os_interface/linux/drm_mock.h
+++ b/opencl/test/unit_test/os_interface/linux/drm_mock.h
@@ -106,6 +106,7 @@ class DrmMock : public Drm {
         I915_SCHEDULER_CAP_PRIORITY |
         I915_SCHEDULER_CAP_PREEMPTION;
     int StoredExecSoftPin = 0;
+    int StoredRetValForVmId = 1;
 
     bool disableSomeTopology = false;
 
diff --git a/opencl/test/unit_test/os_interface/linux/drm_tests.cpp b/opencl/test/unit_test/os_interface/linux/drm_tests.cpp
index 43d223c685..4b19259464 100644
--- a/opencl/test/unit_test/os_interface/linux/drm_tests.cpp
+++ b/opencl/test/unit_test/os_interface/linux/drm_tests.cpp
@@ -390,3 +390,36 @@ TEST(DrmTest, givenDrmWithPerContextVMRequiredWhenCreatingOsContextsThenImplicit
     OsContextLinux osContext2(drmMock, 5u, 1, aub_stream::ENGINE_RCS, PreemptionMode::Disabled, false, false, false);
     EXPECT_EQ(0u, drmMock.receivedCreateContextId);
 }
+
+TEST(DrmTest, givenDrmWithPerContextVMRequiredWhenCreatingOsContextsThenImplicitVmIdPerContextIsQueriedAndStored) {
+    auto &rootEnv = *platform()->peekExecutionEnvironment()->rootDeviceEnvironments[0];
+    rootEnv.executionEnvironment.setPerContextMemorySpace();
+
+    DrmMock drmMock;
+    EXPECT_TRUE(drmMock.requirePerContextVM);
+
+    drmMock.StoredRetValForVmId = 20;
+
+    OsContextLinux osContext(drmMock, 0u, 1, aub_stream::ENGINE_RCS, PreemptionMode::Disabled, false, false, false);
+    EXPECT_EQ(0u, drmMock.receivedCreateContextId);
+    EXPECT_EQ(1u, drmMock.receivedContextParamRequestCount);
+
+    auto &drmVmIds = osContext.getDrmVmIds();
+    EXPECT_EQ(1u, drmVmIds.size());
+
+    EXPECT_EQ(20u, drmVmIds[0]);
+}
+
+TEST(DrmTest, givenNoPerContextVmsDrmWhenCreatingOsContextsThenVmIdIsNotQueriedAndStored) {
+    DrmMock drmMock;
+    EXPECT_FALSE(drmMock.requirePerContextVM);
+
+    drmMock.StoredRetValForVmId = 1;
+
+    OsContextLinux osContext(drmMock, 0u, 1, aub_stream::ENGINE_RCS, PreemptionMode::Disabled, false, false, false);
+    EXPECT_EQ(0u, drmMock.receivedCreateContextId);
+    EXPECT_EQ(1u, drmMock.receivedContextParamRequestCount);
+
+    auto &drmVmIds = osContext.getDrmVmIds();
+    EXPECT_EQ(0u, drmVmIds.size());
+}
diff --git a/shared/source/os_interface/linux/drm_neo.cpp b/shared/source/os_interface/linux/drm_neo.cpp
index f5a892ac78..f1c6692f02 100644
--- a/shared/source/os_interface/linux/drm_neo.cpp
+++ b/shared/source/os_interface/linux/drm_neo.cpp
@@ -243,6 +243,21 @@ void Drm::destroyDrmVirtualMemory(uint32_t drmVmId) {
     UNRECOVERABLE_IF(ret != 0);
 }
 
+uint32_t Drm::queryVmId(uint32_t drmContextId) {
+    drm_i915_gem_context_param param{};
+    param.ctx_id = drmContextId;
+    param.value = 0;
+    param.param = I915_CONTEXT_PARAM_VM;
+    auto retVal = this->ioctl(DRM_IOCTL_I915_GEM_CONTEXT_GETPARAM, &param);
+    DEBUG_BREAK_IF(retVal != 0);
+    (void)retVal;
+
+    auto vmId = static_cast<uint32_t>(param.value);
+    DEBUG_BREAK_IF(vmId == 0);
+
+    return vmId;
+}
+
 int Drm::getEuTotal(int &euTotal) {
     return getParamIoctl(I915_PARAM_EU_TOTAL, &euTotal);
 }
diff --git a/shared/source/os_interface/linux/drm_neo.h b/shared/source/os_interface/linux/drm_neo.h
index 30896eb631..c4761716f2 100644
--- a/shared/source/os_interface/linux/drm_neo.h
+++ b/shared/source/os_interface/linux/drm_neo.h
@@ -72,6 +72,7 @@ class Drm {
     void destroyDrmVirtualMemory(uint32_t drmVmId);
     uint32_t createDrmContext(uint32_t drmVmId);
     void destroyDrmContext(uint32_t drmContextId);
+    uint32_t queryVmId(uint32_t drmContextId);
     void setLowPriorityContextParam(uint32_t drmContextId);
 
     unsigned int bindDrmContext(uint32_t drmContextId, uint32_t deviceIndex, aub_stream::EngineType engineType);
diff --git a/shared/source/os_interface/linux/os_context_linux.cpp b/shared/source/os_interface/linux/os_context_linux.cpp
index ed2963fd17..28392948aa 100644
--- a/shared/source/os_interface/linux/os_context_linux.cpp
+++ b/shared/source/os_interface/linux/os_context_linux.cpp
@@ -41,6 +41,11 @@ OsContextLinux::OsContextLinux(Drm &drm, uint32_t contextId, DeviceBitfield devi
             }
             this->engineFlag = drm.bindDrmContext(drmContextId, deviceIndex, engineType);
             this->drmContextIds.push_back(drmContextId);
+
+            if (drm.isPerContextVMRequired()) {
+                drmVmId = drm.queryVmId(drmContextId);
+                this->drmVmIds.push_back(drmVmId);
+            }
         }
     }
 }
diff --git a/shared/source/os_interface/linux/os_context_linux.h b/shared/source/os_interface/linux/os_context_linux.h
index d949b2c77a..c1a9d0c77a 100644
--- a/shared/source/os_interface/linux/os_context_linux.h
+++ b/shared/source/os_interface/linux/os_context_linux.h
@@ -24,11 +24,13 @@ class OsContextLinux : public OsContext {
 
     unsigned int getEngineFlag() const { return engineFlag; }
     const std::vector<uint32_t> &getDrmContextIds() const { return drmContextIds; }
+    const std::vector<uint32_t> &getDrmVmIds() const { return drmVmIds; }
     Drm &getDrm() const;
 
   protected:
     unsigned int engineFlag = 0;
     std::vector<uint32_t> drmContextIds;
+    std::vector<uint32_t> drmVmIds;
     Drm &drm;
 };
 } // namespace NEO