Skip to content

chore(cuda): add nsys & ncu in run.sh #548

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
gaxiom wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

chore(cuda): add nsys & ncu in run.sh #548

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
48ebd6f
add nsys & ncu
gaxiom Jan 2, 2026
eacc3d5
-logs
gaxiom Jan 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 42 additions & 31 deletions run.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,11 @@
# --mode <MODE> Set the proving mode (default: prove-app)
# Valid modes: prove-app, prove-stark
# --cuda Force CUDA acceleration (auto-detected if nvidia-smi available)
#
# --nsys Run with nsys profiling
# --block <N> Set the block number to prove (default: 23992138)
# --ncu <KERNEL> Run with ncu profiling for specified kernel
# --launch-skip N Skip N kernel launches (for ncu)
# --launch-count N Profile N kernel launches (for ncu)
# Examples:
# ./run.sh # Run with defaults
# ./run.sh --mode prove-stark # Run in prove-stark mode
Expand Down Expand Up @@ -98,22 +102,19 @@ fi
MODE_OVERRIDE=""
USE_CUDA=false
CUDA_REASON=""
USE_NSYS=false
USE_NCU=false

while [[ $# -gt 0 ]]; do
case $1 in
--mode)
MODE_OVERRIDE="$2"
shift 2
;;
--cuda)
USE_CUDA=true
CUDA_REASON="requested via script argument"
shift
;;
*)
echo "Unknown argument: $1"
exit 1
;;
--mode) MODE_OVERRIDE="$2"; shift 2 ;;
--block) BLOCK_NUMBER_OVERRIDE="$2"; shift 2 ;;
--cuda) USE_CUDA=true; CUDA_REASON="requested via --cuda"; shift ;;
--nsys) USE_NSYS=true; USE_CUDA=true; CUDA_REASON="requested via --nsys"; shift ;;
--ncu) USE_NCU=true; USE_CUDA=true; CUDA_REASON="requested via --ncu"; ncu_kernel="$2"; shift 2 ;;
--launch-skip) launch_skip="$2"; shift 2 ;;
--launch-count) launch_count="$2"; shift 2 ;;
*) echo "Unknown argument: $1"; exit 1 ;;
esac
done

Expand All @@ -126,8 +127,10 @@ if [ "$USE_CUDA" = "true" ]; then
echo "Using CUDA acceleration ($CUDA_REASON)."
fi

if [ "$NVIDIA_SMI_READY" = "true" ]; then
if [ "$NVIDIA_SMI_READY" = "true" ] && [ "$USE_NSYS" = "false" ]; then
start_gpu_monitor "$GPU_MONITOR_INTERVAL"
elif [ "$USE_NSYS" = "true" ]; then
echo "GPU memory monitoring disabled for nsys profiling."
else
echo "nvidia-smi not detected; GPU memory monitoring disabled."
fi
Expand All @@ -151,7 +154,7 @@ cd "$WORKDIR"

PROFILE="release"
FEATURES="metrics,jemalloc,unprotected,nightly-features"
BLOCK_NUMBER=23992138
BLOCK_NUMBER="${BLOCK_NUMBER_OVERRIDE:-23992138}"
# switch to +nightly-2025-08-19 if using tco
TOOLCHAIN="+nightly-2025-08-19" # "+stable"
BIN_NAME="openvm-reth-benchmark-bin"
Expand Down Expand Up @@ -182,7 +185,9 @@ echo "Unsupported architecture: $arch"
exit 1
;;
esac
export JEMALLOC_SYS_WITH_MALLOC_CONF="retain:true,background_thread:true,metadata_thp:always,dirty_decay_ms:10000,muzzy_decay_ms:10000,abort_conf:true"
if [ "$USE_NSYS" = "false" ]; then
export JEMALLOC_SYS_WITH_MALLOC_CONF="retain:true,background_thread:true,metadata_thp:always,dirty_decay_ms:10000,muzzy_decay_ms:10000,abort_conf:true"
fi
RUSTFLAGS=$RUSTFLAGS cargo $TOOLCHAIN build --bin $BIN_NAME --profile=$PROFILE --no-default-features --features=$FEATURES
PARAMS_DIR="$HOME/.openvm/params/"

Expand All @@ -193,17 +198,23 @@ else
TARGET_DIR="$PROFILE"
fi

RUST_LOG="info,p3_=warn" OUTPUT_PATH="metrics.json" VPMM_PAGES=$VPMM_PAGES VPMM_PAGE_SIZE=$VPMM_PAGE_SIZE ./target/$TARGET_DIR/$BIN_NAME \
--kzg-params-dir $PARAMS_DIR \
--mode $MODE \
--block-number $BLOCK_NUMBER \
--rpc-url $RPC_1 \
--cache-dir rpc-cache \
--app-log-blowup 1 \
--leaf-log-blowup 1 \
--internal-log-blowup 2 \
--root-log-blowup 3 \
--max-segment-length $MAX_SEGMENT_LENGTH \
--segment-max-cells $SEGMENT_MAX_CELLS \
--num-children-leaf 1 \
--num-children-internal 3
BIN=./target/$TARGET_DIR/$BIN_NAME
BIN_ARGS="--kzg-params-dir $PARAMS_DIR --mode $MODE --block-number $BLOCK_NUMBER --rpc-url $RPC_1 --cache-dir rpc-cache \
--app-log-blowup 1 --leaf-log-blowup 1 --internal-log-blowup 2 --root-log-blowup 3 \
--max-segment-length $MAX_SEGMENT_LENGTH --segment-max-cells $SEGMENT_MAX_CELLS --num-children-leaf 1 --num-children-internal 3"
export VPMM_PAGES=$VPMM_PAGES VPMM_PAGE_SIZE=$VPMM_PAGE_SIZE

if [ "$USE_NSYS" = "true" ]; then
NSYS_OUTPUT="reth.nsys-rep"
nsys profile --trace=cuda --force-overwrite=true -o "$NSYS_OUTPUT" $BIN $BIN_ARGS
nsys stats --report cuda_gpu_kern_sum "$NSYS_OUTPUT"
nsys stats --report cuda_gpu_mem_time_sum "$NSYS_OUTPUT"
nsys stats --report cuda_gpu_mem_size_sum "$NSYS_OUTPUT"
elif [ "$USE_NCU" = "true" ]; then
NCU_OUTPUT="reth-${ncu_kernel}.ncu-rep"
sudo env PATH=$PATH ncu --target-processes all --kernel-name "$ncu_kernel" -f -o "$NCU_OUTPUT" \
--launch-skip "${launch_skip:-0}" --launch-count "${launch_count:-4}" --set full $BIN $BIN_ARGS
else
export RUST_LOG="info,p3_=warn" OUTPUT_PATH="metrics.json"
$BIN $BIN_ARGS
fi
Loading