run_all_attention.sh

#!/bin/bash
###############################################################################
# run_all_attention.sh - sweep every entry in model_configs.json's
# attention_models through bench_attention.py, across backends / phases /
# batch sizes / seq lens.
#
# Mirrors run_all_models.sh's docker bring-up, but attention is a single-GPU
# op so it runs `python3 bench_attention.py` directly (no torchrun / EP).
#
# Env knobs (all optional):
#   GPU_ID             default 0 (which visible device to run on)
#   MODELS             comma-separated subset of attention_models keys (default: all)
#   BACKENDS           default 'triton'   (e.g. 'triton,aiter')
#   PHASES             default 'extend,decode'
#   MLA_MODE           default 'materialized'  ('absorbed' = compressed-latent decode)
#   DTYPE              default 'bf16'
#   BATCH_SIZES        default '1,8,32,128'
#   SEQ_LENS           default '1024,4096,8192,16384'
#   WARMUP / ITERS     default 5 / 20
#   KINETO_NUM_TESTS   default 15
#   DOCKER_IMAGE       default lmsysorg/sglang:v0.5.11-rocm720-mi35x
#   CONTAINER_NAME     default bench_attn_sweep
#   CONFIG_FILE        default $SCRIPT_DIR/model_configs.json
###############################################################################
set -euo pipefail

SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)

# NOTE: this image is the only prebuilt one carrying
# sglang.test.kits.attention_unittest (required by bench_attention.py).
DOCKER_IMAGE=${DOCKER_IMAGE:-lmsysorg/sglang-rocm:v0.5.12.post1-rocm720-mi35x-20260601}
CONTAINER_NAME=${CONTAINER_NAME:-bench_attn_sweep}

GPU_ID=${GPU_ID:-0}
CONFIG_FILE=${CONFIG_FILE:-${SCRIPT_DIR}/model_configs.json}
# Attention backend + MLA mode are NOT swept: bench_attention.py lets SGLang
# auto-select the backend (aiter/triton by arch & head count) and auto-derive
# the MLA mode per phase. We only enable AITER via the env var below.
PHASES=${PHASES:-extend,decode}
# Force every model's attention backend to this (default aiter): MLA models with
# kv-head count not in {16,128} (e.g. Kimi-K2) and DSA models (GLM-5) would
# otherwise be benched on triton. Set FORCE_BACKEND='' to restore auto-select.
FORCE_BACKEND=${FORCE_BACKEND-aiter}
DTYPE=${DTYPE:-bf16}
# KV-cache storage dtype, independent of DTYPE (q/k/v compute). 'auto' reuses
# DTYPE (original bf16 behaviour); 'fp8_e4m3'/'fp8_e5m2' bench an fp8 KV cache.
KV_CACHE_DTYPE=${KV_CACHE_DTYPE:-auto}
# DSA (DeepSeek Sparse Attention) decode kernel for DSA models (GLM-5). Empty =
# use each model's serve.dsa_decode_backend from model_configs.json (GLM-5 ->
# tilelang per the AMD cookbook). Set DSA_DECODE_BACKEND=none to force the dense
# absorbed-MLA triton approximation, or =aiter/tilelang to override globally.
DSA_DECODE_BACKEND=${DSA_DECODE_BACKEND:-}
# APPLY_MODEL_ENV=1 (default) exports each model's serve.env from the config
# (e.g. Kimi-K2 needs SGLANG_ROCM_FUSED_DECODE_MLA=0). Set 0 to skip.
APPLY_MODEL_ENV=${APPLY_MODEL_ENV:-1}
BATCH_SIZES=${BATCH_SIZES:-1,8,32,128}
SEQ_LENS=${SEQ_LENS:-1024,4096,8192,16384}
WARMUP=${WARMUP:-5}
ITERS=${ITERS:-20}
KINETO_NUM_TESTS=${KINETO_NUM_TESTS:-15}

if [[ ! -f "$CONFIG_FILE" ]]; then
    echo "ERROR: $CONFIG_FILE not found" >&2
    exit 1
fi

if [[ -n "${MODELS:-}" ]]; then
    IFS=',' read -ra MODEL_LIST <<< "$MODELS"
else
    mapfile -t MODEL_LIST < <(python3 -c "
import json
with open('$CONFIG_FILE') as fh:
    doc = json.load(fh)
for name in doc.get('attention_models', {}):
    print(name)
")
fi

if [[ ${#MODEL_LIST[@]} -eq 0 ]]; then
    echo "ERROR: no attention models found (CONFIG_FILE=$CONFIG_FILE MODELS='${MODELS:-}')" >&2
    exit 1
fi

LOG_DIR=${LOG_DIR:-${SCRIPT_DIR}/logs}
mkdir -p "$LOG_DIR"
STAMP=$(date +%Y%m%d_%H%M%S)
RUN_LOG="$LOG_DIR/attn_sweep_${STAMP}.log"
RESULTS_DIR="$LOG_DIR/attn_results_${STAMP}"
TRACE_DIR="$LOG_DIR/attn_traces_${STAMP}"
mkdir -p "$RESULTS_DIR" "$TRACE_DIR"

cat <<EOF | tee "$RUN_LOG"
====================================================================
[run_all_attention] DOCKER_IMAGE   = $DOCKER_IMAGE
[run_all_attention] CONTAINER_NAME = $CONTAINER_NAME
[run_all_attention] GPU_ID         = $GPU_ID
[run_all_attention] MODELS (${#MODEL_LIST[@]}) = ${MODEL_LIST[*]}
[run_all_attention] BACKEND        = ${FORCE_BACKEND:-SGLang auto-select} (SGLANG_USE_AITER=1)
[run_all_attention] PHASES         = $PHASES
[run_all_attention] MLA_MODE       = auto-derived per phase
[run_all_attention] DTYPE          = $DTYPE
[run_all_attention] KV_CACHE_DTYPE = $KV_CACHE_DTYPE
[run_all_attention] BATCH_SIZES    = $BATCH_SIZES
[run_all_attention] SEQ_LENS       = $SEQ_LENS
[run_all_attention] WARMUP/ITERS   = $WARMUP / $ITERS
[run_all_attention] RESULTS_DIR    = $RESULTS_DIR
[run_all_attention] TRACE_DIR      = $TRACE_DIR
====================================================================
EOF

docker rm -f "$CONTAINER_NAME" 2>/dev/null || true

docker run -d \
    --name "$CONTAINER_NAME" \
    --ipc=host \
    --network=host \
    --device=/dev/kfd \
    --device=/dev/dri \
    --cap-add=SYS_PTRACE \
    --group-add video \
    --privileged \
    --shm-size=64g \
    -v /home:/io \
    -v /apps:/apps \
    -w "$SCRIPT_DIR" \
    "$DOCKER_IMAGE" sleep infinity 2>&1 | tee -a "$RUN_LOG"

if [[ "${ATTACH:-0}" == "1" ]]; then
    docker exec -it "$CONTAINER_NAME" /bin/bash
    exit $?
fi

_safe() { echo "${1//[^A-Za-z0-9._-]/_}"; }

# Backend + MLA mode are auto-selected by bench_attention.py (SGLang's own
# logic), so we no longer sweep them: one run per model, tagged by model only.
for model in "${MODEL_LIST[@]}"; do
    model_safe=$(_safe "$model")
    tag="${model_safe}"
    out_json="$RESULTS_DIR/${tag}.json"
    sub_trace="$TRACE_DIR/$tag"
    mkdir -p "$sub_trace"

    # Per-model serve params from model_configs.json: env vars to export (e.g.
    # Kimi-K2 SGLANG_ROCM_FUSED_DECODE_MLA=0) and the DSA decode backend default
    # (e.g. GLM-5 -> tilelang). A non-empty global DSA_DECODE_BACKEND overrides.
    serve_eval=$(python3 -c "
import json, shlex
d = json.load(open('$CONFIG_FILE'))['attention_models'].get('$model', {})
s = d.get('serve', {}) or {}
env = (s.get('env', {}) or {}) if '$APPLY_MODEL_ENV' == '1' else {}
print('MODEL_ENVFLAGS=' + shlex.quote(' '.join(f'-e {k}={v}' for k, v in env.items())))
print('SERVE_DSA=' + shlex.quote(str(s.get('dsa_decode_backend', '') or '')))
")
    eval "$serve_eval"
    dsa_be="${DSA_DECODE_BACKEND:-$SERVE_DSA}"
    dsa_flag=""
    if [[ -n "$dsa_be" && "$dsa_be" != "none" ]]; then
        dsa_flag="--dsa-decode-backend $dsa_be"
    fi

    echo "===== model=$model  backend=AUTO  dsa_decode=${dsa_be:-(none)}  env=[${MODEL_ENVFLAGS}]  phases=$PHASES =====" | tee -a "$RUN_LOG"
    docker exec \
        -e SGLANG_USE_AITER=1 \
        -e BENCH_FORCE_BACKEND="$FORCE_BACKEND" \
        -e HIP_VISIBLE_DEVICES="$GPU_ID" \
        -e CUDA_VISIBLE_DEVICES="$GPU_ID" \
        -e OMP_NUM_THREADS=8 \
        ${MODEL_ENVFLAGS} \
        "$CONTAINER_NAME" bash -lc "
set -euo pipefail
cd $SCRIPT_DIR
python3 bench_attention.py \\
    --config-file '$CONFIG_FILE' \\
    --model-name '$model' \\
    --dtype $DTYPE \\
    --kv-cache-dtype $KV_CACHE_DTYPE \\
    ${FORCE_BACKEND:+--force-backend $FORCE_BACKEND} \\
    $dsa_flag \\
    --phases $PHASES \\
    --batch-sizes $BATCH_SIZES \\
    --seq-lens $SEQ_LENS \\
    --warmup $WARMUP --iters $ITERS \\
    --kineto-num-tests $KINETO_NUM_TESTS \\
    --trace-dir '$sub_trace' \\
    --output '$out_json'
" 2>&1 | tee -a "$RUN_LOG" || echo "WARN: model=$model failed" | tee -a "$RUN_LOG"
done

if [[ "${KEEP_CONTAINER:-0}" != "1" ]]; then
    docker rm -f "$CONTAINER_NAME" >/dev/null 2>&1 || true
fi

echo "[run_all_attention] done. log=$RUN_LOG results=$RESULTS_DIR traces=$TRACE_DIR"