diff --git a/run_tb.sh b/run_tb.sh
index 0ccc2f0..f7ed68e 100755
--- a/run_tb.sh
+++ b/run_tb.sh
@@ -1,10 +1,18 @@
 #!/bin/bash
 # run_tb.sh - Run Vivado XSIM testbench for a module
 #
-# Usage: ./run_tb.sh <module>
-#        ./run_tb.sh mod_add
+# Usage: ./run_tb.sh <module> [K] [CASE]
+#        ./run_tb.sh mod_add              # run the module's full xsim_run.tcl
+#        ./run_tb.sh top                  # ML-KEM KeyGen: all K, all cases
+#        ./run_tb.sh top 2                # only K=2 (ML-KEM-512), all its cases
+#        ./run_tb.sh top 2 3              # only K=2, only CASE=3
+#        ./run_tb.sh top 4 0              # only K=4 (ML-KEM-1024), CASE=0
 #        ./run_tb.sh --list
 #
+# For the 'top' module, K (2/3/4) and CASE select a single KAT run so you can
+# iterate quickly instead of elaborating 3 snapshots and running all 11 cases.
+# Other modules ignore the extra args and run their xsim_run.tcl verbatim.
+#
 # Prerequisites:
 #   Vivado 2019.2 at /opt/Xilinx/Vivado/2019.2/
 
@@ -31,6 +39,8 @@ if [ -z "$1" ]; then
 fi
 
 MODULE="$1"
+SEL_K="$2"          # optional: 2/3/4 (top module only)
+SEL_CASE="$3"       # optional: KAT case index (top module only)
 TB_DIR="$SCRIPT_DIR/sync_rtl/$MODULE/TB"
 TCL_FILE="$TB_DIR/xsim_run.tcl"
 
@@ -80,4 +90,46 @@ execute_tcl() {
     rm -f "$tmp_script"
 }
 
+# Fast single-K / single-case path for the 'top' (ML-KEM KeyGen) module.
+# Compiles the same RTL as the tcl, but elaborates ONLY the requested K and
+# runs ONLY the requested CASE(s). Avoids the 3-snapshot + 11-run full sweep.
+run_top_selected() {
+    local tcl_file="$1" k="$2" csel="$3"
+    set +e   # we do our own error handling / result parsing here
+
+    rm -rf xsim.dir .Xil
+
+    # Compile every xvlog line from the tcl, verbatim.
+    while read -r cmd; do
+        eval "$cmd" || { echo "COMPILE FAILED: $cmd"; return 1; }
+    done < <(grep -E '^xvlog ' "$tcl_file")
+
+    # Cases per K: K=2 has 0..4, K=3/4 have 0..2.
+    local cases
+    if [ "$k" = "2" ]; then cases="0 1 2 3 4"; else cases="0 1 2"; fi
+    if [ -n "$csel" ]; then cases="$csel"; fi
+
+    echo "  xelab tb_mlkem_kg_katK_xsim -generic_top KP=$k -s mlkem_kg_k$k --timescale 1ns/1ps"
+    xelab tb_mlkem_kg_katK_xsim -generic_top KP=$k -s mlkem_kg_k$k --timescale 1ns/1ps \
+        || { echo "ELAB FAILED for K=$k"; return 1; }
+
+    local fail=0
+    for c in $cases; do
+        local log="/tmp/run_tb_top_k${k}_c${c}.log"
+        echo "  xsim mlkem_kg_k$k -R -testplusarg CASE=$c   (log: $log)"
+        xsim "mlkem_kg_k$k" -R -testplusarg "CASE=$c" > "$log" 2>&1
+        local pf nf
+        pf=$(grep -oE 'PASS|FAIL' "$log" | tail -1)
+        nf=$(grep -c 'cannot be opened' "$log")
+        echo "    K=$k CASE=$c -> ${pf:-NORESULT}  (file-not-found=$nf)"
+        { [ "$pf" = "PASS" ] && [ "$nf" -eq 0 ]; } || fail=1
+    done
+    return $fail
+}
+
+if [ "$MODULE" = "top" ] && [ -n "$SEL_K" ]; then
+    run_top_selected "$TCL_FILE" "$SEL_K" "$SEL_CASE"
+    exit $?
+fi
+
 execute_tcl "$TCL_FILE"