From 8af5e26717793b58078c23cadfededa337d7d3cb Mon Sep 17 00:00:00 2001
From: Qianfeng Zhang <Qianfeng.Zhang@amd.com>
Date: Thu, 5 Feb 2026 14:56:31 +0000
Subject: [PATCH] Add softmax selection to two of the testing scripts

---
 .../scripts/test_jagged_causal_mattn0_full0.sh   | 16 +++++++++++-----
 .../test_jagged_causal_mattn256_full256.sh       | 16 +++++++++++-----
 2 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/example/ck_tile/18_hstu_attention/scripts/test_jagged_causal_mattn0_full0.sh b/example/ck_tile/18_hstu_attention/scripts/test_jagged_causal_mattn0_full0.sh
index 99de954415..85f636a828 100644
--- a/example/ck_tile/18_hstu_attention/scripts/test_jagged_causal_mattn0_full0.sh
+++ b/example/ck_tile/18_hstu_attention/scripts/test_jagged_causal_mattn0_full0.sh
@@ -1,14 +1,20 @@
 #!/bin/bash
 
-ndist=0
+set +x
 
+ndist=0
+BUILD=build
+
+USE_SOFTMAX=0
 if [ $# -ge 1 ]; then
-    ndist=$1
+    USE_SOFTMAX=$1
 fi
 
-set +x
-BUILD=build
-EXE=$BUILD/bin/tile_example_hstu_attention
+if [ $USE_SOFTMAX -eq 1 ]; then
+    EXE="$BUILD/bin/tile_example_hstu_attention -softmax=1"
+else
+    EXE="$BUILD/bin/tile_example_hstu_attention"
+fi
 
 dtype="bf16"
 
diff --git a/example/ck_tile/18_hstu_attention/scripts/test_jagged_causal_mattn256_full256.sh b/example/ck_tile/18_hstu_attention/scripts/test_jagged_causal_mattn256_full256.sh
index 605891f815..2667ef8c96 100644
--- a/example/ck_tile/18_hstu_attention/scripts/test_jagged_causal_mattn256_full256.sh
+++ b/example/ck_tile/18_hstu_attention/scripts/test_jagged_causal_mattn256_full256.sh
@@ -1,14 +1,20 @@
 #!/bin/bash
 
-ndist=0
+set +x
 
+ndist=0
+BUILD=build
+
+USE_SOFTMAX=0
 if [ $# -ge 1 ]; then
-    ndist=$1
+    USE_SOFTMAX=$1
 fi
 
-set +x
-BUILD=build
-EXE=$BUILD/bin/tile_example_hstu_attention
+if [ $USE_SOFTMAX -eq 1 ]; then
+    EXE="$BUILD/bin/tile_example_hstu_attention -softmax=1"
+else
+    EXE="$BUILD/bin/tile_example_hstu_attention"
+fi
 
 dtype="bf16"