diff --git a/Dockerfile_trtllm b/Dockerfile_trtllm
index 68a6e563..7213a2b5 100644
--- a/Dockerfile_trtllm
+++ b/Dockerfile_trtllm
@@ -1,7 +1,7 @@
-ARG CUDA_ARCH_LIST="75-real;80-real;86-real;89-real;90-real"
-ARG OMPI_VERSION="4.1.7rc1"
-ARG BUILD_TYPE=release
-ARG IS_GHA_BUILD=false
+ARG cuda_arch_list="75-real;80-real;86-real;89-real;90-real"
+ARG ompi_version="4.1.7rc1"
+ARG build_type=release
+ARG is_gha_build=false
 
 # Build dependencies resolver stage
 FROM lukemathwalker/cargo-chef:latest-rust-1.84.0 AS chef
@@ -48,9 +48,9 @@ ENV TENSORRT_INSTALL_PREFIX=/usr/local/tensorrt
 
 # Install OpenMPI
 FROM cuda-builder AS mpi-builder
-ARG OMPI_VERSION
+ARG ompi_version
 
-ENV OMPI_TARBALL_FILENAME="openmpi-$OMPI_VERSION.tar.bz2"
+ENV OMPI_TARBALL_FILENAME="openmpi-$ompi_version.tar.bz2"
 RUN wget "https://download.open-mpi.org/release/open-mpi/v4.1/$OMPI_TARBALL_FILENAME" -P /opt/src && \
     mkdir /usr/src/mpi && \
     tar -xf "/opt/src/$OMPI_TARBALL_FILENAME" -C /usr/src/mpi --strip-components=1 && \
@@ -79,7 +79,7 @@ RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | bash -s -- -y &&
 # Build actual TGI
 ARG build_type
 ARG cuda_arch_list
-ARG is_gha_build
+ARG is_gha_build="false"
 
 # SCCACHE Specifics args - before finding a better, more generic, way...
 ARG aws_access_key_id
@@ -91,6 +91,7 @@ ARG sscache_region
 
 ENV CMAKE_PREFIX_PATH="/usr/local/mpi:/usr/local/tensorrt:$CMAKE_PREFIX_PATH"
 ENV CUDA_ARCH_LIST=${cuda_arch_list}
+ENV IS_GHA_BUILD ${is_gha_build}
 ENV LD_LIBRARY_PATH="/usr/local/mpi/lib:$LD_LIBRARY_PATH"
 ENV PKG_CONFIG_PATH="/usr/local/mpi/lib/pkgconfig:$PKG_CONFIG_PATH"
 ENV USE_LLD_LINKER=ON
diff --git a/backends/trtllm/CMakeLists.txt b/backends/trtllm/CMakeLists.txt
index ad53cb75..55c61f5b 100644
--- a/backends/trtllm/CMakeLists.txt
+++ b/backends/trtllm/CMakeLists.txt
@@ -85,6 +85,10 @@ endif ()
 #### Unit Tests ####
 if (${TGI_TRTLLM_BACKEND_BUILD_TESTS})
     message(STATUS "Building tests")
+    option(TGI_TRTLLM_BACKEND_ENABLE_ASAN "Enable AddressSanitizer")
+    option(TGI_TRTLLM_BACKEND_ENABLE_UBSAN "Enable UndefinedSanitizer")
+    option(TGI_TRTLLM_BACKEND_ENABLE_MSAN "Enable MemorySanitizer")
+
     FetchContent_Declare(
             Catch2
             URL https://github.com/catchorg/Catch2/archive/refs/tags/v3.7.1.tar.gz
@@ -112,32 +116,20 @@ if (${TGI_TRTLLM_BACKEND_BUILD_TESTS})
     target_link_libraries(tgi_trtllm_backend_tests PRIVATE tensorrt_llm nvinfer_plugin_tensorrt_llm tensorrt_llm_nvrtc_wrapper)
 
     if (CMAKE_BUILD_TYPE MATCHES "Debug")
-        if (NOT "${ASan_${CMAKE_C_COMPILER_ID}_FLAGS}" STREQUAL "")
+        if (${TGI_TRTLLM_BACKEND_ENABLE_ASAN})
+            message(STATUS "Enabled AddressSanitizer")
             target_link_options(tgi_trtllm_backend_tests BEFORE PUBLIC -fsanitize=address)
         endif ()
 
-        if (NOT "${UBSan_${CMAKE_C_COMPILER_ID}_FLAGS}" STREQUAL "")
+        if (${TGI_TRTLLM_BACKEND_ENABLE_UBSAN})
+            message(STATUS "Enabled UndefinedSanitizer")
             target_link_options(tgi_trtllm_backend_tests BEFORE PUBLIC -fsanitize=undefined)
         endif ()
 
-        if (NOT "${MSan_${CMAKE_C_COMPILER_ID}_FLAGS}" STREQUAL "")
+        if (${TGI_TRTLLM_BACKEND_ENABLE_MSAN})
+            message(STATUS "Enabled MemorySanitizer")
             target_link_options(tgi_trtllm_backend_tests BEFORE PUBLIC -fsanitize=memory)
         endif ()
-
-        #        execute_process(COMMAND CMAKE_C_COMPILER --print-file-name=libasan.so OUTPUT_VARIABLE LIBASAN_PATH OUTPUT_STRIP_TRAILING_WHITESPACE RESULT_VARIABLE LIBASAN_FOUND)
-        #        execute_process(COMMAND CMAKE_C_COMPILER --print-file-name=libubsan.so OUTPUT_VARIABLE LIBUBSAN_PATH OUTPUT_STRIP_TRAILING_WHITESPACE RESULT_VARIABLE LIBUBSAN_FOUND)
-
-        #        if (LIBASAN_FOUND)
-        #            set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS}  -fsanitize=address")
-        #            set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fsanitize=address")
-        #            target_link_options(tgi_trtllm_backend_tests BEFORE PUBLIC -fsanitize=address)
-        #        endif ()
-
-        #        if (LIBUBSAN_FOUND)
-        #            set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS}  -fsanitize=undefined")
-        #            set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fsanitize=undefined")
-        #            target_link_options(tgi_trtllm_backend_tests BEFORE PUBLIC -fsanitize=undefined)
-        #        endif ()
     endif ()
 
     if(CMAKE_BUILD_TYPE MATCHES "Debug")