diff --git a/Makefile b/Makefile
index 618933b..3c2de1d 100644
--- a/Makefile
+++ b/Makefile
@@ -14,11 +14,14 @@ CUDA_OUTPUT_FILE = -o $@
 # NVCC flags
 # -t=0 is short for --threads, 0 = number of CPUs on the machine
 NVCC_FLAGS = -O3 -t=0 --use_fast_math
-NVCC_LDFLAGS = -lcublas -lcublasLt -lcudnn
-NVCC_INCLUDES = -I../cudnn-frontend/include # TODO: Don't hardcode path
+NVCC_LDFLAGS = -lcublas -lcublasLt
+NVCC_INCLUDES =
 NVCC_LDLIBS =
 NCLL_INCUDES =
 
+# autodect a lot of various supports on current platform
+$(info ---------------------------------------------)
+
 ifneq ($(OS), Windows_NT)
   NVCC := $(shell which nvcc 2>/dev/null)
 
@@ -59,6 +62,21 @@ else
   endif
 endif
 
+# Check and include cudnn if available
+# Currently hard-coding a bunch of stuff here for Linux, todo make this better/nicer
+ifeq ($(SHELL_UNAME), Linux)
+  ifeq ($(shell [ -d ../cudnn-frontend/include ] && echo "exists"), exists)
+    $(info ✓ cuDNN found, will run with flash-attention)
+    NVCC_INCLUDES += -I../cudnn-frontend/include
+    NVCC_LDFLAGS += -lcudnn
+    CFLAGS += -DENABLE_CUDNN
+  else
+    $(info ✗ cuDNN not found, disabling flash-attention)
+    $(info ---> You need cuDNN: https://developer.nvidia.com/cudnn)
+    $(info ---> And the cuDNN front-end: https://github.com/NVIDIA/cudnn-frontend/tree/main)
+  endif
+endif
+
 # Check if OpenMP is available
 # This is done by attempting to compile an empty file with OpenMP flags
 # OpenMP makes the code a lot faster so I advise installing it
@@ -79,46 +97,47 @@ else
         LDFLAGS += -L/opt/homebrew/opt/libomp/lib
         LDLIBS += -lomp
         INCLUDES += -I/opt/homebrew/opt/libomp/include
-        $(info OpenMP found, compiling with OpenMP support)
+        $(info ✓ OpenMP found)
       else ifeq ($(shell [ -d /usr/local/opt/libomp/lib ] && echo "exists"), exists)
         # macOS with Homebrew on Intel
         CFLAGS += -Xclang -fopenmp -DOMP
         LDFLAGS += -L/usr/local/opt/libomp/lib
         LDLIBS += -lomp
         INCLUDES += -I/usr/local/opt/libomp/include
-        $(info OpenMP found, compiling with OpenMP support)
+        $(info ✓ OpenMP found)
       else
-        $(warning OpenMP not found, skipping OpenMP support)
+        $(info ✗ OpenMP not found)
       endif
     else
       # Check for OpenMP support in GCC or Clang on Linux
       ifeq ($(shell echo | $(CC) -fopenmp -x c -E - > /dev/null 2>&1; echo $$?), 0)
         CFLAGS += -fopenmp -DOMP
         LDLIBS += -lgomp
-        $(info OpenMP found, compiling with OpenMP support)
+        $(info ✓ OpenMP found)
       else
-        $(warning OpenMP not found, skipping OpenMP support)
+        $(info ✗ OpenMP not found)
       endif
     endif
   endif
 endif
 
+# Check if OpenMPI and NCCL are available, include them if so, for multi-GPU training
 ifeq ($(NO_MULTI_GPU), 1)
-  $(info Multi-GPU (OpenMPI + NCCL) is manually disabled)
+  $(info → Multi-GPU (OpenMPI + NCCL) is manually disabled)
 else
   ifneq ($(OS), Windows_NT)
     # Detect if running on macOS or Linux
     ifeq ($(SHELL_UNAME), Darwin)
-      $(warning Multi-GPU on CUDA on Darwin is not supported, skipping OpenMPI + NCCL support)
+      $(info ✗ Multi-GPU on CUDA on Darwin is not supported, skipping OpenMPI + NCCL support)
     else ifeq ($(shell [ -d /usr/lib/x86_64-linux-gnu/openmpi/lib/ ] && [ -d /usr/lib/x86_64-linux-gnu/openmpi/include/ ] && echo "exists"), exists)
-      $(info OpenMPI found, adding support)
+      $(info ✓ OpenMPI found, OK to train with multiple GPUs)
       NVCC_INCLUDES += -I/usr/lib/x86_64-linux-gnu/openmpi/include
       NVCC_LDFLAGS += -L/usr/lib/x86_64-linux-gnu/openmpi/lib/
       NVCC_LDLIBS += -lmpi -lnccl
       NVCC_FLAGS += -DMULTI_GPU
     else
-      $(warning OpenMPI is not found, disabling multi-GPU support)
-      $(warning On Linux you can try install OpenMPI with `sudo apt install openmpi-bin openmpi-doc libopenmpi-dev`)
+      $(info ✗ OpenMPI is not found, disabling multi-GPU support)
+      $(info ---> On Linux you can try install OpenMPI with `sudo apt install openmpi-bin openmpi-doc libopenmpi-dev`)
     endif
   endif
 endif
@@ -145,12 +164,14 @@ TARGETS = train_gpt2 test_gpt2
 
 # Conditional inclusion of CUDA targets
 ifeq ($(NVCC),)
-    $(info nvcc not found, skipping CUDA builds)
+    $(info ✗ nvcc not found, skipping GPU/CUDA builds)
 else
-    $(info nvcc found, including CUDA builds)
+    $(info ✓ nvcc found, including GPU/CUDA support)
     TARGETS += train_gpt2cu test_gpt2cu train_gpt2fp32cu test_gpt2fp32cu
 endif
 
+$(info ---------------------------------------------)
+
 all: $(TARGETS)
 
 train_gpt2: train_gpt2.c
diff --git a/train_gpt2.cu b/train_gpt2.cu
index ee26322..b74b59c 100644
--- a/train_gpt2.cu
+++ b/train_gpt2.cu
@@ -32,7 +32,6 @@ make train_gpt2cu PRECISION=FP32
 This reads & runs in fp32, B=4, T=64, LR=1e-4, val/sample never (200),
 -a 1 is "overfit single batch", -x 10 is 10 iterations, and -f 0 disables tf32
 */
-#define ENABLE_CUDNN // can be enabled via nvcc "-DENABLE_CUDNN"
 
 #include <stdio.h>
 #include <stdlib.h>
@@ -2474,11 +2473,11 @@ int main(int argc, char *argv[]) {
     cuda_arch_major = deviceProp.major;
     cuda_arch_minor = deviceProp.minor;
 
-    // setup cuBLAS and cuBLASLt
+    // set up cuBLAS and cuBLASLt
     cublasCheck(cublasCreate(&cublas_handle));
     cublasCheck(cublasLtCreate(&cublaslt_handle));
     cudaCheck(cudaMalloc(&cublaslt_workspace, cublaslt_workspace_size));
-
+    // setup compute precision settings for cublas
     // TF32 precision is equivalent to torch.set_float32_matmul_precision('high')
     int enable_tf32 = cuda_arch_major >= 8 ? 1 : 0;
     if (override_enable_tf32 == 0) { enable_tf32 = 0; } // force to zero via arg
@@ -2487,6 +2486,7 @@ int main(int argc, char *argv[]) {
     cublasCheck(cublasSetMathMode(cublas_handle, cublas_math_mode));
     if(cublas_compute_type); // unused in BF16 mode, avoid warning
 
+    // set up cuDNN
     #ifdef ENABLE_CUDNN
     checkCudnnErr(cudnnCreate(&cudnn_handle));
     #endif
@@ -2648,7 +2648,7 @@ int main(int argc, char *argv[]) {
     // add a total average, for optimizations that are only mild improvements (excluding 1st batch as warmup)
     printf0("total average iteration time: %f ms\n", total_sum_iteration_time_s / (train_num_batches-1) * 1000);
 
-    // free
+    // free and destroy everything
     dataloader_free(&train_loader);
     dataloader_free(&val_loader);
     tokenizer_free(&tokenizer);
@@ -2656,13 +2656,10 @@ int main(int argc, char *argv[]) {
     free(cpu_logits_raw);
     free(cpu_logits);
     free(gen_tokens);
-
     #ifdef ENABLE_CUDNN
-    if (cudnn_workspace != NULL) {
-        cudaCheck(cudaFree(cudnn_workspace));
-    }
-    #endif
+    if (cudnn_workspace != NULL) { cudaCheck(cudaFree(cudnn_workspace)); }
     checkCudnnErr(cudnnDestroy(cudnn_handle));
+    #endif
     cudaCheck(cudaFree(cublaslt_workspace));
     cublasCheck(cublasDestroy(cublas_handle));
     cublasCheck(cublasLtDestroy(cublaslt_handle));