diff --git a/runtime/kernel/kernel.cpp b/runtime/kernel/kernel.cpp
index bd4300aa90..f12025cbd8 100644
--- a/runtime/kernel/kernel.cpp
+++ b/runtime/kernel/kernel.cpp
@@ -1392,9 +1392,10 @@ cl_int Kernel::setArgImageWithMipLevel(uint32_t argIndex,
         patch<uint32_t, uint32_t>(kernelArgInfo.offsetHeap, crossThreadData, kernelArgInfo.offsetObjectId);
         patch<uint32_t, cl_uint>(imageDesc.num_mip_levels, crossThreadData, kernelArgInfo.offsetNumMipLevels);
 
+        auto pixelSize = pImage->getSurfaceFormatInfo().ImageElementSizeInBytes;
         patch<uint64_t, uint64_t>(pImage->getGraphicsAllocation()->getGpuAddress(), crossThreadData, kernelArgInfo.offsetFlatBaseOffset);
-        patch<uint32_t, size_t>(imageDesc.image_width - 1, crossThreadData, kernelArgInfo.offsetFlatWidth);
-        patch<uint32_t, size_t>(imageDesc.image_height - 1, crossThreadData, kernelArgInfo.offsetFlatHeight);
+        patch<uint32_t, size_t>((imageDesc.image_width * pixelSize) - 1, crossThreadData, kernelArgInfo.offsetFlatWidth);
+        patch<uint32_t, size_t>((imageDesc.image_height * pixelSize) - 1, crossThreadData, kernelArgInfo.offsetFlatHeight);
         patch<uint32_t, size_t>(imageDesc.image_row_pitch - 1, crossThreadData, kernelArgInfo.offsetFlatPitch);
 
         addAllocationToCacheFlushVector(argIndex, pImage->getGraphicsAllocation());
diff --git a/unit_tests/kernel/kernel_image_arg_tests.cpp b/unit_tests/kernel/kernel_image_arg_tests.cpp
index c80be5b3ff..d6083ee8d5 100644
--- a/unit_tests/kernel/kernel_image_arg_tests.cpp
+++ b/unit_tests/kernel/kernel_image_arg_tests.cpp
@@ -60,15 +60,16 @@ TEST_F(KernelImageArgTest, givenKernelWithFlatImageTokensWhenArgIsSetThenPatchAl
 
     pKernel->setArg(0, sizeof(memObj), &memObj);
     auto crossThreadData = reinterpret_cast<uint32_t *>(pKernel->getCrossThreadData());
+    auto pixelSize = image->getSurfaceFormatInfo().ImageElementSizeInBytes;
 
     auto offsetFlatBaseOffset = ptrOffset(crossThreadData, pKernel->getKernelInfo().kernelArgInfo[0].offsetFlatBaseOffset);
     EXPECT_EQ(imageBaseAddress, *reinterpret_cast<uint64_t *>(offsetFlatBaseOffset));
 
     auto offsetFlatWidth = ptrOffset(crossThreadData, pKernel->getKernelInfo().kernelArgInfo[0].offsetFlatWidth);
-    EXPECT_EQ(imageWidth - 1, *offsetFlatWidth);
+    EXPECT_EQ(static_cast<uint32_t>((imageWidth * pixelSize) - 1), *offsetFlatWidth);
 
     auto offsetFlatHeight = ptrOffset(crossThreadData, pKernel->getKernelInfo().kernelArgInfo[0].offsetFlatHeight);
-    EXPECT_EQ(imageHeight - 1, *offsetFlatHeight);
+    EXPECT_EQ(static_cast<uint32_t>((imageHeight * pixelSize) - 1), *offsetFlatHeight);
 
     auto offsetFlatPitch = ptrOffset(crossThreadData, pKernel->getKernelInfo().kernelArgInfo[0].offsetFlatPitch);
     EXPECT_EQ(imageRowPitch - 1, *offsetFlatPitch);