diff --git a/.gitignore b/.gitignore index 6246c25..c58219b 100644 --- a/.gitignore +++ b/.gitignore @@ -76,14 +76,21 @@ build/ sglang_sm120_cache/ vllm_sm120_cache/ *_sm120_cache/ +sglang_nightly_cu13_cache/ # 一次性 bench 输出 bench-output/ +# 单数 dataset/（可重新下载，与 datasets/ 同源） +dataset/ + # 逐请求原始日志（体积大；汇总见 results.json / report.md） experiments/**/raw_outputs/ **/dummy_sharegpt.json +# pro6000 统一 bench CLI 的一次性测试输出（ops_ 前缀 run_id） +experiments/pro6000/*/results/ops_*/ + # Virtual environments (keep only docs under envs/) envs/* !envs/README.md @@ -96,3 +103,6 @@ skills-lock.json # 临时备份文件 *.bak.glm52orig +*.bak +*.bak.* +*.tmp diff --git a/deploy/profiles/pro6000/dsv4_pro6000_sglang_tp16.env b/deploy/profiles/pro6000/dsv4_pro6000_sglang_tp16.env new file mode 100644 index 0000000..365629a --- /dev/null +++ b/deploy/profiles/pro6000/dsv4_pro6000_sglang_tp16.env @@ -0,0 +1,34 @@ +# DeepSeek-V4-Pro SGLang multi-node TP=16 EP=2 deployment profile (2x RTX 6000D). +# Nodes: 10.101.0.11 (rank 0) + 10.101.0.13 (rank 1), 16x RTX 6000D (85GB). +# Key: SGLANG_SHARED_EXPERT_TP1=1 + --ep-size 2 -> TP=16 works with FP8 block=128. +# Model-team only. Ops only run `python -m sskj.bench` against the served URL. + +PLATFORM=pro6000 +EXPERIMENT=dsv4_pro6000_sglang_tp16 +MODEL_NAME=DeepSeek-V4-Pro +ENGINE=sglang +RUNTIME=docker +DOCKER_IMAGE=lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45 +CONTAINER_NAME=${EXPERIMENT}_node${NODE_RANK} +MODEL_PATH=/data/hf_models/DeepSeek-V4-Pro +SERVED_MODEL_NAME=default +PORT=30000 +HEALTH_PATH=/health +HEALTH_HOST=10.101.0.11 +HEALTH_WAIT_S=2400 +CONTAINER_PYTHON=python3 + +# ---- Multi-node topology (rank order; rank 0 exposes the HTTP API) ---- +NNODES=2 +NODE_HOSTS="10.101.0.11 10.101.0.13" +NODE_SSH_USER=root +MASTER_IP=10.101.0.11 +DIST_PORT=20000 + +DEVICE_VARS="CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7" +ENGINE_ENV="NCCL_SOCKET_IFNAME=eth0 NCCL_DEBUG=WARN SGLANG_SHARED_EXPERT_TP1=1 PYTHONUNBUFFERED=1 HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1" + +DOCKER_FLAGS="--gpus all --network host --ipc=host --ulimit memlock=-1 --shm-size=20g" +VOLUMES="${MODEL_PATH}:${MODEL_PATH}:ro" + +LAUNCH_ARGS="--model-path ${MODEL_PATH} --tp-size ${TP} --ep-size 2 --nnodes ${NNODES} --node-rank ${NODE_RANK} --dist-init-addr ${MASTER_IP}:${DIST_PORT} --trust-remote-code --host 0.0.0.0 --port ${PORT} --mem-fraction-static 0.9 --cuda-graph-max-bs-decode 64 --max-running-requests 256" diff --git a/deploy/profiles/pro6000/dsv4_pro6000_sglang_tp16_eagle.env b/deploy/profiles/pro6000/dsv4_pro6000_sglang_tp16_eagle.env new file mode 100644 index 0000000..ff37c53 --- /dev/null +++ b/deploy/profiles/pro6000/dsv4_pro6000_sglang_tp16_eagle.env @@ -0,0 +1,32 @@ +# DeepSeek-V4-Pro SGLang multi-node TP=16 EP=2 + EAGLE speculative decoding +# (2x RTX 6000D). Same topology as dsv4_pro6000_sglang_tp16, plus EAGLE flags. + +PLATFORM=pro6000 +EXPERIMENT=dsv4_pro6000_sglang_tp16_eagle +MODEL_NAME=DeepSeek-V4-Pro +ENGINE=sglang +RUNTIME=docker +DOCKER_IMAGE=lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45 +CONTAINER_NAME=${EXPERIMENT}_node${NODE_RANK} +MODEL_PATH=/data/hf_models/DeepSeek-V4-Pro +SERVED_MODEL_NAME=default +PORT=30000 +HEALTH_PATH=/health +HEALTH_HOST=10.101.0.11 +HEALTH_WAIT_S=2400 +CONTAINER_PYTHON=python3 + +# ---- Multi-node topology (rank order; rank 0 exposes the HTTP API) ---- +NNODES=2 +NODE_HOSTS="10.101.0.11 10.101.0.13" +NODE_SSH_USER=root +MASTER_IP=10.101.0.11 +DIST_PORT=20000 + +DEVICE_VARS="CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7" +ENGINE_ENV="NCCL_SOCKET_IFNAME=eth0 NCCL_DEBUG=WARN SGLANG_SHARED_EXPERT_TP1=1 PYTHONUNBUFFERED=1 HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1" + +DOCKER_FLAGS="--gpus all --network host --ipc=host --ulimit memlock=-1 --shm-size=20g" +VOLUMES="${MODEL_PATH}:${MODEL_PATH}:ro" + +LAUNCH_ARGS="--model-path ${MODEL_PATH} --tp-size ${TP} --ep-size 2 --nnodes ${NNODES} --node-rank ${NODE_RANK} --dist-init-addr ${MASTER_IP}:${DIST_PORT} --trust-remote-code --host 0.0.0.0 --port ${PORT} --mem-fraction-static 0.9 --cuda-graph-max-bs-decode 64 --max-running-requests 256 --speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4" diff --git a/deploy/profiles/pro6000/dsv4_pro6000_sglang_tp_dp_matrix.env b/deploy/profiles/pro6000/dsv4_pro6000_sglang_tp_dp_matrix.env new file mode 100644 index 0000000..f172f72 --- /dev/null +++ b/deploy/profiles/pro6000/dsv4_pro6000_sglang_tp_dp_matrix.env @@ -0,0 +1,24 @@ +# DeepSeek-V4-Flash SGLang TPxDP matrix profile on a single RTX 6000D node (8 GPUs). +# TP=2/DP=4 is excluded because Marlin weight loading OOMs on this machine. + +PLATFORM=pro6000 +EXPERIMENT=dsv4_pro6000_sglang_tp_dp_matrix +MODEL_NAME=DeepSeek-V4-Flash +ENGINE=sglang +RUNTIME=docker +DOCKER_IMAGE=lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45 +CONTAINER_NAME=${EXPERIMENT}_sglang_tp${TP}_dp${DP} +MODEL_PATH=/data/6000D/DeepSeek-V4-Flash +SERVED_MODEL_NAME=deepseek-v4-flash +PORT=30031 +HEALTH_PATH=/health +HEALTH_WAIT_S=2400 +CONTAINER_PYTHON=python3 + +DEVICE_VARS="CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7" +ENGINE_ENV="PYTHONUNBUFFERED=1 HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True" + +DOCKER_FLAGS="--gpus all --privileged --ipc=host --network host --ulimit memlock=-1 --ulimit stack=67108864" +VOLUMES="${MODEL_PATH}:${MODEL_PATH}:ro" + +LAUNCH_ARGS="--model-path ${MODEL_PATH} --trust-remote-code --tp-size ${TP} --moe-runner-backend auto --mem-fraction-static 0.9 --context-length 131072 --max-running-requests 64 --host 0.0.0.0 --port ${PORT}" diff --git a/deploy/profiles/pro6000/dsv4_pro6000_vllm_tp_dp_matrix.env b/deploy/profiles/pro6000/dsv4_pro6000_vllm_tp_dp_matrix.env new file mode 100644 index 0000000..1c0de38 --- /dev/null +++ b/deploy/profiles/pro6000/dsv4_pro6000_vllm_tp_dp_matrix.env @@ -0,0 +1,24 @@ +# DeepSeek-V4-Flash vLLM TPxDP matrix profile on a single RTX 6000D node (8 GPUs). + +PLATFORM=pro6000 +EXPERIMENT=dsv4_pro6000_vllm_tp_dp_matrix +MODEL_NAME=DeepSeek-V4-Flash +ENGINE=vllm +SERVER_CMD=vllm serve +RUNTIME=docker +DOCKER_IMAGE=vllm-sm120-dsv4:0.25.1-fi0.6.14 +CONTAINER_NAME=${EXPERIMENT}_vllm_tp${TP}_dp${DP} +MODEL_PATH=/data/6000D/DeepSeek-V4-Flash +SERVED_MODEL_NAME=deepseek-v4-flash +PORT=30030 +HEALTH_PATH=/health +HEALTH_WAIT_S=2400 + +DEVICE_VARS="CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7" +ENGINE_ENV="PYTHONUNBUFFERED=1 HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1" + +DOCKER_FLAGS="--gpus all --privileged --ipc=host --network host --ulimit memlock=-1 --ulimit stack=67108864" +VOLUMES="${MODEL_PATH}:${MODEL_PATH}:ro" + +LAUNCH_ARGS="--model ${MODEL_PATH} --trust-remote-code --kv-cache-dtype fp8 --block-size 256 --tensor-parallel-size ${TP} --gpu-memory-utilization 0.9 --max-model-len 131072 --max-num-seqs 128 --host 0.0.0.0 --port ${PORT}" +DP_FLAG="--data-parallel-size ${DP}" diff --git a/deploy/profiles/pro6000/glm52_pro6000_sglang_tp16.env b/deploy/profiles/pro6000/glm52_pro6000_sglang_tp16.env new file mode 100644 index 0000000..8945aa6 --- /dev/null +++ b/deploy/profiles/pro6000/glm52_pro6000_sglang_tp16.env @@ -0,0 +1,37 @@ +# GLM-5.2-FP8 SGLang multi-node TP=16 deployment profile (2x RTX 6000D). +# node0 = 174.1.51.5 (pro6000D.1, HTTP API), node1 = 174.1.51.7 (pro6000D.3, +# compute-only). Model ~700GB fp8 so 2-node TP=16 is mandatory. +# Requires the tilelang/dsa patch mounts below (experiment-dir sources). + +PLATFORM=pro6000 +EXPERIMENT=glm52_pro6000_sglang_multinode_tp16 +MODEL_NAME=GLM-5.2-FP8 +ENGINE=sglang +RUNTIME=docker +DOCKER_IMAGE=lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45 +CONTAINER_NAME=${EXPERIMENT}_sglang_tp${TP}_dp${DP}_node${NODE_RANK} +MODEL_PATH=/data/hf_models/GLM-5.2-FP8 +SERVED_MODEL_NAME=GLM-5.2-FP8 +PORT=30031 +HEALTH_PATH=/health +HEALTH_HOST=174.1.51.5 +HEALTH_WAIT_S=2400 +CONTAINER_PYTHON=python3 +TP=16 +DP=1 + +# ---- Multi-node topology (rank order; rank 0 exposes the HTTP API) ---- +NNODES=2 +NODE_HOSTS="174.1.51.5 174.1.51.7" +NODE_SSH_USER=root +MASTER_IP=174.1.51.5 +DIST_PORT=50000 + +DEVICE_VARS="CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7" +ENGINE_ENV="PYTHONUNBUFFERED=1 HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:256 NCCL_IB_HCA=mlx5_0,mlx5_3 NCCL_MIN_NCHANNELS=8 NCCL_IB_QPS_PER_CONNECTION=4 NCCL_NET_GDR_LEVEL=PHB NCCL_SOCKET_IFNAME=eth1 NCCL_IB_RETRY_CNT=14 NCCL_IB_TIMEOUT=16" + +DOCKER_FLAGS="--gpus all --privileged --ipc=host --network host --ulimit memlock=-1 --ulimit stack=67108864" +VOLUMES="${MODEL_PATH}:${MODEL_PATH}:ro ${ROOT_DIR}/sglang_nightly_cu13_cache:/root/.cache" +PATCH_MOUNTS="${ROOT_DIR}/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/tilelang_kernel_sm120.py:/sgl-workspace/sglang/python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py:ro ${ROOT_DIR}/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/dsa_backend.py:/sgl-workspace/sglang/python/sglang/srt/layers/attention/dsa_backend.py:ro ${ROOT_DIR}/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/dsa_indexer.py:/sgl-workspace/sglang/python/sglang/srt/layers/attention/dsa/dsa_indexer.py:ro" + +LAUNCH_ARGS="--model-path ${MODEL_PATH} --trust-remote-code --tp-size ${TP} --dp-size ${DP} --quantization fp8 --kv-cache-dtype bfloat16 --moe-runner-backend auto --mem-fraction-static 0.80 --context-length 131072 --max-running-requests 64 --cuda-graph-backend-decode disabled --dsa-prefill-backend tilelang --dsa-decode-backend tilelang --host 0.0.0.0 --port ${PORT} --dist-init-addr ${MASTER_IP}:${DIST_PORT} --nnodes ${NNODES} --node-rank ${NODE_RANK}" diff --git a/deploy/profiles/pro6000/qwen3_235b_pro6000_sglang_tp8.env b/deploy/profiles/pro6000/qwen3_235b_pro6000_sglang_tp8.env new file mode 100644 index 0000000..e05b538 --- /dev/null +++ b/deploy/profiles/pro6000/qwen3_235b_pro6000_sglang_tp8.env @@ -0,0 +1,26 @@ +# Qwen3-235B-A22B SGLang TP=8 profile on a single RTX 6000D node (8 GPUs). +# NVIDIA counterpart of p800/qwen3_235b_p800_sglang_tp8 (same workload for +# comparability: isl=2048 osl=2048 c=16). + +PLATFORM=pro6000 +EXPERIMENT=qwen3_235b_pro6000_sglang_tp8 +MODEL_NAME=Qwen3-235B-A22B +ENGINE=sglang +RUNTIME=docker +DOCKER_IMAGE=sglang-sm120-dsv4:0.5.15.post1-fi0.6.14-sm120fix1 +CONTAINER_NAME=qwen3_235b_pro6000_sglang_tp8 +MODEL_PATH=/data/6000D/Qwen3-235B-A22B +SERVED_MODEL_NAME=/data/6000D/Qwen3-235B-A22B +PORT=30010 +HEALTH_PATH=/health +HEALTH_WAIT_S=1800 +CONTAINER_PYTHON=python3 +TP=8 + +DEVICE_VARS="CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7" +ENGINE_ENV="PYTHONUNBUFFERED=1 HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True" + +DOCKER_FLAGS="--gpus all --privileged --ipc=host --network host --ulimit memlock=-1 --ulimit stack=67108864" +VOLUMES="/data:/data" + +LAUNCH_ARGS="--model-path ${MODEL_PATH} --trust-remote-code --tp-size ${TP} --mem-fraction-static 0.9 --context-length 8192 --max-running-requests 16 --host 0.0.0.0 --port ${PORT}" diff --git a/experiments/pro6000/dsv4_pro6000_sglang_tp16/matrix.json b/experiments/pro6000/dsv4_pro6000_sglang_tp16/matrix.json new file mode 100644 index 0000000..6d6482c --- /dev/null +++ b/experiments/pro6000/dsv4_pro6000_sglang_tp16/matrix.json @@ -0,0 +1,128 @@ +{ + "comment": "ISL/OSL matrix for dsv4_pro6000_sglang_tp16 (multi-node TP16 EP2). Y=must test, P=optional, N=skip. Extracted from run_batch.sh legacy inline matrix; concurrency per ISL mirrors get_conc_for().", + "mode": "Y", + "matrix": { + "1024": { + "128": "Y", + "256": "Y", + "512": "Y", + "1024": "Y", + "2048": "Y", + "4096": "Y" + }, + "4096": { + "128": "Y", + "256": "Y", + "512": "Y", + "1024": "Y", + "2048": "Y", + "4096": "Y" + }, + "8192": { + "128": "Y", + "256": "Y", + "512": "Y", + "1024": "Y", + "2048": "Y", + "4096": "Y" + }, + "16384": { + "128": "Y", + "256": "Y", + "512": "Y", + "1024": "Y", + "2048": "Y", + "4096": "P" + }, + "32768": { + "128": "Y", + "256": "Y", + "512": "Y", + "1024": "Y", + "2048": "Y", + "4096": "P" + }, + "65536": { + "128": "Y", + "256": "Y", + "512": "Y", + "1024": "Y", + "2048": "P", + "4096": "N" + }, + "131072": { + "128": "Y", + "256": "Y", + "512": "Y", + "1024": "P", + "2048": "N", + "4096": "N" + }, + "262144": { + "128": "Y", + "256": "Y", + "512": "P", + "1024": "N", + "2048": "N", + "4096": "N" + }, + "524288": { + "128": "Y", + "256": "P", + "512": "N", + "1024": "N", + "2048": "N", + "4096": "N" + }, + "1048576": { + "128": "Y", + "256": "P", + "512": "N", + "1024": "N", + "2048": "N", + "4096": "N" + } + }, + "concurrency": { + "1024": { + "low": 1, + "high": 128 + }, + "4096": { + "low": 1, + "high": 128 + }, + "8192": { + "low": 1, + "high": 32 + }, + "16384": { + "low": 1, + "high": 32 + }, + "32768": { + "low": 1, + "high": 8 + }, + "65536": { + "low": 1, + "high": 8 + }, + "131072": { + "low": 1, + "high": 4 + }, + "262144": { + "low": 1, + "high": 4 + }, + "524288": { + "low": 1, + "high": 4 + }, + "1048576": { + "low": 1, + "high": 4 + } + } +} \ No newline at end of file diff --git a/experiments/pro6000/dsv4_pro6000_sglang_tp16/run_batch.sh b/experiments/pro6000/dsv4_pro6000_sglang_tp16/run_batch.sh index 4868c6a..b0fb6f8 100755 --- a/experiments/pro6000/dsv4_pro6000_sglang_tp16/run_batch.sh +++ b/experiments/pro6000/dsv4_pro6000_sglang_tp16/run_batch.sh @@ -152,9 +152,9 @@ check_and_recover_server() { fi log "WARN: Server health failed. Checking..." local hok=$(ssh -o StrictHostKeyChecking=no -o ConnectTimeout=5 "${HEAD_NODE}" \ - "docker ps --filter name=${EXPERIMENT}_head --format '{{.Names}}' 2>/dev/null" || echo "") + "docker ps --filter name=${EXPERIMENT}_node0 --format '{{.Names}}' 2>/dev/null" || echo "") local wok=$(ssh -o StrictHostKeyChecking=no -o ConnectTimeout=5 "${WORKER_NODE}" \ - "docker ps --filter name=${EXPERIMENT}_worker --format '{{.Names}}' 2>/dev/null" || echo "") + "docker ps --filter name=${EXPERIMENT}_node1 --format '{{.Names}}' 2>/dev/null" || echo "") if [[ -z "$hok" || -z "$wok" ]]; then log "ERROR: Containers died. Restarting server..." bash "${SCRIPT_DIR}/start_sglang_multinode.sh" || { log "FATAL: Restart failed"; return 1; } diff --git a/experiments/pro6000/dsv4_pro6000_sglang_tp16/start_sglang_multinode.sh b/experiments/pro6000/dsv4_pro6000_sglang_tp16/start_sglang_multinode.sh index 8436d46..e6b726a 100755 --- a/experiments/pro6000/dsv4_pro6000_sglang_tp16/start_sglang_multinode.sh +++ b/experiments/pro6000/dsv4_pro6000_sglang_tp16/start_sglang_multinode.sh @@ -1,84 +1,40 @@ #!/usr/bin/env bash -# ============================================================ -# Start SGLang multi-node TP=16 EP=2 across 2 nodes. -# Key: SGLANG_SHARED_EXPERT_TP1=1 + --ep-size 2 -# Usage: ./start_sglang_multinode.sh [--dry-run] -# ============================================================ +# Start SGLang multi-node TP=16 EP=2 across 2 nodes through the shared +# deployment layer (`python -m sskj.deploy start`, NODE_HOSTS orchestration). +# Usage: start_sglang_multinode.sh [--dry-run] set -Eeuo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# shellcheck source=/dev/null source "${SCRIPT_DIR}/config.env" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/platform.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" + +RUNTIME_BASE="${RUNTIME_BASE:-${SCRIPT_DIR}/runtime}" +mkdir -p "${RUNTIME_BASE}/logs" DRY_RUN="${DRY_RUN:-0}" -HEAD_CONTAINER="${EXPERIMENT}_head" -WORKER_CONTAINER="${EXPERIMENT}_worker" +DEPLOY_PROFILE="${DEPLOY_PROFILE:-pro6000/dsv4_pro6000_sglang_tp16}" -log() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*"; } +log "=== SGLang multi-node TP=${TP_SIZE} EP=${EP_SIZE} via deploy profile ${DEPLOY_PROFILE} ===" -build_docker_cmd() { - local node_rank="$1" container_name="$2" - cat </dev/null; -docker run -d --name ${container_name} \\ - --gpus all --network host --ipc=host --ulimit memlock=-1 --shm-size=20g \\ - -v ${MODEL_PATH}:${MODEL_PATH}:ro \\ - -e NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME} \\ - -e NCCL_DEBUG=${NCCL_DEBUG} \\ - -e SGLANG_SHARED_EXPERT_TP1=1 \\ - ${DOCKER_IMAGE} \\ - python3 -m sglang.launch_server \\ - --model-path ${MODEL_PATH} \\ - --tp-size ${TP_SIZE} --ep-size ${EP_SIZE} --nnodes ${NNODES} \\ - --node-rank ${node_rank} --dist-init-addr ${HEAD_IP}:${DIST_INIT_PORT} \\ - --trust-remote-code --host 0.0.0.0 --port ${SGLANG_PORT} \\ - --mem-fraction-static ${MEM_FRACTION_STATIC} \\ - --cuda-graph-max-bs-decode ${CUDA_GRAPH_MAX_BS_DECODE} \\ - --max-running-requests ${MAX_RUNNING_REQUESTS} -DOCKER_CMD -} - -main() { - log "=== SGLang multi-node TP=${TP_SIZE} EP=${EP_SIZE} ===" - - if [[ "$DRY_RUN" == "0" ]]; then - ssh -o StrictHostKeyChecking=no "${HEAD_NODE}" "docker rm -f ${HEAD_CONTAINER} 2>/dev/null" || true - ssh -o StrictHostKeyChecking=no "${WORKER_NODE}" "docker rm -f ${WORKER_CONTAINER} 2>/dev/null" || true - sleep 2 - fi - - log "Starting WORKER (${WORKER_NODE}, node_rank=1) ..." - local cmd; cmd="$(build_docker_cmd 1 "${WORKER_CONTAINER}")" - if [[ "$DRY_RUN" == "1" ]]; then log "[DRY] $cmd"; else - ssh -o StrictHostKeyChecking=no "${WORKER_NODE}" "bash -c '${cmd}'" >/dev/null - log "Worker started." - fi - sleep 5 - - log "Starting HEAD (${HEAD_NODE}, node_rank=0) ..." - cmd="$(build_docker_cmd 0 "${HEAD_CONTAINER}")" - if [[ "$DRY_RUN" == "1" ]]; then log "[DRY] $cmd"; else - ssh -o StrictHostKeyChecking=no "${HEAD_NODE}" "bash -c '${cmd}'" >/dev/null - log "Head started." - fi - - # Health check - if [[ "$DRY_RUN" == "1" ]]; then log "[DRY] Skip health check"; return 0; fi - log "Waiting for health on ${HEAD_IP}:${SGLANG_PORT} ..." - for (( i=1; i<=${HEALTH_CHECK_RETRIES}; i++ )); do - if curl --fail --silent --max-time 5 "http://${HEAD_IP}:${SGLANG_PORT}/health" >/dev/null 2>&1; then - log "Server healthy! (${i}s)" - return 0 - fi - local hok=$(ssh -o StrictHostKeyChecking=no "${HEAD_NODE}" \ - "docker ps --filter name=${HEAD_CONTAINER} --format '{{.Names}}' 2>/dev/null" || true) - local wok=$(ssh -o StrictHostKeyChecking=no "${WORKER_NODE}" \ - "docker ps --filter name=${WORKER_CONTAINER} --format '{{.Names}}' 2>/dev/null" || true) - if [[ -z "$hok" ]]; then log "ERROR: Head died"; ssh "${HEAD_NODE}" "docker logs --tail 30 ${HEAD_CONTAINER}" || true; return 1; fi - if [[ -z "$wok" ]]; then log "ERROR: Worker died"; ssh "${WORKER_NODE}" "docker logs --tail 30 ${WORKER_CONTAINER}" || true; return 1; fi - if (( i % 12 == 0 )); then log " waiting... ${i}s"; fi - sleep "${HEALTH_CHECK_INTERVAL_S}" - done - log "ERROR: Timeout"; return 1 -} - -main "$@" +if [[ "$DRY_RUN" == "1" ]]; then + PYTHONPATH="${ROOT_DIR}/src" "${DEPLOY_PYTHON:-python3}" -m sskj.deploy start \ + --profile "$(deploy_profile_abs "$DEPLOY_PROFILE")" \ + --tp "$TP_SIZE" --dp "1" \ + --port "${SGLANG_PORT:-30000}" \ + --model-path "$MODEL_PATH" \ + --log-dir "${RUNTIME_BASE}/logs" \ + --dry-run +else + deploy_start_multinode \ + "$DEPLOY_PROFILE" "$TP_SIZE" "1" \ + "${RUNTIME_BASE}/logs" \ + "${SGLANG_PORT:-30000}" \ + "$MODEL_PATH" \ + "$EXPERIMENT" +fi diff --git a/experiments/pro6000/dsv4_pro6000_sglang_tp16/stop_sglang_multinode.sh b/experiments/pro6000/dsv4_pro6000_sglang_tp16/stop_sglang_multinode.sh index fd76361..7144f09 100755 --- a/experiments/pro6000/dsv4_pro6000_sglang_tp16/stop_sglang_multinode.sh +++ b/experiments/pro6000/dsv4_pro6000_sglang_tp16/stop_sglang_multinode.sh @@ -1,14 +1,36 @@ #!/usr/bin/env bash -# Stop SGLang multi-node server and clean up. -set -euo pipefail +# Stop SGLang multi-node server (both nodes) through the shared deployment layer. +# Usage: stop_sglang_multinode.sh [--dry-run] +set -Eeuo pipefail + SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# shellcheck source=/dev/null source "${SCRIPT_DIR}/config.env" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/platform.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" -log() { echo "[$(date '+%H:%M:%S')] $*"; } +DRY_RUN="${DRY_RUN:-0}" +DEPLOY_PROFILE="${DEPLOY_PROFILE:-pro6000/dsv4_pro6000_sglang_tp16}" -log "Stopping ${EXPERIMENT} ..." -ssh -o StrictHostKeyChecking=no "${HEAD_NODE}" \ - "docker rm -f ${EXPERIMENT}_head 2>/dev/null && echo 'Head stopped' || echo 'Head already stopped'" -ssh -o StrictHostKeyChecking=no "${WORKER_NODE}" \ - "docker rm -f ${EXPERIMENT}_worker 2>/dev/null && echo 'Worker stopped' || echo 'Worker already stopped'" +log "=== Stopping SGLang multi-node via deploy profile ${DEPLOY_PROFILE} ===" + +if [[ "$DRY_RUN" == "1" ]]; then + PYTHONPATH="${ROOT_DIR}/src" "${DEPLOY_PYTHON:-python3}" -m sskj.deploy stop \ + --profile "$(deploy_profile_abs "$DEPLOY_PROFILE")" \ + --tp "$TP_SIZE" --dp "1" \ + --port "${SGLANG_PORT:-30000}" \ + --model-path "$MODEL_PATH" \ + --container-name "$EXPERIMENT" \ + --dry-run +else + deploy_stop \ + "$DEPLOY_PROFILE" "$TP_SIZE" "1" \ + "${SGLANG_PORT:-30000}" \ + "$MODEL_PATH" \ + "$EXPERIMENT" +fi log "Done." diff --git a/experiments/pro6000/dsv4_pro6000_sglang_tp16_eagle/run_batch.sh b/experiments/pro6000/dsv4_pro6000_sglang_tp16_eagle/run_batch.sh index 4868c6a..b0fb6f8 100755 --- a/experiments/pro6000/dsv4_pro6000_sglang_tp16_eagle/run_batch.sh +++ b/experiments/pro6000/dsv4_pro6000_sglang_tp16_eagle/run_batch.sh @@ -152,9 +152,9 @@ check_and_recover_server() { fi log "WARN: Server health failed. Checking..." local hok=$(ssh -o StrictHostKeyChecking=no -o ConnectTimeout=5 "${HEAD_NODE}" \ - "docker ps --filter name=${EXPERIMENT}_head --format '{{.Names}}' 2>/dev/null" || echo "") + "docker ps --filter name=${EXPERIMENT}_node0 --format '{{.Names}}' 2>/dev/null" || echo "") local wok=$(ssh -o StrictHostKeyChecking=no -o ConnectTimeout=5 "${WORKER_NODE}" \ - "docker ps --filter name=${EXPERIMENT}_worker --format '{{.Names}}' 2>/dev/null" || echo "") + "docker ps --filter name=${EXPERIMENT}_node1 --format '{{.Names}}' 2>/dev/null" || echo "") if [[ -z "$hok" || -z "$wok" ]]; then log "ERROR: Containers died. Restarting server..." bash "${SCRIPT_DIR}/start_sglang_multinode.sh" || { log "FATAL: Restart failed"; return 1; } diff --git a/experiments/pro6000/dsv4_pro6000_sglang_tp16_eagle/start_sglang_multinode.sh b/experiments/pro6000/dsv4_pro6000_sglang_tp16_eagle/start_sglang_multinode.sh index aa4b1f2..32a09d3 100755 --- a/experiments/pro6000/dsv4_pro6000_sglang_tp16_eagle/start_sglang_multinode.sh +++ b/experiments/pro6000/dsv4_pro6000_sglang_tp16_eagle/start_sglang_multinode.sh @@ -1,82 +1,40 @@ #!/usr/bin/env bash -# ============================================================ -# Start SGLang multi-node TP=16 EP=2 + EAGLE speculative decoding -# Usage: ./start_sglang_multinode.sh [--dry-run] -# ============================================================ +# Start SGLang multi-node TP=16 EP=2 across 2 nodes through the shared +# deployment layer (`python -m sskj.deploy start`, NODE_HOSTS orchestration). +# Usage: start_sglang_multinode.sh [--dry-run] set -Eeuo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# shellcheck source=/dev/null source "${SCRIPT_DIR}/config.env" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/platform.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" + +RUNTIME_BASE="${RUNTIME_BASE:-${SCRIPT_DIR}/runtime}" +mkdir -p "${RUNTIME_BASE}/logs" DRY_RUN="${DRY_RUN:-0}" -HEAD_CONTAINER="${EXPERIMENT}_head" -WORKER_CONTAINER="${EXPERIMENT}_worker" +DEPLOY_PROFILE="${DEPLOY_PROFILE:-pro6000/dsv4_pro6000_sglang_tp16_eagle}" -log() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*"; } +log "=== SGLang multi-node TP=${TP_SIZE} EP=${EP_SIZE} via deploy profile ${DEPLOY_PROFILE} ===" -build_docker_cmd() { - local node_rank="$1" container_name="$2" - cat </dev/null; -docker run -d --name ${container_name} \\ - --gpus all --network host --ipc=host --ulimit memlock=-1 --shm-size=20g \\ - -v ${MODEL_PATH}:${MODEL_PATH}:ro \\ - -e NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME} \\ - -e NCCL_DEBUG=${NCCL_DEBUG} \\ - -e SGLANG_SHARED_EXPERT_TP1=1 \\ - ${DOCKER_IMAGE} \\ - python3 -m sglang.launch_server \\ - --model-path ${MODEL_PATH} \\ - --tp-size ${TP_SIZE} --ep-size ${EP_SIZE} --nnodes ${NNODES} \\ - --node-rank ${node_rank} --dist-init-addr ${HEAD_IP}:${DIST_INIT_PORT} \\ - --trust-remote-code --host 0.0.0.0 --port ${SGLANG_PORT} \\ - --mem-fraction-static ${MEM_FRACTION_STATIC} \\ - --cuda-graph-max-bs-decode ${CUDA_GRAPH_MAX_BS_DECODE} \\ - --max-running-requests ${MAX_RUNNING_REQUESTS} \\ - --speculative-algorithm ${SPECULATIVE_ALGORITHM} \\ - --speculative-num-steps ${SPECULATIVE_NUM_STEPS} \\ - --speculative-eagle-topk ${SPECULATIVE_EAGLE_TOPK} \\ - --speculative-num-draft-tokens ${SPECULATIVE_NUM_DRAFT_TOKENS} -DOCKER_CMD -} - -main() { - log "=== SGLang TP=${TP_SIZE} EP=${EP_SIZE} + EAGLE ===" - - if [[ "$DRY_RUN" == "0" ]]; then - ssh -o StrictHostKeyChecking=no "${HEAD_NODE}" "docker rm -f ${HEAD_CONTAINER} 2>/dev/null" || true - ssh -o StrictHostKeyChecking=no "${WORKER_NODE}" "docker rm -f ${WORKER_CONTAINER} 2>/dev/null" || true - sleep 2 - fi - - log "Starting WORKER (${WORKER_NODE}, node_rank=1) ..." - local cmd; cmd="$(build_docker_cmd 1 "${WORKER_CONTAINER}")" - if [[ "$DRY_RUN" == "1" ]]; then log "[DRY] $cmd"; else - ssh -o StrictHostKeyChecking=no "${WORKER_NODE}" "bash -c '${cmd}'" >/dev/null - log "Worker started." - fi - sleep 5 - - log "Starting HEAD (${HEAD_NODE}, node_rank=0) ..." - cmd="$(build_docker_cmd 0 "${HEAD_CONTAINER}")" - if [[ "$DRY_RUN" == "1" ]]; then log "[DRY] $cmd"; else - ssh -o StrictHostKeyChecking=no "${HEAD_NODE}" "bash -c '${cmd}'" >/dev/null - log "Head started." - fi - - if [[ "$DRY_RUN" == "1" ]]; then log "[DRY] Skip health check"; return 0; fi - log "Waiting for health on ${HEAD_IP}:${SGLANG_PORT} ..." - for (( i=1; i<=${HEALTH_CHECK_RETRIES}; i++ )); do - if curl --fail --silent --max-time 5 "http://${HEAD_IP}:${SGLANG_PORT}/health" >/dev/null 2>&1; then - log "Server healthy! (${i}s)"; return 0 - fi - local hok=$(ssh -o StrictHostKeyChecking=no "${HEAD_NODE}" "docker ps --filter name=${HEAD_CONTAINER} --format '{{.Names}}' 2>/dev/null" || true) - local wok=$(ssh -o StrictHostKeyChecking=no "${WORKER_NODE}" "docker ps --filter name=${WORKER_CONTAINER} --format '{{.Names}}' 2>/dev/null" || true) - if [[ -z "$hok" ]]; then log "ERROR: Head died"; ssh "${HEAD_NODE}" "docker logs --tail 30 ${HEAD_CONTAINER}" || true; return 1; fi - if [[ -z "$wok" ]]; then log "ERROR: Worker died"; ssh "${WORKER_NODE}" "docker logs --tail 30 ${WORKER_CONTAINER}" || true; return 1; fi - sleep "${HEALTH_CHECK_INTERVAL_S}" - done - log "ERROR: Timeout"; return 1 -} - -main "$@" +if [[ "$DRY_RUN" == "1" ]]; then + PYTHONPATH="${ROOT_DIR}/src" "${DEPLOY_PYTHON:-python3}" -m sskj.deploy start \ + --profile "$(deploy_profile_abs "$DEPLOY_PROFILE")" \ + --tp "$TP_SIZE" --dp "1" \ + --port "${SGLANG_PORT:-30000}" \ + --model-path "$MODEL_PATH" \ + --log-dir "${RUNTIME_BASE}/logs" \ + --dry-run +else + deploy_start_multinode \ + "$DEPLOY_PROFILE" "$TP_SIZE" "1" \ + "${RUNTIME_BASE}/logs" \ + "${SGLANG_PORT:-30000}" \ + "$MODEL_PATH" \ + "$EXPERIMENT" +fi diff --git a/experiments/pro6000/dsv4_pro6000_sglang_tp16_eagle/stop_sglang_multinode.sh b/experiments/pro6000/dsv4_pro6000_sglang_tp16_eagle/stop_sglang_multinode.sh index f9bf14b..b392739 100755 --- a/experiments/pro6000/dsv4_pro6000_sglang_tp16_eagle/stop_sglang_multinode.sh +++ b/experiments/pro6000/dsv4_pro6000_sglang_tp16_eagle/stop_sglang_multinode.sh @@ -1,9 +1,36 @@ #!/usr/bin/env bash -set -euo pipefail +# Stop SGLang multi-node server (both nodes) through the shared deployment layer. +# Usage: stop_sglang_multinode.sh [--dry-run] +set -Eeuo pipefail + SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# shellcheck source=/dev/null source "${SCRIPT_DIR}/config.env" -log() { echo "[$(date '+%H:%M:%S')] $*"; } -log "Stopping ${EXPERIMENT} ..." -ssh -o StrictHostKeyChecking=no "${HEAD_NODE}" "docker rm -f ${EXPERIMENT}_head 2>/dev/null && echo 'Head stopped' || echo 'Head already stopped'" -ssh -o StrictHostKeyChecking=no "${WORKER_NODE}" "docker rm -f ${EXPERIMENT}_worker 2>/dev/null && echo 'Worker stopped' || echo 'Worker already stopped'" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/platform.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" + +DRY_RUN="${DRY_RUN:-0}" +DEPLOY_PROFILE="${DEPLOY_PROFILE:-pro6000/dsv4_pro6000_sglang_tp16_eagle}" + +log "=== Stopping SGLang multi-node via deploy profile ${DEPLOY_PROFILE} ===" + +if [[ "$DRY_RUN" == "1" ]]; then + PYTHONPATH="${ROOT_DIR}/src" "${DEPLOY_PYTHON:-python3}" -m sskj.deploy stop \ + --profile "$(deploy_profile_abs "$DEPLOY_PROFILE")" \ + --tp "$TP_SIZE" --dp "1" \ + --port "${SGLANG_PORT:-30000}" \ + --model-path "$MODEL_PATH" \ + --container-name "$EXPERIMENT" \ + --dry-run +else + deploy_stop \ + "$DEPLOY_PROFILE" "$TP_SIZE" "1" \ + "${SGLANG_PORT:-30000}" \ + "$MODEL_PATH" \ + "$EXPERIMENT" +fi log "Done." diff --git a/experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/run_adaptive_concurrency.sh b/experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/run_adaptive_concurrency.sh index bf9f543..88f7827 100755 --- a/experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/run_adaptive_concurrency.sh +++ b/experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/run_adaptive_concurrency.sh @@ -15,6 +15,10 @@ source "${SCRIPT_DIR}/config.env" source "${SCRIPT_DIR}/adaptive_config.env" # shellcheck source=/dev/null source "${SCRIPT_DIR}/../../../scripts/common/adaptive_bench_lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" + +DEPLOY_PROFILE="${DEPLOY_PROFILE:-pro6000/dsv4_pro6000_sglang_tp_dp_matrix}" ENGINE="sglang" ENGINE_PORT="$SGLANG_PORT" @@ -37,29 +41,9 @@ engine_is_healthy() { engine_stop_server() { local tp="$1" local dp="$2" - local pid_file="${RUNTIME_BASE}/${EXPERIMENT}_sglang_tp${tp}_dp${dp}.pid" - if [[ -f "$pid_file" ]]; then - local pid - pid="$(cat "$pid_file")" - if [[ -n "${CONTAINER_NAME:-}" ]]; then - if docker exec "$CONTAINER_NAME" kill -0 "$pid" 2>/dev/null; then - log "stopping sglang in persistent container pid=${pid} tp=${tp} dp=${dp}" - docker exec "$CONTAINER_NAME" kill "$pid" 2>/dev/null || true - sleep 5 - docker exec "$CONTAINER_NAME" kill -9 "$pid" 2>/dev/null || true - fi - elif kill -0 "$pid" 2>/dev/null; then - log "stopping sglang server pid=${pid} tp=${tp} dp=${dp}" - kill "$pid" 2>/dev/null || true - sleep 5 - kill -9 "$pid" 2>/dev/null || true - fi - rm -f "$pid_file" - fi - if [[ -z "${CONTAINER_NAME:-}" ]]; then - docker rm -f "${EXPERIMENT}_sglang_tp${tp}_dp${dp}" >/dev/null 2>&1 || true - fi + log "stopping sglang server tp=${tp} dp=${dp} via deploy profile" + bash "${SCRIPT_DIR}/stop_sglang_docker.sh" "$tp" "$dp" >> "${ADAPTIVE_LOG_DIR}/sglang_tp${tp}_dp${dp}.server.outer.log" 2>&1 || true ACTIVE_ENGINE_SERVER_LOG="" sleep 2 } @@ -67,21 +51,7 @@ engine_stop_server() { engine_build_server_args() { local tp="$1" local dp="$2" - local -a args=( - python3 -m sglang.launch_server --model-path "$MODEL_PATH" - --trust-remote-code - --tp-size "$tp" - --moe-runner-backend "$MOE_RUNNER_BACKEND" - --mem-fraction-static "$MEM_FRACTION_STATIC" - --context-length "$CONTEXT_LENGTH" - --max-running-requests "$MAX_RUNNING_REQUESTS" - --host 0.0.0.0 - --port "$ENGINE_PORT" - ) - if (( dp > 1 )); then - args+=(--dp-size "$dp") - fi - printf '%q ' "${args[@]}" + deploy_render_args "$DEPLOY_PROFILE" "$tp" "$dp" "$ENGINE_PORT" "$MODEL_PATH" } engine_start_server() { diff --git a/experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/run_adaptive_concurrency_add16.sh b/experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/run_adaptive_concurrency_add16.sh index bf1ea37..e192412 100755 --- a/experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/run_adaptive_concurrency_add16.sh +++ b/experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/run_adaptive_concurrency_add16.sh @@ -15,6 +15,10 @@ source "${SCRIPT_DIR}/config.env" source "${SCRIPT_DIR}/adaptive_config.env" # shellcheck source=/dev/null source "${SCRIPT_DIR}/../../../scripts/common/adaptive_bench_lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" + +DEPLOY_PROFILE="${DEPLOY_PROFILE:-pro6000/dsv4_pro6000_sglang_tp_dp_matrix}" ENGINE="sglang" ENGINE_PORT="$SGLANG_PORT" @@ -37,29 +41,9 @@ engine_is_healthy() { engine_stop_server() { local tp="$1" local dp="$2" - local pid_file="${RUNTIME_BASE}/${EXPERIMENT}_sglang_tp${tp}_dp${dp}.pid" - if [[ -f "$pid_file" ]]; then - local pid - pid="$(cat "$pid_file")" - if [[ -n "${CONTAINER_NAME:-}" ]]; then - if docker exec "$CONTAINER_NAME" kill -0 "$pid" 2>/dev/null; then - log "stopping sglang in persistent container pid=${pid} tp=${tp} dp=${dp}" - docker exec "$CONTAINER_NAME" kill "$pid" 2>/dev/null || true - sleep 5 - docker exec "$CONTAINER_NAME" kill -9 "$pid" 2>/dev/null || true - fi - elif kill -0 "$pid" 2>/dev/null; then - log "stopping sglang server pid=${pid} tp=${tp} dp=${dp}" - kill "$pid" 2>/dev/null || true - sleep 5 - kill -9 "$pid" 2>/dev/null || true - fi - rm -f "$pid_file" - fi - if [[ -z "${CONTAINER_NAME:-}" ]]; then - docker rm -f "${EXPERIMENT}_sglang_tp${tp}_dp${dp}" >/dev/null 2>&1 || true - fi + log "stopping sglang server tp=${tp} dp=${dp} via deploy profile" + bash "${SCRIPT_DIR}/stop_sglang_docker.sh" "$tp" "$dp" >> "${ADAPTIVE_LOG_DIR}/sglang_tp${tp}_dp${dp}.server.outer.log" 2>&1 || true ACTIVE_ENGINE_SERVER_LOG="" sleep 2 } @@ -67,21 +51,7 @@ engine_stop_server() { engine_build_server_args() { local tp="$1" local dp="$2" - local -a args=( - python3 -m sglang.launch_server --model-path "$MODEL_PATH" - --trust-remote-code - --tp-size "$tp" - --moe-runner-backend "$MOE_RUNNER_BACKEND" - --mem-fraction-static "$MEM_FRACTION_STATIC" - --context-length "$CONTEXT_LENGTH" - --max-running-requests "$MAX_RUNNING_REQUESTS" - --host 0.0.0.0 - --port "$ENGINE_PORT" - ) - if (( dp > 1 )); then - args+=(--dp-size "$dp") - fi - printf '%q ' "${args[@]}" + deploy_render_args "$DEPLOY_PROFILE" "$tp" "$dp" "$ENGINE_PORT" "$MODEL_PATH" } engine_start_server() { @@ -181,6 +151,8 @@ export -f engine_run_bench export ENGINE_PORT MODEL_PATH RESULT_BASE DOCKER_IMAGE USE_DOCKER_CLIENT export BENCH_DATASET_NAME DATASET_PATH RANDOM_RANGE_RATIO BENCH_WARMUP_MAX_REQUESTS PYTHON SGLANG_BENCH_MODULE +adaptive_main "$@" + export SEARCH_START_CONCURRENCY=16 export SEARCH_ADDEND=16 # If the initial concurrency violates the TTFT SLO, search downward. Stop at @@ -189,5 +161,3 @@ export SEARCH_INITIAL_BACKOFF_CONCURRENCIES="8 1" # When concurrency 1 still has a severely excessive TTFT, stop the remaining # shapes in this TP/DP group. Zero disables this rule. export TTFT_GROUP_SKIP_MS="${TTFT_GROUP_SKIP_MS:-8000}" - -adaptive_main "$@" diff --git a/experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/run_bench.sh b/experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/run_bench.sh index 2b886e4..ec67a45 100755 --- a/experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/run_bench.sh +++ b/experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/run_bench.sh @@ -11,6 +11,8 @@ source "${SCRIPT_DIR}/../../../scripts/common/lib.sh" source "${SCRIPT_DIR}/../../../scripts/common/platform.sh" # shellcheck source=/dev/null source "${SCRIPT_DIR}/config.env" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" RUN_ID="${RUN_ID:-$(date '+%Y%m%d-%H%M%S')}" RESULT_BASE="${SCRIPT_DIR}/results" @@ -51,33 +53,8 @@ stop_server() { local tp="$1" local dp="$2" - local pid_file="${RUNTIME_BASE}/${EXPERIMENT}_sglang_tp${tp}_dp${dp}.pid" - if [[ -f "$pid_file" ]]; then - local pid - pid="$(cat "$pid_file")" - # In container-reuse mode, pid is the server PID inside the container. - if [[ -n "${CONTAINER_NAME:-}" ]]; then - if docker exec "$CONTAINER_NAME" kill -0 "$pid" 2>/dev/null; then - log "stopping sglang server inside container (tp=${tp}, dp=${dp}, pid=${pid})" - docker exec "$CONTAINER_NAME" kill "$pid" 2>/dev/null || true - sleep 5 - docker exec "$CONTAINER_NAME" kill -9 "$pid" 2>/dev/null || true - fi - else - if kill -0 "$pid" 2>/dev/null; then - log "stopping sglang server pid=${pid} (tp=${tp}, dp=${dp})" - kill "$pid" 2>/dev/null || true - sleep 5 - kill -9 "$pid" 2>/dev/null || true - fi - fi - rm -f "$pid_file" - fi - - # Fallback: remove any Docker container started by this experiment (legacy mode). - if [[ -z "${CONTAINER_NAME:-}" ]]; then - docker rm -f "${EXPERIMENT}_sglang_tp${tp}_dp${dp}" >/dev/null 2>&1 || true - fi + log "stopping sglang server tp=${tp} dp=${dp} via deploy profile" + bash "${SCRIPT_DIR}/stop_sglang_docker.sh" "$tp" "$dp" >> "${log_dir_global}/sglang_tp${tp}_dp${dp}.server.outer.log" 2>&1 || true # Fallback: kill any SGLang launch processes for this model. pkill -9 -f "sglang.launch_server.*${MODEL_NAME}" 2>/dev/null || true @@ -89,23 +66,11 @@ build_server_args() { local tp="$1" local dp="$2" - local args=( - "python3 -m sglang.launch_server" --model-path "$MODEL_PATH" - --trust-remote-code - --tp-size "$tp" - --moe-runner-backend "$MOE_RUNNER_BACKEND" - --mem-fraction-static "$MEM_FRACTION_STATIC" - --context-length "$CONTEXT_LENGTH" - --max-running-requests "$MAX_RUNNING_REQUESTS" - --host 0.0.0.0 - --port "$SGLANG_PORT" - ) - if [[ "$dp" -gt 1 ]]; then - args+=( - --dp-size "$dp" - ) - fi - printf '%s ' "${args[@]}" + deploy_render_args \ + "${DEPLOY_PROFILE:-pro6000/dsv4_pro6000_sglang_tp_dp_matrix}" \ + "$tp" "$dp" \ + "${SGLANG_PORT:-30031}" \ + "$MODEL_PATH" } start_server() { diff --git a/experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/start_sglang_docker.sh b/experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/start_sglang_docker.sh index 6df6177..409b605 100755 --- a/experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/start_sglang_docker.sh +++ b/experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/start_sglang_docker.sh @@ -1,105 +1,33 @@ #!/usr/bin/env bash -# Start SGLang server in Docker for a given TPxDP configuration. +# Start the SGLang TPxDP server through the shared deployment layer. # Usage: start_sglang_docker.sh -# -# Uses the verified SGLang SM120 image and keeps its JIT cache on -# persistent storage. The container is removed automatically on stop. -set -e +set -Eeuo pipefail -TP="${1}" -DP="${2}" +TP="${1:-}" +DP="${2:-}" +if [[ -z "$TP" || -z "$DP" ]]; then + echo "Usage: $0 " + exit 1 +fi SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" # shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/platform.sh" +# shellcheck source=/dev/null source "${SCRIPT_DIR}/config.env" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" RUNTIME_BASE="${RUNTIME_BASE:-${SCRIPT_DIR}/runtime}" -mkdir -p "${RUNTIME_BASE}/logs" "${RUNTIME_BASE}/tmp" "$CACHE_DIR" +mkdir -p "${RUNTIME_BASE}/logs" "${RUNTIME_BASE}/tmp" -IMAGE="${DOCKER_IMAGE:-lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45}" -PORT="${SGLANG_PORT:-30031}" -NAME="${EXPERIMENT}_sglang_tp${TP}_dp${DP}" -PID_FILE="${RUNTIME_BASE}/${EXPERIMENT}_sglang_tp${TP}_dp${DP}.pid" - -LOG="${RUNTIME_BASE}/logs/${EXPERIMENT}_sglang_docker_tp${TP}_dp${DP}_$(date +%Y%m%d_%H%M%S).log" -rm -f "$PID_FILE" - -# Clean up any stale container with the same name. -docker rm -f "$NAME" >/dev/null 2>&1 || true - -SERVER_ARGS=( - -m sglang.launch_server - --model-path "$MODEL_PATH" - --trust-remote-code - --tp-size "$TP" - --moe-runner-backend "$MOE_RUNNER_BACKEND" - --mem-fraction-static "$MEM_FRACTION_STATIC" - --context-length "$CONTEXT_LENGTH" - --max-running-requests "$MAX_RUNNING_REQUESTS" - --host 0.0.0.0 - --port "$PORT" -) - -if [[ "$DP" -gt 1 ]]; then - SERVER_ARGS+=( - --dp-size "$DP" - ) -fi - -SERVER_ARGS_STR="${SERVER_ARGS[*]}" - -echo "=== Starting SGLang server in Docker (TP=${TP}, DP=${DP}) ===" -echo "Image: $IMAGE" -echo "Model: $MODEL_PATH" -echo "Container name: $NAME" -echo "Host port: $PORT" -echo "Command: python3 ${SERVER_ARGS_STR}" -echo "Log: $LOG" - -# Run docker in the foreground so that killing the host process stops the -# container (the --rm flag ensures cleanup). nohup lets us background it and -# capture the host PID in the same way as the bare-metal start script. -nohup docker run --rm \ - --name "$NAME" \ - --gpus all \ - --privileged \ - --ipc=host \ - --network host \ - --ulimit memlock=-1 \ - --ulimit stack=67108864 \ - --entrypoint python3 \ - -v "${MODEL_PATH}:${MODEL_PATH}:ro" \ - -v "${CACHE_DIR}:/root/.cache" \ - -v "${RUNTIME_BASE}/tmp:/tmp" \ - -e CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" \ - -e PYTHONUNBUFFERED=1 \ - -e HF_HUB_OFFLINE=1 \ - -e TRANSFORMERS_OFFLINE=1 \ - -e PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" \ - "$IMAGE" \ - "${SERVER_ARGS[@]}" \ - > "$LOG" 2>&1 & - -PID=$! -echo $PID > "$PID_FILE" -echo "PID: $PID" -echo "Waiting for health on port ${PORT}..." - -for i in $(seq 1 240); do - if curl --fail --silent --show-error --max-time 5 "http://127.0.0.1:${PORT}/health" >/dev/null 2>&1; then - echo "SGLang server is ready at http://127.0.0.1:${PORT}" - echo "Log: $LOG" - exit 0 - fi - if ! kill -0 $PID 2>/dev/null; then - echo "ERROR: Docker SGLang server exited early" - tail -200 "$LOG" - exit 1 - fi - echo "Waiting... ($i/240)" - sleep 5 -done - -echo "ERROR: Docker SGLang server not healthy after 240 retries" -tail -200 "$LOG" -exit 1 +log "starting SGLang server tp=${TP} dp=${DP} via deploy profile" +deploy_start \ + "${DEPLOY_PROFILE:-pro6000/dsv4_pro6000_sglang_tp_dp_matrix}" \ + "$TP" "$DP" \ + "${RUNTIME_BASE}/logs" \ + "${SGLANG_PORT:-30031}" \ + "$MODEL_PATH" \ + "${EXPERIMENT}_sglang_tp${TP}_dp${DP}" diff --git a/experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/stop_sglang_docker.sh b/experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/stop_sglang_docker.sh new file mode 100644 index 0000000..fd31f74 --- /dev/null +++ b/experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/stop_sglang_docker.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +# Stop the SGLang TPxDP server through the shared deployment layer. +# Usage: stop_sglang_docker.sh +set -Eeuo pipefail + +TP="${1:-}" +DP="${2:-}" +if [[ -z "$TP" || -z "$DP" ]]; then + echo "Usage: $0 " + exit 1 +fi + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/platform.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/config.env" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" + +log "stopping SGLang server tp=${TP} dp=${DP} via deploy profile" +deploy_stop \ + "${DEPLOY_PROFILE:-pro6000/dsv4_pro6000_sglang_tp_dp_matrix}" \ + "$TP" "$DP" \ + "${SGLANG_PORT:-30031}" \ + "$MODEL_PATH" \ + "${EXPERIMENT}_sglang_tp${TP}_dp${DP}" diff --git a/experiments/pro6000/dsv4_pro6000_vllm_tp_dp_matrix/run_adaptive_concurrency.sh b/experiments/pro6000/dsv4_pro6000_vllm_tp_dp_matrix/run_adaptive_concurrency.sh index ebe4b5e..a5ee5f1 100755 --- a/experiments/pro6000/dsv4_pro6000_vllm_tp_dp_matrix/run_adaptive_concurrency.sh +++ b/experiments/pro6000/dsv4_pro6000_vllm_tp_dp_matrix/run_adaptive_concurrency.sh @@ -15,6 +15,10 @@ source "${SCRIPT_DIR}/config.env" source "${SCRIPT_DIR}/adaptive_config.env" # shellcheck source=/dev/null source "${SCRIPT_DIR}/../../../scripts/common/adaptive_bench_lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" + +DEPLOY_PROFILE="${DEPLOY_PROFILE:-pro6000/dsv4_pro6000_vllm_tp_dp_matrix}" ENGINE="vllm" ENGINE_PORT="$VLLM_PORT" @@ -38,20 +42,9 @@ engine_is_healthy() { engine_stop_server() { local tp="$1" local dp="$2" - local pid_file="${RUNTIME_BASE}/${EXPERIMENT}_vllm_tp${tp}_dp${dp}.pid" - if [[ -f "$pid_file" ]]; then - local pid - pid="$(cat "$pid_file")" - if kill -0 "$pid" 2>/dev/null; then - log "stopping vllm server pid=${pid} tp=${tp} dp=${dp}" - kill "$pid" 2>/dev/null || true - sleep 5 - kill -9 "$pid" 2>/dev/null || true - fi - rm -f "$pid_file" - fi - docker rm -f "${EXPERIMENT}_vllm_tp${tp}_dp${dp}" >/dev/null 2>&1 || true + log "stopping vllm server tp=${tp} dp=${dp} via deploy profile" + bash "${SCRIPT_DIR}/stop_vllm_docker.sh" "$tp" "$dp" >> "${ADAPTIVE_LOG_DIR}/vllm_tp${tp}_dp${dp}.server.outer.log" 2>&1 || true ACTIVE_ENGINE_SERVER_LOG="" sleep 2 } @@ -59,22 +52,7 @@ engine_stop_server() { engine_build_server_args() { local tp="$1" local dp="$2" - local -a args=( - vllm serve "$MODEL_PATH" - --trust-remote-code - --kv-cache-dtype "$KV_CACHE_DTYPE" - --block-size "$BLOCK_SIZE" - --tensor-parallel-size "$tp" - --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" - --max-model-len "$MAX_MODEL_LEN" - --max-num-seqs "$MAX_NUM_SEQS" - --host 0.0.0.0 - --port "$ENGINE_PORT" - ) - if (( dp > 1 )); then - args+=(--data-parallel-size "$dp") - fi - printf '%q ' "${args[@]}" + deploy_render_args "$DEPLOY_PROFILE" "$tp" "$dp" "$ENGINE_PORT" "$MODEL_PATH" } engine_start_server() { diff --git a/experiments/pro6000/dsv4_pro6000_vllm_tp_dp_matrix/run_adaptive_concurrency_add16.sh b/experiments/pro6000/dsv4_pro6000_vllm_tp_dp_matrix/run_adaptive_concurrency_add16.sh index 854e994..4fa313f 100755 --- a/experiments/pro6000/dsv4_pro6000_vllm_tp_dp_matrix/run_adaptive_concurrency_add16.sh +++ b/experiments/pro6000/dsv4_pro6000_vllm_tp_dp_matrix/run_adaptive_concurrency_add16.sh @@ -15,6 +15,10 @@ source "${SCRIPT_DIR}/config.env" source "${SCRIPT_DIR}/adaptive_config.env" # shellcheck source=/dev/null source "${SCRIPT_DIR}/../../../scripts/common/adaptive_bench_lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" + +DEPLOY_PROFILE="${DEPLOY_PROFILE:-pro6000/dsv4_pro6000_vllm_tp_dp_matrix}" ENGINE="vllm" ENGINE_PORT="$VLLM_PORT" @@ -38,20 +42,9 @@ engine_is_healthy() { engine_stop_server() { local tp="$1" local dp="$2" - local pid_file="${RUNTIME_BASE}/${EXPERIMENT}_vllm_tp${tp}_dp${dp}.pid" - if [[ -f "$pid_file" ]]; then - local pid - pid="$(cat "$pid_file")" - if kill -0 "$pid" 2>/dev/null; then - log "stopping vllm server pid=${pid} tp=${tp} dp=${dp}" - kill "$pid" 2>/dev/null || true - sleep 5 - kill -9 "$pid" 2>/dev/null || true - fi - rm -f "$pid_file" - fi - docker rm -f "${EXPERIMENT}_vllm_tp${tp}_dp${dp}" >/dev/null 2>&1 || true + log "stopping vllm server tp=${tp} dp=${dp} via deploy profile" + bash "${SCRIPT_DIR}/stop_vllm_docker.sh" "$tp" "$dp" >> "${ADAPTIVE_LOG_DIR}/vllm_tp${tp}_dp${dp}.server.outer.log" 2>&1 || true ACTIVE_ENGINE_SERVER_LOG="" sleep 2 } @@ -59,22 +52,7 @@ engine_stop_server() { engine_build_server_args() { local tp="$1" local dp="$2" - local -a args=( - vllm serve "$MODEL_PATH" - --trust-remote-code - --kv-cache-dtype "$KV_CACHE_DTYPE" - --block-size "$BLOCK_SIZE" - --tensor-parallel-size "$tp" - --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" - --max-model-len "$MAX_MODEL_LEN" - --max-num-seqs "$MAX_NUM_SEQS" - --host 0.0.0.0 - --port "$ENGINE_PORT" - ) - if (( dp > 1 )); then - args+=(--data-parallel-size "$dp") - fi - printf '%q ' "${args[@]}" + deploy_render_args "$DEPLOY_PROFILE" "$tp" "$dp" "$ENGINE_PORT" "$MODEL_PATH" } engine_start_server() { @@ -163,6 +141,8 @@ export -f engine_run_bench export ENGINE_PORT MODEL_PATH RESULT_BASE DOCKER_CLIENT_IMAGE USE_DOCKER_CLIENT export BENCH_DATASET_NAME DATASET_PATH RANDOM_RANGE_RATIO BENCH_WARMUP_MAX_REQUESTS PYTHON SGLANG_BENCH_MODULE +adaptive_main "$@" + export SEARCH_START_CONCURRENCY=16 export SEARCH_ADDEND=16 # If the initial concurrency violates the TTFT SLO, search downward. Stop at @@ -171,5 +151,3 @@ export SEARCH_INITIAL_BACKOFF_CONCURRENCIES="8 1" # When concurrency 1 still has a severely excessive TTFT, stop the remaining # shapes in this TP/DP group. Zero disables this rule. export TTFT_GROUP_SKIP_MS="${TTFT_GROUP_SKIP_MS:-8000}" - -adaptive_main "$@" diff --git a/experiments/pro6000/dsv4_pro6000_vllm_tp_dp_matrix/start_vllm_docker.sh b/experiments/pro6000/dsv4_pro6000_vllm_tp_dp_matrix/start_vllm_docker.sh index cd07631..0dd8374 100755 --- a/experiments/pro6000/dsv4_pro6000_vllm_tp_dp_matrix/start_vllm_docker.sh +++ b/experiments/pro6000/dsv4_pro6000_vllm_tp_dp_matrix/start_vllm_docker.sh @@ -1,114 +1,33 @@ #!/usr/bin/env bash -# Start vLLM server in Docker for a given TPxDP configuration. +# Start the vLLM TPxDP server through the shared deployment layer. # Usage: start_vllm_docker.sh -# -# Uses the verified vLLM SM120 image and keeps its cache on persistent -# storage. The container is removed automatically on stop. -set -e +set -Eeuo pipefail -TP="${1}" -DP="${2}" +TP="${1:-}" +DP="${2:-}" +if [[ -z "$TP" || -z "$DP" ]]; then + echo "Usage: $0 " + exit 1 +fi SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" # shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/platform.sh" +# shellcheck source=/dev/null source "${SCRIPT_DIR}/config.env" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" RUNTIME_BASE="${RUNTIME_BASE:-${SCRIPT_DIR}/runtime}" -mkdir -p "${RUNTIME_BASE}/logs" "${RUNTIME_BASE}/tmp" "$CACHE_DIR" +mkdir -p "${RUNTIME_BASE}/logs" "${RUNTIME_BASE}/tmp" -IMAGE="${DOCKER_IMAGE:-vllm-sm120-dsv4:0.25.1-fi0.6.14}" -PORT="${VLLM_PORT:-30030}" -NAME="${EXPERIMENT}_vllm_tp${TP}_dp${DP}" -PID_FILE="${RUNTIME_BASE}/${EXPERIMENT}_vllm_tp${TP}_dp${DP}.pid" - -LOG="${RUNTIME_BASE}/logs/${EXPERIMENT}_vllm_docker_tp${TP}_dp${DP}_$(date +%Y%m%d_%H%M%S).log" -rm -f "$PID_FILE" - -# Clean up any stale container with the same name. -docker rm -f "$NAME" >/dev/null 2>&1 || true - -SERVER_ARGS=( - serve "$MODEL_PATH" - --trust-remote-code - --kv-cache-dtype "$KV_CACHE_DTYPE" - --block-size "$BLOCK_SIZE" - --tensor-parallel-size "$TP" - --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" - --max-model-len "$MAX_MODEL_LEN" - --max-num-seqs "$MAX_NUM_SEQS" - --host 0.0.0.0 - --port "$PORT" -) - -if [[ "$DP" -gt 1 ]]; then - SERVER_ARGS+=( - --data-parallel-size "$DP" - ) -fi - -SERVER_ARGS_STR="${SERVER_ARGS[*]}" - -echo "=== Starting vLLM server in Docker (TP=${TP}, DP=${DP}) ===" -echo "Image: $IMAGE" -echo "Model: $MODEL_PATH" -echo "Container name: $NAME" -echo "Host port: $PORT" -echo "Command: vllm ${SERVER_ARGS_STR}" -echo "Log: $LOG" - -# Run docker in the foreground so that killing the host process stops the -# container (the --rm flag ensures cleanup). nohup lets us background it and -# capture the host PID in the same way as the bare-metal start script. -nohup docker run --rm \ - --name "$NAME" \ - --gpus all \ - --privileged \ - --ipc=host \ - --network host \ - --ulimit memlock=-1 \ - --ulimit stack=67108864 \ - --entrypoint vllm \ - -v "${MODEL_PATH}:${MODEL_PATH}:ro" \ - -v "${CACHE_DIR}:/root/.cache" \ - -v "${RUNTIME_BASE}/tmp:/tmp" \ - -e CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" \ - -e PYTHONUNBUFFERED=1 \ - -e HF_HUB_OFFLINE=1 \ - -e TRANSFORMERS_OFFLINE=1 \ - "$IMAGE" \ - "${SERVER_ARGS[@]}" \ - > "$LOG" 2>&1 & - -PID=$! -echo $PID > "$PID_FILE" -echo "PID: $PID" -echo "Waiting for health on port ${PORT}..." - -for i in $(seq 1 600); do - if curl --fail --silent --show-error --max-time 5 "http://127.0.0.1:${PORT}/health" >/dev/null 2>&1; then - echo "vLLM server is ready at http://127.0.0.1:${PORT}" - echo "Log: $LOG" - exit 0 - fi - # Check if the Docker container is still running (not the nohup PID). - # Allow a brief grace period for the container to appear in docker ps. - container_running=0 - for _ in $(seq 1 3); do - if docker ps --filter "name=${NAME}" --format '{{.Names}}' | grep -q "^${NAME}$"; then - container_running=1 - break - fi - sleep 1 - done - if [[ "$container_running" -eq 0 ]]; then - echo "ERROR: Docker vLLM container exited early" - tail -200 "$LOG" - exit 1 - fi - echo "Waiting... ($i/600)" - sleep 5 -done - -echo "ERROR: Docker vLLM server not healthy after 600 retries" -tail -200 "$LOG" -exit 1 +log "starting vLLM server tp=${TP} dp=${DP} via deploy profile" +deploy_start \ + "${DEPLOY_PROFILE:-pro6000/dsv4_pro6000_vllm_tp_dp_matrix}" \ + "$TP" "$DP" \ + "${RUNTIME_BASE}/logs" \ + "${VLLM_PORT:-30030}" \ + "$MODEL_PATH" \ + "${EXPERIMENT}_vllm_tp${TP}_dp${DP}" diff --git a/experiments/pro6000/dsv4_pro6000_vllm_tp_dp_matrix/stop_vllm_docker.sh b/experiments/pro6000/dsv4_pro6000_vllm_tp_dp_matrix/stop_vllm_docker.sh new file mode 100644 index 0000000..8224d78 --- /dev/null +++ b/experiments/pro6000/dsv4_pro6000_vllm_tp_dp_matrix/stop_vllm_docker.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +# Stop the vLLM TPxDP server through the shared deployment layer. +# Usage: stop_vllm_docker.sh +set -Eeuo pipefail + +TP="${1:-}" +DP="${2:-}" +if [[ -z "$TP" || -z "$DP" ]]; then + echo "Usage: $0 " + exit 1 +fi + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/platform.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/config.env" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" + +log "stopping vLLM server tp=${TP} dp=${DP} via deploy profile" +deploy_stop \ + "${DEPLOY_PROFILE:-pro6000/dsv4_pro6000_vllm_tp_dp_matrix}" \ + "$TP" "$DP" \ + "${VLLM_PORT:-30030}" \ + "$MODEL_PATH" \ + "${EXPERIMENT}_vllm_tp${TP}_dp${DP}" diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/README.md b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/README.md new file mode 100644 index 0000000..ff5a7bb --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/README.md @@ -0,0 +1,130 @@ +# GLM-5.2-FP8 多机部署（SGLang TP=16，2×RTX 6000D 节点） + +在 pro6000D.1 + pro6000D.3 两台机器（各 8×RTX 6000D）上多机部署 GLM-5.2-FP8， +并对齐 H20 / DSV4 实验的 benchmark 方法做 TP×DP 矩阵测试。 + +## 为什么必须多机 + +GLM-5.2 是 MoE 模型：**总参数约 735B，FP8 权重约 700GB**（256 路由专家 + 1 共享 +专家，每 token 激活 8 个，DeepSeek-V3 架构）。单机 8 卡 RTX 6000D 显存共 685GB， +**装不下 700GB 权重**，因此 2 节点 16 卡是必需项，不是优化项。 + +## 架构 + +``` + HTTP 请求 (bench_serving) + │ + ▼ + ┌──────────────┐ NCCL/RoCE (41.6GB/s) ┌──────────────┐ + │ pro6000D.1 │ ◄══════════════════════► │ pro6000D.3 │ + │ node-rank 0 │ 计算网 10.101/10.102 │ node-rank 1 │ + │ 持有 1/2 权重 │ │ 持有 1/2 权重 │ + │ HTTP :30031 │ │ 不暴露 HTTP │ + └──────────────┘ └──────────────┘ +``` + +- **node0（pro6000D.1）**：主节点，对外暴露 HTTP API，接收所有请求。 +- **node1（pro6000D.3）**：从节点，只做计算，不对外服务。 +- 两机通过 NCCL 连接：`--dist-init-addr --nnodes 2 --node-rank <0|1>`。 +- NCCL bootstrap 走管理网（174.1.51.x），数据面走 RoCE 计算网（mlx5_0/mlx5_3）。 +- benchmark 客户端只打 node0 的 HTTP，感知不到背后是两台机器。 + +## 前置条件（两台机器都要满足） + +1. **模型文件**：两机都要有完整的 GLM-5.2-FP8，路径统一为 + `/data/hf_models/GLM-5.2-FP8`（141 个 safetensors + index.json）。 + - pro6000D.1：下载中（截至编写时 112/141 分片）。 + - pro6000D.3：**当前没有，需补齐**（从 .1 rsync 传过去或重新下载）。 +2. **Docker 镜像**：两机都要有 + `lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45`（sglang dev，原生支持 + GLM-5.2 的 `glm_moe_dsa` 架构）。 + - pro6000D.1：已有。 + - pro6000D.3：下载中。 +3. **SSH 互通**：node0 能 `ssh pro6000D.3` 免密登录（已配置）。 +4. **NCCL 跨机已验证**：双 NIC 调优后 all_reduce busbw 41.6 GB/s（见带宽报告）。 + +## 文件说明 + +| 文件 | 作用 | +|------|------| +| `config.env` | 全局配置：模型路径、TP=16、上下文 256K、NCCL 调优、节点拓扑 | +| `start_sglang_node.sh ` | 单节点启动（参数化 rank），node0/node1 共用 | +| `start_sglang_multinode.sh ` | 多机编排：sync→起 node1→起 node0→等 health | +| `stop_sglang_multinode.sh ` | 多机停止：停 node0（本地）+ node1（ssh） | +| `start_sglang_dp.sh ` | run_bench 入口，委托给 multinode 编排 | +| `run_bench.sh` | TP×DP 矩阵 benchmark（复用 DSV4 模板，改 stop 为多机） | +| `matrix.json` | ISL×OSL 测试矩阵（复用 DSV4） | +| `adaptive_config.env` | 自适应并发搜索配置（复用 DSV4） | + +## NCCL 调优（关键，不可省） + +跨机必须带以下环境变量（已固化在 config.env，注入容器），否则带宽从 41.6 +跌到 21 GB/s： + +```bash +NCCL_IB_HCA=mlx5_0,mlx5_3 # 启用两张 RoCE 网卡 +NCCL_MIN_NCHANNELS=8 # 强制 8 channel 分流到两卡 +NCCL_IB_QPS_PER_CONNECTION=4 # 每连接 4 QP 提升 RoCE 并行 +NCCL_NET_GDR_LEVEL=PHB # GPUDirect RDMA +NCCL_SOCKET_IFNAME=eth1 # bootstrap 走管理网网卡 +``` + +## 用法 + +### 冒烟测试（单场景验证端到端） + +```bash +cd /data/yy/sskj/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16 +GRID_LIMIT=1 DRY_RUN=0 bash run_bench.sh +``` + +### 全矩阵（low/high 两个并发点） + +```bash +CONCURRENCY_SAMPLES=2 bash run_bench.sh +``` + +### 手动起停（不走 bench） + +```bash +bash start_sglang_dp.sh 16 1 # 起两机，等 node0 health +bash stop_sglang_multinode.sh 16 1 # 停两机 +``` + +### 手动单独起 node1（调试） + +```bash +ssh pro6000D.3 'cd /data/yy/sskj/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16 && bash start_sglang_node.sh 1' +``` + +## 启动流程说明 + +`start_sglang_multinode.sh` 做的事： + +1. **rsync 实验目录到 node1**（让 node1 有 config.env + start_sglang_node.sh） +2. **ssh 后台起 node1**：`ssh pro6000D.3 '... start_sglang_node.sh 1'`，node1 + 启动容器后进入 rendezvous 等待窗口，等 node0 连接 +3. **本地起 node0**：`start_sglang_node.sh 0`，node0 发起 NCCL 连接，两机会合后 + 一起加载模型 +4. **轮询 node0 /health**：就绪后返回（最多等 480×5s=40 分钟，多机加载慢） + +node1 的日志在远端 `/tmp/glm52_node1_inner.log` 和本实验 `runtime/logs/` 下。 + +## 与 H20 基线对比 + +- H20 基线：`experiments/h20/glm_h20_vllm_tp_dp_matrix`，单机 8 卡 vLLM TP=8。 +- 本实验：2 机 16 卡 SGLang TP=16。 +- **框架不同**（vLLM vs SGLang）、**并行不同**（单机 TP=8 vs 多机 TP=16）， + 对比时需注明。但 benchmark 方法一致（sglang bench_serving + 同一 matrix.json + + 自适应并发），吞吐/延迟指标可对照。 +- 用 `compare.py` 对比两边 results。 + +## 注意事项 + +- **单机 TP=8 装不下**：735GB 权重 ÷ 8 卡 ≈ 92GB/卡 > 85.6GB 单卡上限，会 OOM。 + 必须多机。 +- **CUDA 版本不一致**（.1 是 12.8、.3 是 13.0）不影响：两机都在 cu130 容器内跑， + 宿主机 CUDA 版本无关。 +- **node1 不暴露 HTTP**：所有请求打 node0，node1 纯计算。 +- **加载慢**：700GB 权重 + NCCL 建连，首次加载预计 10-20 分钟，health 轮询 + 上限设了 40 分钟。 diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/adaptive_config.env b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/adaptive_config.env new file mode 100644 index 0000000..3cf7dd7 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/adaptive_config.env @@ -0,0 +1,56 @@ +# Adaptive concurrency search settings. +# +# For each fixed (TP, DP, ISL, OSL), probe: +# C = start, start * multiplier, ... up to max +# and stop after Total TPS has less than TPS_MIN_GAIN_PCT meaningful growth for +# PLATEAU_PATIENCE consecutive points. + +SEARCH_START_CONCURRENCY="${SEARCH_START_CONCURRENCY:-1}" +SEARCH_MAX_CONCURRENCY="${SEARCH_MAX_CONCURRENCY:-64}" + +# At the add16 initial probe, restart and retry C=8 then C=1 after an OOM. +ENABLE_INITIAL_OOM_BACKOFF="${ENABLE_INITIAL_OOM_BACKOFF:-1}" +SEARCH_MULTIPLIER="${SEARCH_MULTIPLIER:-2}" +NUM_PROMPTS_MULTIPLIER="${NUM_PROMPTS_MULTIPLIER:-5}" + +# A gain below 2% is treated as throughput saturation. Two consecutive +# low-gain points prevent one noisy measurement from stopping the search. +TPS_MIN_GAIN_PCT="${TPS_MIN_GAIN_PCT:-2.0}" +PLATEAU_PATIENCE="${PLATEAU_PATIENCE:-2}" + +# Stop a shape when p95 TTFT exceeds the SLO; keep group skipping disabled. +TTFT_SLO_MS="${TTFT_SLO_MS:-4000}" +ENABLE_TTFT_SLO_STOP="${ENABLE_TTFT_SLO_STOP:-1}" + +# Keep the same random workload semantics as the fixed matrix baseline. +# DATASET_PATH must contain at least SEARCH_MAX_CONCURRENCY times +# NUM_PROMPTS_MULTIPLIER valid two-turn conversations. Set this explicitly to +# random-ids to use generated token IDs without a ShareGPT seed dataset. +BENCH_DATASET_NAME="${BENCH_DATASET_NAME:-random}" +# SGLang interprets 0.0 as Uniform[1, requested_len]. Use 1.0 for fixed +# ISL/OSL points; lower values intentionally benchmark a length distribution. +RANDOM_RANGE_RATIO="${RANDOM_RANGE_RATIO:-1.0}" +# Before each measured point, warm up with the same concurrency so lazy kernel +# compilation and CUDA graph capture are excluded from TTFT/TPS. 0 means no +# cap; set a positive cap only when very high-concurrency warmup is impractical. +BENCH_WARMUP_MAX_REQUESTS="${BENCH_WARMUP_MAX_REQUESTS:-0}" + +# Reject a point if the completed request count or actual token lengths do not +# match the requested workload. +INPUT_LENGTH_TOLERANCE_PCT="${INPUT_LENGTH_TOLERANCE_PCT:-5.0}" +OUTPUT_LENGTH_TOLERANCE_PCT="${OUTPUT_LENGTH_TOLERANCE_PCT:-10.0}" + +MAX_POINT_RETRIES="${MAX_POINT_RETRIES:-1}" +SERVER_RESTART_COOLDOWN_S="${SERVER_RESTART_COOLDOWN_S:-10}" +SCENARIO_TIMEOUT_S="${SCENARIO_TIMEOUT_S:-1800}" +GPU_MEM_SAMPLE_INTERVAL_S="${GPU_MEM_SAMPLE_INTERVAL_S:-1}" + +# Optional space-separated filters, useful for smoke tests: +# TP_LIST="8" ISL_LIST="1024" OSL_LIST="128" +TP_LIST="${TP_LIST:-}" +ISL_LIST="${ISL_LIST:-}" +OSL_LIST="${OSL_LIST:-}" + +DRY_RUN="${DRY_RUN:-0}" +# Counts ISL/OSL shapes per TP/DP config, not individual concurrency probes. +GRID_LIMIT="${GRID_LIMIT:-0}" diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/compare.py b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/compare.py new file mode 100755 index 0000000..981609a --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/compare.py @@ -0,0 +1,162 @@ +#!/usr/bin/env python3 +"""Cross TP×DP configuration comparison for dsv4_h200_sglang_tp_dp_matrix. + +Usage: + python3 compare.py --run-root results/ [--output comparison.md] +""" +import argparse +import json +import re +from collections import defaultdict +from pathlib import Path + + +def load_result(result_root: Path) -> dict: + path = result_root / "results.json" + with open(path, "r", encoding="utf-8") as f: + return json.load(f) + + +def slo_status(ttft_p95_ms: float, tpot_mean_ms: float, + ttft_limit_ms: float = 3000.0, tpot_limit_ms: float = 50.0) -> str: + ttft_ok = ttft_p95_ms < ttft_limit_ms + tpot_ok = tpot_mean_ms < tpot_limit_ms + if ttft_ok and tpot_ok: + return "PASS" + if ttft_ok or tpot_ok: + return "PARTIAL" + return "FAIL" + + +def gpu_memory_str(gpu: dict | None) -> str: + if not gpu: + return "-" + peak = gpu.get("peak_used_mb", 0) + total = gpu.get("memory_total_mb", 0) + if total: + return f"{peak:.0f}/{total:.0f} ({100*peak/total:.1f}%)" + return f"{peak:.0f}" + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--run-root", type=Path, required=True) + parser.add_argument("-o", "--output", type=Path, default=Path("comparison.md")) + parser.add_argument("--ttft-limit", type=float, default=3000.0) + parser.add_argument("--tpot-limit", type=float, default=50.0) + args = parser.parse_args() + + # Discover configurations: tp*_dp* directories. + configs = [] + for subdir in sorted(args.run_root.iterdir()): + if not subdir.is_dir(): + continue + name = subdir.name + if not (name.startswith("tp") and "_dp" in name): + continue + results_json = subdir / "results.json" + if not results_json.exists(): + continue + configs.append((name, load_result(subdir))) + + if not configs: + print(f"No tp*_dp* results found under {args.run_root}") + return + + model = configs[0][1].get("metadata", {}).get("model", "unknown") + hardware = configs[0][1].get("metadata", {}).get("hardware", "unknown") + + # Group by scenario name. + by_scenario: dict[str, dict[str, dict]] = defaultdict(dict) + skipped: dict[str, dict[str, str]] = defaultdict(dict) + for label, data in configs: + for s in data.get("scenarios", []): + key = s["name"] + if s.get("status") == "skipped_oom": + skipped[key][label] = s.get("note", "skipped") + else: + by_scenario[key][label] = s + + with open(args.output, "w", encoding="utf-8") as f: + f.write(f"# SGLang TP×DP matrix comparison ({hardware})\n\n") + f.write("## Summary\n\n") + f.write(f"- Model: `{model}`\n") + f.write(f"- Hardware: {hardware}\n") + f.write("- Backend: SGLang (Docker)\n") + f.write("- Benchmark client: `sglang.benchmark.serving`\n") + f.write(f"- SLO reference: TTFT P95 < {args.ttft_limit}ms, TPOT mean < {args.tpot_limit}ms\n\n") + + # Configuration overview. + f.write("### Configurations\n\n") + f.write("| Config | TP | DP | GPUs/replica | Notes |\n") + f.write("|---|---:|---:|---:|---|\n") + for label, data in configs: + cfg = data.get("config", {}) + tp = cfg.get("tp", "?") + dp = cfg.get("dp", "?") + f.write(f"| {label} | {tp} | {dp} | {tp} | server args recorded per ISL in results.json |\n") + f.write("\n") + + # Side-by-side table. + f.write("## Side-by-side results\n\n") + headers = [ + "Scenario", "ISL", "DSL", "Config", "Conc", "Req/s", "OutTok/s", + "TTFT P95(ms)", "TTFT P99(ms)", "TPOT Mean(ms)", "TPOT P95(ms)", + "TPOT P99(ms)", "E2E P99(ms)", "Peak GPU mem", "SLO" + ] + f.write("| " + " | ".join(headers) + " |\n") + f.write("|" + "|".join(["---"] * len(headers)) + "|\n") + + for scenario_name in sorted(by_scenario.keys(), key=lambda x: tuple(map(int, re.findall(r"\d+", x)))): + _, isl, dsl = re.findall(r"\d+", scenario_name) + for label, data in configs: + s = by_scenario[scenario_name].get(label) + if s is None: + if scenario_name in skipped and label in skipped[scenario_name]: + note = skipped[scenario_name][label] + f.write(f"| {scenario_name} | {isl} | {dsl} | {label} | - | - | - | - | - | - | - | - | - | - | {note} |\n") + continue + cfg = s["config"] + m = s["metrics"] + status = slo_status(m["ttft_ms"]["p95"], m["tpot_ms"]["mean"], args.ttft_limit, args.tpot_limit) + gpu = m.get("gpu_memory") + f.write( + f"| {scenario_name} | {isl} | {dsl} | {label} | {cfg['concurrency']} | " + f"{m['request_throughput']:.2f} | {m['output_token_throughput']:.2f} | " + f"{m['ttft_ms']['p95']:.2f} | {m['ttft_ms']['p99']:.2f} | " + f"{m['tpot_ms']['mean']:.2f} | {m['tpot_ms']['p95']:.2f} | {m['tpot_ms']['p99']:.2f} | " + f"{m['e2e_ms']['p99']:.2f} | {gpu_memory_str(gpu)} | {status} |\n" + ) + + # Best throughput per ISL/DSL. + f.write("\n## Best throughput per (ISL, DSL)\n\n") + f.write("| ISL | DSL | Best Config | Concurrency | OutTok/s | TTFT P95(ms) | TPOT Mean(ms) | SLO |\n") + f.write("|---:|---:|---|---:|---:|---:|---:|---:|\n") + best_by_shape: dict[tuple[int, int], tuple[float, str, dict]] = {} + for scenario_name, backends in by_scenario.items(): + _, isl, dsl = re.findall(r"\d+", scenario_name) + isl_i, dsl_i = int(isl), int(dsl) + for label, s in backends.items(): + m = s["metrics"] + out_tok = m["output_token_throughput"] + if (isl_i, dsl_i) not in best_by_shape or out_tok > best_by_shape[(isl_i, dsl_i)][0]: + best_by_shape[(isl_i, dsl_i)] = (out_tok, label, s) + for (isl_i, dsl_i), (out_tok, label, s) in sorted(best_by_shape.items()): + m = s["metrics"] + status = slo_status(m["ttft_ms"]["p95"], m["tpot_ms"]["mean"], args.ttft_limit, args.tpot_limit) + f.write( + f"| {isl_i} | {dsl_i} | {label} | {s['config']['concurrency']} | " + f"{out_tok:.2f} | {m['ttft_ms']['p95']:.2f} | {m['tpot_ms']['mean']:.2f} | {status} |\n" + ) + + f.write("\n## Notes\n\n") + f.write("- SLO check uses TTFT P95 and TPOT mean.\n") + f.write("- A PARTIAL indicates one of the two metrics is out of target; FAIL indicates both are out.\n") + f.write("- `Peak GPU mem` shows peak used / total MB and utilization percentage.\n") + f.write("- Optional (P) combinations that failed are marked as skipped/OOM and do not break the run.\n") + + print(f"Wrote comparison to {args.output}") + + +if __name__ == "__main__": + main() diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/config.env b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/config.env new file mode 100755 index 0000000..3fca577 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/config.env @@ -0,0 +1,145 @@ +#!/usr/bin/env bash +# Multi-node SGLang deployment for GLM-5.2-FP8 on 2x RTX 6000D nodes (16 GPUs). +# +# Topology: +# node0 = pro6000D.1 (gpu1, 174.1.51.5, 10.101.0.11 / 10.102.0.11) +# node1 = pro6000D.3 (gpu3, 174.1.51.7, 10.101.0.13 / 10.102.0.13) +# Each node has 8x RTX 6000D (85.6GB) = 685GB; GLM-5.2-FP8 is ~700GB, so a +# single node cannot hold the weights -> 2-node TP=16 is mandatory. +# +# SGLang multi-node: every node runs its own `python3 -m sglang.launch_server` +# and they connect via NCCL using --dist-init-addr/--nnodes/--node-rank. +# node0 exposes the HTTP API; node1 is compute-only. The benchmark client only +# talks to node0. + +# --- Path bootstrap ------------------------------------------------------- +# run_bench.sh sources platform.sh (which sets ROOT_DIR) before this file, but +# the start_* scripts source this file directly. Provide a fallback so +# ${ROOT_DIR} / ${SCRIPT_DIR} expand correctly in both paths. +SCRIPT_DIR="${SCRIPT_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")" && pwd)}" +ROOT_DIR="${ROOT_DIR:-$(cd "${SCRIPT_DIR}/../../.." 2>/dev/null && pwd)}" +# If the relative climb failed (e.g. dir moved), fall back to the script dir. +[[ -z "$ROOT_DIR" || ! -d "$ROOT_DIR/platforms" ]] && ROOT_DIR="$SCRIPT_DIR" + +EXPERIMENT="glm52_pro6000_sglang_multinode_tp16" +MODEL_NAME="GLM-5.2-FP8" +MODEL_PATH="/data/hf_models/GLM-5.2-FP8" +SERVED_MODEL_NAME="GLM-5.2-FP8" + +SGLANG_PORT="${SGLANG_PORT:-30031}" + +# Python interpreter for orchestration scripts (parse_backend.py, compare.py, +# etc.) and the benchmark client. Defaults to the system python3 if the sglang +# venv does not exist on the host. +VENV_CLIENT="${VENV_CLIENT:-/root/.miniconda3/envs/sglang}" + +# Run the benchmark client natively (0) or inside Docker (1). +USE_DOCKER_CLIENT="${USE_DOCKER_CLIENT:-1}" + +# All 8 GPUs per node participate in TP=16 (8 per node x 2 nodes). +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}" + +# Runtime working directory for logs, pid files, and tmp. Defaults to a local +# directory under this experiment so the benchmark is self-contained. +RUNTIME_BASE="${RUNTIME_BASE:-${SCRIPT_DIR}/runtime}" + +# --------------------------------------------------------------------------- +# Multi-node topology +# --------------------------------------------------------------------------- +# Management network is used for the NCCL bootstrap (dist-init-addr). The data +# plane (NCCL/RoCE) is selected automatically by NCCL inside the container via +# the NCCL_* env vars below. +NODE0_HOST="${NODE0_HOST:-pro6000D.1}" +NODE1_HOST="${NODE1_HOST:-pro6000D.3}" +NODE0_IP="${NODE0_IP:-174.1.51.5}" +NODE1_IP="${NODE1_IP:-174.1.51.7}" +NNODES="${NNODES:-2}" +# Port for NCCL bootstrap (dist-init-addr). Must be free on node0. +DIST_INIT_PORT="${DIST_INIT_PORT:-50000}" +# SSH alias for reaching node1 from the control machine (node0). +NODE1_SSH="${NODE1_SSH:-root@174.1.51.7}" + +# --------------------------------------------------------------------------- +# Parallel configuration. TP=16 across 2 nodes (8 GPUs/node). +# Constraint (sglang): tp_size * pp_size must be divisible by nnodes. +# 16 * 1 = 16, 16 % 2 == 0 OK +# --------------------------------------------------------------------------- +TP_SIZE="${TP_SIZE:-16}" +DP_SIZE="${DP_SIZE:-1}" +# PARALLEL_CONFIGS is kept for compatibility with the matrix runner; the single +# multi-node config is "16 1". +declare -a PARALLEL_CONFIGS=( + "16 1" +) + +# --------------------------------------------------------------------------- +# SGLang server settings (aligned with dsv4_pro6000_sglang_tp_dp_matrix). +# --------------------------------------------------------------------------- +MEM_FRACTION_STATIC="${MEM_FRACTION_STATIC:-0.80}" +MOE_RUNNER_BACKEND="${MOE_RUNNER_BACKEND:-auto}" +# TileLang DSA on SM120 currently requires BF16 KV cache. Keep the smoke-test +# context at 128K so the doubled KV-cache element size fits beside FP8 weights. +CONTEXT_LENGTH="${CONTEXT_LENGTH:-131072}" +MAX_RUNNING_REQUESTS="${MAX_RUNNING_REQUESTS:-64}" +# GLM-5.2 is shipped as FP8 weights; tell SGLang explicitly. +QUANTIZATION="${QUANTIZATION:-fp8}" +KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-bfloat16}" + +# --------------------------------------------------------------------------- +# NCCL tuning for cross-node RoCE (mandatory, otherwise bandwidth collapses). +# Verified in the bandwidth report: dual-NIC tuned all_reduce = 41.6 GB/s, +# vs 21 GB/s without these knobs. +# --------------------------------------------------------------------------- +export NCCL_IB_HCA="${NCCL_IB_HCA:-mlx5_0,mlx5_3}" +export NCCL_MIN_NCHANNELS="${NCCL_MIN_NCHANNELS:-8}" +export NCCL_IB_QPS_PER_CONNECTION="${NCCL_IB_QPS_PER_CONNECTION:-4}" +export NCCL_NET_GDR_LEVEL="${NCCL_NET_GDR_LEVEL:-PHB}" +# Bind the NCCL socket (bootstrap/OOB) to the management NIC eth1 (174.1.51.x). +# eth0/eth3 are the RoCE data-plane NICs; eth1 carries the mgmt IP. +export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-eth1}" +export NCCL_IB_RETRY_CNT="${NCCL_IB_RETRY_CNT:-14}" +export NCCL_IB_TIMEOUT="${NCCL_IB_TIMEOUT:-16}" + +# --------------------------------------------------------------------------- +# Deployment switch. 0 = native sglang venv, 1 = Docker. +# --------------------------------------------------------------------------- +USE_DOCKER="${USE_DOCKER:-1}" +DOCKER_IMAGE="${DOCKER_IMAGE:-lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45}" + +# To use ShareGPT, set BENCH_DATASET_NAME=random and DATASET_PATH explicitly. +BENCH_DATASET_NAME="${BENCH_DATASET_NAME:-random}" +DATASET_PATH="${DATASET_PATH:-${ROOT_DIR}/dataset/ShareGPT_V3_unfiltered_cleaned_split.json}" +SGLANG_BENCH_MODULE="${SGLANG_BENCH_MODULE:-sglang.benchmark.serving}" +CACHE_DIR="${CACHE_DIR:-${ROOT_DIR}/sglang_nightly_cu13_cache}" + +# Matrix and concurrency rules are defined in matrix.json by default. +MATRIX_FILE="${MATRIX_FILE:-${SCRIPT_DIR:-.}/matrix.json}" +MATRIX_MODE="${MATRIX_MODE:-Y}" + +# Sampling density for concurrency. +# 0 = use the default heuristic in generate_scenarios.py (6-8 points). +# 2 = only test the low and high endpoints. +export CONCURRENCY_SAMPLES="${CONCURRENCY_SAMPLES:-2}" + +# Per-scenario timeout to avoid hangs (seconds). Multi-node load is slower, so +# the default is raised. +SCENARIO_TIMEOUT_S="${SCENARIO_TIMEOUT_S:-2400}" + +# GPU memory sampling interval (seconds). +GPU_MEM_SAMPLE_INTERVAL_S="${GPU_MEM_SAMPLE_INTERVAL_S:-1}" + +# Dry-run mode: if 1, only log the server args and scenario plan without +# starting any server or sending requests. +DRY_RUN="${DRY_RUN:-0}" + +# Per-config scenario limit for quick smoke tests. 0 = run all generated +# scenarios. +GRID_LIMIT="${GRID_LIMIT:-0}" + +# PyTorch CUDA allocator setting for the SGLang server. expandable_segments +# reduces fragmentation from the GiB-scale indexer temporaries that OOM the +# SM120 torch fallback (fp8_paged_mqa_logits_torch_sm120) at ISL >= 4096. +PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True,max_split_size_mb:256}" + +# Seconds to wait for node1 to come up and register before starting node0. +NODE1_READINESS_WAIT_S="${NODE1_READINESS_WAIT_S:-15}" diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/dsa_backend.py b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/dsa_backend.py new file mode 100644 index 0000000..8e19ced --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/dsa_backend.py @@ -0,0 +1,3321 @@ +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, + Dict, + List, + Literal, + Optional, + Tuple, + TypeAlias, +) + +import torch + +from sglang.srt.configs.model_config import get_dsa_index_topk, is_deepseek_dsa +from sglang.srt.runtime_context import get_parallel + +logger = logging.getLogger(__name__) +from sglang.kernels.ops.attention.dsa.dequant_k_cache import ( + dequantize_k_cache_paged, + gather_dequant_requant_fp8_paged, +) +from sglang.kernels.ops.attention.dsa.quant_k_cache import quantize_k_cache +from sglang.kernels.ops.attention.dsa.transform_index import ( + transform_index_page_table_decode, + transform_index_page_table_prefill, +) +from sglang.kernels.ops.attention.utils import ( + concat_mla_absorb_q_general, + mla_quantize_and_rope_for_fp8, + seqlens_expand_triton, +) +from sglang.kernels.ops.kvcache.cache_ops import concat_and_cast_q_fp8_pad +from sglang.srt.environ import envs +from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.dsa.dsa_backend_mtp_precompute import ( + DeepseekSparseAttnBackendMTPPrecomputeMixin, + PrecomputedMetadata, + compute_cu_seqlens, +) +from sglang.srt.layers.attention.dsa.dsa_indexer import BaseIndexerMetadata +from sglang.srt.layers.attention.dsa.dsa_topk_backend import ( + DSATopKBackend, + TopkTransformMethod, +) +from sglang.srt.layers.attention.dsa.utils import ( + can_dsa_prefill_cp_round_robin_split, + compute_dsa_seqlens, + dsa_cp_round_robin_split_data, + dsa_cp_round_robin_split_q_seqs, + dsa_use_prefill_cp, + is_dsa_enable_prefill_cp, + is_dsa_prefill_cp_in_seq_split, + pad_dsa_cache_seqlens, + should_use_dsa_fused_topk, +) +from sglang.srt.layers.utils.cp_utils import ( + cp_all_gather_rerange_output, + cp_split_and_rebuild_position, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.runtime_context import get_buffer +from sglang.srt.utils import ( + get_bool_env_var, + is_cuda, + is_gfx95_supported, + is_hip, + is_sm100_supported, + is_sm120_supported, + print_warning_once, +) + +# Opt-in (default off): route the fp8 sparse-MLA prefill path through the Triton +# per-query flash kernel instead of TileLang. Validated on gfx950 (GLM-5.1 @ +# TP4: 16 heads, d_v=512, tail=64). Reads q_nope/q_rope directly (skips the +# concat). Enable with SGLANG_DSA_TRITON_PREFILL=1. Decode stays on TileLang. +_DSA_TRITON_PREFILL = get_bool_env_var("SGLANG_DSA_TRITON_PREFILL") +_IS_GFX95 = is_gfx95_supported() + +if is_cuda(): + import deep_gemm + +if TYPE_CHECKING: + from sglang.srt.layers.radix_attention import RadixAttention + from sglang.srt.model_executor.model_runner import ModelRunner + from sglang.srt.speculative.spec_info import SpecInput + + +def _all_gather_dsa_trtllm_fp8_kv( + forward_batch: ForwardBatch, + k: torch.Tensor, + k_rope: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + kv_lora_rank = k.shape[-1] + qk_rope_head_dim = k_rope.shape[-1] + kv_dtype = k.dtype + kv = torch.cat((k, k_rope), dim=-1).view(torch.uint8) + kv = cp_all_gather_rerange_output( + kv, + get_parallel().attn_cp_size, + forward_batch, + torch.cuda.current_stream(), + ).view(kv_dtype) + return kv.split((kv_lora_rank, qk_rope_head_dim), dim=-1) + + +_is_hip = is_hip() + +if _is_hip: + from sglang.kernels.ops.attention.dsa.triton_kernel import get_valid_kv_indices + from sglang.kernels.ops.quantization.fp8_kernel import fp8_dtype + + try: + from aiter import ( # noqa: F401 + flash_attn_varlen_func, + get_mla_metadata_info_v1, + get_mla_metadata_v1, + mha_batch_prefill_func, + paged_attention_ragged, + ) + from aiter.mla import mla_decode_fwd, mla_prefill_fwd # noqa: F401 + except ImportError: + print( + "aiter is AMD specific kernel library. Please make sure aiter is installed on your AMD device." + ) +else: + from sglang.jit_kernel.flash_attention import ( + flash_attn_varlen_func, + flash_attn_with_kvcache, + ) + + +def _to_2d_context_lens(seqlens_32: torch.Tensor, batch_size: int) -> torch.Tensor: + # Always normalize to (N_total, 1) layout, to avoid deadlock at deep_gemm.fp8_paged_mqa_logits + if seqlens_32.dim() == 2: + if seqlens_32.size(1) == 1: + return seqlens_32 + # Fall through and re-flatten if the caller already gave us a (bs, next_n) + # view — we want (N_total, 1) regardless. + seqlens_32 = seqlens_32.reshape(-1) + return seqlens_32.contiguous().view(-1, 1) + + +@dataclass(frozen=True) +class DSAFlashMLAMetadata: + """Metadata only needed by FlashMLA""" + + flashmla_metadata: torch.Tensor + num_splits: torch.Tensor + + def slice(self, sli): + return DSAFlashMLAMetadata( + flashmla_metadata=self.flashmla_metadata, + num_splits=self.num_splits[sli], + ) + + def copy_(self, other: DSAFlashMLAMetadata): + self.flashmla_metadata.copy_(other.flashmla_metadata) + self.num_splits.copy_(other.num_splits) + + +@dataclass(frozen=True) +class DSAMetadata: + page_size: int + + # Sequence lengths for the forward batch + cache_seqlens_int32: torch.Tensor + # Maximum sequence length for query + max_seq_len_q: int + # Maximum sequence length for key + max_seq_len_k: int + # Cumulative sequence lengths for query + cu_seqlens_q: torch.Tensor + # Cumulative sequence lengths for key + cu_seqlens_k: torch.Tensor + # Page table, the index of KV Cache Tables/Blocks + # this table is always with page_size = 1. + # None for fused-decode CUDA graphs where the wide [bs, max_ctx_len] table is + # never read (attention uses topk_indices, indexer uses real_page_table); the + # graph then only materializes the compact real_page_table. See + # `dsa_drop_wide_page_table`. + page_table_1: Optional[torch.Tensor] + + # NOTE(dark): This will property be used in: + # 1. dense decode/prefill, we use paged flash attention, need real_page_table + # 2. sparse decode/prefill, indexer need real_page_table to compute the score + real_page_table: torch.Tensor + + # DSA metadata (dsa prefill are expanded) + dsa_cache_seqlens_int32: torch.Tensor # this seqlens is clipped to `topk` + dsa_cu_seqlens_q: torch.Tensor # must be arange(0, len(dsa_cu_seqlens_k)) + dsa_cu_seqlens_k: torch.Tensor # cumsum of `dsa_cache_seqlens_int32` + dsa_extend_seq_lens_list: List[int] + dsa_seqlens_expanded: torch.Tensor # expanded, unclipped `seqlens` + dsa_max_seqlen_q: Literal[1] = 1 # always 1 for decode, variable for extend + + flashmla_metadata: Optional[DSAFlashMLAMetadata] = None + # DeepGEMM schedule metadata for paged MQA logits (decode/target_verify/draft_extend only). + # Precomputed once per forward batch and reused across layers. + paged_mqa_schedule_metadata: Optional[torch.Tensor] = None + # 2D context_lens used to build the schedule above; the indexer reuses it + # as DG's `context_lens` arg so the broadcast doesn't rebuild per layer. + paged_mqa_ctx_lens_2d: Optional[torch.Tensor] = None + # Precomputed once per forward batch and reused across layers: the + # DeepSeek-V4 top-k v2 plan (cluster-threshold metadata) for the folded + # decode top-k transform. None unless SGLANG_OPT_USE_TOPK_V2 and decode. + topk_v2_plan: Optional[torch.Tensor] = None + # The sum of sequence lengths for key, prefill only + seq_lens_sum: Optional[int] = None + # The flattened 1D page table with shape (seq_lens_sum,), prefill only + # this table is always with page_size = 1 + page_table_1_flattened: Optional[torch.Tensor] = None + # The offset of topk indices in ragged kv, prefill only + # shape: (seq_lens_sum,) + topk_indices_offset: Optional[torch.Tensor] = None + + # k_start and k_end in kv cache for each token. + indexer_k_start_end: Optional[Tuple[torch.Tensor, torch.Tensor]] = None + # seq lens for each batch. + indexer_seq_lens_cpu: Optional[torch.Tensor] = None + # seq lens for each batch. + indexer_seq_lens: Optional[torch.Tensor] = None + # batch index for each token. + token_to_batch_idx: Optional[torch.Tensor] = None + + +@torch.compile +def _compiled_cat(tensors: list[torch.Tensor], dim: int = -1) -> torch.Tensor: + return torch.cat(tensors, dim=dim) + + +def _cat(tensors: list[torch.Tensor], dim: int = -1) -> torch.Tensor: + """ + Concatenate two tensors along the last dimension. + Use this function to concatenate q_nope and q_rope or k_nope and k_rope. + """ + assert len(tensors) == 2 + + qk_nope, qk_rope = tensors + assert qk_nope.ndim == 3 and qk_rope.ndim == 3 + + torch._dynamo.mark_dynamic(qk_nope, 0) + torch._dynamo.mark_dynamic(qk_rope, 0) + + return _compiled_cat([qk_nope, qk_rope], dim=dim) + + +@dataclass(frozen=True) +class DSAIndexerMetadata(BaseIndexerMetadata): + attn_metadata: DSAMetadata + topk_transform_method: TopkTransformMethod + topk_backend: DSATopKBackend = DSATopKBackend.SGL_KERNEL + paged_mqa_schedule_metadata: Optional[torch.Tensor] = None + paged_mqa_ctx_lens_2d: Optional[torch.Tensor] = None + force_unfused_topk: bool = False + + def get_seqlens_int32(self) -> torch.Tensor: + return self.attn_metadata.cache_seqlens_int32 + + def get_page_table_64(self) -> torch.Tensor: + return self.attn_metadata.real_page_table + + def get_page_table_1(self) -> torch.Tensor: + return self.attn_metadata.page_table_1 + + def get_seqlens_expanded(self) -> torch.Tensor: + return self.attn_metadata.dsa_seqlens_expanded + + def get_cu_seqlens_k(self) -> torch.Tensor: + return self.attn_metadata.cu_seqlens_k + + def get_indexer_kvcache_range(self) -> Tuple[torch.Tensor, torch.Tensor]: + return self.attn_metadata.indexer_k_start_end + + def get_indexer_seq_len(self) -> torch.Tensor: + return self.attn_metadata.indexer_seq_lens + + def get_indexer_seq_len_cpu(self) -> torch.Tensor: + return self.attn_metadata.indexer_seq_lens_cpu + + def get_dsa_extend_len_cpu(self) -> List[int]: + return self.attn_metadata.dsa_extend_seq_lens_list + + def get_token_to_batch_idx(self) -> torch.Tensor: + return self.attn_metadata.token_to_batch_idx + + def topk_transform( + self, + logits: torch.Tensor, + topk: int, + ks: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + ke_offset: Optional[torch.Tensor] = None, + batch_idx_list: Optional[List[int]] = None, + topk_indices_offset_override: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if topk_indices_offset_override is not None: + cu_topk_indices_offset = topk_indices_offset_override + cu_seqlens_q_topk = None + elif cu_seqlens_q is not None: + cu_seqlens_q = cu_seqlens_q.to(torch.int32) + cu_seqlens_q_topk = compute_cu_seqlens(cu_seqlens_q) + cu_topk_indices_offset = torch.repeat_interleave( + cu_seqlens_q_topk[:-1], + cu_seqlens_q, + ) + else: + cu_seqlens_q_topk = self.attn_metadata.cu_seqlens_q + cu_topk_indices_offset = self.attn_metadata.topk_indices_offset + if ke_offset is not None: + seq_lens_topk = ke_offset + else: + seq_lens_topk = self.get_seqlens_expanded() + return self.topk_backend.topk_transform( + logits=logits, + lengths=seq_lens_topk, + topk=topk, + topk_transform_method=self.topk_transform_method, + attn_metadata=self.attn_metadata, + cu_seqlens_q_topk=cu_seqlens_q_topk, + topk_indices_offset=cu_topk_indices_offset, + row_starts=ks, + batch_idx_list=batch_idx_list, + force_unfused_topk=self.force_unfused_topk, + ) + + +_DSA_IMPL_T: TypeAlias = Literal[ + "flashmla_sparse", "flashmla_sparse_q8", "flashmla_kv", "fa3", "tilelang", "trtllm" +] + + +class DeepseekSparseAttnBackend( + DeepseekSparseAttnBackendMTPPrecomputeMixin, AttentionBackend +): + # Decode/verify/draft graph replay rebuilds metadata from static buffers + # (page-table width) and never reads seq_lens_cpu / seq_lens_sum; opt out of + # the D2H sync. The eager fallback derives lengths from GPU seq_lens. + needs_cpu_seq_lens: bool = False + + def __init__( + self, + model_runner: ModelRunner, + skip_prefill: bool = False, + speculative_step_id=0, + topk=0, + speculative_num_steps=0, + seed_dsa_topk_from_draft_extend: bool = False, + ): + super().__init__() + self.forward_metadata: DSAMetadata + self.device = model_runner.device + assert isinstance(model_runner.page_size, int) + self.real_page_size = model_runner.page_size + self.num_splits = ( + 1 if model_runner.server_args.enable_deterministic_inference else 0 + ) + self.use_dsa = is_deepseek_dsa(model_runner.model_config.hf_config) + assert self.use_dsa, "DSA backend only supports DeepSeek DSA" + self.dsa_kv_cache_store_fp8 = ( + model_runner.token_to_kv_pool.dsa_kv_cache_store_fp8 + ) + self.dsa_index_topk = get_dsa_index_topk(model_runner.model_config.hf_config) + self.max_context_len = model_runner.model_config.context_len + self.num_q_heads = ( + model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size + ) + self.kv_cache_dim = model_runner.token_to_kv_pool.kv_cache_dim + self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim + self.kv_lora_rank = model_runner.model_config.kv_lora_rank + self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim + + assert model_runner.req_to_token_pool is not None + self.req_to_token_pool = model_runner.req_to_token_pool + self.token_to_kv_pool = model_runner.token_to_kv_pool + self.hisparse_coordinator = model_runner.hisparse_coordinator + self.req_to_token = model_runner.req_to_token_pool.req_to_token + + self.use_mha: bool = False + self.dsa_prefill_impl: _DSA_IMPL_T = ( + model_runner.server_args.dsa_prefill_backend + ) + self.dsa_decode_impl: _DSA_IMPL_T = model_runner.server_args.dsa_decode_backend + self.dsa_topk_backend: DSATopKBackend = DSATopKBackend( + model_runner.server_args.dsa_topk_backend + ) + if self.num_q_heads <= 64: + self.flashmla_kv_num_q_heads = 64 + elif self.num_q_heads <= 128: + self.flashmla_kv_num_q_heads = 128 + else: + # Keep original head count if it exceeds current padded variants. + self.flashmla_kv_num_q_heads = self.num_q_heads + self.enable_auto_select_prefill_impl = self.dsa_prefill_impl == "flashmla_auto" + + self._arange_buf = torch.arange(16384, device=self.device, dtype=torch.int32) + + if _is_hip: + max_bs = model_runner.req_to_token_pool.size + + self.kv_indptr = torch.zeros( + (max_bs + 1,), dtype=torch.int32, device=model_runner.device + ) + + self.kv_indices = torch.zeros( + max_bs * self.dsa_index_topk, + dtype=torch.int32, + device=self.device, + ) + # Aiter mla_decode_fwd supports num_heads multiples of 16 in range [16, 128]. + # For models with fewer heads per GPU (e.g. GLM-5 64 heads / TP8 = 8), need to pad the heads to 16. + self.need_pad_heads = self.num_q_heads < 16 + self.head_repeat_factor = ( + 16 // self.num_q_heads if self.num_q_heads < 16 else 1 + ) + self.num_head_padded = self.num_q_heads * self.head_repeat_factor + self.aiter_dsa_max_split_per_batch = 64 + self.aiter_dsa_metadata_capacity = 0 + self.aiter_dsa_metadata_max_seqlen_q = 0 + self.aiter_dsa_metadata_q_dtype = None + self.aiter_dsa_metadata_kv_dtype = None + self.aiter_dsa_kv_last_page_lens = None + self.aiter_dsa_work_metadata = None + + if ( + self.dsa_prefill_impl == "aiter" or self.dsa_decode_impl == "aiter" + ) and model_runner.kv_cache_dtype == fp8_dtype: + self._ensure_aiter_dsa_decode_metadata_buffer( + max_seqlen_q=1, + batch_size=max_bs, + q_dtype=torch.bfloat16, + kv_dtype=fp8_dtype, + ) + + # Speculative decoding + self.topk = model_runner.server_args.speculative_eagle_topk or 0 + self.speculative_num_steps = speculative_num_steps + self.speculative_num_draft_tokens = ( + model_runner.server_args.speculative_num_draft_tokens + ) + self.speculative_step_id = speculative_step_id + self.use_fused_topk = should_use_dsa_fused_topk( + model_runner.server_args, seed_dsa_topk_from_draft_extend + ) + if envs.SGLANG_DSA_FUSE_TOPK.get() and not self.use_fused_topk: + print_warning_once( + "Disabling fused DSA top-k for IndexShare under PD disaggregation." + ) + + self.device_capability = torch.cuda.get_device_capability() + self.device_sm_major = self.device_capability[0] + self.kv_cache_dtype = model_runner.kv_cache_dtype + + # `flashmla_sparse_q8` = the native FP8 SM90 sparse-prefill kernel. It always + # runs FP8 (requires fp8_e4m3 KV) and is SM90-only, so validate both at + # construction: an unsupported config must fail at launch rather than + # mid-forward. `flashmla_sparse` remains the bf16 path with no such + # requirement. + if self.dsa_prefill_impl == "flashmla_sparse_q8": + if self.kv_cache_dtype != torch.float8_e4m3fn: + raise ValueError( + "--dsa-prefill-backend flashmla_sparse_q8 is native FP8 and requires " + f"--kv-cache-dtype fp8_e4m3 (got kv_cache_dtype={self.kv_cache_dtype}); " + "use --dsa-prefill-backend flashmla_sparse for the bf16 path." + ) + if self.device_sm_major != 9: + raise ValueError( + "--dsa-prefill-backend flashmla_sparse_q8 is SM90-only; got compute " + f"capability sm_{self.device_sm_major}x." + ) + + # `flashmla_sparse_q8` is prefill-only (FP8 decode goes through + # `flashmla_kv`); reject it as a decode backend, since argparse accepts it + # via the shared DSA_CHOICES list. + if self.dsa_decode_impl == "flashmla_sparse_q8": + raise ValueError( + "--dsa-decode-backend flashmla_sparse_q8 is not supported: " + "flashmla_sparse_q8 is a prefill-only backend. For FP8, use " + "--dsa-prefill-backend flashmla_sparse_q8 together with " + "--dsa-decode-backend flashmla_kv." + ) + + # Q8KV8 per-call device-tensor caches, populated lazily on the first + # Q8KV8 dispatch (no-ops for other backends). + self._q8kv8_identity_scale: Optional[torch.Tensor] = None + self._q8kv8_qpad_buf: Optional[torch.Tensor] = None + + # Allocate global workspace buffer for TRT-LLM kernels (ragged attention on SM100/B200, or trtllm decode) + if self.device_sm_major >= 10 or self.dsa_decode_impl == "trtllm": + self.workspace_buffer = get_buffer( + "dsa_trtllm_workspace", + lambda: torch.empty( + envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.get(), + dtype=torch.uint8, + device=model_runner.device, + ), + ) + else: + self.workspace_buffer = None + + def _make_aiter_dsa_decode_metadata_buffer( + self, + max_seqlen_q: int, + batch_size: int, + q_dtype: torch.dtype, + kv_dtype: torch.dtype, + ): + ( + (work_metadata_size, work_metadata_type), + (work_indptr_size, work_indptr_type), + (work_info_set_size, work_info_set_type), + (reduce_indptr_size, reduce_indptr_type), + (reduce_final_map_size, reduce_final_map_type), + (reduce_partial_map_size, reduce_partial_map_type), + ) = get_mla_metadata_info_v1( + batch_size, + max_seqlen_q, + self.num_head_padded, + q_dtype, + kv_dtype, + is_sparse=True, + fast_mode=False, + num_kv_splits=self.aiter_dsa_max_split_per_batch, + intra_batch_mode=True, + ) + + return ( + torch.empty( + work_metadata_size, dtype=work_metadata_type, device=self.device + ), + torch.empty(work_indptr_size, dtype=work_indptr_type, device=self.device), + torch.empty( + work_info_set_size, dtype=work_info_set_type, device=self.device + ), + torch.empty( + reduce_indptr_size, dtype=reduce_indptr_type, device=self.device + ), + torch.empty( + reduce_final_map_size, dtype=reduce_final_map_type, device=self.device + ), + torch.empty( + reduce_partial_map_size, + dtype=reduce_partial_map_type, + device=self.device, + ), + ) + + def _ensure_aiter_dsa_decode_metadata_buffer( + self, + max_seqlen_q: int, + batch_size: int, + q_dtype: torch.dtype, + kv_dtype: torch.dtype, + ) -> None: + if ( + self.aiter_dsa_work_metadata is not None + and self.aiter_dsa_metadata_capacity >= batch_size + and self.aiter_dsa_metadata_max_seqlen_q == max_seqlen_q + and self.aiter_dsa_metadata_q_dtype == q_dtype + and self.aiter_dsa_metadata_kv_dtype == kv_dtype + ): + return + + ( + self.aiter_dsa_work_metadata, + self.aiter_dsa_work_indptr, + self.aiter_dsa_work_info_set, + self.aiter_dsa_reduce_indptr, + self.aiter_dsa_reduce_final_map, + self.aiter_dsa_reduce_partial_map, + ) = self._make_aiter_dsa_decode_metadata_buffer( + max_seqlen_q=max_seqlen_q, + batch_size=batch_size, + q_dtype=q_dtype, + kv_dtype=kv_dtype, + ) + self.aiter_dsa_kv_last_page_lens = torch.ones( + (batch_size,), dtype=torch.int32, device=self.device + ) + self.aiter_dsa_metadata_capacity = batch_size + self.aiter_dsa_metadata_max_seqlen_q = max_seqlen_q + self.aiter_dsa_metadata_q_dtype = q_dtype + self.aiter_dsa_metadata_kv_dtype = kv_dtype + + def _prepare_aiter_dsa_decode_metadata( + self, + qo_indptr: torch.Tensor, + kv_indptr: torch.Tensor, + bs: int, + max_seqlen_q: int, + q_dtype: torch.dtype, + kv_dtype: torch.dtype, + ) -> dict: + self._ensure_aiter_dsa_decode_metadata_buffer( + max_seqlen_q=max_seqlen_q, + batch_size=bs, + q_dtype=q_dtype, + kv_dtype=kv_dtype, + ) + self.aiter_dsa_kv_last_page_lens[:bs].fill_(1) + kv_last_page_lens = self.aiter_dsa_kv_last_page_lens[:bs] + + get_mla_metadata_v1( + qo_indptr, + kv_indptr, + kv_last_page_lens, + self.num_head_padded, + 1, + False, + self.aiter_dsa_work_metadata, + self.aiter_dsa_work_info_set, + self.aiter_dsa_work_indptr, + self.aiter_dsa_reduce_indptr, + self.aiter_dsa_reduce_final_map, + self.aiter_dsa_reduce_partial_map, + page_size=1, + kv_granularity=16, + max_seqlen_qo=max_seqlen_q, + uni_seqlen_qo=max_seqlen_q, + fast_mode=False, + topk=self.dsa_index_topk, + max_split_per_batch=self.aiter_dsa_max_split_per_batch, + intra_batch_mode=True, + dtype_q=q_dtype, + dtype_kv=kv_dtype, + ) + + return { + "kv_last_page_lens": kv_last_page_lens, + "work_meta_data": self.aiter_dsa_work_metadata, + "work_indptr": self.aiter_dsa_work_indptr, + "work_info_set": self.aiter_dsa_work_info_set, + "reduce_indptr": self.aiter_dsa_reduce_indptr, + "reduce_final_map": self.aiter_dsa_reduce_final_map, + "reduce_partial_map": self.aiter_dsa_reduce_partial_map, + "intra_batch_mode": True, + "num_kv_splits": self.aiter_dsa_max_split_per_batch, + } + + def _build_paged_mqa_schedule_2d_ctx_lens( + self, + forward_mode: ForwardMode, + cache_seqlens_int32: torch.Tensor, + seqlens_expanded: torch.Tensor, + batch_size: int, + ) -> torch.Tensor: + # target_verify with next_n>=2 uses DG-native q=[B,next_n,H,D] which + # needs a [B, next_n] schedule; everything else stays per-token. + # TODO: SM90 supports DG-native next_n in {1,2} too — enable once + # validated; for now DG-native is SM100+ only. + next_n = self.speculative_num_draft_tokens + if ( + forward_mode.is_target_verify() + and next_n + and next_n >= 2 + and is_sm100_supported() + ): + return cache_seqlens_int32.view(-1, 1).expand(-1, next_n).contiguous() + if forward_mode.is_target_verify() or forward_mode.is_draft_extend_v2(): + return _to_2d_context_lens(seqlens_expanded, batch_size) + return _to_2d_context_lens(cache_seqlens_int32, batch_size) + + def _refresh_paged_mqa_schedule_metadata( + self, + metadata: DSAMetadata, + seqlens_32_2d: torch.Tensor, + ) -> None: + if is_sm120_supported(): + return + new_schedule = deep_gemm.get_paged_mqa_logits_metadata( + seqlens_32_2d, 64, deep_gemm.get_num_sms() + ) + if metadata.paged_mqa_schedule_metadata is None: + object.__setattr__(metadata, "paged_mqa_schedule_metadata", new_schedule) + else: + metadata.paged_mqa_schedule_metadata.copy_(new_schedule) + + def _build_topk_v2_plan( + self, seqlens_expanded: torch.Tensor + ) -> Optional[torch.Tensor]: + # Preprocess the folded top-k v2 plan once per forward (shared across + # layers), at metadata-build time, from the same seqlens the transform + # receives as `lengths` (dsa_seqlens_expanded). This must cover EVERY shape + # that dispatches to `_topk_transform_v2_paged` -- decode AND MTP + # target-verify / draft-extend, whose expanded row count is exactly what v2 + # sees -- otherwise the helper's plan-present assertion fires. None only + # when the fold is disabled; such metadata is never dispatched to v2. + if not envs.SGLANG_OPT_USE_TOPK_V2.get(): + return None + from sglang.jit_kernel.dsv4.topk import plan_topk_v2 + + return plan_topk_v2(seqlens_expanded) + + def _refresh_topk_v2_plan(self, metadata: DSAMetadata) -> None: + # Refresh the plan in-place under CUDA graph replay so the captured + # read sees fresh cluster metadata for the replay's decode seq lengths. + # `copy_` preserves the buffer's data_ptr captured by the graph. None + # means it was not built (fold disabled / non-decode shape), and such a + # metadata object is never dispatched to the v2 helper, so there is + # nothing to refresh. + if metadata.topk_v2_plan is None: + return + from sglang.jit_kernel.dsv4.topk import plan_topk_v2 + + metadata.topk_v2_plan.copy_(plan_topk_v2(metadata.dsa_seqlens_expanded)) + + def _get_fused_topk_page_table(self, topk_indices: torch.Tensor) -> torch.Tensor: + if ( + self.dsa_topk_backend.is_sgl_kernel() + or self.dsa_topk_backend.is_flashinfer() + ): + return topk_indices + raise RuntimeError( + f"Unsupported {self.dsa_topk_backend = } for SGLANG_DSA_FUSE_TOPK." + ) + + def get_device_int32_arange(self, length: int) -> torch.Tensor: + if length > len(self._arange_buf): + next_pow_of_2 = 1 << (length - 1).bit_length() + self._arange_buf = torch.arange( + next_pow_of_2, device=self.device, dtype=torch.int32 + ) + return self._arange_buf[:length] + + def _graph_page_table_width(self, metadata: DSAMetadata) -> int: + """Column count to scan req_to_token during graph replay. Reads the wide + page_table_1 width when present, else req_to_token's width (the wide table + is dropped for fused decode graphs, see `dsa_drop_wide_page_table`).""" + if metadata.page_table_1 is not None: + return metadata.page_table_1.shape[1] + return self.req_to_token.shape[1] + + def _transform_table_1_to_real(self, page_table: torch.Tensor) -> torch.Tensor: + page_size = self.real_page_size + if page_size == 1: + return page_table + max_seqlen_k = page_table.shape[1] + strided_indices = torch.arange( + 0, max_seqlen_k, page_size, device=page_table.device, dtype=torch.int32 + ) + return page_table[:, strided_indices] // page_size + + def init_forward_metadata_out_graph( + self, + forward_batch: ForwardBatch, + in_capture: bool = False, + ): + seq_lens_cpu = ( + forward_batch.seq_lens.cpu() if in_capture else forward_batch.seq_lens_cpu + ) + self._apply_cuda_graph_metadata( + bs=forward_batch.batch_size, + req_pool_indices=forward_batch.req_pool_indices, + seq_lens=forward_batch.seq_lens, + seq_lens_cpu=seq_lens_cpu, + forward_mode=forward_batch.forward_mode, + spec_info=forward_batch.spec_info, + out_cache_loc=getattr(forward_batch, "out_cache_loc", None), + actual_forward_mode=getattr(forward_batch, "actual_forward_mode", None), + ) + + def init_forward_metadata(self, forward_batch: ForwardBatch): + """Init the metadata for a forward pass.""" + batch_size = forward_batch.batch_size + device = forward_batch.seq_lens.device + + if forward_batch.forward_mode.is_target_verify(): + draft_token_num = self.speculative_num_draft_tokens + else: + draft_token_num = 0 + + cache_seqlens_int32 = (forward_batch.seq_lens + draft_token_num).to(torch.int32) + cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32) + if forward_batch.seq_lens_cpu is not None: + max_seqlen_k = int( + forward_batch.seq_lens_cpu.max().item() + draft_token_num + ) + else: + # needs_cpu_seq_lens=False nulls the host mirror for spec-v2 relay + # batches; graph replay uses the static page-table width, so only this + # eager (e.g. over-capture-bs) fallback needs a length here. + max_seqlen_k = int(forward_batch.seq_lens.max().item()) + draft_token_num + # [b, max_seqlen_k] + page_table = self.req_to_token_pool.req_to_token[ + forward_batch.req_pool_indices, :max_seqlen_k + ] + + page_table_1_flattened = None + topk_indices_offset = None + + # Centralized dispatch: decide all strategies for this batch + self.set_dsa_prefill_impl(forward_batch) + dsa_impl_for_batch = ( + self.dsa_decode_impl + if ( + forward_batch.forward_mode.is_decode_or_idle() + or forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend_v2() + ) + else self.dsa_prefill_impl + ) + use_flashmla_kv = (not self.use_mha) and dsa_impl_for_batch == "flashmla_kv" + topk_transform_method = self.get_topk_transform_method( + forward_batch.forward_mode + ) + # Batch indices selected when cp enabled: After splitting multiple sequences, + # a certain cp rank may not have some of these sequences. + # We use bs_idx_cpu to mark which sequences are finally selected by the current cp rank, + # a default value of None indicates that all sequences are selected. + bs_idx_cpu = None + # seq_len_cpu of selected sequences + indexer_seq_lens_cpu = forward_batch.seq_lens_cpu + indexer_seq_lens = forward_batch.seq_lens + + if forward_batch.forward_mode.is_decode_or_idle(): + extend_seq_lens_cpu = [1] * batch_size + max_seqlen_q = 1 + cu_seqlens_q = self.get_device_int32_arange(batch_size + 1) + seqlens_expanded = cache_seqlens_int32 + elif forward_batch.forward_mode.is_target_verify(): + max_seqlen_q = 1 + cu_seqlens_q = torch.arange( + 0, + batch_size * self.speculative_num_draft_tokens + 1, + 1, + dtype=torch.int32, + device=device, + ) + extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * batch_size + forward_batch.extend_seq_lens_cpu = extend_seq_lens_cpu + + seqlens_expanded = seqlens_expand_triton( + torch.tensor(extend_seq_lens_cpu, dtype=torch.int32, device=device), + cache_seqlens_int32, + self.speculative_num_draft_tokens * batch_size, + self.speculative_num_draft_tokens, + ) + page_table = torch.repeat_interleave( + page_table, repeats=self.speculative_num_draft_tokens, dim=0 + ) + elif forward_batch.forward_mode.is_draft_extend_v2(): + if forward_batch.extend_prefix_lens_cpu is None: + assert forward_batch.extend_prefix_lens is not None + forward_batch.extend_prefix_lens_cpu = ( + forward_batch.extend_prefix_lens.cpu().tolist() + ) + if forward_batch.seq_lens_cpu is None: + forward_batch.seq_lens_cpu = forward_batch.seq_lens.cpu() + forward_batch.seq_lens_sum = int(forward_batch.seq_lens_cpu.sum()) + assert ( + forward_batch.extend_seq_lens_cpu is not None + and forward_batch.extend_seq_lens is not None + and forward_batch.extend_prefix_lens_cpu is not None + ), "All of them must not be None" + + extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu + assert forward_batch.extend_seq_lens is not None + + max_seqlen_q = 1 + cu_seqlens_q = torch.arange( + 0, + forward_batch.extend_num_tokens + 1, + 1, + dtype=torch.int32, + device=device, + ) + + seqlens_expanded = seqlens_expand_triton( + forward_batch.extend_seq_lens, + cache_seqlens_int32, + sum(extend_seq_lens_cpu), + self.speculative_num_draft_tokens, + ) + if forward_batch.forward_mode.is_draft_extend_v2(): + # DRAFT_EXTEND_V2: V2 worker pre-fills draft KV cache with ALL speculated + # tokens upfront. All requests extend by the same fixed + # (speculative_num_draft_tokens). Use scalar to avoid GPU sync. + page_table = torch.repeat_interleave( + page_table, repeats=self.speculative_num_draft_tokens, dim=0 + ) + else: + # DRAFT_EXTEND: the draft worker extends by (num_correct_drafts + 1) + # per request after verification. Lengths vary per request based on + # how many tokens were accepted. + page_table = torch.repeat_interleave( + page_table, repeats=forward_batch.extend_seq_lens, dim=0 + ) + elif forward_batch.forward_mode.is_extend(): + assert ( + forward_batch.extend_seq_lens_cpu is not None + and forward_batch.extend_seq_lens is not None + and forward_batch.extend_prefix_lens_cpu is not None + ), "All of them must not be None" + extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu + assert forward_batch.extend_seq_lens is not None + extend_seq_lens = forward_batch.extend_seq_lens + + seqlens_expanded = torch.cat( + [ + torch.arange( + kv_len - qo_len + 1, + kv_len + 1, + dtype=torch.int32, + device=device, + ) + for qo_len, kv_len in zip( + forward_batch.extend_seq_lens_cpu, + forward_batch.seq_lens_cpu.tolist(), + strict=True, + ) + ] + ) + + if can_dsa_prefill_cp_round_robin_split(forward_batch): + seqlens_expanded = dsa_cp_round_robin_split_data(seqlens_expanded) + extend_seq_lens_cpu, extend_seq_lens, bs_idx_cpu, bs_idx = ( + dsa_cp_round_robin_split_q_seqs( + extend_seq_lens_cpu, extend_seq_lens + ) + ) + indexer_seq_lens_cpu = indexer_seq_lens_cpu[bs_idx_cpu] + indexer_seq_lens = indexer_seq_lens[bs_idx] + cache_seqlens_int32 = cache_seqlens_int32[bs_idx] + cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32) + max_seqlen_k = ( + int(indexer_seq_lens_cpu.max().item() + draft_token_num) + if len(indexer_seq_lens_cpu) != 0 + else 0 + ) + page_table = page_table[bs_idx, :max_seqlen_k] + + if any(forward_batch.extend_prefix_lens_cpu) or bs_idx_cpu is not None: + max_seqlen_q = ( + max(extend_seq_lens_cpu) if len(extend_seq_lens_cpu) != 0 else 1 + ) + cu_seqlens_q = compute_cu_seqlens(extend_seq_lens.to(torch.int32)) + else: + max_seqlen_q = max_seqlen_k + cu_seqlens_q = cu_seqlens_k + + # Check if MHA FP8 dequantization is needed + mha_dequantize_needed = ( + self.use_mha and self.token_to_kv_pool.dtype == torch.float8_e4m3fn + ) + forward_batch.using_mha_one_shot_fp8_dequant = mha_dequantize_needed + + # page_table_1_flattened is only used when prefix sharing is enabled: + has_prefix_sharing = any(forward_batch.extend_prefix_lens_cpu) + if has_prefix_sharing and ( + topk_transform_method == TopkTransformMethod.RAGGED + or mha_dequantize_needed + ): + page_table_1_flattened = torch.cat( + [ + page_table[i, :kv_len] + for i, kv_len in enumerate( + indexer_seq_lens_cpu.tolist(), + ) + ] + ) + assert page_table_1_flattened.shape[0] == sum( + indexer_seq_lens_cpu + ), f"{page_table_1_flattened.shape[0] = } must be the same as {sum(indexer_seq_lens_cpu) = }" + + # Validate indices when logical tokens exceed physical capacity + # This is likely to be triggered by PP with high kv reuse & parallelism + kv_cache_capacity = ( + self.token_to_kv_pool.size + self.token_to_kv_pool.page_size + ) + if forward_batch.seq_lens_sum > kv_cache_capacity: + max_idx = page_table_1_flattened.max().item() + assert max_idx < kv_cache_capacity, ( + f"Invalid page table index: max={max_idx}, " + f"kv_cache_capacity={kv_cache_capacity}" + ) + + if topk_transform_method == TopkTransformMethod.RAGGED: + topk_indices_offset = torch.repeat_interleave( + cu_seqlens_k[:-1], + extend_seq_lens, + ) + else: + assert False, f"Unsupported {forward_batch.forward_mode = }" + + indexer_k_start_end, token_to_batch_idx = self._cal_indexer_k_start_end( + forward_batch, bs_idx_cpu + ) + # 1D, expanded seqlens (1D means cheap to compute, so always compute it) + dsa_cache_seqlens_int32 = compute_dsa_seqlens( + original_seq_lens=seqlens_expanded, + dsa_index_topk=self.dsa_index_topk, + ) + dsa_cache_seqlens_int32 = pad_dsa_cache_seqlens( + forward_batch, dsa_cache_seqlens_int32 + ) + dsa_cu_seqlens_k = compute_cu_seqlens(dsa_cache_seqlens_int32) + dsa_cu_seqlens_q = self.get_device_int32_arange(len(dsa_cu_seqlens_k)) + + paged_mqa_schedule_metadata = None + paged_mqa_ctx_lens_2d = None + if is_cuda() and ( + forward_batch.forward_mode.is_decode_or_idle() + or forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend_v2() + ): + paged_mqa_ctx_lens_2d = self._build_paged_mqa_schedule_2d_ctx_lens( + forward_batch.forward_mode, + cache_seqlens_int32, + seqlens_expanded, + forward_batch.batch_size, + ) + # NOTE: block_kv arg must be 64 here — DG computes SPLIT_KV = + # block_kv * 4 and both DG's and the indexer's compute kernels + # require SPLIT_KV = 256; this is independent of the cache page size. + if not is_sm120_supported(): + paged_mqa_schedule_metadata = ( + deep_gemm.get_paged_mqa_logits_metadata( + paged_mqa_ctx_lens_2d, 64, deep_gemm.get_num_sms() + ) + ) + + metadata = DSAMetadata( + page_size=self.real_page_size, + cache_seqlens_int32=cache_seqlens_int32, + max_seq_len_q=max_seqlen_q, + max_seq_len_k=max_seqlen_k, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + seq_lens_sum=forward_batch.seq_lens_sum, + page_table_1=page_table, + page_table_1_flattened=page_table_1_flattened, + flashmla_metadata=( + self._compute_flashmla_metadata( + cache_seqlens=dsa_cache_seqlens_int32, + seq_len_q=1, + ) + if use_flashmla_kv + else None + ), + paged_mqa_schedule_metadata=paged_mqa_schedule_metadata, + paged_mqa_ctx_lens_2d=paged_mqa_ctx_lens_2d, + dsa_cache_seqlens_int32=dsa_cache_seqlens_int32, + dsa_cu_seqlens_q=dsa_cu_seqlens_q, + dsa_cu_seqlens_k=dsa_cu_seqlens_k, + dsa_seqlens_expanded=seqlens_expanded, + dsa_extend_seq_lens_list=extend_seq_lens_cpu, + real_page_table=self._transform_table_1_to_real(page_table), + dsa_max_seqlen_q=1, + topk_indices_offset=topk_indices_offset, + indexer_k_start_end=indexer_k_start_end, + indexer_seq_lens_cpu=indexer_seq_lens_cpu, + indexer_seq_lens=indexer_seq_lens, + token_to_batch_idx=token_to_batch_idx, + topk_v2_plan=self._build_topk_v2_plan(seqlens_expanded), + ) + self.forward_metadata = metadata + + def _cal_indexer_k_start_end( + self, + forward_batch: ForwardBatch, + bs_idx: Optional[List[int]] = None, + ): + if not forward_batch.forward_mode.is_extend_without_speculative(): + return None, None + if forward_batch.batch_size == 0 or (bs_idx is not None and len(bs_idx) == 0): + empty_t = torch.empty(0, dtype=torch.int32, device=self.device) + return (empty_t, empty_t), empty_t + + # Suppose there are two requests, with extend_seq_len = [3, 2] + # and seq_lens = [10, 4] + # The logits matrix looks like this, with * representing the valid logits + # and - representing the invalid logits: + # + # ********--|---- + # *********-|---- + # **********|---- + # ----------|***- + # ----------|**** + # + # ks = [0, 0, 0, 10, 10] + # ke = [8, 9, 10, 13, 14] + ks_list = [] + ke_list = [] + token_to_batch_idx = [] + + q_offset = 0 + k_offset = 0 + + assert ( + forward_batch.seq_lens_cpu is not None + and forward_batch.extend_seq_lens_cpu is not None + ) + for i in range(forward_batch.batch_size): + seq_len = forward_batch.seq_lens_cpu[i].item() + assert isinstance(seq_len, int) + extend_seq_len = forward_batch.extend_seq_lens_cpu[i] + ks = torch.full( + (extend_seq_len,), k_offset, dtype=torch.int32, device=self.device + ) + kv_len = seq_len + if forward_batch.forward_mode.is_target_verify(): + kv_len += self.speculative_num_draft_tokens + seq_lens_expanded = torch.arange( + kv_len - extend_seq_len + 1, + kv_len + 1, + dtype=torch.int32, + device=self.device, + ) + ke = ks + seq_lens_expanded + ks_list.append(ks) + ke_list.append(ke) + + # bi: The index within the selected batch bs_idx. Entries that were not selected are ignored. + bi = bs_idx.index(i) if (bs_idx is not None and i in bs_idx) else i + tb = torch.full( + (extend_seq_len,), bi, dtype=torch.int32, device=self.device + ) + token_to_batch_idx.append(tb) + + if bs_idx is None or i in bs_idx: # skip batch not included in bs_idx + q_offset += extend_seq_len + k_offset += seq_len + + ks = torch.cat(ks_list, dim=0) + ke = torch.cat(ke_list, dim=0) + token_to_batch_idx = torch.cat(token_to_batch_idx, dim=0) + if bs_idx is not None: + assert can_dsa_prefill_cp_round_robin_split(forward_batch) + ks = dsa_cp_round_robin_split_data(ks) + ke = dsa_cp_round_robin_split_data(ke) + token_to_batch_idx = dsa_cp_round_robin_split_data(token_to_batch_idx) + return (ks, ke), token_to_batch_idx + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + """Initialize CUDA graph state for the attention backend. + + Args: + max_bs (int): Maximum batch size to support in CUDA graphs + + This creates fixed-size tensors that will be reused during CUDA graph replay + to avoid memory allocations. + """ + # Whether we can skip the wide [max_num_tokens, max_ctx_len] page_size=1 + # page table in the decode CUDA graph. It is dead weight there only when the + # decode top-k routes to the fused v2 kernel: attention reads topk_indices + # and the indexer reads the compact real_page_table, so nothing reads the + # page_size=1 table. This MUST match the exact condition under which + # `DSATopKBackend.topk_transform` dispatches decode PAGED to + # `_topk_transform_v2_paged` -- otherwise the legacy transform would read a + # dropped (None) table. Hence: fused top-k AND v2 enabled AND index_topk in + # the kernel's supported range, on CUDA with page_size>1. Excludes HIP (its + # indexer reads page_table_1), hisparse (needs page_size=1 loc translation), + # and spec decoding (MTP precompute fast-path + target-verify/draft-extend + # still consume the wide table). Computed once from stable config; the graph + # is captured once per process. + self.dsa_drop_wide_page_table = ( + is_cuda() + and not _is_hip + and self.real_page_size > 1 + and self.hisparse_coordinator is None + and not self.speculative_num_draft_tokens + and self.use_fused_topk + and envs.SGLANG_OPT_USE_TOPK_V2.get() + and self.dsa_index_topk is not None + and self.dsa_index_topk <= 2048 + ) + + max_ctx_len = self.req_to_token.shape[1] + self.decode_cuda_graph_metadata: Dict = { + "cache_seqlens": torch.ones( + max_num_tokens, dtype=torch.int32, device=self.device + ), + "cu_seqlens_q": torch.arange( + 0, max_bs + 1, dtype=torch.int32, device=self.device + ), + "cu_seqlens_k": torch.zeros( + max_bs + 1, dtype=torch.int32, device=self.device + ), + # fake page_table for sparse_prefill + # Match req_to_token's width exactly. It is over-allocated beyond + # context_len because spec decoding lets seq_len transiently overshoot. + # When dropping the wide table (fused decode), allocate only the compact + # page_size=64 real table; else allocate the wide page_size=1 table and + # derive real from it per batch size. + "real_page_table": ( + torch.zeros( + max_num_tokens, + (max_ctx_len + self.real_page_size - 1) // self.real_page_size, + dtype=torch.int32, + device=self.device, + ) + if self.dsa_drop_wide_page_table + else None + ), + "page_table": ( + None + if self.dsa_drop_wide_page_table + else torch.zeros( + max_num_tokens, + max_ctx_len, + dtype=torch.int32, + device=self.device, + ) + ), + "flashmla_metadata": ( + self._compute_flashmla_metadata( + cache_seqlens=torch.ones( + max_num_tokens, dtype=torch.int32, device=self.device + ), + seq_len_q=1, + ) + if self.dsa_decode_impl == "flashmla_kv" + else None + ), + } + + def _build_forward_metadata_cuda_graph( + self, + bs: int, + num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: Optional[torch.Tensor], + forward_mode: ForwardMode, + spec_info: Optional[SpecInput], + out_cache_loc: Optional[torch.Tensor] = None, + actual_forward_mode: Optional[ForwardMode] = None, + ): + """Create and store DSAMetadata for a new batch size during CUDA graph capture.""" + self.set_dsa_prefill_impl(forward_batch=None) + + if forward_mode.is_decode_or_idle(): + # Normal Decode + # Get sequence information + cache_seqlens_int32 = seq_lens.to(torch.int32) + cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32) + + # Use max context length for seq_len_k + real_rows = bs + if self.dsa_drop_wide_page_table: + page_table_1 = None + max_seqlen_k = self.req_to_token.shape[1] + else: + page_table_1 = self.decode_cuda_graph_metadata["page_table"][:bs, :] + max_seqlen_k = page_table_1.shape[1] + max_seqlen_q = 1 + + # Precompute page table + # Precompute cumulative sequence lengths + + # NOTE(dark): this is always arange, since we are decoding + cu_seqlens_q = self.decode_cuda_graph_metadata["cu_seqlens_q"][: bs + 1] + dsa_cache_seqlens_int32 = compute_dsa_seqlens( + cache_seqlens_int32, dsa_index_topk=self.dsa_index_topk + ) + + seqlens_expanded = cache_seqlens_int32 + dsa_extend_seq_lens_list = [1] * bs + if self.dsa_decode_impl == "flashmla_kv": + flashmla_metadata = self.decode_cuda_graph_metadata[ + "flashmla_metadata" + ].slice(slice(0, bs + 1)) + flashmla_metadata.copy_( + self._compute_flashmla_metadata( + cache_seqlens=dsa_cache_seqlens_int32, + seq_len_q=1, + ) + ) + else: + flashmla_metadata = None + elif forward_mode.is_target_verify() or forward_mode.is_draft_extend_v2(): + cache_seqlens_int32 = (seq_lens + self.speculative_num_draft_tokens).to( + torch.int32 + ) + cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32) + max_seqlen_q = 1 + real_rows = bs * self.speculative_num_draft_tokens + if self.dsa_drop_wide_page_table: + page_table_1 = None + max_seqlen_k = self.req_to_token.shape[1] + else: + page_table_1 = self.decode_cuda_graph_metadata["page_table"][ + :real_rows, : + ] + max_seqlen_k = page_table_1.shape[1] + + cu_seqlens_q = torch.arange( + 0, + bs * self.speculative_num_draft_tokens + 1, + 1, + dtype=torch.int32, + device=self.device, + ) + + extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * bs + + seqlens_int32_cpu = [ + self.speculative_num_draft_tokens + kv_len + for kv_len in seq_lens.tolist() + ] + seqlens_expanded = torch.cat( + [ + torch.arange( + kv_len - qo_len + 1, + kv_len + 1, + dtype=torch.int32, + device=self.device, + ) + for qo_len, kv_len in zip( + extend_seq_lens_cpu, + seqlens_int32_cpu, + strict=True, + ) + ] + ) + dsa_cache_seqlens_int32 = compute_dsa_seqlens( + seqlens_expanded, dsa_index_topk=self.dsa_index_topk + ) + dsa_extend_seq_lens_list = [1] * bs * self.speculative_num_draft_tokens + + if self.dsa_decode_impl == "flashmla_kv": + flashmla_metadata = self.decode_cuda_graph_metadata[ + "flashmla_metadata" + ].slice(slice(0, bs * self.speculative_num_draft_tokens + 1)) + + flashmla_metadata.copy_( + self._compute_flashmla_metadata( + cache_seqlens=dsa_cache_seqlens_int32, + seq_len_q=1, + ) + ) + else: + flashmla_metadata = None + + dsa_cu_seqlens_k = compute_cu_seqlens(dsa_cache_seqlens_int32) + dsa_cu_seqlens_q = self.get_device_int32_arange(len(dsa_cu_seqlens_k)) + if self.dsa_drop_wide_page_table: + # Compact page_size=64 static buffer; filled per-replay by the fused + # metadata kernel straight from req_to_token (no wide table needed). + real_page_table = self.decode_cuda_graph_metadata["real_page_table"][ + :real_rows, : + ] + else: + real_page_table = self._transform_table_1_to_real(page_table_1) + + paged_mqa_schedule_metadata = None + paged_mqa_ctx_lens_2d = None + if is_cuda() and ( + forward_mode.is_decode_or_idle() + or forward_mode.is_target_verify() + or forward_mode.is_draft_extend_v2() + ): + paged_mqa_ctx_lens_2d = self._build_paged_mqa_schedule_2d_ctx_lens( + forward_mode, cache_seqlens_int32, seqlens_expanded, bs + ) + if not is_sm120_supported(): + paged_mqa_schedule_metadata = ( + deep_gemm.get_paged_mqa_logits_metadata( + paged_mqa_ctx_lens_2d, 64, deep_gemm.get_num_sms() + ) + ) + + metadata = DSAMetadata( + page_size=self.real_page_size, + cache_seqlens_int32=cache_seqlens_int32, + max_seq_len_q=max_seqlen_q, + max_seq_len_k=max_seqlen_k, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + page_table_1=page_table_1, + flashmla_metadata=flashmla_metadata, + paged_mqa_schedule_metadata=paged_mqa_schedule_metadata, + paged_mqa_ctx_lens_2d=paged_mqa_ctx_lens_2d, + dsa_cache_seqlens_int32=dsa_cache_seqlens_int32, + dsa_cu_seqlens_q=dsa_cu_seqlens_q, + dsa_cu_seqlens_k=dsa_cu_seqlens_k, + dsa_seqlens_expanded=seqlens_expanded, + real_page_table=real_page_table, + dsa_extend_seq_lens_list=dsa_extend_seq_lens_list, + topk_v2_plan=self._build_topk_v2_plan(seqlens_expanded), + ) + self.decode_cuda_graph_metadata[bs] = metadata + self.forward_metadata = metadata + + def _apply_cuda_graph_metadata( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + forward_mode: ForwardMode, + spec_info: Optional[SpecInput], + out_cache_loc: Optional[torch.Tensor] = None, + actual_forward_mode: Optional[ForwardMode] = None, + ): + """Shared capture+replay body for the cuda-graph init path. + + Public entry: :py:meth:`init_forward_metadata_out_graph`. Spec runners + also call this directly via _apply_cuda_graph_metadata when they + need to pass out_cache_loc / actual_forward_mode explicitly. + """ + if bs not in self.decode_cuda_graph_metadata: + self._build_forward_metadata_cuda_graph( + bs, + None, + req_pool_indices, + seq_lens, + seq_lens_cpu, + forward_mode, + spec_info, + out_cache_loc, + actual_forward_mode, + ) + return + + self.set_dsa_prefill_impl(forward_batch=None) + + seq_lens = seq_lens[:bs] + req_pool_indices = req_pool_indices[:bs] + + # Normal Decode + metadata: DSAMetadata = self.decode_cuda_graph_metadata[bs] + used_fused_metadata_generation = False + target_verify_ctx_lens_written = False + if forward_mode.is_decode_or_idle(): + # Normal Decode + max_len = self._graph_page_table_width(metadata) + + if is_cuda() and not _is_hip: + from sglang.kernels.ops.attention.dsa_metadata import ( + fused_dsa_decode_metadata, + ) + + fused_dsa_decode_metadata( + seq_lens=seq_lens, + req_pool_indices=req_pool_indices, + req_to_token=self.req_to_token, + cache_seqlens=metadata.cache_seqlens_int32, + cu_seqlens_k=metadata.cu_seqlens_k, + page_table_1=metadata.page_table_1, + dsa_cache_seqlens=metadata.dsa_cache_seqlens_int32, + dsa_cu_seqlens_k=metadata.dsa_cu_seqlens_k, + real_page_table=metadata.real_page_table, + bs=bs, + max_len=max_len, + dsa_index_topk=self.dsa_index_topk, + real_page_size=self.real_page_size, + ) + cache_seqlens = metadata.cache_seqlens_int32 + dsa_cache_seqlens = metadata.dsa_cache_seqlens_int32 + seqlens_expanded = cache_seqlens + page_indices = None + used_fused_metadata_generation = True + + if not used_fused_metadata_generation: + cache_seqlens = seq_lens.to(torch.int32) + metadata.cache_seqlens_int32.copy_(cache_seqlens) + metadata.cu_seqlens_k[1:].copy_( + torch.cumsum(cache_seqlens, dim=0, dtype=torch.int32) + ) + page_indices = self.req_to_token[req_pool_indices, :max_len] + metadata.page_table_1[:, :max_len].copy_(page_indices) + dsa_cache_seqlens = compute_dsa_seqlens( + cache_seqlens, dsa_index_topk=self.dsa_index_topk + ) + metadata.dsa_cache_seqlens_int32.copy_(dsa_cache_seqlens) + seqlens_expanded = cache_seqlens + elif forward_mode.is_target_verify(): + max_seqlen_k = self._graph_page_table_width(metadata) + + if is_cuda() and not _is_hip: + from sglang.kernels.ops.attention.dsa_metadata import ( + fused_dsa_target_verify_metadata, + ) + + paged_mqa_ctx_lens_2d = None + if ( + self.speculative_num_draft_tokens >= 2 + and is_sm100_supported() + and metadata.paged_mqa_ctx_lens_2d is not None + and metadata.paged_mqa_ctx_lens_2d.dim() == 2 + and metadata.paged_mqa_ctx_lens_2d.size(0) == bs + and metadata.paged_mqa_ctx_lens_2d.size(1) + == self.speculative_num_draft_tokens + ): + paged_mqa_ctx_lens_2d = metadata.paged_mqa_ctx_lens_2d + + fused_dsa_target_verify_metadata( + seq_lens=seq_lens, + req_pool_indices=req_pool_indices, + req_to_token=self.req_to_token, + cache_seqlens=metadata.cache_seqlens_int32, + cu_seqlens_k=metadata.cu_seqlens_k, + page_table_1=metadata.page_table_1, + seqlens_expanded=metadata.dsa_seqlens_expanded, + dsa_cache_seqlens=metadata.dsa_cache_seqlens_int32, + dsa_cu_seqlens_k=metadata.dsa_cu_seqlens_k, + real_page_table=metadata.real_page_table, + bs=bs, + max_seqlen_k=max_seqlen_k, + dsa_index_topk=self.dsa_index_topk, + real_page_size=self.real_page_size, + next_n=self.speculative_num_draft_tokens, + paged_mqa_ctx_lens_2d=paged_mqa_ctx_lens_2d, + ) + target_verify_ctx_lens_written = paged_mqa_ctx_lens_2d is not None + cache_seqlens = metadata.cache_seqlens_int32 + seqlens_expanded = metadata.dsa_seqlens_expanded[ + : self.speculative_num_draft_tokens * bs + ] + dsa_cache_seqlens = metadata.dsa_cache_seqlens_int32[ + : self.speculative_num_draft_tokens * bs + ] + page_indices = None + used_fused_metadata_generation = True + + if not used_fused_metadata_generation: + cache_seqlens = (seq_lens + self.speculative_num_draft_tokens).to( + torch.int32 + ) + metadata.cache_seqlens_int32.copy_(cache_seqlens) + metadata.cu_seqlens_k[1:].copy_( + torch.cumsum(cache_seqlens, dim=0, dtype=torch.int32) + ) + page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k] + page_indices = torch.repeat_interleave( + page_indices, repeats=self.speculative_num_draft_tokens, dim=0 + ) + metadata.page_table_1[:, :max_seqlen_k].copy_(page_indices) + + # Fill the constant per-req qo lengths on-device; torch.tensor(list, + # device=cuda) does a pageable H2D copy that blocks the host. + extend_seq_lens = torch.full( + (bs,), + self.speculative_num_draft_tokens, + dtype=torch.int32, + device=self.device, + ) + seqlens_expanded = seqlens_expand_triton( + extend_seq_lens, + cache_seqlens, + self.speculative_num_draft_tokens * bs, + self.speculative_num_draft_tokens, + ) + metadata.dsa_seqlens_expanded.copy_(seqlens_expanded) + dsa_cache_seqlens = compute_dsa_seqlens( + seqlens_expanded, self.dsa_index_topk + ) + metadata.dsa_cache_seqlens_int32.copy_(dsa_cache_seqlens) + elif forward_mode.is_draft_extend_v2(): + # V2 draft-extend processes the full padded tree width + # (speculative_num_draft_tokens) per req -- a static shape, like + # target-verify -- so graph replay stays host-sync-free. seq_lens + # already includes the draft KV written by prepare_for_draft_extend; + # the per-req accept length is handled downstream by output + # selection, not by reshaping the page table here. + max_seqlen_k = self._graph_page_table_width(metadata) + total_extend_len = self.speculative_num_draft_tokens * bs + + # See target-verify note: fill on-device to avoid the blocking + # pageable H2D from torch.tensor(list, device=cuda). + extend_seq_lens = torch.full( + (bs,), + self.speculative_num_draft_tokens, + dtype=torch.int32, + device=self.device, + ) + + if is_cuda() and not _is_hip: + from sglang.kernels.ops.attention.dsa_metadata import ( + fused_dsa_draft_extend_metadata, + ) + + fused_dsa_draft_extend_metadata( + seq_lens=seq_lens, + extend_seq_lens=extend_seq_lens, + req_pool_indices=req_pool_indices, + req_to_token=self.req_to_token, + cache_seqlens=metadata.cache_seqlens_int32, + cu_seqlens_k=metadata.cu_seqlens_k, + page_table_1=metadata.page_table_1, + seqlens_expanded=metadata.dsa_seqlens_expanded, + dsa_cache_seqlens=metadata.dsa_cache_seqlens_int32, + dsa_cu_seqlens_k=metadata.dsa_cu_seqlens_k, + real_page_table=metadata.real_page_table, + bs=bs, + total_len=total_extend_len, + max_seqlen_k=max_seqlen_k, + dsa_index_topk=self.dsa_index_topk, + real_page_size=self.real_page_size, + max_extend_len=self.speculative_num_draft_tokens, + max_total_len=bs * self.speculative_num_draft_tokens, + static_extend_len=True, + ) + cache_seqlens = metadata.cache_seqlens_int32 + seqlens_expanded = metadata.dsa_seqlens_expanded[:total_extend_len] + dsa_cache_seqlens = metadata.dsa_cache_seqlens_int32[:total_extend_len] + page_indices = None + used_fused_metadata_generation = True + + if not used_fused_metadata_generation: + cache_seqlens = seq_lens.to(torch.int32) + metadata.cache_seqlens_int32.copy_(cache_seqlens) + metadata.cu_seqlens_k[1:].copy_( + torch.cumsum(cache_seqlens, dim=0, dtype=torch.int32) + ) + + page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k] + page_indices = torch.repeat_interleave( + page_indices, repeats=self.speculative_num_draft_tokens, dim=0 + ) + metadata.page_table_1[:, :max_seqlen_k].copy_(page_indices) + + seqlens_expanded = seqlens_expand_triton( + extend_seq_lens, + cache_seqlens, + total_extend_len, + self.speculative_num_draft_tokens, + ) + metadata.dsa_seqlens_expanded[: seqlens_expanded.shape[0]].copy_( + seqlens_expanded + ) + dsa_cache_seqlens = compute_dsa_seqlens( + seqlens_expanded, self.dsa_index_topk + ) + metadata.dsa_cache_seqlens_int32.copy_(dsa_cache_seqlens) + + # Update DeepGEMM paged MQA schedule metadata outside the captured graph. + if is_cuda() and ( + forward_mode.is_decode_or_idle() + or forward_mode.is_target_verify() + or forward_mode.is_draft_extend_v2() + ): + if forward_mode.is_draft_extend_v2(): + schedule_seqlens_expanded = metadata.dsa_seqlens_expanded + else: + schedule_seqlens_expanded = seqlens_expanded + if target_verify_ctx_lens_written: + seqlens_32_2d = metadata.paged_mqa_ctx_lens_2d + else: + seqlens_32_2d = self._build_paged_mqa_schedule_2d_ctx_lens( + forward_mode, + metadata.cache_seqlens_int32, + schedule_seqlens_expanded, + bs, + ) + self._refresh_paged_mqa_schedule_metadata(metadata, seqlens_32_2d) + self._refresh_topk_v2_plan(metadata) + # `copy_` preserves the buffer's data_ptr that the captured graph captured. + if not target_verify_ctx_lens_written: + if metadata.paged_mqa_ctx_lens_2d is None: + object.__setattr__(metadata, "paged_mqa_ctx_lens_2d", seqlens_32_2d) + else: + metadata.paged_mqa_ctx_lens_2d.copy_(seqlens_32_2d) + seqlens_expanded_size = seqlens_expanded.shape[0] + assert ( + metadata.dsa_cache_seqlens_int32 is not None + and metadata.dsa_cu_seqlens_k is not None + and self.dsa_index_topk is not None + ) + + if not used_fused_metadata_generation: + metadata.dsa_cu_seqlens_k[1 : 1 + seqlens_expanded_size].copy_( + torch.cumsum(dsa_cache_seqlens, dim=0, dtype=torch.int32) + ) + # NOTE(dark): (dsa-) cu_seqlens_q is always arange, no need to copy + + assert self.real_page_size == metadata.page_size + if self.real_page_size > 1: + if not used_fused_metadata_generation: + real_table = self._transform_table_1_to_real(page_indices) + new_rows = real_table.shape[0] + new_cols = real_table.shape[1] + metadata.real_page_table[:new_rows, :new_cols].copy_(real_table) + else: + assert metadata.real_page_table is metadata.page_table_1 + + if self.dsa_decode_impl == "flashmla_kv": + flashmla_metadata = metadata.flashmla_metadata.slice( + slice(0, seqlens_expanded_size + 1) + ) + flashmla_metadata.copy_( + self._compute_flashmla_metadata( + cache_seqlens=dsa_cache_seqlens, + seq_len_q=1, + ) + ) + + self.forward_metadata = metadata + + def init_forward_metadata_replay_cuda_graph_from_precomputed( + self, + bs: int, + precomputed: PrecomputedMetadata, + forward_mode: ForwardMode, + ): + """Fast path: copy precomputed metadata to this backend's metadata. + + This function only performs copy operations, no computation. + + Args: + bs: Batch size + precomputed: Precomputed metadata to copy from + forward_mode: Forward mode + """ + self.set_dsa_prefill_impl(forward_batch=None) + + metadata = self.decode_cuda_graph_metadata[bs] + + # Track whether fused kernel succeeded + fused_kernel_succeeded = False + + # Use fused CUDA kernel for all copy operations + if not _is_hip: + try: + from sglang.jit_kernel.fused_metadata_copy import ( + fused_metadata_copy_cuda, + ) + + # Map forward_mode to integer enum + if forward_mode.is_decode_or_idle(): + mode_int = 0 # DECODE + elif forward_mode.is_target_verify(): + mode_int = 1 # TARGET_VERIFY + else: + raise ValueError(f"Unsupported forward_mode: {forward_mode}") + + # Prepare FlashMLA tensors if needed + flashmla_num_splits_src = None + flashmla_num_splits_dst = None + flashmla_metadata_src = None + flashmla_metadata_dst = None + if precomputed.flashmla_metadata is not None: + flashmla_num_splits_src = precomputed.flashmla_metadata.num_splits + flashmla_num_splits_dst = metadata.flashmla_metadata.num_splits + flashmla_metadata_src = ( + precomputed.flashmla_metadata.flashmla_metadata + ) + flashmla_metadata_dst = metadata.flashmla_metadata.flashmla_metadata + + # Call fused kernel + fused_metadata_copy_cuda( + # Source tensors + precomputed.cache_seqlens, + precomputed.cu_seqlens_k, + precomputed.page_indices, + precomputed.dsa_cache_seqlens, + precomputed.seqlens_expanded, + precomputed.dsa_cu_seqlens_k, + precomputed.real_page_table, + flashmla_num_splits_src, + flashmla_metadata_src, + # Destination tensors + metadata.cache_seqlens_int32, + metadata.cu_seqlens_k, + metadata.page_table_1, + metadata.dsa_cache_seqlens_int32, + metadata.dsa_seqlens_expanded, + metadata.dsa_cu_seqlens_k, + ( + metadata.real_page_table + if precomputed.real_page_table is not None + else None + ), + flashmla_num_splits_dst, + flashmla_metadata_dst, + # Parameters + mode_int, + bs, + precomputed.max_len, + precomputed.max_seqlen_k, + precomputed.seqlens_expanded_size, + ) + + # Successfully used fused kernel + fused_kernel_succeeded = True + + except ImportError: + print( + "Warning: Fused metadata copy kernel not available, falling back to individual copies." + ) + except Exception as e: + print( + f"Warning: Fused metadata copy kernel failed with error: {e}, falling back to individual copies." + ) + + # Fallback to individual copy operations if the fused kernel is unavailable + # or fails at runtime. + if not fused_kernel_succeeded: + # Copy basic seqlens + metadata.cache_seqlens_int32.copy_(precomputed.cache_seqlens) + metadata.cu_seqlens_k[1:].copy_(precomputed.cu_seqlens_k[1:]) + + # Mode-specific copy logic + if forward_mode.is_decode_or_idle(): + # Decode mode + metadata.page_table_1[:, : precomputed.max_len].copy_( + precomputed.page_indices + ) + metadata.dsa_cache_seqlens_int32.copy_(precomputed.dsa_cache_seqlens) + # seqlens_expanded is same as cache_seqlens (already copied) + + elif forward_mode.is_target_verify(): + # Target verify mode + metadata.page_table_1[:, : precomputed.max_seqlen_k].copy_( + precomputed.page_indices + ) + metadata.dsa_seqlens_expanded.copy_(precomputed.seqlens_expanded) + metadata.dsa_cache_seqlens_int32.copy_(precomputed.dsa_cache_seqlens) + + # Copy DSA cu_seqlens + size = precomputed.seqlens_expanded_size + metadata.dsa_cu_seqlens_k[1 : 1 + size].copy_( + precomputed.dsa_cu_seqlens_k[1 : 1 + size] + ) + + # Copy real page table + if precomputed.real_page_table is not None: + rows, cols = precomputed.real_page_table.shape + metadata.real_page_table[:rows, :cols].copy_( + precomputed.real_page_table + ) + + # Copy FlashMLA metadata in fallback path + if precomputed.flashmla_metadata is not None: + size = precomputed.seqlens_expanded_size + flashmla_metadata = metadata.flashmla_metadata.slice(slice(0, size + 1)) + flashmla_metadata.copy_(precomputed.flashmla_metadata) + + # Refresh DeepGEMM paged MQA schedule metadata for the actual seqlens of + # this replay (the captured graph holds stale data otherwise, which can + # deadlock the kernel when the runtime work decomposition diverges from + # the captured one). + if is_cuda(): + if forward_mode.is_decode_or_idle(): + seqlens_32_2d = _to_2d_context_lens(metadata.cache_seqlens_int32, bs) + else: + seqlens_32_2d = self._build_paged_mqa_schedule_2d_ctx_lens( + forward_mode, + metadata.cache_seqlens_int32, + metadata.dsa_seqlens_expanded, + bs, + ) + self._refresh_paged_mqa_schedule_metadata(metadata, seqlens_32_2d) + self._refresh_topk_v2_plan(metadata) + if metadata.paged_mqa_ctx_lens_2d is None: + object.__setattr__(metadata, "paged_mqa_ctx_lens_2d", seqlens_32_2d) + else: + metadata.paged_mqa_ctx_lens_2d.copy_(seqlens_32_2d) + + self.forward_metadata = metadata + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + # For multi-head latent attention + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, + topk_indices: Optional[torch.Tensor] = None, + cos_sin_cache: Optional[torch.Tensor] = None, + is_neox: Optional[bool] = False, + llama_4_scaling: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + + causal = not layer.is_cross_attention + metadata = self.forward_metadata + assert causal, "DSA is causal only" + + dsa_impl = ( + self.dsa_decode_impl + if ( + forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend_v2() + ) + else self.dsa_prefill_impl + ) + + if dsa_impl == "trtllm" and not self.use_mha: + return self._forward_trtllm( + q, + k, + v, + layer, + forward_batch, + metadata.dsa_cache_seqlens_int32, + save_kv_cache, + q_rope, + k_rope, + topk_indices, + cos_sin_cache, + is_neox, + llama_4_scaling, + is_prefill=True, + ) + + if k is not None: + assert v is not None + if save_kv_cache: + cache_loc = ( + forward_batch.out_cache_loc + if not layer.is_cross_attention + else forward_batch.encoder_out_cache_loc + ) + self.token_to_kv_pool.set_mla_kv_buffer( # type: ignore + layer, + cache_loc, + k, + k_rope, + ) + + # Use MHA kernel if in MHA_ONE_SHOT mode + if self.use_mha: + assert k is not None and v is not None + assert q_rope is None, "MHA_ONE_SHOT path should not pass q_rope" + assert ( + layer.tp_k_head_num == layer.tp_q_head_num > 1 + ), "MHA_ONE_SHOT requires dense multi-head config" + return self._forward_standard_mha( + q=q, + k=k, + v=v, + layer=layer, + forward_batch=forward_batch, + metadata=metadata, + ) + + # Do absorbed multi-latent attention (MLA path) + assert q_rope is not None + kv_cache = self.token_to_kv_pool.get_key_buffer(layer.layer_id) + + if q_rope is not None: + q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) + q_rope = q_rope.view( + -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim + ) + else: + q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) + q_nope = q_all[:, :, : layer.v_head_dim] + q_rope = q_all[:, :, layer.v_head_dim :] + + # NOTE(dark): here, we use page size = 1 + topk_transform_method = self.get_topk_transform_method( + forward_batch.forward_mode + ) + + if self.use_fused_topk: + if topk_indices is not None: + topk_indices = self._pad_topk_indices(topk_indices, q_nope.shape[0]) + + page_table_1 = self._get_fused_topk_page_table(topk_indices) + else: + if topk_transform_method == TopkTransformMethod.RAGGED: + if topk_indices is not None: + topk_indices = self._pad_topk_indices(topk_indices, q_nope.shape[0]) + topk_indices_offset = metadata.topk_indices_offset + assert topk_indices_offset is not None + mask = topk_indices != -1 + topk_indices_offset = ( + topk_indices_offset.unsqueeze(1) + if topk_indices_offset.ndim == 1 + else topk_indices_offset + ) + topk_indices = torch.where( + mask, topk_indices + topk_indices_offset, topk_indices + ) + elif topk_transform_method == TopkTransformMethod.PAGED: + assert metadata.dsa_extend_seq_lens_list is not None + page_table_1 = transform_index_page_table_prefill( + page_table=metadata.page_table_1, + topk_indices=topk_indices, + extend_lens_cpu=metadata.dsa_extend_seq_lens_list, + page_size=1, + output_num_tokens=q_nope.shape[0], + page_table_is_expanded=( + forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend_v2() + ), + cu_seqlens_q=metadata.cu_seqlens_q, + ) + + # todo hisparse: to cover more backends + if self.hisparse_coordinator is not None: + # flash_mla_sparse_fwd / tilelang require int32 page indices. + page_table_1 = self.token_to_kv_pool.translate_loc_to_hisparse_device( + page_table_1 + ).to(torch.int32) + + if dsa_impl == "tilelang": + if q_rope is not None: + # Triton prefill kernel reads q_nope/q_rope directly, skipping + # the concat (it splits q into main/tail internally anyway). + # Gated to gfx950 + the validated shape (16 heads, d_v=512, + # tail=64, topk=2048); everything else uses TileLang. + if ( + _DSA_TRITON_PREFILL + and _IS_GFX95 + and kv_cache.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz) + and layer.tp_q_head_num == 16 + and layer.v_head_dim == 512 + and (layer.head_dim - layer.v_head_dim) == 64 + and page_table_1.shape[-1] == 2048 + and q_nope.shape[0] >= 512 + ): + from sglang.kernels.ops.attention.dsa.triton_sparse_mla import ( + triton_sparse_mla_fwd, + ) + + return triton_sparse_mla_fwd( + q_nope=q_nope, + q_rope=q_rope, + kv=kv_cache, + indices=page_table_1.unsqueeze(1), + sm_scale=layer.scaling, + d_v=layer.v_head_dim, + ) + q_all = concat_mla_absorb_q_general(q_nope, q_rope) + return self._forward_tilelang( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + sm_scale=layer.scaling, + v_head_dim=layer.v_head_dim, + ) + elif dsa_impl in ("flashmla_sparse", "flashmla_sparse_q8"): + if topk_transform_method == TopkTransformMethod.RAGGED: + _has_prefix = any(forward_batch.extend_prefix_lens_cpu) + page_table_1 = topk_indices + + # `flashmla_sparse_q8` = native FP8 sparse prefill (constructor + # guarantees fp8_e4m3 KV + SM90). The helper consumes q_nope/q_rope + # directly (fusing the concat with the bf16->fp8 cast), so no bf16 + # q_all is materialized on this path. The prefix path hands over the + # paged fp8 KV as-is; the non-prefix path passes the gathered bf16 KV. + if dsa_impl == "flashmla_sparse_q8": + if _has_prefix: + page_table_1_flattened = ( + self.forward_metadata.page_table_1_flattened + ) + assert page_table_1_flattened is not None + return self._forward_flashmla_sparse_q8kv8( + q_nope=q_nope, + q_rope=q_rope, + kv_bf16=None, + paged_kv_cache=kv_cache, + page_table_1_flattened=page_table_1_flattened, + page_table_1=page_table_1, + sm_scale=layer.scaling, + v_head_dim=layer.v_head_dim, + ) + kv_cache = _cat([k, k_rope], dim=-1) + return self._forward_flashmla_sparse_q8kv8( + q_nope=q_nope, + q_rope=q_rope, + kv_bf16=kv_cache, + paged_kv_cache=None, + page_table_1_flattened=None, + page_table_1=page_table_1, + sm_scale=layer.scaling, + v_head_dim=layer.v_head_dim, + ) + + # bf16 path (dsa_impl == "flashmla_sparse"). + if _has_prefix: + page_table_1_flattened = ( + self.forward_metadata.page_table_1_flattened + ) + assert page_table_1_flattened is not None + kv_cache = dequantize_k_cache_paged( + kv_cache, page_table_1_flattened + ) + else: + kv_cache = _cat([k, k_rope], dim=-1) + + if q_rope is not None: + q_all = concat_mla_absorb_q_general(q_nope, q_rope) + return self._forward_flashmla_sparse( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + sm_scale=layer.scaling, + v_head_dim=layer.v_head_dim, + ) + elif dsa_impl == "flashmla_kv": + if q_rope is not None: + q_all = concat_mla_absorb_q_general(q_nope, q_rope) + return self._forward_flashmla_kv( + q_all=q_all, + kv_cache=kv_cache, + sm_scale=layer.scaling, + v_head_dim=layer.v_head_dim, + # TODO optimize args + layer=layer, + metadata=metadata, + page_table_1=page_table_1, + ) + elif dsa_impl == "fa3": + return self._forward_fa3( + q_rope=q_rope, + kv_cache=kv_cache, + v_head_dim=layer.v_head_dim, + q_nope=q_nope, + page_table=page_table_1, + cache_seqlens=metadata.dsa_cache_seqlens_int32, + cu_seqlens_q=metadata.dsa_cu_seqlens_q, + cu_seqlens_k=metadata.dsa_cu_seqlens_k, + max_seqlen_q=metadata.dsa_max_seqlen_q, + sm_scale=layer.scaling, + logit_cap=layer.logit_cap, + page_size=1, + ) + elif dsa_impl == "aiter": + if q_rope is not None: + q_all = torch.cat([q_nope, q_rope], dim=-1) + return self._forward_aiter_extend( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + layer=layer, + ) + else: + raise ValueError( + f"Unsupported {dsa_impl = } for forward_extend. Consider using an other attention backend." + ) + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache=True, + # For multi-head latent attention + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, + topk_indices: Optional[torch.Tensor] = None, + cos_sin_cache: Optional[torch.Tensor] = None, + is_neox: Optional[bool] = False, + llama_4_scaling: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + + causal = not layer.is_cross_attention + metadata = self.forward_metadata + assert causal, "DSA is causal only" + + if self.dsa_decode_impl == "trtllm": + return self._forward_trtllm( + q, + k, + v, + layer, + forward_batch, + metadata.cache_seqlens_int32, + save_kv_cache, + q_rope, + k_rope, + topk_indices, + cos_sin_cache, + is_neox, + llama_4_scaling, + ) + + if k is not None: + assert v is not None + if save_kv_cache: + cache_loc = ( + forward_batch.out_cache_loc + if not layer.is_cross_attention + else forward_batch.encoder_out_cache_loc + ) + self.token_to_kv_pool.set_mla_kv_buffer( # type: ignore + layer, + cache_loc, + k, + k_rope, + ) + + # Do absorbed multi-latent attention + kv_cache = self.token_to_kv_pool.get_key_buffer(layer.layer_id) + if q_rope is not None: + q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) + q_rope = q_rope.view( + -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim + ) + # Caller passed split q_nope / q_rope; we'll need to concat below if + # the chosen impl wants q_all. + q_all = None + else: + # Caller passed already-concatenated q (q_all = q). Reuse it directly + # via a zero-copy view; the impl-specific blocks below will skip the + # otherwise redundant concat_mla_absorb_q_general call. + q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) + q_nope = q_all[:, :, : layer.v_head_dim] + q_rope = q_all[:, :, layer.v_head_dim :] + + # Align topk_indices with q dimensions + if topk_indices is not None: + topk_indices = self._pad_topk_indices(topk_indices, q_nope.shape[0]) + + if self.hisparse_coordinator is not None: + page_table_1 = self.hisparse_coordinator.swap_in_selected_pages( + forward_batch.req_pool_indices, + forward_batch.seq_lens, + topk_indices, + layer.layer_id, + ) + elif self.use_fused_topk: + page_table_1 = self._get_fused_topk_page_table(topk_indices) + else: + page_table_1 = transform_index_page_table_decode( + page_table=metadata.page_table_1, + topk_indices=topk_indices, + page_size=1, + ) + + if self.dsa_decode_impl == "flashmla_sparse": + if q_rope is not None: + q_all = concat_mla_absorb_q_general(q_nope, q_rope) + return self._forward_flashmla_sparse( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + sm_scale=layer.scaling, + v_head_dim=layer.v_head_dim, + ) + elif self.dsa_decode_impl == "flashmla_kv": + if q_rope is not None: + q_all = concat_mla_absorb_q_general(q_nope, q_rope) + return self._forward_flashmla_kv( + q_all=q_all, + kv_cache=kv_cache, + sm_scale=layer.scaling, + v_head_dim=layer.v_head_dim, + # TODO optimize args + layer=layer, + metadata=metadata, + page_table_1=page_table_1, + ) + elif self.dsa_decode_impl == "tilelang": + # Cat-skip (HIP-only): when caller passes q_rope=None on HIP, q_all + # has already been set to a zero-copy view of q in the else branch + # above and we can reuse it directly. The `not _is_hip` clause keeps + # CUDA / MUSA paths byte-identical to pre-patch by always re-cat. + if q_all is None or not _is_hip: + q_all = concat_mla_absorb_q_general(q_nope, q_rope) + return self._forward_tilelang( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + sm_scale=layer.scaling, + v_head_dim=layer.v_head_dim, + ) + elif self.dsa_decode_impl == "fa3": + return self._forward_fa3( + q_rope=q_rope, + kv_cache=kv_cache, + v_head_dim=layer.v_head_dim, + q_nope=q_nope, + page_table=page_table_1, + cache_seqlens=metadata.dsa_cache_seqlens_int32, + cu_seqlens_q=metadata.dsa_cu_seqlens_q, + cu_seqlens_k=metadata.dsa_cu_seqlens_k, + max_seqlen_q=metadata.dsa_max_seqlen_q, + sm_scale=layer.scaling, + logit_cap=layer.logit_cap, + page_size=1, + ) + elif self.dsa_decode_impl == "aiter": + if q_all is None or not _is_hip: + q_all = torch.cat([q_nope, q_rope], dim=-1) + return self._forward_aiter( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + layer=layer, + metadata=metadata, + bs=forward_batch.batch_size, + ) + + else: + assert False, f"Unsupported {self.dsa_decode_impl = }" + + def _forward_fa3( + self, + q_rope: torch.Tensor, + kv_cache: torch.Tensor, + v_head_dim: int, + q_nope: torch.Tensor, + page_table: torch.Tensor, + cache_seqlens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + sm_scale: float, + logit_cap: float, + page_size: int, + ) -> torch.Tensor: + k_rope_cache = kv_cache[:, :, v_head_dim:] + c_kv_cache = kv_cache[:, :, :v_head_dim] + qk_rope_dim = k_rope_cache.shape[-1] + k_rope_cache = k_rope_cache.view(-1, page_size, 1, qk_rope_dim) + c_kv_cache = c_kv_cache.view(-1, page_size, 1, v_head_dim) + o = flash_attn_with_kvcache( + q=q_rope, + k_cache=k_rope_cache, + v_cache=c_kv_cache, + qv=q_nope, + page_table=page_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k_new=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + softmax_scale=sm_scale, + causal=True, + softcap=logit_cap, + return_softmax_lse=False, + num_splits=self.num_splits, + ) + return o # type: ignore + + def _forward_flashmla_sparse( + self, + q_all: torch.Tensor, + kv_cache: torch.Tensor, + v_head_dim: int, + page_table_1: torch.Tensor, + sm_scale: float, + ) -> torch.Tensor: + from sgl_kernel.flash_mla import flash_mla_sparse_fwd + + # FlashMLA sparse kernel requires num_heads to be a multiple of 64 (Hopper) or 128 (Blackwell) + # When using TP, num_heads might be smaller (e.g., 256//8=32) + num_tokens, num_heads, head_dim = q_all.shape + + # Determine required padding based on GPU architecture (use cached value) + required_padding = 128 if self.device_sm_major >= 10 else 64 + + need_padding = num_heads % required_padding != 0 + + if need_padding: + assert required_padding % num_heads == 0, ( + f"num_heads {num_heads} cannot be padded to {required_padding}. " + f"TP size may be too large for this model." + ) + + # Pad q to required size + q_padded = q_all.new_zeros((num_tokens, required_padding, head_dim)) + q_padded[:, :num_heads, :] = q_all + q_input = q_padded + else: + q_input = q_all + + # indices shape must be (s_q, h_kv=1, topk), keep h_kv=1 unchanged + indices_input = page_table_1.unsqueeze(1) + + o, _, _ = flash_mla_sparse_fwd( + q=q_input, + kv=kv_cache, + indices=indices_input, + sm_scale=sm_scale, + d_v=v_head_dim, + ) + + # Trim output back to original num_heads if we padded + if need_padding: + o = o[:, :num_heads, :] + + return o + + def _forward_flashmla_sparse_q8kv8( + self, + q_nope: torch.Tensor, + q_rope: torch.Tensor, + kv_bf16: Optional[torch.Tensor], + v_head_dim: int, + page_table_1: torch.Tensor, + sm_scale: float, + paged_kv_cache: Optional[torch.Tensor] = None, + page_table_1_flattened: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Native FP8 (q8 x kv8) sparse-prefill attention (SM90 JIT kernel). + + Same contract as ``_forward_flashmla_sparse`` but executed through the + FP8 ``sparse_mla_q8kv8_prefill_fwd`` kernel. Identity per-tensor + scales (scalar 1.0) are used: a raw bf16->fp8 cast of q/kv is accurate + on real DeepSeek-V3 magnitudes, so no dynamic rescaling is applied. + The kernel runs via its fixed full-topk entry (``attn_sink`` / + ``topk_length`` left None), keeping control flow identical across DP + ranks; -1 topk sentinels are clamped to distinct zero pad rows inside + the kernel. + + Two KV paths: + * non-prefix extend: ``kv_bf16`` (the gathered bf16 KV) is cast into + a zero-padded fp8 buffer. + * prefix extend: ``paged_kv_cache`` (fp8, 656 B/token: nope_fp8 + + per-group scales + rope_bf16) is gathered, dequantized per group, + and requantized to per-tensor fp8 in one fused Triton kernel + (``gather_dequant_requant_fp8_paged``) — no intermediate bf16 + materialization. + """ + from sglang.jit_kernel.sparse_mla_q8kv8_prefill_sm90 import ( + sparse_mla_q8kv8_prefill_fwd, + ) + + num_tokens, num_heads, d_nope = q_nope.shape + head_dim = d_nope + q_rope.shape[-1] + dev = q_nope.device + + # The SM90 kernel requires num_heads % 64 == 0; smaller head counts + # (high-TP splits) are zero-padded up to 64. + required_padding = 64 + need_padding = num_heads % required_padding != 0 + + # Build the fp8 q. concat_and_cast_q_fp8_pad fuses the nope/rope + # concat with the bf16->fp8 cast in one Triton kernel (bit-exact vs + # concat + .to(fp8)); it requires power-of-two head/dim counts (a + # tl.arange constraint), so non-power-of-two head counts fall back to + # the generic concat + cast. + if need_padding: + if required_padding % num_heads != 0: + raise ValueError( + f"num_heads={num_heads} cannot be padded to {required_padding}; " + "this TP size is incompatible with flashmla_sparse_q8." + ) + # Cached zero-padded fp8 q buffer: the pad rows [num_heads:64] are + # zero on first alloc and only ever read by the kernel; the active + # slice is overwritten each forward. Eager-mode DSA runs layers + # sequentially on one stream, so single-buffer reuse is safe. + # Grown on demand. + buf = self._q8kv8_qpad_buf + if buf is None or buf.shape[0] < num_tokens: + buf = torch.zeros( + (num_tokens, required_padding, head_dim), + dtype=torch.float8_e4m3fn, + device=dev, + ) + self._q8kv8_qpad_buf = buf + q_fp8 = buf[:num_tokens] + # head counts that divide 64 are powers of two, so the fused + # concat-cast is always applicable here. + concat_and_cast_q_fp8_pad(q_fp8, q_nope, q_rope, num_heads) + elif (num_heads & (num_heads - 1)) == 0: + q_fp8 = q_nope.new_empty( + (num_tokens, num_heads, head_dim), dtype=torch.float8_e4m3fn + ) + concat_and_cast_q_fp8_pad(q_fp8, q_nope, q_rope, num_heads) + else: + # Generic fallback for non-power-of-two head counts. + q_fp8 = concat_mla_absorb_q_general(q_nope, q_rope).to(torch.float8_e4m3fn) + + # Identity per-tensor scale, cached: creating it per call is a + # host->device copy that synchronizes the stream. + identity_scale = self._q8kv8_identity_scale + if identity_scale is None: + identity_scale = torch.tensor([1.0], dtype=torch.float32, device=dev) + self._q8kv8_identity_scale = identity_scale + + # KV: append `topk` trailing zero rows so the kernel's -1-sentinel + # clamp can map every padded topk slot to a DISTINCT zero row. + # Mapping many slots onto one shared row would serialize the kernel's + # KV gather; distinct zero rows are value-identical (zero KV + # contributes nothing to the softmax-weighted sum) at full speed. + topk = page_table_1.shape[-1] + if paged_kv_cache is not None: + kv_padded = gather_dequant_requant_fp8_paged( + paged_kv_cache, + page_table_1_flattened, + extra_rows=topk, + ).view(-1, 1, head_dim) + else: + kv_padded = kv_bf16.new_zeros( + (kv_bf16.shape[0] + topk, *kv_bf16.shape[1:]), + dtype=torch.float8_e4m3fn, + ) + kv_padded[: kv_bf16.shape[0]].copy_(kv_bf16) + kv_padded = kv_padded.view(-1, 1, head_dim) + + o, _, _ = sparse_mla_q8kv8_prefill_fwd( + q=q_fp8, + kv=kv_padded, + indices=page_table_1.unsqueeze(1), + sm_scale=sm_scale, + q_scale=identity_scale, + kv_scale=identity_scale, + d_v=v_head_dim, + attn_sink=None, + topk_length=None, + ) + + # Trim the output back to the original head count if we padded. + if need_padding: + o = o[:, :num_heads, :] + return o + + def _forward_flashmla_kv( + self, + q_all: torch.Tensor, + kv_cache: torch.Tensor, + v_head_dim: int, + sm_scale: float, + layer, + metadata: DSAMetadata, + page_table_1, + ) -> torch.Tensor: + from sgl_kernel.flash_mla import flash_mla_with_kvcache + + cache_seqlens = metadata.dsa_cache_seqlens_int32 + assert metadata.flashmla_metadata is not None + + # TODO the 2nd dim is seq_len_q, need to be >1 when MTP + q_all = q_all.view(-1, 1, layer.tp_q_head_num, layer.head_dim) + num_q_heads = q_all.shape[2] + target_q_heads = self.flashmla_kv_num_q_heads + if target_q_heads != num_q_heads: + # Pad q heads to match FlashMLA decode supported head-count variants. + q_input = q_all.new_zeros( + q_all.shape[0], q_all.shape[1], target_q_heads, q_all.shape[3] + ) + q_input[:, :, :num_q_heads, :] = q_all + else: + q_input = q_all + + kv_cache = kv_cache.view(-1, self.real_page_size, 1, self.kv_cache_dim) + assert self.real_page_size == 64, "only page size 64 is supported" + + if not self.dsa_kv_cache_store_fp8: + # inefficiently quantize the whole cache + kv_cache = quantize_k_cache(kv_cache) + + indices = page_table_1.unsqueeze(1) + assert ( + indices.shape[-1] == self.dsa_index_topk + ) # requirement of FlashMLA decode kernel + + o, _ = flash_mla_with_kvcache( + q=q_input, + k_cache=kv_cache, + cache_seqlens=cache_seqlens, + head_dim_v=v_head_dim, + tile_scheduler_metadata=metadata.flashmla_metadata.flashmla_metadata, + num_splits=metadata.flashmla_metadata.num_splits, + softmax_scale=sm_scale, + indices=indices, + # doc says it is not used, but if pass in None then error + block_table=torch.empty( + (q_all.shape[0], 0), dtype=torch.int32, device=q_all.device + ), + is_fp8_kvcache=True, + ) + + if target_q_heads != num_q_heads: + o = o[:, :, :num_q_heads, :] + + return o + + def _forward_standard_mha( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + metadata: DSAMetadata, + ) -> torch.Tensor: + """Standard MHA using FlashAttention varlen for MHA_ONE_SHOT mode.""" + q = q.view(-1, layer.tp_q_head_num, layer.head_dim) + k = k.view(-1, layer.tp_k_head_num, layer.head_dim) + v = v.view(-1, layer.tp_v_head_num, layer.v_head_dim) + + # MHA_ONE_SHOT: k/v include all tokens (prefix + current) + cu_seqlens_q = metadata.cu_seqlens_q + cu_seqlens_k = metadata.cu_seqlens_k + max_seqlen_k = metadata.max_seq_len_k + causal = True + + # Verify batch sizes match (length of cu_seqlens should be batch_size + 1) + assert len(cu_seqlens_q) == len(cu_seqlens_k), ( + f"batch_size mismatch: cu_seqlens_q has {len(cu_seqlens_q)-1} requests, " + f"cu_seqlens_k has {len(cu_seqlens_k)-1} requests" + ) + + # Use TRTLLm ragged attention for SM100 (Blackwell/B200) to avoid FA4 accuracy issues + if self.device_sm_major >= 10: + import flashinfer + + seq_lens = metadata.cache_seqlens_int32 + return flashinfer.prefill.trtllm_ragged_attention_deepseek( + query=q, + key=k, + value=v, + workspace_buffer=self.workspace_buffer, + seq_lens=seq_lens, + max_q_len=metadata.max_seq_len_q, + max_kv_len=max_seqlen_k, + bmm1_scale=layer.scaling, + bmm2_scale=1.0, + o_sf_scale=1.0, + batch_size=forward_batch.batch_size, + window_left=-1, + cum_seq_lens_q=cu_seqlens_q, + cum_seq_lens_kv=cu_seqlens_k, + enable_pdl=False, + is_causal=causal, + return_lse=False, + skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_PREFILL_THRESHOLD_SCALE_FACTOR.get(), + ) + + # Use FA3 for SM90 (Hopper/H200) + return flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=metadata.max_seq_len_q, + max_seqlen_k=max_seqlen_k, + softmax_scale=layer.scaling, + causal=causal, + ) + + def _forward_tilelang( + self, + q_all: torch.Tensor, + kv_cache: torch.Tensor, + v_head_dim: int, + page_table_1: torch.Tensor, + sm_scale: float, + ) -> torch.Tensor: + from sglang.kernels.ops.attention.dsa.tilelang_kernel import tilelang_sparse_fwd + + return tilelang_sparse_fwd( + q=q_all, + kv=kv_cache, + indices=page_table_1.unsqueeze(1), + sm_scale=sm_scale, + d_v=v_head_dim, + ) + + def _forward_aiter( + self, + q_all: torch.Tensor, + kv_cache: torch.Tensor, + page_table_1: torch.Tensor, + layer: RadixAttention, + metadata: DSAMetadata, + bs: int, + ) -> torch.Tensor: + q = q_all.reshape(-1, layer.tp_q_head_num * layer.head_dim) + + if layer.head_dim != layer.v_head_dim: + o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim)) + else: + o = torch.empty_like(q) + + if self.need_pad_heads: + q_kernel = q.view( + -1, layer.tp_q_head_num, layer.head_dim + ).repeat_interleave(self.head_repeat_factor, dim=1) + o_kernel = q.new_empty( + ( + q.shape[0], + layer.tp_q_head_num * self.head_repeat_factor, + layer.v_head_dim, + ) + ) + else: + q_kernel = q.view(-1, layer.tp_q_head_num, layer.head_dim) + o_kernel = o.view(-1, layer.tp_q_head_num, layer.v_head_dim) + + q_scale = None + kv_scale = None + aiter_persistent_kwargs = {} + if kv_cache.dtype == fp8_dtype: + kv_scale = torch.ones((), dtype=torch.float32, device=q_kernel.device) + + kv_indptr = self.kv_indptr + + non_minus1_mask = page_table_1 != -1 + non_minus1_counts = non_minus1_mask.sum(dim=1) + kv_indptr[1 : bs + 1] = torch.cumsum(non_minus1_counts, dim=0) + + kv_indices = self.kv_indices + get_valid_kv_indices(page_table_1, kv_indptr, kv_indices, bs) + + kv_last_page_lens = metadata.cu_seqlens_q + if kv_cache.dtype == fp8_dtype: + aiter_persistent_kwargs = self._prepare_aiter_dsa_decode_metadata( + metadata.cu_seqlens_q, + kv_indptr, + bs, + metadata.max_seq_len_q, + q_kernel.dtype, + kv_cache.dtype, + ) + kv_last_page_lens = aiter_persistent_kwargs.pop("kv_last_page_lens") + + mla_decode_fwd( + q_kernel, + kv_cache.view(-1, 1, 1, layer.head_dim), + o_kernel, + metadata.cu_seqlens_q, + kv_indptr, + kv_indices, + kv_last_page_lens, + metadata.max_seq_len_q, + sm_scale=layer.scaling, + logit_cap=layer.logit_cap, + q_scale=q_scale, + kv_scale=kv_scale, + **aiter_persistent_kwargs, + ) + + if self.need_pad_heads: + o = o_kernel[:, :: self.head_repeat_factor, :] + + return o + + def _forward_aiter_extend( + self, + q_all: torch.Tensor, + kv_cache: torch.Tensor, + page_table_1: torch.Tensor, + layer: RadixAttention, + ) -> torch.Tensor: + num_tokens = q_all.shape[0] + q = q_all.reshape(-1, layer.tp_q_head_num * layer.head_dim) + + if layer.head_dim != layer.v_head_dim: + o = q.new_empty((num_tokens, layer.tp_q_head_num * layer.v_head_dim)) + else: + o = torch.empty_like(q) + + if self.need_pad_heads: + q_kernel = q.view( + -1, layer.tp_q_head_num, layer.head_dim + ).repeat_interleave(self.head_repeat_factor, dim=1) + o_kernel = q.new_empty( + ( + num_tokens, + layer.tp_q_head_num * self.head_repeat_factor, + layer.v_head_dim, + ) + ) + else: + q_kernel = q.view(-1, layer.tp_q_head_num, layer.head_dim) + o_kernel = o.view(-1, layer.tp_q_head_num, layer.v_head_dim) + + q_scale = None + kv_scale = None + aiter_persistent_kwargs = {} + if kv_cache.dtype == fp8_dtype: + kv_scale = torch.ones((), dtype=torch.float32, device=q_kernel.device) + + non_minus1_mask = page_table_1 != -1 + non_minus1_counts = non_minus1_mask.sum(dim=1) + + kv_indptr = torch.zeros(num_tokens + 1, dtype=torch.int32, device=self.device) + kv_indptr[1:] = torch.cumsum(non_minus1_counts, dim=0) + + # Allocate kv_indices with upper-bound size (num_tokens * topk) + topk = page_table_1.shape[1] + kv_indices = torch.zeros( + num_tokens * topk, dtype=torch.int32, device=self.device + ) + + # Use get_valid_kv_indices kernel to extract valid indices + get_valid_kv_indices(page_table_1, kv_indptr, kv_indices, num_tokens) + + # Build cu_seqlens_q for extend: each token is treated as seq_len_q=1 + cu_seqlens_q = torch.arange( + 0, num_tokens + 1, dtype=torch.int32, device=self.device + ) + kv_last_page_lens = cu_seqlens_q + if kv_cache.dtype == fp8_dtype: + aiter_persistent_kwargs = self._prepare_aiter_dsa_decode_metadata( + cu_seqlens_q, + kv_indptr, + num_tokens, + 1, + q_kernel.dtype, + kv_cache.dtype, + ) + kv_last_page_lens = aiter_persistent_kwargs.pop("kv_last_page_lens") + + # TODO support more forward_mode + mla_decode_fwd( + q_kernel, + kv_cache.view(-1, 1, 1, layer.head_dim), + o_kernel, + cu_seqlens_q, + kv_indptr, + kv_indices, + kv_last_page_lens, + 1, # max_seq_len_q = 1 for per-token attention + sm_scale=layer.scaling, + logit_cap=layer.logit_cap, + q_scale=q_scale, + kv_scale=kv_scale, + **aiter_persistent_kwargs, + ) + + if self.need_pad_heads: + o = o_kernel[:, :: self.head_repeat_factor, :] + + return o + + def _forward_trtllm( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: RadixAttention, + forward_batch: ForwardBatch, + seq_lens: torch.Tensor, + save_kv_cache=True, + # For multi-head latent attention + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, + topk_indices: Optional[torch.Tensor] = None, + cos_sin_cache: Optional[torch.Tensor] = None, + is_neox: Optional[bool] = False, + llama_4_scaling: Optional[torch.Tensor] = None, + is_prefill: bool = False, + ) -> torch.Tensor: + """Forward using TRT-LLM sparse MLA kernel.""" + import flashinfer.decode + + metadata = self.forward_metadata + + merge_query = q_rope is not None + if self.kv_cache_dtype == torch.float8_e4m3fn: + # For FP8 path, we quantize the query and rope parts and merge them into a single tensor + # Note: rope application in deepseek_v2.py:forward_absorb_prepare is skipped for FP8 decode path of this trtllm_mla backend + assert q_rope is not None, "For FP8 path q_rope should not be None." + assert k_rope is not None, "For FP8 path k_rope should not be None." + assert ( + cos_sin_cache is not None + ), "For FP8 path cos_sin_cache should not be None." + + rope_positions = forward_batch.positions + if dsa_use_prefill_cp(forward_batch): + rope_positions = cp_split_and_rebuild_position( + forward_batch, rope_positions + ) + + q, k, k_rope = mla_quantize_and_rope_for_fp8( + q, + q_rope, + k.squeeze(1), + k_rope.squeeze(1), + rope_positions, + cos_sin_cache, + is_neox, + self.kv_lora_rank, + self.qk_rope_head_dim, + ) + if save_kv_cache and dsa_use_prefill_cp(forward_batch): + k, k_rope = _all_gather_dsa_trtllm_fp8_kv(forward_batch, k, k_rope) + merge_query = False + + # Save KV cache if requested + if save_kv_cache: + assert ( + k is not None and k_rope is not None + ), "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None." + cache_loc = ( + forward_batch.out_cache_loc + if not layer.is_cross_attention + else forward_batch.encoder_out_cache_loc + ) + self.token_to_kv_pool.set_mla_kv_buffer(layer, cache_loc, k, k_rope) + + k_cache = self.token_to_kv_pool.get_key_buffer(layer.layer_id) + kv_cache = k_cache.view(-1, self.real_page_size, self.kv_cache_dim).unsqueeze(1) + + if merge_query: + q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim) + q_rope_reshaped = q_rope.view( + -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim + ) + q_all = concat_mla_absorb_q_general(q_nope, q_rope_reshaped) + else: + q_all = q.view(-1, layer.tp_q_head_num, layer.head_dim) + + if self.use_fused_topk: + if topk_indices is not None: + topk_indices = self._pad_topk_indices(topk_indices, q.shape[0]) + page_table_1 = self._get_fused_topk_page_table(topk_indices) + elif is_prefill: + page_table_1 = transform_index_page_table_prefill( + page_table=metadata.page_table_1, + topk_indices=topk_indices, + extend_lens_cpu=metadata.dsa_extend_seq_lens_list, + page_size=1, + output_num_tokens=q.shape[0], + page_table_is_expanded=( + forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend_v2() + ), + cu_seqlens_q=metadata.cu_seqlens_q, + ) + else: + if topk_indices is not None: + topk_indices = self._pad_topk_indices(topk_indices, q.shape[0]) + page_table_1 = transform_index_page_table_decode( + page_table=metadata.page_table_1, + topk_indices=topk_indices, + page_size=1, + ) + + q_scale = 1.0 + k_scale = ( + layer.k_scale_float + if getattr(layer, "k_scale_float", None) is not None + else 1.0 + ) + bmm1_scale = q_scale * k_scale * layer.scaling + + batch_size = page_table_1.shape[0] + _, num_heads, head_dim = q_all.shape + + q = q_all.view(batch_size, 1, num_heads, head_dim) + kv = kv_cache.view(-1, 1, self.real_page_size, self.kv_cache_dim) + block_tables = page_table_1.unsqueeze(1) + seq_lens = metadata.cache_seqlens_int32 if seq_lens is None else seq_lens + + if ( + dsa_use_prefill_cp(forward_batch) + and is_dsa_prefill_cp_in_seq_split() + and forward_batch.attn_cp_metadata is not None + ): + cp_meta = forward_batch.attn_cp_metadata + seq_chunks = list(torch.split(seq_lens, cp_meta.split_list, dim=0)) + seq_lens = torch.cat([seq_chunks[i] for i in cp_meta.zigzag_index], dim=0) + + out = flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla( + query=q, + kv_cache=kv, + workspace_buffer=self.workspace_buffer, + qk_nope_head_dim=self.qk_nope_head_dim, + kv_lora_rank=self.kv_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + block_tables=block_tables, + seq_lens=seq_lens, + max_seq_len=metadata.max_seq_len_k, + sparse_mla_top_k=self.dsa_index_topk, + bmm1_scale=bmm1_scale, + backend="trtllm-gen", + skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(), + ) + + return out + + def _pad_topk_indices( + self, topk_indices: torch.Tensor, num_tokens: int + ) -> torch.Tensor: + current_tokens = topk_indices.shape[0] + if current_tokens == num_tokens: + return topk_indices + + assert current_tokens <= num_tokens, ( + f"topk_indices rows ({current_tokens}) > num_tokens ({num_tokens}); " + "this indicates a mismatch between indexer output and q layout." + ) + + pad_size = num_tokens - current_tokens + padding = torch.full( + (pad_size, topk_indices.shape[1]), + -1, + dtype=topk_indices.dtype, + device=topk_indices.device, + ) + return torch.cat([topk_indices, padding], dim=0) + + def get_cuda_graph_seq_len_fill_value(self): + """Get the fill value for sequence length in CUDA graph.""" + return 1 + + def set_dsa_prefill_impl(self, forward_batch: Optional[ForwardBatch] = None): + """ + Decide all attention prefill dispatch strategies for this batch. + """ + from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( + is_in_breakable_cuda_graph, + ) + from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( + is_in_tc_piecewise_cuda_graph, + ) + from sglang.srt.utils import get_device_sm, is_blackwell + + # Decide MHA vs MLA + if is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph(): + # Can't branch on seq_lens_cpu in graph replay, force MHA off to + # guarantee correctness. + self.use_mha = False + elif ( + forward_batch and forward_batch.forward_mode.is_extend_without_speculative() + ): + # Check if sequence meets criteria for MHA_ONE_SHOT + assert forward_batch.seq_lens_cpu is not None + max_kv_len = forward_batch.seq_lens_cpu.max().item() + sum_seq_lens = sum(forward_batch.seq_lens_cpu) + device_sm = get_device_sm() + + # Requirements: H200/B200, short sequences, supported dtype, fits in chunk + self.use_mha = ( + ( + device_sm == 90 or (device_sm >= 100 and device_sm < 110) + ) # SM90/SM100 only + and max_kv_len + <= envs.SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD.get() # Short enough for MHA + and self.token_to_kv_pool.dtype in [torch.bfloat16, torch.float8_e4m3fn] + and sum_seq_lens + <= forward_batch.get_max_chunk_capacity() # Fits in chunk + and (not is_dsa_enable_prefill_cp()) # CP not enabled + and (self.hisparse_coordinator is None) + ) + else: + self.use_mha = False # Decode/verify always use MLA + + # Set MLA implementation only if not using MHA + if not self.use_mha and self.enable_auto_select_prefill_impl: + if self.dsa_kv_cache_store_fp8: + if ( + is_blackwell() + and forward_batch is not None + and forward_batch.forward_mode == ForwardMode.EXTEND + ): + total_kv_tokens = forward_batch.seq_lens_sum + total_q_tokens = forward_batch.extend_num_tokens + # Heuristic based on benchmarking flashmla_kv vs flashmla_sparse + dequantize_k_cache_paged + if total_kv_tokens < total_q_tokens * 512: + self.dsa_prefill_impl = "flashmla_sparse" + return + self.dsa_prefill_impl = "flashmla_kv" + else: + # bf16 kv cache + self.dsa_prefill_impl = "flashmla_sparse" + + def get_topk_transform_method( + self, forward_mode: Optional[ForwardMode] = None + ) -> TopkTransformMethod: + """ + SGLANG_DSA_FUSE_TOPK controls whether to fuse the topk transform into the topk kernel. + This method is used to select the topk transform method which can be fused or unfused. + """ + if ( + # disable for MTP + self.dsa_kv_cache_store_fp8 + # flashmla_sparse_q8 shares flashmla_sparse's RAGGED prefill routing — the q8 + # dispatch lives inside the RAGGED branch of forward_extend; without this the + # transform is PAGED, the q8 path is skipped, and the bf16 kernel crashes on + # fp8 KV ("kv must have dtype kBFloat16"). + and self.dsa_prefill_impl in ("flashmla_sparse", "flashmla_sparse_q8") + and forward_mode == ForwardMode.EXTEND + ): + topk_transform_method = TopkTransformMethod.RAGGED + else: + topk_transform_method = TopkTransformMethod.PAGED + return topk_transform_method + + def get_indexer_metadata( + self, layer_id: int, forward_batch: ForwardBatch + ) -> DSAIndexerMetadata: + force_unfused = not self.use_fused_topk or ( + self.hisparse_coordinator is not None + and forward_batch.forward_mode.is_decode_or_idle() + ) + return DSAIndexerMetadata( + attn_metadata=self.forward_metadata, + topk_transform_method=self.get_topk_transform_method( + forward_batch.forward_mode + ), + topk_backend=self.dsa_topk_backend, + paged_mqa_schedule_metadata=self.forward_metadata.paged_mqa_schedule_metadata, + paged_mqa_ctx_lens_2d=self.forward_metadata.paged_mqa_ctx_lens_2d, + force_unfused_topk=force_unfused, + ) + + def _compute_flashmla_metadata(self, cache_seqlens: torch.Tensor, seq_len_q: int): + from sgl_kernel.flash_mla import get_mla_metadata + + num_heads_q = self.flashmla_kv_num_q_heads + + flashmla_metadata, num_splits = get_mla_metadata( + cache_seqlens=cache_seqlens, + # TODO doc says `num_q_tokens_per_q_seq * num_heads_q // num_heads_k` + # but the name looks like need seq_len_q? + num_q_tokens_per_head_k=seq_len_q * num_heads_q // 1, + num_heads_k=1, + num_heads_q=num_heads_q, + is_fp8_kvcache=True, + topk=self.dsa_index_topk, + ) + + return DSAFlashMLAMetadata( + flashmla_metadata=flashmla_metadata, + num_splits=num_splits, + ) + + +class DeepseekSparseAttnMultiStepBackend: + + # Per-step draft decode replays from precomputed GPU metadata; opt out so + # decide_needs_cpu_seq_lens' OR over the backends stays False. + needs_cpu_seq_lens: bool = False + + def __init__( + self, + model_runner: ModelRunner, + topk: int, + speculative_num_steps: int, + seed_dsa_topk_from_draft_extend: bool = False, + ): + self.topk = topk + self.speculative_num_steps = speculative_num_steps + self.attn_backends = [] + for i in range(self.speculative_num_steps - 1): + self.attn_backends.append( + DeepseekSparseAttnBackend( + model_runner, + speculative_step_id=i, + topk=self.topk, + speculative_num_steps=self.speculative_num_steps, + seed_dsa_topk_from_draft_extend=seed_dsa_topk_from_draft_extend, + ) + ) + + def init_forward_metadata(self, forward_batch: ForwardBatch): + for i in range(self.speculative_num_steps - 1): + self.attn_backends[i].init_forward_metadata(forward_batch) + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + for i in range(self.speculative_num_steps - 1): + self.attn_backends[i].init_cuda_graph_state(max_bs, max_num_tokens) + + def init_forward_metadata_out_graph( + self, + forward_batch: ForwardBatch, + in_capture: bool = False, + ): + from sglang.srt.model_executor.forward_batch_info import build_inner_fb_view + + if in_capture: + inner_fb = build_inner_fb_view( + forward_batch, + bs=forward_batch.batch_size, + forward_mode=ForwardMode.DECODE, + ) + for i in range(self.speculative_num_steps - 1): + self.attn_backends[i].init_forward_metadata_out_graph( + inner_fb, in_capture=True + ) + return + + bs = forward_batch.batch_size + # Precompute metadata once (shared across all backends) + precomputed = self.attn_backends[0]._precompute_replay_metadata( + bs=bs, + req_pool_indices=forward_batch.req_pool_indices, + seq_lens=forward_batch.seq_lens, + seq_lens_cpu=forward_batch.seq_lens_cpu, + forward_mode=ForwardMode.DECODE, + ) + + # Use multi-backend fused copy when we have 3 or more backends + # This is 3x faster than calling the single-backend copy 3 times + if self.speculative_num_steps > 3: + try: + from sglang.jit_kernel.fused_metadata_copy import ( + fused_metadata_copy_multi_cuda, + ) + + metadata0 = self.attn_backends[0].decode_cuda_graph_metadata[bs] + metadata1 = self.attn_backends[1].decode_cuda_graph_metadata[bs] + metadata2 = self.attn_backends[2].decode_cuda_graph_metadata[bs] + + # Set dsa_prefill_impl for first 3 backends (required by the method) + for i in range(3): + self.attn_backends[i].set_dsa_prefill_impl(forward_batch=None) + + # Prepare FlashMLA tensors if needed + flashmla_num_splits_src = None + flashmla_metadata_src = None + flashmla_num_splits_dst0 = None + flashmla_num_splits_dst1 = None + flashmla_num_splits_dst2 = None + flashmla_metadata_dst0 = None + flashmla_metadata_dst1 = None + flashmla_metadata_dst2 = None + + if precomputed.flashmla_metadata is not None: + flashmla_num_splits_src = precomputed.flashmla_metadata.num_splits + flashmla_metadata_src = ( + precomputed.flashmla_metadata.flashmla_metadata + ) + flashmla_num_splits_dst0 = metadata0.flashmla_metadata.num_splits + flashmla_num_splits_dst1 = metadata1.flashmla_metadata.num_splits + flashmla_num_splits_dst2 = metadata2.flashmla_metadata.num_splits + flashmla_metadata_dst0 = ( + metadata0.flashmla_metadata.flashmla_metadata + ) + flashmla_metadata_dst1 = ( + metadata1.flashmla_metadata.flashmla_metadata + ) + flashmla_metadata_dst2 = ( + metadata2.flashmla_metadata.flashmla_metadata + ) + + # Call the multi-backend fused kernel for first 3 backends + fused_metadata_copy_multi_cuda( + # Source tensors + precomputed.cache_seqlens, + precomputed.cu_seqlens_k, + precomputed.page_indices, + precomputed.dsa_cache_seqlens, + precomputed.dsa_cu_seqlens_k, + precomputed.real_page_table, + flashmla_num_splits_src, + flashmla_metadata_src, + # Destination tensors for backend 0 + metadata0.cache_seqlens_int32, + metadata0.cu_seqlens_k, + metadata0.page_table_1, + metadata0.dsa_cache_seqlens_int32, + metadata0.dsa_cu_seqlens_k, + ( + metadata0.real_page_table + if precomputed.real_page_table is not None + else None + ), + flashmla_num_splits_dst0, + flashmla_metadata_dst0, + # Destination tensors for backend 1 + metadata1.cache_seqlens_int32, + metadata1.cu_seqlens_k, + metadata1.page_table_1, + metadata1.dsa_cache_seqlens_int32, + metadata1.dsa_cu_seqlens_k, + ( + metadata1.real_page_table + if precomputed.real_page_table is not None + else None + ), + flashmla_num_splits_dst1, + flashmla_metadata_dst1, + # Destination tensors for backend 2 + metadata2.cache_seqlens_int32, + metadata2.cu_seqlens_k, + metadata2.page_table_1, + metadata2.dsa_cache_seqlens_int32, + metadata2.dsa_cu_seqlens_k, + ( + metadata2.real_page_table + if precomputed.real_page_table is not None + else None + ), + flashmla_num_splits_dst2, + flashmla_metadata_dst2, + # Parameters + bs, + precomputed.max_len, + precomputed.seqlens_expanded_size, + ) + + # Copy remaining backends one by one (if > 3 backends) + for i in range(3, self.speculative_num_steps - 1): + self.attn_backends[ + i + ].init_forward_metadata_replay_cuda_graph_from_precomputed( + bs=bs, + precomputed=precomputed, + forward_mode=ForwardMode.DECODE, + ) + except (ImportError, Exception) as e: + # Fallback to loop if multi-backend kernel not available or fails + if isinstance(e, ImportError): + print( + "Warning: Multi-backend fused metadata copy kernel not available, falling back to loop." + ) + else: + print( + f"Warning: Multi-backend fused metadata copy kernel failed with error: {e}, falling back to loop." + ) + for i in range(self.speculative_num_steps - 1): + self.attn_backends[ + i + ].init_forward_metadata_replay_cuda_graph_from_precomputed( + bs=bs, + precomputed=precomputed, + forward_mode=ForwardMode.DECODE, + ) + else: + # Less than 3 backends: copy to each backend individually + for i in range(self.speculative_num_steps - 1): + self.attn_backends[ + i + ].init_forward_metadata_replay_cuda_graph_from_precomputed( + bs=bs, + precomputed=precomputed, + forward_mode=ForwardMode.DECODE, + ) + + def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch) -> None: + for i in range(self.speculative_num_steps - 1): + self.attn_backends[i].init_forward_metadata_in_graph(forward_batch) + + +# Backward-compat aliases (deprecated: use DSA class names) +DeepseekSparseAttnBackend = DeepseekSparseAttnBackend +DeepseekSparseAttnMultiStepBackend = DeepseekSparseAttnMultiStepBackend +DSAMetadata = DSAMetadata +DSAFlashMLAMetadata = DSAFlashMLAMetadata +DSAIndexerMetadata = DSAIndexerMetadata diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/dsa_indexer.py b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/dsa_indexer.py new file mode 100644 index 0000000..71186cf --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/dsa_indexer.py @@ -0,0 +1,2540 @@ +from __future__ import annotations + +import contextlib +import logging +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union + +import torch +from einops import rearrange + +from sglang.jit_kernel.fused_store_index_cache import ( + can_use_dsa_fused_store, + fused_store_index_k_cache, +) +from sglang.kernels.ops.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz +from sglang.srt.compilation.compilation_config import register_split_op +from sglang.srt.environ import envs +from sglang.srt.layers.attention.dsa.paged_mqa_logits_backend import ( + DSAPagedMQALogitsBackend, +) +from sglang.srt.layers.attention.dsa.utils import ( + aiter_can_use_preshuffle_paged_mqa, + is_dsa_enable_prefill_cp, + is_dsa_prefill_cp_in_seq_split, + is_graph_dsa_split_op_surface, +) +from sglang.srt.layers.dp_attention import attn_tp_all_gather_into_tensor +from sglang.srt.layers.layernorm import LayerNorm, RMSNorm +from sglang.srt.layers.utils import MultiPlatformOp +from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( + eager_on_graph, +) +from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( + is_in_breakable_cuda_graph, +) +from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( + get_tc_piecewise_forward_context, + is_in_tc_piecewise_cuda_graph, +) +from sglang.srt.runtime_context import get_parallel, get_server_args +from sglang.srt.state_capturer.indexer_topk import ( + maybe_capture_indexer_topk, +) +from sglang.srt.utils import ( + add_prefix, + ceil_align, + get_bool_env_var, + is_cuda, + is_gfx95_supported, + is_hip, + is_npu, + is_sm120_supported, + is_xpu, +) +from sglang.srt.utils.custom_op import register_custom_op + +logger = logging.getLogger(__name__) + +global _use_multi_stream +_is_cuda = is_cuda() +_is_hip = is_hip() +_is_npu = is_npu() +_is_xpu = is_xpu() + +if not _is_npu: + from sglang.jit_kernel.dsa import ( + aiter_paged_mqa_logits, + cutedsl_paged_mqa_logits, + deepgemm_paged_mqa_logits_native, + deepgemm_paged_mqa_logits_split, + ) +else: + aiter_paged_mqa_logits = None + cutedsl_paged_mqa_logits = None + deepgemm_paged_mqa_logits_native = None + deepgemm_paged_mqa_logits_split = None + +if _is_cuda: + from sglang.jit_kernel.dsa import pick_dsl_expand +else: + pick_dsl_expand = None + +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +_is_fp8_fnuz = is_fp8_fnuz() +_is_gfx95_supported = is_gfx95_supported() +# Whether the aiter preshuffle paged-MQA path (page_size=64 + Preshuffle=True + +# KVBlockSize=64) can be used. Falls back to the legacy page_size=1 / KVBlockSize=1 +# path when the gluon kernel is unavailable (Triton<3.5 and no AOT bundle). +_use_aiter_preshuffle = aiter_can_use_preshuffle_paged_mqa() +if _use_aiter and not _use_aiter_preshuffle: + logger.warning( + "ROCm DSA indexer: aiter preshuffle paged-MQA path is unavailable " + "(needs Triton>=3.5.0 or AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS=1); " + "falling back to legacy page_size=1 / KVBlockSize=1 path." + ) +if _is_cuda: + try: + import deep_gemm + except ImportError as e: + deep_gemm = e + +if _use_aiter: + from aiter.ops.cache import indexer_k_quant_and_cache + +if is_npu(): + import torch_npu + from sglang.srt.hardware_backend.npu.utils import get_indexer_weight_stream + +from sglang.srt.distributed import ( + get_attn_tp_group, +) +from sglang.srt.distributed.parallel_state import get_pp_group +from sglang.srt.layers import deep_gemm_wrapper +from sglang.srt.layers.communicator import ScatterMode +from sglang.srt.layers.linear import ReplicatedLinear +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.rotary_embedding import get_rope_wrapper +from sglang.srt.layers.utils.cp_utils import cp_all_gather_rerange_output +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_executor.forward_context import ( + get_attn_backend, + get_req_to_token_pool, + get_token_to_kv_pool, +) +from sglang.srt.model_executor.runner import get_is_capture_mode +from sglang.srt.runtime_context import get_server_args + +_use_ag_after_qlora = envs.SGLANG_USE_AG_AFTER_QLORA.get() +if TYPE_CHECKING: + from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool + + +DUAL_STREAM_TOKEN_THRESHOLD = 1024 if _is_cuda else 0 +GRAPH_WEIGHTS_PROJ_LORA_ERROR = ( + "DSA indexer weights_proj LoRA is incompatible with " + "piecewise/breakable CUDA graph; remove the explicit " + "prefill cuda-graph backend override or drop " + "indexer.weights_proj from the LoRA target modules." +) + + +def _is_in_piecewise_or_breakable_cuda_graph() -> bool: + return is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph() + + +def _uses_dsa_attention_backend(forward_batch: ForwardBatch) -> bool: + attn_backend = get_attn_backend() + server_args = get_server_args() + prefill_backend, decode_backend = server_args.get_attention_backends() + prefill_backend = ( + getattr(attn_backend, "prefill_attention_backend_str", None) or prefill_backend + ) + decode_backend = ( + getattr(attn_backend, "decode_attention_backend_str", None) or decode_backend + ) + + if forward_batch.forward_mode.is_decode_or_idle(): + backend_name = decode_backend + elif ( + forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend_v2() + ): + backend_name = ( + decode_backend + if server_args.speculative_attention_mode == "decode" + else prefill_backend + ) + else: + backend_name = prefill_backend + + return backend_name in ("dsa", "nsa") + + +if _is_cuda: + from sglang.jit_kernel.dsv4 import fused_q_indexer_rope_first_quant + from sglang.jit_kernel.dsv32 import ( + fused_k_indexer_norm_rope, + fused_k_indexer_norm_rope_store, + ) + + def _scale_head_gate_graph_fake_impl( + weights_raw: torch.Tensor, + n_heads_inv_sqrt: float, + softmax_scale: float, + q_scale: torch.Tensor, + ) -> torch.Tensor: + return torch.empty( + (weights_raw.shape[0], weights_raw.shape[1], q_scale.shape[-1]), + dtype=torch.float32, + device=weights_raw.device, + ) + + # In-graph (PCG/BCG) head gate for the fused path: weights_proj is folded + # into wk_weights_proj, so weights_raw is precomputed and there is no GEMM. + @register_custom_op(fake_impl=_scale_head_gate_graph_fake_impl) + def scale_head_gate_graph( + weights_raw: torch.Tensor, + n_heads_inv_sqrt: float, + softmax_scale: float, + q_scale: torch.Tensor, + ) -> torch.Tensor: + weights = weights_raw * n_heads_inv_sqrt + return weights.unsqueeze(-1) * q_scale * softmax_scale + + def _logits_head_gate_graph_fake_impl( + x: torch.Tensor, + weight: torch.Tensor, + n_heads_inv_sqrt: float, + softmax_scale: float, + q_scale: torch.Tensor, + ) -> torch.Tensor: + return torch.empty( + (x.shape[0], weight.shape[0], q_scale.shape[-1]), + dtype=torch.float32, + device=x.device, + ) + + # In-graph (PCG/BCG) head gate for the NON-prefill path + @register_custom_op(fake_impl=_logits_head_gate_graph_fake_impl) + def logits_head_gate_graph( + x: torch.Tensor, + weight: torch.Tensor, + n_heads_inv_sqrt: float, + softmax_scale: float, + q_scale: torch.Tensor, + ) -> torch.Tensor: + out = torch.mm(x, weight.t(), out_dtype=torch.float32) + weights = out * n_heads_inv_sqrt + weights = weights.unsqueeze(-1) * q_scale * softmax_scale + return weights + + @register_custom_op(mutates_args=["topk_indices"]) + @register_split_op() + def broadcast_indexer_topk_from_rank0_(topk_indices: torch.Tensor) -> None: + _broadcast_indexer_topk_from_rank0_impl(topk_indices) + + +def _broadcast_indexer_topk_from_rank0_impl(topk_indices: torch.Tensor) -> None: + group = get_attn_tp_group() + if group.world_size == 1: + return + + if topk_indices.device.type == "cuda" and torch.cuda.is_current_stream_capturing(): + if group.pynccl_comm is None: + raise RuntimeError( + "SGLANG_DSA_TOPK_BROADCAST requires PyNCCL during CUDA graph capture." + ) + with group.pynccl_comm.change_state(enable=True): + group.pynccl_comm.broadcast(topk_indices, src=0) + else: + group.broadcast(topk_indices, src=0) + + +def _broadcast_indexer_topk_from_rank0( + topk_indices: Optional[torch.Tensor], +) -> Optional[torch.Tensor]: + # Sync only the finalized indexer output. Internal topk_transform calls can + # be chunked differently across ranks, which would make collectives diverge. + if topk_indices is None or not envs.SGLANG_DSA_TOPK_BROADCAST.get(): + return topk_indices + + if is_in_tc_piecewise_cuda_graph(): + broadcast_indexer_topk_from_rank0_(topk_indices) + else: + _broadcast_indexer_topk_from_rank0_impl(topk_indices) + return topk_indices + + +class BaseIndexerMetadata(ABC): + @abstractmethod + def get_seqlens_int32(self) -> torch.Tensor: + """ + Return: (batch_size,) int32 tensor + """ + + @abstractmethod + def get_page_table_64(self) -> torch.Tensor: + """ + Return: (batch_size, num_blocks) int32, page table. + The page size of the table is 64. + """ + + @abstractmethod + def get_page_table_1(self) -> torch.Tensor: + """ + Return: (batch_size, num_blocks) int32, page table. + The page size of the table is 1. + """ + + @abstractmethod + def get_seqlens_expanded(self) -> torch.Tensor: + """ + Return: (sum_extend_seq_len,) int32 tensor + """ + + def get_indexer_kvcache_range(self) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Return: (tokens, ), (tokens, ) int32, k_start and k_end in kv cache(token,xxx) for each token. + """ + + def get_indexer_seq_len_cpu(self) -> torch.Tensor: + """ + Return: seq lens for each batch. + """ + + def get_indexer_seq_len(self) -> torch.Tensor: + """ + Return: seq lens for each batch. + """ + + def get_dsa_extend_len_cpu(self) -> List[int]: + """ + Return: extend seq lens for each batch. + """ + + def get_token_to_batch_idx(self) -> torch.Tensor: + """ + Return: batch idx for each token. + """ + + @abstractmethod + def topk_transform( + self, + logits: torch.Tensor, + topk: int, + ) -> torch.Tensor: + """ + Perform topk selection on the logits and possibly transform the result. + + NOTE that attention backend may override this function to do some + transformation, which means the result of this topk_transform may not + be the topk indices of the input logits. + + Return: Anything, since it will be passed to the attention backend + for further processing on sparse attention computation. + Don't assume it is the topk indices of the input logits. + """ + + +def rotate_activation(x: torch.Tensor) -> torch.Tensor: + # from sgl_kernel import hadamard_transform + if _is_hip: + from fast_hadamard_transform import hadamard_transform + elif _is_xpu: + from sgl_kernel import hadamard_transform + else: + from sglang.jit_kernel.hadamard import hadamard_transform + + hidden_size = x.size(-1) + assert ( + hidden_size & (hidden_size - 1) + ) == 0, "Hidden size must be a power of 2 for Hadamard transform." + return hadamard_transform(x, scale=hidden_size**-0.5) + + +class Indexer(MultiPlatformOp): + _MQA_LOGITS_BYTES_PER_ELEM = 4 + _MQA_LOGITS_STATIC_SKIP_ELEMS = 8_000_000 + _MQA_LOGITS_TOTAL_MEM_FRACTION = 0.3 + _mqa_logits_budget_bytes: Dict[int, int] = {} + + @staticmethod + def _mqa_logits_free_mem_fraction() -> float: + return envs.SGLANG_DSA_MQA_LOGITS_FREE_MEM_FRACTION.get() + + def __init__( + self, + hidden_size: int, + index_n_heads: int, + index_head_dim: int, + rope_head_dim: int, + index_topk: int, + q_lora_rank: int, + max_position_embeddings: int, + rope_theta: float, + layer_id: int, + scale_fmt: Optional[str], + block_size: int = 128, + rope_scaling: Optional[Dict[str, Any]] = None, + is_neox_style: bool = True, + prefix: str = "", + quant_config: Optional[QuantizationConfig] = None, + alt_stream: Optional[torch.cuda.Stream] = None, + config=None, + ): + super().__init__() + self.hidden_size = hidden_size + self.n_heads = index_n_heads + self.head_dim = index_head_dim + self.rope_head_dim = rope_head_dim + self.index_topk = index_topk + self.q_lora_rank = q_lora_rank + self.layer_id = layer_id + self.use_dsa_indexer_fusion = ( + _is_cuda + and not envs.SGLANG_DISABLE_DSA_INDEXER_FUSION.get() + and not is_neox_style + ) + self.alt_stream = alt_stream + self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() + if self.dsa_enable_prefill_cp: + self.cp_size = get_parallel().attn_cp_size + self.cp_rank = get_parallel().attn_cp_rank + else: + self.cp_size = None + self.cp_rank = None + if _is_cuda: + self.sm_count = deep_gemm.get_num_sms() + self.half_device_sm_count = ceil_align(self.sm_count // 2, 8) + pp_size = get_server_args().pp_size + self.logits_with_pp_recv = pp_size > 1 and not get_pp_group().is_last_rank + else: + self.logits_with_pp_recv = False + + self.wq_b = ReplicatedLinear( + self.q_lora_rank, + self.n_heads * self.head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("wq_b", prefix), + ) + + if self.use_dsa_indexer_fusion: + self.wk_weights_proj = ReplicatedLinear( + self.hidden_size, + self.head_dim + self.n_heads, + bias=False, + params_dtype=torch.bfloat16, + prefix=add_prefix("wk_weights_proj", prefix), + ) + else: + self.wk = ReplicatedLinear( + self.hidden_size, + self.head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("wk", prefix), + ) + self.weights_proj = ReplicatedLinear( + self.hidden_size, + self.n_heads, + bias=False, + params_dtype=torch.bfloat16, + prefix=add_prefix("weights_proj", prefix), + ) + if ( + config is not None + and getattr(config, "index_k_norm_type", "layer") == "rms" + ): + self.k_norm = RMSNorm(self.head_dim) + else: + self.k_norm = LayerNorm( + self.head_dim, dtype=torch.bfloat16 if _use_aiter else torch.float32 + ) + self.rotary_emb = get_rope_wrapper( + rope_head_dim, + rotary_dim=rope_head_dim, + max_position=max_position_embeddings, + base=rope_theta, # type: ignore + rope_scaling=rope_scaling, + is_neox_style=is_neox_style, + device=get_server_args().device, + ) + self.block_size = block_size + self.scale_fmt = scale_fmt + self.softmax_scale = self.head_dim**-0.5 + self.num_init_tokens = self.num_local_tokens = 0 + if config is not None: + self.num_init_tokens = getattr(config, "index_init_tokens", 0) + self.num_local_tokens = getattr(config, "index_local_tokens", 0) + + self.paged_mqa_logits_backend = DSAPagedMQALogitsBackend.resolve( + get_server_args().dsa_paged_mqa_logits_backend + ) + + @contextlib.contextmanager + def _with_real_sm_count(self): + # When pipeline parallelism is enabled, each PP rank initiates a recv operation after the _pp_launch_batch + # request to receive the PP proxy tensor or output from the previous stage, occupying one SM resource. + # Model execution runs in parallel with the recv operation, so the SMs available to the indexer must be reduced + # by 1. Currently, the last rank starts the send result + recv request only after waiting for execution results. + if self.logits_with_pp_recv: + pp_recv_sm_count = 1 + with deep_gemm_wrapper.configure_deep_gemm_num_sms( + self.sm_count - pp_recv_sm_count + ): + yield + else: + yield + + @property + def _indexer_cos_sin_cache(self) -> torch.Tensor: + return self.rotary_emb.cos_sin_cache + + def _weights_proj_bf16_in_fp32_out( + self, x: Union[torch.Tensor, Tuple[torch.Tensor, ...]] + ) -> torch.Tensor: + # aiter (ROCm gfx95): extract the passthrough bf16 tensor from the + # 3-tuple (fp8, scale, bf16) produced by fused_rms_fp8_group_quant, + # avoiding an expensive FP8-to-bf16 dequantization. + if _use_aiter and _is_gfx95_supported and isinstance(x, tuple) and len(x) == 3: + x = x[2] + if _is_cuda: + return torch.mm(x, self.weights_proj.weight.t(), out_dtype=torch.float32) + + weights, _ = self.weights_proj(x) + if _is_hip: + # Return bf16; multiplying with q_scale promotes back to fp32. + return weights + return weights.float() + + @torch.compile(dynamic=True) + def _project_and_scale_head_gates( + self, x: Union[torch.Tensor, Tuple[torch.Tensor, ...]] + ): + weights = self._weights_proj_bf16_in_fp32_out(x) + weights = weights * self.n_heads**-0.5 + return weights + + @torch.compile(dynamic=True) + def _get_logits_head_gate( + self, x: Union[torch.Tensor, Tuple[torch.Tensor, ...]], q_scale: torch.Tensor + ): + weights = self._weights_proj_bf16_in_fp32_out(x) + weights = weights * self.n_heads**-0.5 + weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale + return weights + + @torch.compile(dynamic=True) + def _apply_q_scale_and_softmax_scale( + self, weights: torch.Tensor, q_scale: torch.Tensor + ): + return weights.unsqueeze(-1) * q_scale * self.softmax_scale + + @torch.compile(dynamic=True) + def _scale_head_gates(self, weights_raw: torch.Tensor, q_scale: torch.Tensor): + weights = weights_raw * self.n_heads**-0.5 + return weights.unsqueeze(-1) * q_scale * self.softmax_scale + + def _fused_k_weights(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + kw, _ = self.wk_weights_proj(x) + return kw.split([self.head_dim, self.n_heads], dim=-1) + + def _maybe_rotate(self, x: torch.Tensor) -> torch.Tensor: + # Fusion drops the (logit-preserving) Hadamard rotation; without it the + # index-K cache here matches the fused path that decode reads back. + return x if self.use_dsa_indexer_fusion else rotate_activation(x) + + def _should_skip_logits_computation(self, forward_batch: ForwardBatch) -> bool: + if ( + forward_batch.forward_mode.is_extend_without_speculative() + and forward_batch.seq_lens_cpu is not None + ): + max_kv_len = forward_batch.seq_lens_cpu.max().item() + return max_kv_len <= self.index_topk + return False + + def _get_q_k_bf16( + self, + q_lora: torch.Tensor, + x: torch.Tensor, + positions: torch.Tensor, + enable_dual_stream: bool, + forward_batch: ForwardBatch, + ): + weights_raw = None + if enable_dual_stream: + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + + with deep_gemm_wrapper.configure_deep_gemm_num_sms( + self.half_device_sm_count + ): + query, _ = self.wq_b(q_lora) + query = rearrange(query, "l (h d) -> l h d", d=self.head_dim) + q_rope, _ = torch.split( + query, + [self.rope_head_dim, self.head_dim - self.rope_head_dim], + dim=-1, + ) + with torch.cuda.stream(self.alt_stream): + # TODO we should also put DeepGEMM half SM here? + if self.use_dsa_indexer_fusion: + key, weights_raw = self._fused_k_weights(x) + else: + key, _ = self.wk(x) + key = self.k_norm(key) + + k_rope, _ = torch.split( + key, + [self.rope_head_dim, self.head_dim - self.rope_head_dim], + dim=-1, + ) + + current_stream.wait_stream(self.alt_stream) + else: + query, _ = self.wq_b(q_lora) + query = rearrange(query, "l (h d) -> l h d", d=self.head_dim) + q_rope, _ = torch.split( + query, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 + ) + if self.use_dsa_indexer_fusion: + key, weights_raw = self._fused_k_weights(x) + else: + key, _ = self.wk(x) + key = self.k_norm(key) + k_rope, _ = torch.split( + key, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 + ) + + q_rope, k_rope = self.rotary_emb(positions, q_rope, k_rope) + + self._update_rope_guarded(query[..., : self.rope_head_dim], q_rope) + self._update_rope_guarded(key[..., : self.rope_head_dim], k_rope) + + if enable_dual_stream: + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + query = self._maybe_rotate(query) + + with torch.cuda.stream(self.alt_stream): + key = self._maybe_rotate(key) + current_stream.wait_stream(self.alt_stream) + elif ( + self.alt_stream is not None + and forward_batch.attn_cp_metadata is not None + and self.dsa_enable_prefill_cp + ): + key = self._maybe_rotate(key) + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + query = self._maybe_rotate(query) + + with torch.cuda.stream(self.alt_stream): + key = cp_all_gather_rerange_output( + key.contiguous(), + self.cp_size, + forward_batch, + torch.cuda.current_stream(), + ) + current_stream.wait_stream(self.alt_stream) + return query, key, weights_raw + else: + query = self._maybe_rotate(query) + key = self._maybe_rotate(key) + + # allgather+rerrange + if forward_batch.attn_cp_metadata is not None and self.dsa_enable_prefill_cp: + key = cp_all_gather_rerange_output( + key.contiguous(), + self.cp_size, + forward_batch, + torch.cuda.current_stream(), + ) + return query, key, weights_raw + + def _get_k_bf16( + self, + x: torch.Tensor, + positions: torch.Tensor, + enable_dual_stream: bool, + ): + # Non-fusion path only; self.wk does not exist when fusion is on. + key, _ = self.wk(x) + key = self.k_norm(key) + k_rope, _ = torch.split( + key, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 + ) + + _, k_rope = self.rotary_emb(positions, k_rope, k_rope) + self._update_rope_guarded(key[..., : self.rope_head_dim], k_rope) + key = rotate_activation(key) + + return key + + def _fused_k_prepare_and_store( + self, + key_raw: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + layer_id: int, + act_quant, + out_cache_loc: Optional[torch.Tensor] = None, + ) -> None: + if out_cache_loc is None: + out_cache_loc = forward_batch.out_cache_loc + pool = get_token_to_kv_pool() + page_size = pool.page_size + if hasattr(pool, "invalidate_index_buffer_for_layer"): + pool.invalidate_index_buffer_for_layer(layer_id) + if hasattr(pool, "_is_layer_owned") and not pool._is_layer_owned(layer_id): + return + if ( + not _is_fp8_fnuz + and out_cache_loc is not None + and can_use_dsa_fused_store(torch.bfloat16, out_cache_loc.dtype, page_size) + ): + fused_k_indexer_norm_rope_store( + key_raw, + pool.get_index_k_with_scale_buffer(layer_id=layer_id), + out_cache_loc, + self.k_norm.weight, + self.k_norm.bias, + self.k_norm.variance_epsilon, + self._indexer_cos_sin_cache, + positions, + page_size, + ) + return + + # Fallback: separate K kernel + store kernel. + key = fused_k_indexer_norm_rope( + key_raw, + self.k_norm.weight, + self.k_norm.bias, + self.k_norm.variance_epsilon, + self._indexer_cos_sin_cache, + positions, + ) + self._store_index_k_cache( + forward_batch=forward_batch, + layer_id=layer_id, + key=key, + act_quant=act_quant, + out_cache_loc=out_cache_loc, + ) + + def _fused_q_prepare_and_store( + self, + x: torch.Tensor, + q_lora: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + layer_id: int, + act_quant, + *, + num_tokens: Optional[int] = None, + enable_dual_stream: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + # num_tokens (graph split-op contract) slices q/k/positions/out_cache_loc + # to the unpadded count; the returned q_fp8/weights are sliced to match. + q_scale_gate = self.softmax_scale * self.n_heads**-0.5 + out_cache_loc = forward_batch.out_cache_loc + if num_tokens is not None: + positions = positions[:num_tokens] + out_cache_loc = out_cache_loc[:num_tokens] + + if self.alt_stream is None or not enable_dual_stream: + kw, _ = self.wk_weights_proj(x) + key, weights_raw = kw.split([self.head_dim, self.n_heads], dim=-1) + if num_tokens is not None: + key = key[:num_tokens] + weights_raw = weights_raw[:num_tokens] + self._fused_k_prepare_and_store( + key, + positions, + forward_batch, + layer_id, + act_quant, + out_cache_loc=out_cache_loc, + ) + q = self.wq_b(q_lora)[0].view(-1, self.n_heads, self.head_dim) + if num_tokens is not None: + q = q[:num_tokens] + return fused_q_indexer_rope_first_quant( + q.contiguous(), + weights_raw, + q_scale_gate, + self._indexer_cos_sin_cache, + positions, + ) + + # Two overlap stages: wq_b GEMM (alt) || wk_weights_proj GEMM (current), + # then fused Q kernel (current) || fused K kernel + cache store (alt). + # wait_stream calls are ordered by issue position so each side waits only + # on the GEMMs it consumes, not on the other side's fused kernel. + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + with torch.cuda.stream(self.alt_stream): + q = self.wq_b(q_lora)[0].view(-1, self.n_heads, self.head_dim) + if num_tokens is not None: + q = q[:num_tokens] + + kw, _ = self.wk_weights_proj(x) + key, weights_raw = kw.split([self.head_dim, self.n_heads], dim=-1) + if num_tokens is not None: + key = key[:num_tokens] + weights_raw = weights_raw[:num_tokens] + + current_stream.wait_stream(self.alt_stream) + self.alt_stream.wait_stream(current_stream) + q_fp8, weights = fused_q_indexer_rope_first_quant( + q.contiguous(), + weights_raw, + q_scale_gate, + self._indexer_cos_sin_cache, + positions, + ) + with torch.cuda.stream(self.alt_stream): + self._fused_k_prepare_and_store( + key, + positions, + forward_batch, + layer_id, + act_quant, + out_cache_loc=out_cache_loc, + ) + + current_stream.wait_stream(self.alt_stream) + return q_fp8, weights + + @staticmethod + def _update_rope_guarded(dst: torch.Tensor, src: torch.Tensor) -> None: + # On AMD with in-place RoPE kernels, self-aliasing can occur; + # skip write-back when src/dst tensors point to a single memory. + # data_ptr() is not comparable inside torch.compile, so skip the guard there. + if not torch.compiler.is_compiling() and src.data_ptr() == dst.data_ptr(): + return + dst.copy_(src) + + @staticmethod + def _get_index_k_read_buffer(pool, layer_id: int) -> torch.Tensor: + # Read path: prefer the owner-broadcast scratch buffer under DSA cache + # layer split; fall back to the owned buffer for plain pools. Stores go + # through get_index_k_with_scale_buffer() (owned buffer) instead. + if hasattr(pool, "get_broadcastable_index_k_with_scale_buffer"): + return pool.get_broadcastable_index_k_with_scale_buffer(layer_id) + return pool.get_index_k_with_scale_buffer(layer_id=layer_id) + + @staticmethod + def _pad_heads_for_deep_gemm(q_fp8, weights): + """Pad q and weights to 32 heads when num_heads < 32, + so that block_q = 128/num_heads doesn't exceed seq_len_alignment(4).""" + num_heads = q_fp8.shape[1] + if num_heads >= 32: + return q_fp8, weights, num_heads + target_heads = 32 + q_fp8 = torch.nn.functional.pad(q_fp8, (0, 0, 0, target_heads - num_heads)) + weights = torch.nn.functional.pad(weights, (0, target_heads - num_heads)) + return q_fp8, weights, num_heads + + def _mask_init_and_local_tokens( + self, + logits: torch.Tensor, + lengths: torch.Tensor, + row_starts: Optional[torch.Tensor] = None, + ): + if self.num_init_tokens == 0 and self.num_local_tokens == 0: + return logits + if row_starts is None: + row_starts = lengths.new_zeros(lengths.shape[0]) + num_init_tokens = self.num_init_tokens + num_local_tokens = self.num_local_tokens + if num_init_tokens > 0: + init_idxs = ( + torch.arange( + num_init_tokens, dtype=lengths.dtype, device=lengths.device + )[None, :] + + row_starts[:, None] + ) + init_idxs.clamp_max_(logits.shape[-1] - 1) + logits.scatter_(dim=1, index=init_idxs, value=float("inf")) + if num_local_tokens > 0: + local_idxs = ( + lengths[:, None] + - 1 + + row_starts[:, None] + - torch.arange( + num_local_tokens, dtype=lengths.dtype, device=lengths.device + )[None, :] + ) + local_idxs.clamp_min_(0) + logits.scatter_(dim=1, index=local_idxs, value=float("inf")) + return logits + + def _get_topk_paged( + self, + forward_batch: ForwardBatch, + layer_id: int, + q_fp8: torch.Tensor, + weights: torch.Tensor, + metadata: BaseIndexerMetadata, + ) -> torch.Tensor: + if TYPE_CHECKING: + assert isinstance(get_token_to_kv_pool(), DSATokenToKVPool) + + page_size = get_token_to_kv_pool().page_size + # NOTE(dark): blocksize = 64 is hardcoded in deep_gemm + if _is_hip: + if _use_aiter_preshuffle: + assert ( + page_size % 16 == 0 + ), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}" + else: + assert ( + page_size == 1 + ), f"HIP legacy DSA path requires page_size == 1, got {page_size}" + else: + assert page_size == 64, "only support page size 64" + # NOTE(dark): this support extend/decode/decode+graph + if _is_hip and not _use_aiter_preshuffle: + block_tables = metadata.get_page_table_1() + else: + block_tables = metadata.get_page_table_64() + + max_seq_len = block_tables.shape[1] * page_size + kv_cache_fp8 = self._get_index_k_read_buffer(get_token_to_kv_pool(), layer_id) + + blocksize = page_size + if ( + forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend_v2() + ): + seqlens_32 = metadata.get_seqlens_expanded() + else: + seqlens_32 = metadata.get_seqlens_int32() + # Reuse pre-computed schedule metadata if available (from init_forward_metadata), + # otherwise fall back to computing it here. + schedule_metadata = getattr(metadata, "paged_mqa_schedule_metadata", None) + assert len(q_fp8.shape) == 3 + # attn_tp_size > 1 or MAX_LEN padding mode can leave padding in the + # hidden states; q_offset is the real (unpadded) q length. + q_offset = sum(metadata.get_dsa_extend_len_cpu()) + + B = metadata.get_seqlens_int32().shape[0] + next_n = q_offset // B if B > 0 else 0 + use_cute_dsl = ( + self.paged_mqa_logits_backend.is_cutedsl() + and not forward_batch.forward_mode.is_draft_extend_v2() + ) + dsl_expand_factor, dsl_atom = 1, 1 + if ( + use_cute_dsl + and forward_batch.forward_mode.is_target_verify() + and next_n >= 2 + ): + assert pick_dsl_expand is not None, "CuTe DSL paged MQA is CUDA-only." + dsl_expand_factor, dsl_atom = pick_dsl_expand( + next_n, + batch_size=B, + max_ctx=max_seq_len, + num_sms=self.sm_count, + kernel_atoms=(1, 2, 3, 4), + num_heads=self.n_heads, + ) + ctx_2d = getattr(metadata, "paged_mqa_ctx_lens_2d", None) + use_dg_native = ( + not use_cute_dsl + and _is_cuda + and forward_batch.forward_mode.is_target_verify() + and next_n >= 2 + and ctx_2d is not None + and ctx_2d.shape == (B, next_n) + ) + + if use_dg_native: + seqlens_32_2d = ctx_2d + elif seqlens_32.dim() == 2: + seqlens_32_2d = seqlens_32 + else: + seqlens_32_2d = seqlens_32.unsqueeze(-1) + if _is_cuda and not is_sm120_supported(): + if schedule_metadata is None: + schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata( + seqlens_32_2d, blocksize, self.sm_count + ) + + assert len(kv_cache_fp8.shape) == 2 + block_kv = page_size + num_heads_kv = 1 + head_dim_with_sf = 132 + kv_cache_fp8 = kv_cache_fp8.view( + kv_cache_fp8.shape[0], block_kv, num_heads_kv, head_dim_with_sf + ) + assert len(weights.shape) == 3 + weights = weights.squeeze(2) + + if _is_cuda and is_sm120_supported(): + from sglang.srt.layers.attention.dsv4.indexer import ( + fp8_paged_mqa_logits_torch_sm120, + ) + + logits = fp8_paged_mqa_logits_torch_sm120( + q_fp8[:q_offset].unsqueeze(1), + kv_cache_fp8, + weights[:q_offset], + seqlens_32_2d, + block_tables, + schedule_metadata, + max_seq_len, + clean_logits=False, + ) + elif self.paged_mqa_logits_backend.is_aiter(): + logits = aiter_paged_mqa_logits( + q_fp8, + kv_cache_fp8, + weights, + seqlens_32, + block_tables, + max_seq_len, + preshuffle=_use_aiter_preshuffle, + kv_block_size=block_kv, + ) + elif use_cute_dsl: + logits = cutedsl_paged_mqa_logits( + q_fp8, + kv_cache_fp8, + weights, + metadata.get_seqlens_int32(), + block_tables, + schedule_metadata, + max_seq_len, + q_offset=q_offset, + B=B, + next_n=next_n, + is_target_verify=forward_batch.forward_mode.is_target_verify(), + dsl_expand_factor=dsl_expand_factor, + dsl_atom=dsl_atom, + blocksize=blocksize, + sm_count=self.sm_count, + get_paged_mqa_logits_metadata_fn=deep_gemm.get_paged_mqa_logits_metadata, + ) + elif use_dg_native: + logits = deepgemm_paged_mqa_logits_native( + deep_gemm.fp8_paged_mqa_logits, + q_fp8, + kv_cache_fp8, + weights, + seqlens_32_2d, + block_tables, + schedule_metadata, + max_seq_len, + q_offset=q_offset, + B=B, + next_n=next_n, + ) + else: + logits = deepgemm_paged_mqa_logits_split( + deep_gemm.fp8_paged_mqa_logits, + q_fp8, + kv_cache_fp8, + weights, + seqlens_32_2d, + block_tables, + schedule_metadata, + max_seq_len, + q_offset=q_offset, + ) + + # NOTE(dark): logits should be cleaned in topk_transform + self._mask_init_and_local_tokens(logits, seqlens_32) + topk_result = metadata.topk_transform(logits, self.index_topk) + # Restore possible padding exist in the hidden states. + if not _is_hip and q_offset < q_fp8.shape[0]: + pad_len = q_fp8.shape[0] - q_offset + padding = torch.full( + (pad_len, topk_result.shape[1]), + -1, + dtype=topk_result.dtype, + device=topk_result.device, + ) + topk_result = torch.cat([topk_result, padding], dim=0) + return topk_result + + def _get_mqa_logits_budget_bytes(self, device_index: int) -> int: + free_mem_fraction = self._mqa_logits_free_mem_fraction() + cached_budget = self._mqa_logits_budget_bytes.get(device_index) + if cached_budget is not None: + return cached_budget + + total_mem = torch.cuda.get_device_properties(device_index).total_memory + + total_mem_budget = int(total_mem * self._MQA_LOGITS_TOTAL_MEM_FRACTION) + mem_fraction_static = get_server_args().mem_fraction_static + if mem_fraction_static is None: + static_budget = total_mem_budget + else: + static_free_mem = int(total_mem * max(0.0, 1.0 - mem_fraction_static)) + static_budget = min( + int(static_free_mem * free_mem_fraction), + total_mem_budget, + ) + static_budget = max(1, static_budget) + + # Keep the static serving-memory guard during CUDA graph capture without + # caching it. The first non-capture prefill path will cache the real + # free-memory budget below. + if get_is_capture_mode(): + return static_budget + + # Match the original free-memory guard: logits_bytes * 2 > free_mem. + # torch.cuda.mem_get_info synchronizes the host, so cache the result, + # capped by the workload-independent serving-memory headroom. + free_mem, _ = torch.cuda.mem_get_info(device_index) + budget_bytes = min(int(free_mem * free_mem_fraction), static_budget) + + budget_bytes = max(1, budget_bytes) + self._mqa_logits_budget_bytes[device_index] = budget_bytes + return budget_bytes + + def _should_chunk_mqa_logits( + self, num_q: int, num_k: int, device_index: int + ) -> Tuple[bool, int]: + """ + Detect whether we need to chunk the MQA logits computation to avoid OOM + Return: (need_chunk, logits_budget_bytes) + """ + # Quick static check for normal batches + if num_q * num_k < self._MQA_LOGITS_STATIC_SKIP_ELEMS: + return False, 0 + + logits_bytes = num_q * num_k * self._MQA_LOGITS_BYTES_PER_ELEM + logits_budget_bytes = self._get_mqa_logits_budget_bytes(device_index) + + need_chunk = logits_bytes > logits_budget_bytes + return need_chunk, logits_budget_bytes + + def _get_topk_ragged( + self, + enable_dual_stream: bool, + forward_batch: ForwardBatch, + layer_id: int, + q_fp8: torch.Tensor, + weights: torch.Tensor, + metadata: BaseIndexerMetadata, + topk_result: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if TYPE_CHECKING: + assert isinstance(get_token_to_kv_pool(), DSATokenToKVPool) + + assert forward_batch.forward_mode.is_extend_without_speculative() + + page_size = get_token_to_kv_pool().page_size + if _is_hip: + if _use_aiter_preshuffle: + assert ( + page_size % 16 == 0 + ), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}" + else: + assert ( + page_size == 1 + ), f"HIP legacy DSA path requires page_size == 1, got {page_size}" + else: + assert page_size == 64, "only support page size 64" + + assert len(weights.shape) == 3 + assert ( + forward_batch.seq_lens_cpu is not None + and forward_batch.extend_seq_lens_cpu is not None + ) + weights = weights.squeeze(-1) + + if _is_hip and not _use_aiter_preshuffle: + block_tables = metadata.get_page_table_1() + else: + block_tables = metadata.get_page_table_64() + + assert ( + forward_batch.seq_lens_cpu is not None + and forward_batch.extend_seq_lens_cpu is not None + ) + + batch_size = len(block_tables) + token_nums, _, _ = q_fp8.shape + device = q_fp8.device + device_index = device.index + assert device_index is not None, "q_fp8 must be on an indexed CUDA device" + + if topk_result is None: + topk_result = torch.full( + (token_nums, self.index_topk), -1, device=device, dtype=torch.int32 + ) + if batch_size == 0: + return topk_result + + ks, ke = metadata.get_indexer_kvcache_range() + + indexer_seq_lens_cpu = metadata.get_indexer_seq_len_cpu() + seq_len_sum = torch.sum(indexer_seq_lens_cpu).item() + max_seq_len = torch.max(indexer_seq_lens_cpu).item() + k_fp8, k_scale = get_token_to_kv_pool().get_index_k_scale_buffer( + layer_id, + metadata.get_indexer_seq_len(), + block_tables, + seq_len_sum, + max_seq_len, + ) + if _is_fp8_fnuz: + k_fp8 = k_fp8.view(torch.float8_e4m3fnuz) + else: + k_fp8 = k_fp8.view(torch.float8_e4m3fn) + + k_scale = k_scale.view(torch.float32).squeeze(-1) + kv_fp8 = (k_fp8, k_scale) + + # Check if we need to chunk to avoid OOM + seq_lens_expanded = metadata.get_seqlens_expanded() + token_to_batch_idx = metadata.get_token_to_batch_idx() + q_offset = ks.shape[0] + k_offset = k_fp8.shape[0] + need_chunk, logits_budget_bytes = self._should_chunk_mqa_logits( + q_offset, k_offset, device_index + ) + + if not need_chunk: + assert q_fp8[:q_offset].shape[0] != 0 + with self._with_real_sm_count(): + if _is_hip: + from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits + + kv, scale = kv_fp8 + # Match the CUDA deep_gemm path (clean_logits=False): the topk + # transform masks invalid positions via ks/ke/lengths, so the + # -inf pre-fill of the [tokens x seq_len_kv] logits buffer is + # redundant and grows quadratically with context length. + logits = fp8_mqa_logits( + q_fp8[:q_offset], + kv, + scale, + weights[:q_offset], + ks, + ke, + clean_logits=False, + ) + else: + q_padded, w_padded, _ = self._pad_heads_for_deep_gemm( + q_fp8[:q_offset], weights[:q_offset] + ) + logits = deep_gemm.fp8_mqa_logits( + q_padded, + kv_fp8, + w_padded, + ks, + ke, + clean_logits=False, + ) + assert logits.shape[0] == len(seq_lens_expanded) + assert logits.shape[1] == k_offset + + self._mask_init_and_local_tokens(logits, seq_lens_expanded, ks) + raw_topk_result = metadata.topk_transform(logits, self.index_topk, ks=ks) + topk_result[:q_offset] = raw_topk_result + return topk_result + + bytes_per_row = k_offset * self._MQA_LOGITS_BYTES_PER_ELEM + max_rows = max(1, int(logits_budget_bytes // max(bytes_per_row, 1))) + max_rows = min(max_rows, q_offset) + + global_topk_offset = metadata.attn_metadata.topk_indices_offset + cu_seqlens_q_full = None + if global_topk_offset is None: + cu_seqlens_q_full = torch.ones(q_offset, dtype=torch.int32, device=device) + + assert ( + seq_lens_expanded.shape[0] == q_offset + ), f"seq_lens_expanded length mismatch: {seq_lens_expanded.shape[0]} != {q_offset}" + if global_topk_offset is not None: + assert ( + global_topk_offset.shape[0] >= q_offset + ), f"topk_indices_offset too short: {global_topk_offset.shape[0]} < {q_offset}" + + start = 0 + while start < q_offset: + end = min(start + max_rows, q_offset) + + with self._with_real_sm_count(): + if _is_hip: + from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits + + kv, scale = kv_fp8 + # clean_logits=False: topk transform handles masking (see above) + logits_chunk = fp8_mqa_logits( + q_fp8[start:end], + kv, + scale, + weights[start:end], + ks[start:end], + ke[start:end], + clean_logits=False, + ) + else: + q_padded, w_padded, _ = self._pad_heads_for_deep_gemm( + q_fp8[start:end], weights[start:end] + ) + logits_chunk = deep_gemm.fp8_mqa_logits( + q_padded, + kv_fp8, + w_padded, + ks[start:end], + ke[start:end], + clean_logits=False, + ) + + lengths_chunk = seq_lens_expanded[start:end] + self._mask_init_and_local_tokens(logits_chunk, lengths_chunk, ks[start:end]) + + # RAGGED: use global offset; PAGED: construct local cu_seqlens_q per chunk + if global_topk_offset is not None: + # RAGGED path + topk_offset_chunk = global_topk_offset[start:end] + cu_seqlens_q_chunk = None + batch_idx_chunk = None + else: + # PAGED path: treat each token as a length-1 sequence + topk_offset_chunk = None + cu_seqlens_q_chunk = cu_seqlens_q_full[start:end] + batch_idx_chunk = token_to_batch_idx[start:end] + + raw_topk_chunk = metadata.topk_transform( + logits_chunk, + self.index_topk, + ks=ks[start:end], + cu_seqlens_q=cu_seqlens_q_chunk, + ke_offset=lengths_chunk, + batch_idx_list=batch_idx_chunk, + topk_indices_offset_override=topk_offset_chunk, + ) + topk_result[start:end] = raw_topk_chunk + start = end + + return topk_result + + def _forward_cuda_k_only( + self, + x: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + layer_id: int, + act_quant, + enable_dual_stream: bool, + metadata: BaseIndexerMetadata, + return_indices: bool = True, + *, + num_tokens: Optional[int] = None, + topk_result: Optional[torch.Tensor] = None, + ) -> Optional[torch.Tensor]: + # Shared by the eager path and the graph DSA split-op dispatch. The two + # keyword args carry the graph contract and default to the eager behavior: + # - num_tokens: slice key/out_cache_loc to the unpadded count (the graph + # runs at a static padded shape). None => full (eager) shape. + # - topk_result: pre-allocated padded buffer to fill in place (a downstream + # captured graph reads it at a fixed address). None => return a fresh, + # naturally-sized tensor. + assert forward_batch.forward_mode.is_extend_without_speculative() + x_meta = x[0] if isinstance(x, tuple) else x + + # Fast path: only compute and store k cache, skip all q and weights ops. + # num_tokens (graph contract) slices to the unpadded count. + out_cache_loc = None + if num_tokens is not None: + assert num_tokens <= forward_batch.out_cache_loc.shape[0] + out_cache_loc = forward_batch.out_cache_loc[:num_tokens] + elif not forward_batch.out_cache_loc.is_contiguous(): + forward_batch.out_cache_loc = forward_batch.out_cache_loc.contiguous() + + # Write the same K representation the decode path reads back: fused + # (no-Hadamard) when fusion is on, else the legacy Hadamard path. + if self.use_dsa_indexer_fusion: + key_raw, _ = self._fused_k_weights(x) + if num_tokens is not None: + assert num_tokens <= key_raw.shape[0] + key_raw = key_raw[:num_tokens] + positions = positions[:num_tokens] + self._fused_k_prepare_and_store( + key_raw, + positions, + forward_batch, + layer_id, + act_quant, + out_cache_loc=out_cache_loc, + ) + else: + key = self._get_k_bf16(x, positions, enable_dual_stream) + if num_tokens is not None: + assert num_tokens <= key.shape[0] + key = key[:num_tokens] + self._store_index_k_cache( + forward_batch=forward_batch, + layer_id=layer_id, + key=key, + act_quant=act_quant, + out_cache_loc=out_cache_loc, + ) + + # MHA doesn't need topk_indices + if not return_indices: + return None + + # MLA: use dummy logits with topk kernel's fast path to generate indices + # When length <= 2048, naive_topk_cuda directly generates [0,1,...,length-1,-1,...] + seq_lens_expanded = metadata.get_seqlens_expanded() + dummy_logits = torch.zeros( + seq_lens_expanded.shape[0], + self.index_topk, + dtype=torch.float32, + device=x_meta.device, + ) + raw_topk_result = metadata.topk_transform(dummy_logits, self.index_topk) + if topk_result is not None: + # PCG/BCG: fill the valid prefix of the padded static buffer and + # leave padded rows at the -1 sentinel. + topk_result[: raw_topk_result.shape[0]] = raw_topk_result + return None + return raw_topk_result + + def _get_topk_ragged_with_cp( + self, + forward_batch: ForwardBatch, + layer_id: int, + q_fp8: torch.Tensor, + weights: torch.Tensor, + metadata: BaseIndexerMetadata, + kv_len: int, + actual_seq_q: int, + cp_index: List[Tuple[int, int, int]] = None, + ) -> torch.Tensor: + assert not _is_in_piecewise_or_breakable_cuda_graph(), ( + "DSA context parallel (_get_topk_ragged_with_cp) not supported under " + "piecewise/breakable CUDA graph" + ) + if TYPE_CHECKING: + assert isinstance(get_token_to_kv_pool(), DSATokenToKVPool) + + page_size = get_token_to_kv_pool().page_size + assert page_size == 64, "only support page size 64" + assert len(weights.shape) == 3 + weights = weights.squeeze(-1) + k_fp8_list = [] + k_scale_list = [] + ks_list = [] + ke_offset_list = [] + offset = 0 + actual_seq_q_list = [] + batch_idx_list = [] + + block_tables = metadata.get_page_table_64() + + assert ( + forward_batch.seq_lens_cpu is not None + and forward_batch.extend_seq_lens_cpu is not None + ) + if cp_index is not None: + # TODO Multi-batch support has accuracy issues + for batch_idx, start_seq_position, end_seq_position in cp_index: + pre_chunk_offset = ( + forward_batch.seq_lens_cpu[batch_idx].item() + - forward_batch.extend_seq_lens_cpu[batch_idx] + ) + start_seq_position += pre_chunk_offset + end_seq_position += pre_chunk_offset + if offset == 0 and batch_idx != 0: + offset += forward_batch.extend_seq_lens_cpu[batch_idx - 1] + k_fp8 = get_token_to_kv_pool().get_index_k_continuous( + layer_id, + end_seq_position, + block_tables[batch_idx], + ) + k_scale = get_token_to_kv_pool().get_index_k_scale_continuous( + layer_id, + end_seq_position, + block_tables[batch_idx], + ) + + extend_seq_len = end_seq_position - start_seq_position + ks = torch.full( + (extend_seq_len,), offset, dtype=torch.int32, device="cuda" + ) + k_fp8_list.append(k_fp8) + k_scale_list.append(k_scale) + ks_list.append(ks) + ke_offset = torch.arange( + start_seq_position + 1, + end_seq_position + 1, + dtype=torch.int32, + device="cuda", + ) + ke_offset_list.append(ke_offset) + actual_seq_q = torch.tensor( + [extend_seq_len], dtype=torch.int32, device="cuda" + ) + actual_seq_q_list.append(actual_seq_q) + batch_idx_list.append(batch_idx) + + k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fn) + k_scale = torch.cat(k_scale_list, dim=0).view(torch.float32).squeeze(-1) + kv_fp8 = (k_fp8, k_scale) + ks = torch.cat(ks_list, dim=0) + ke_offset = torch.cat(ke_offset_list, dim=0) + ke = ks + ke_offset + actual_seq_q = torch.cat(actual_seq_q_list, dim=0) + with self._with_real_sm_count(): + q_padded, w_padded, _ = self._pad_heads_for_deep_gemm(q_fp8, weights) + logits = deep_gemm.fp8_mqa_logits( + q_padded, + kv_fp8, + w_padded, + ks, + ke, + clean_logits=False, + ) + topk_result = metadata.topk_transform( + logits, + self.index_topk, + ks=ks, + cu_seqlens_q=actual_seq_q, + ke_offset=ke_offset, + batch_idx_list=batch_idx_list, + ) + else: + kv_len = ( + forward_batch.seq_lens_cpu[0].item() + - forward_batch.extend_seq_lens_cpu[0] + + kv_len + ) + k_fp8 = get_token_to_kv_pool().get_index_k_continuous( + layer_id, + kv_len, + block_tables[0], + ) + k_scale = get_token_to_kv_pool().get_index_k_scale_continuous( + layer_id, + kv_len, + block_tables[0], + ) + + k_fp8 = k_fp8.view(torch.float8_e4m3fn) + k_scale = k_scale.view(torch.float32).squeeze(-1) + kv_fp8 = (k_fp8, k_scale) + ks = torch.full((actual_seq_q,), offset, dtype=torch.int32, device="cuda") + ke_offset = torch.arange( + (kv_len - actual_seq_q) + 1, + kv_len + 1, + dtype=torch.int32, + device="cuda", + ) + ke = ks + ke_offset + + with self._with_real_sm_count(): + q_padded, w_padded, _ = self._pad_heads_for_deep_gemm(q_fp8, weights) + logits = deep_gemm.fp8_mqa_logits( + q_padded, + kv_fp8, + w_padded, + ks, + ke, + clean_logits=False, + ) + actual_seq_q = torch.tensor([actual_seq_q], dtype=torch.int32).to( + device="cuda", non_blocking=True + ) + topk_result = metadata.topk_transform( + logits, + self.index_topk, + ks=ks, + cu_seqlens_q=actual_seq_q, + ke_offset=ke_offset, + ) + + return topk_result + + def forward_indexer( + self, + q_fp8: torch.Tensor, + weights: torch.Tensor, + forward_batch: ForwardBatch, + topk: int, + layer_id: int, + ) -> Optional[torch.Tensor]: + assert not _is_in_piecewise_or_breakable_cuda_graph(), ( + "DSA forward_indexer (non-CUDA loop path) not supported under " + "piecewise/breakable CUDA graph" + ) + if not _is_npu: + from sglang.kernels.ops.attention.dsa.tilelang_kernel import fp8_index + + page_size = get_token_to_kv_pool().page_size + assert page_size == 64, "only support page size 64" + + assert len(weights.shape) == 3 + weights = weights.squeeze(-1) + + # logits = deep_gemm.fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, ke) + k_fp8_list = [] + k_scale_list = [] + + topk_indices_list = [] + + block_tables = get_req_to_token_pool().req_to_token[ + forward_batch.req_pool_indices, : + ] + strided_indices = torch.arange( + 0, block_tables.shape[-1], page_size, device="cuda" + ) + block_tables = block_tables[:, strided_indices] // page_size + + q_len_start = 0 + + for i in range(forward_batch.batch_size): + seq_len = forward_batch.seq_lens[i].item() + q_len = ( + forward_batch.extend_seq_lens_cpu[i] + if forward_batch.forward_mode.is_extend() + else 1 + ) + q_len_end = q_len_start + q_len + + q_fp8_partial = q_fp8[q_len_start:q_len_end] + q_fp8_partial = q_fp8_partial.unsqueeze(0).contiguous() + + weights_partial = weights[q_len_start:q_len_end] + weights_partial = weights_partial.squeeze(-1).unsqueeze(0).contiguous() + + k_fp8 = get_token_to_kv_pool().get_index_k_continuous( + layer_id, + seq_len, + block_tables[i], + ) + k_scale = get_token_to_kv_pool().get_index_k_scale_continuous( + layer_id, + seq_len, + block_tables[i], + ) + + k_fp8 = k_fp8.view(torch.float8_e4m3fn).unsqueeze(0).contiguous() + k_scale = k_scale.view(torch.float32).squeeze(-1).unsqueeze(0).contiguous() + + index_score = fp8_index( + q_fp8_partial, + weights_partial, + k_fp8, + k_scale, + ) + end_pos = seq_len + topk_indices = index_score.topk(min(topk, end_pos), dim=-1)[1].squeeze(0) + + pad_len = ceil_align(topk_indices.shape[-1], 2048) - topk_indices.shape[-1] + topk_indices = torch.nn.functional.pad( + topk_indices, (0, pad_len), "constant", -1 + ) + + topk_indices_list.append(topk_indices) + + q_len_start = q_len_end + + topk_indices = torch.cat(topk_indices_list, dim=0) + return topk_indices + + def _store_index_k_cache( + self, + forward_batch: ForwardBatch, + layer_id: int, + key: torch.Tensor, + *, + act_quant=None, # fallback only + out_cache_loc: Optional[torch.Tensor] = None, + ) -> None: + """ + Store DSA indexer K cache for current step. + + Preferred: fused_store_index_k_cache(key, cache, out_cache_loc, page_size) + Fallback : act_quant(key) + token_to_kv_pool.set_index_k_scale_buffer(...) + + out_cache_loc will default to forward_batch.out_cache_loc if not provided. + """ + + if out_cache_loc is None: + out_cache_loc = forward_batch.out_cache_loc + + pool = get_token_to_kv_pool() + if hasattr(pool, "invalidate_index_buffer_for_layer"): + pool.invalidate_index_buffer_for_layer(layer_id) + if hasattr(pool, "_is_layer_owned") and not pool._is_layer_owned(layer_id): + return + + if ( + _is_cuda + and (not _is_fp8_fnuz) + and can_use_dsa_fused_store( + key.dtype, + out_cache_loc.dtype, + pool.page_size, + ) + ): + # NOTE: wrapper already normalizes shape/contiguity and asserts dtypes. + buf = pool.get_index_k_with_scale_buffer(layer_id=layer_id) + fused_store_index_k_cache( + key, + buf, + out_cache_loc, + pool.page_size, + ) + return + + # Fast path: AITER fused quant + cache store + # When _use_aiter_preshuffle is True we use the new MFMA 16x16 preshuffle + # layout (page_size>=16). Otherwise we fall back to the legacy row-major + # layout with page_size=1; the same kv_cache.view works for both cases + # because page_size is 1 there. + if _use_aiter: + page_size = pool.page_size + buf = pool.get_index_k_with_scale_buffer(layer_id=layer_id) + kv_cache = buf.view(-1, page_size, 132).view(fp8_dtype) + out_loc = forward_batch.out_cache_loc + if not out_loc.is_contiguous(): + out_loc = out_loc.contiguous() + indexer_k_quant_and_cache( + key, + kv_cache, + out_loc, + self.block_size, + self.scale_fmt, + preshuffle=_use_aiter_preshuffle, + ) + return + + # Fallback: original path + assert act_quant is not None + k_fp8, k_scale = act_quant(key, self.block_size, self.scale_fmt) + + if not out_cache_loc.is_contiguous(): + out_cache_loc = out_cache_loc.contiguous() + + pool.set_index_k_scale_buffer( + layer_id=layer_id, + loc=out_cache_loc, + index_k=k_fp8, + index_k_scale=k_scale, + ) + + def forward_xpu( + self, + x: torch.Tensor, + q_lora: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + layer_id: int, + return_indices: bool = True, + ) -> Optional[torch.Tensor]: + return self.forward_cuda( + x, q_lora, positions, forward_batch, layer_id, return_indices + ) + + def forward_cuda( + self, + x: torch.Tensor, + q_lora: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + layer_id: int, + return_indices: bool = True, + ) -> Optional[torch.Tensor]: + if _is_hip: + from sglang.kernels.ops.attention.dsa.tilelang_kernel import act_quant + elif not _is_npu: + from sglang.kernels.ops.attention.dsa.triton_kernel import act_quant + + if TYPE_CHECKING: + assert isinstance(get_token_to_kv_pool(), DSATokenToKVPool) + + # When upstream uses fused FP8 RMSNorm+quant, activations may be passed as + # a tuple like (x_fp8, x_scale[, y]). Use `x_meta` for shape/device queries. + x_meta = x[0] if isinstance(x, tuple) else x + + in_piecewise_or_breakable_cuda_graph = ( + _is_in_piecewise_or_breakable_cuda_graph() + ) + + # In piecewise/breakable CUDA graph mode, metadata is fetched inside + # custom ops via get_tc_piecewise_forward_context() to prevent Dynamo + # from guarding on forward_metadata identity, which changes each replay + # when init_forward_metadata creates a new ForwardMetadata object. + if not in_piecewise_or_breakable_cuda_graph: + metadata = get_attn_backend().get_indexer_metadata(layer_id, forward_batch) + if metadata is None: + return None + else: + metadata = None + + enable_dual_stream = ( + self.alt_stream is not None + and get_is_capture_mode() + and q_lora.shape[0] > 0 + and q_lora.shape[0] <= DUAL_STREAM_TOKEN_THRESHOLD + ) + + # Determine if should skip topk based on sequence length + # We can only skip the logits computation if cuda graph is not involved + skip_logits_computation = False + if not in_piecewise_or_breakable_cuda_graph: + skip_logits_computation = self._should_skip_logits_computation( + forward_batch + ) + + # Optimization: fast path when skipping topk computation + if skip_logits_computation and (not self.dsa_enable_prefill_cp): + topk_result = self._forward_cuda_k_only( + x, + positions, + forward_batch, + layer_id, + act_quant, + enable_dual_stream, + metadata, + return_indices, + ) + topk_result = _broadcast_indexer_topk_from_rank0(topk_result) + return maybe_capture_indexer_topk(layer_id, topk_result) + + # When weights_proj is LoRA-wrapped, use an eager module call so the + # wrapper owns base+delta and no LoRA kernel runs under torch.compile. + # Fusion folds weights_proj into wk_weights_proj, so weights_proj is + # absent then; short-circuit before touching it. + weights_proj_lora = not self.use_dsa_indexer_fusion and getattr( + self.weights_proj, "set_lora", False + ) + + if ( + self.use_dsa_indexer_fusion + and not in_piecewise_or_breakable_cuda_graph + and forward_batch.attn_cp_metadata is None + ): + q_fp8, weights = self._fused_q_prepare_and_store( + x, q_lora, positions, forward_batch, layer_id, act_quant + ) + elif ( + is_graph_dsa_split_op_surface(forward_batch) + and not self.dsa_enable_prefill_cp + ): + # Default path for non-CP prefill under PCG/BCG: run the whole indexer + # (q/k proj, head gate, k-cache store, topk) as a single eager split op + # instead of capturing it piecemeal in the graph. The split op is + # fusion-aware, so this also covers the fused path here. + if weights_proj_lora: + raise RuntimeError(GRAPH_WEIGHTS_PROJ_LORA_ERROR) + if return_indices: + topk_result = torch.full( + (x.shape[0], self.index_topk), + -1, + device=x.device, + dtype=torch.int32, + ) + else: + topk_result = torch.empty( + (0, self.index_topk), device=x.device, dtype=torch.int32 + ) + graph_dispatch_fn = ( + bcg_dsa_indexer_prefill_split + if is_in_breakable_cuda_graph() + else pcg_dsa_indexer_prefill_split + ) + graph_dispatch_fn( + layer_id=layer_id, + x=x, + q_lora=q_lora, + positions=positions, + topk_result=topk_result, + ) + result = _broadcast_indexer_topk_from_rank0( + topk_result if return_indices else None + ) + return maybe_capture_indexer_topk(layer_id, result) + + elif enable_dual_stream and forward_batch.forward_mode.is_decode_or_idle(): + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + if not self.use_dsa_indexer_fusion: + if weights_proj_lora: + weights = self.weights_proj(x)[0].float() * self.n_heads**-0.5 + else: + weights = self._project_and_scale_head_gates(x) + query, key, weights_raw = self._get_q_k_bf16( + q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch + ) + q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) + with torch.cuda.stream(self.alt_stream): + self._store_index_k_cache( + forward_batch=forward_batch, + layer_id=layer_id, + key=key, + act_quant=act_quant, + ) + current_stream.wait_stream(self.alt_stream) + if self.use_dsa_indexer_fusion: + weights = self._scale_head_gates(weights_raw, q_scale) + else: + weights = self._apply_q_scale_and_softmax_scale(weights, q_scale) + else: + query, key, weights_raw = self._get_q_k_bf16( + q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch + ) + + if enable_dual_stream: + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + + q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) + with torch.cuda.stream(self.alt_stream): + self._store_index_k_cache( + forward_batch=forward_batch, + layer_id=layer_id, + key=key, + act_quant=act_quant, + ) + current_stream.wait_stream(self.alt_stream) + elif not in_piecewise_or_breakable_cuda_graph: + q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) + self._store_index_k_cache( + forward_batch=forward_batch, + layer_id=layer_id, + key=key, + act_quant=act_quant, + ) + else: + # Graph paths not handled by the full DSA indexer split op + # still need q_fp8 for paged topk and q_scale for + # logits_head_gate_graph. K-cache storage is handled by the + # full graph split path when prefill requires it. + q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) + + # aiter (ROCm gfx95): the 3-tuple (fp8, scale, bf16) from + # fused_rms_fp8_group_quant is passed directly to _get_logits_head_gate, + # which extracts the bf16 tensor via _weights_proj_bf16_in_fp32_out, + # completely skipping the FP8 dequantization path below. + if ( + _use_aiter + and _is_gfx95_supported + and isinstance(x, tuple) + and len(x) == 3 + ): + x_for_gate = x + elif isinstance(x, tuple): + assert len(x) in ( + 2, + 3, + ), "For tuple input, only (x, x_s) or (x, x_s, y) formats are accepted" + x_q, x_s = x[0], x[1] + if ( + x_s is not None + and x_q.dim() == 2 + and x_s.dim() == 2 + and x_q.shape[0] == x_s.shape[0] + ): + m, n = x_q.shape + ng = x_s.shape[1] + if ng > 0 and n % ng == 0: + group = n // ng + x_for_gate = ( + x_q.to(torch.float32) + .view(m, ng, group) + .mul_(x_s.to(torch.float32).unsqueeze(-1)) + .view(m, n) + .to(torch.bfloat16) + ) + else: + x_for_gate = x_q.to(torch.bfloat16) + else: + x_for_gate = x_q.to(torch.bfloat16) + else: + x_for_gate = x + + if in_piecewise_or_breakable_cuda_graph: + if self.use_dsa_indexer_fusion: + weights = scale_head_gate_graph( + weights_raw, + self.n_heads**-0.5, + self.softmax_scale, + q_scale, + ) + else: + if weights_proj_lora: + raise RuntimeError(GRAPH_WEIGHTS_PROJ_LORA_ERROR) + weights = logits_head_gate_graph( + x_for_gate, + self.weights_proj.weight, + self.n_heads**-0.5, + self.softmax_scale, + q_scale, + ) + elif self.use_dsa_indexer_fusion: + weights = self._scale_head_gates(weights_raw, q_scale) + elif weights_proj_lora: + weights = self.weights_proj(x_for_gate)[0].float() * self.n_heads**-0.5 + weights = self._apply_q_scale_and_softmax_scale(weights, q_scale) + else: + weights = self._get_logits_head_gate(x_for_gate, q_scale) + + if _is_cuda or _is_hip: + # In piecewise/breakable CUDA graph, any access to seq_lens_cpu + # creates a Dynamo shape guard. These graph modes never have empty + # batches. + if not in_piecewise_or_breakable_cuda_graph: + if forward_batch.seq_lens.numel() == 0: + # this seems b/c max-pad, no worries? + # if x.shape[0] != 0: + # print( + # "HACK: seq_lens empty but x not empty, hackily return all-invalid topk_result" + # ) + topk_result = torch.full( + (x_meta.shape[0], self.index_topk), + -1, + dtype=torch.int, + device=x_meta.device, + ) + topk_result = _broadcast_indexer_topk_from_rank0(topk_result) + return maybe_capture_indexer_topk(layer_id, topk_result) + + if ( + forward_batch.forward_mode.is_decode_or_idle() + or forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend_v2() + ): + topk_result = self._get_topk_paged( + forward_batch, layer_id, q_fp8, weights, metadata + ) + else: + if ( + forward_batch.attn_cp_metadata is not None + and is_dsa_prefill_cp_in_seq_split() + ): + kv_len_prev = forward_batch.attn_cp_metadata.kv_len_prev_list[0] + kv_len_next = forward_batch.attn_cp_metadata.kv_len_next_list[0] + actual_seq_q_prev = ( + forward_batch.attn_cp_metadata.actual_seq_q_prev_list[0] + ) + actual_seq_q_next = ( + forward_batch.attn_cp_metadata.actual_seq_q_next_list[0] + ) + + # TODO support mutil-batch + # cp_batch_seq_index_prev = forward_batch.attn_cp_metadata["cp_batch_seq_index_prev"] + # cp_batch_seq_index_next = forward_batch.attn_cp_metadata["cp_batch_seq_index_next"] + # TODO prev, next, combined into a single call + q_fp8_prev, q_fp8_next = torch.split( + q_fp8, (q_fp8.shape[0] + 1) // 2, dim=0 + ) + weights_prev, weights_next = torch.split( + weights, (weights.shape[0] + 1) // 2, dim=0 + ) + topk_result_prev = self._get_topk_ragged_with_cp( + forward_batch, + layer_id, + q_fp8_prev, + weights_prev, + metadata, + kv_len_prev, + actual_seq_q_prev, + ) + + topk_result_next = self._get_topk_ragged_with_cp( + forward_batch, + layer_id, + q_fp8_next, + weights_next, + metadata, + kv_len_next, + actual_seq_q_next, + ) + topk_result = torch.cat([topk_result_prev, topk_result_next], dim=0) + topk_result = _broadcast_indexer_topk_from_rank0(topk_result) + return maybe_capture_indexer_topk(layer_id, topk_result) + else: + # In-graph (PCG/BCG) non-CP prefill is handled earlier by the + # graph DSA split-op dispatch, so only the eager path reaches + # here. + assert not in_piecewise_or_breakable_cuda_graph, ( + "Internal error: in-graph DSA prefill must go through the " + "graph DSA split-op dispatch" + ) + topk_result = self._get_topk_ragged( + enable_dual_stream, + forward_batch, + layer_id, + q_fp8, + weights, + metadata, + ) + else: + topk_result = self.forward_indexer( + q_fp8.contiguous(), + weights, + forward_batch, + topk=self.index_topk, + layer_id=layer_id, + ) + topk_result = _broadcast_indexer_topk_from_rank0(topk_result) + return maybe_capture_indexer_topk(layer_id, topk_result) + + def forward_npu( + self, + x: torch.Tensor, + q_lora: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + layer_id: int, + layer_scatter_modes=None, + dynamic_scale: torch.Tensor = None, + ) -> torch.Tensor: + if get_attn_backend().forward_metadata.seq_lens_cpu_int is None: + actual_seq_lengths_kv = get_attn_backend().forward_metadata.seq_lens + else: + actual_seq_lengths_kv = get_attn_backend().forward_metadata.seq_lens_cpu_int + is_prefill = ( + forward_batch.forward_mode.is_extend() + and not forward_batch.forward_mode.is_draft_extend_v2() + and not forward_batch.forward_mode.is_target_verify() + ) + + bs = q_lora.shape[0] + + if self.rotary_emb.is_neox_style: + if not hasattr(forward_batch, "npu_indexer_sin_cos_cache"): + cos_sin = self.rotary_emb.cos_sin_cache[positions] + cos, sin = cos_sin.chunk(2, dim=-1) + cos = cos.repeat(1, 2).view(-1, 1, 1, self.rope_head_dim) + sin = sin.repeat(1, 2).view(-1, 1, 1, self.rope_head_dim) + forward_batch.npu_indexer_sin_cos_cache = (sin, cos) + else: + sin, cos = forward_batch.npu_indexer_sin_cos_cache + + if self.alt_stream is not None: + self.alt_stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(self.alt_stream): + q_lora = ( + (q_lora, dynamic_scale) if dynamic_scale is not None else q_lora + ) + q = self.wq_b(q_lora)[ + 0 + ] # [bs, 1536] @ [1536, 64 * 128] = [bs, 64 * 128] + wq_b_event = self.alt_stream.record_event() + q = q.view(bs, self.n_heads, self.head_dim) # [bs, 64, 128] + q_pe, q_nope = torch.split( + q, + [self.rope_head_dim, self.head_dim - self.rope_head_dim], + dim=-1, + ) # [bs, 64, 64 + 64] + q_pe = q_pe.view(bs, self.n_heads, 1, self.rope_head_dim) + q_pe = torch_npu.npu_rotary_mul(q_pe, cos, sin).view( + bs, self.n_heads, self.rope_head_dim + ) # [bs, n, d] + q = torch.cat([q_pe, q_nope], dim=-1) + q.record_stream(self.alt_stream) + q_rope_event = self.alt_stream.record_event() + else: + q_lora = ( + (q_lora, dynamic_scale) if dynamic_scale is not None else q_lora + ) + q = self.wq_b(q_lora)[ + 0 + ] # [bs, 1536] @ [1536, 64 * 128] = [bs, 64 * 128] + q = q.view(bs, self.n_heads, self.head_dim) # [bs, 64, 128] + q_pe, q_nope = torch.split( + q, + [self.rope_head_dim, self.head_dim - self.rope_head_dim], + dim=-1, + ) # [bs, 64, 64 + 64] + q_pe = q_pe.view(bs, self.n_heads, 1, self.rope_head_dim) + q_pe = torch_npu.npu_rotary_mul(q_pe, cos, sin).view( + bs, self.n_heads, self.rope_head_dim + ) # [bs, n, d] + q = torch.cat([q_pe, q_nope], dim=-1) + + if envs.SGLANG_NPU_USE_MULTI_STREAM.get(): + indexer_weight_stream = get_indexer_weight_stream() + indexer_weight_stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(indexer_weight_stream): + x = x.view(-1, self.hidden_size) + weights = self.weights_proj(x.float())[0].to(torch.bfloat16) + weights.record_stream(indexer_weight_stream) + weights_event = indexer_weight_stream.record_event() + else: + x = x.view(-1, self.hidden_size) + weights = self.weights_proj(x.float())[0].to(torch.bfloat16) + + k_proj = self.wk(x)[0] # [b, s, 7168] @ [7168, 128] = [b, s, 128] + k = self.k_norm(k_proj) + if ( + _use_ag_after_qlora + and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED + and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL + ): + k = scattered_to_tp_attn_full(k, forward_batch) + k_pe, k_nope = torch.split( + k, + [self.rope_head_dim, self.head_dim - self.rope_head_dim], + dim=-1, + ) # [bs, 64 + 64] + + k_pe = k_pe.view(-1, 1, 1, self.rope_head_dim) + k_pe = torch.ops.npu.npu_rotary_mul(k_pe, cos, sin).view( + bs, 1, self.rope_head_dim + ) # [bs, 1, d] + k = torch.cat([k_pe, k_nope.unsqueeze(1)], dim=-1) # [bs, 1, 128] + + else: + if envs.SGLANG_NPU_USE_MULTI_STREAM.get(): + indexer_weight_stream = get_indexer_weight_stream() + indexer_weight_stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(indexer_weight_stream): + x = x.view(-1, self.hidden_size) + weights = self.weights_proj(x.float())[0].to(torch.bfloat16) + weights.record_stream(indexer_weight_stream) + weights_event = indexer_weight_stream.record_event() + else: + x = x.view(-1, self.hidden_size) + weights = self.weights_proj(x.float())[0].to(torch.bfloat16) + + q_lora = (q_lora, dynamic_scale) if dynamic_scale is not None else q_lora + q = self.wq_b(q_lora)[0] # [bs, 1536] @ [1536, 64 * 128] = [bs, 64 * 128] + q = q.view(bs, self.n_heads, self.head_dim) # [bs, 64, 128] + q_pe, q_nope = torch.split( + q, + [self.rope_head_dim, self.head_dim - self.rope_head_dim], + dim=-1, + ) # [bs, 64, 64 + 64] + + k_proj = self.wk(x)[0] # [b, s, 7168] @ [7168, 128] = [b, s, 128] + k = self.k_norm(k_proj) + k_pe, k_nope = torch.split( + k, + [self.rope_head_dim, self.head_dim - self.rope_head_dim], + dim=-1, + ) # [bs, 64 + 64] + + k_pe = k_pe.unsqueeze(1) + + if layer_id == 0: + self.rotary_emb.sin_cos_cache = ( + self.rotary_emb.cos_sin_cache.index_select(0, positions) + ) + + q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) + k_pe = k_pe.squeeze(1) + q = torch.cat([q_pe, q_nope], dim=-1) + k = torch.cat([k_pe, k_nope], dim=-1) + + if ( + is_prefill + and self.dsa_enable_prefill_cp + and forward_batch.attn_cp_metadata is not None + ): + k = cp_all_gather_rerange_output( + k.contiguous().view(-1, self.head_dim), + self.cp_size, + forward_batch, + torch.npu.current_stream(), + ) + + get_token_to_kv_pool().set_index_k_buffer( + layer_id, forward_batch.out_cache_loc, k + ) + if is_prefill: + if ( + self.dsa_enable_prefill_cp + and forward_batch.attn_cp_metadata is not None + ): + get_attn_backend().forward_metadata.actual_seq_lengths_q = ( + forward_batch.attn_cp_metadata.actual_seq_q_prev_tensor, + forward_batch.attn_cp_metadata.actual_seq_q_next_tensor, + ) + if sum(forward_batch.extend_prefix_lens_cpu) > 0: + total_kv_len_prev_tensor = ( + forward_batch.attn_cp_metadata.kv_len_prev_tensor + + forward_batch.extend_prefix_lens.squeeze() + ) + total_kv_len_next_tensor = ( + forward_batch.attn_cp_metadata.kv_len_next_tensor + + forward_batch.extend_prefix_lens.squeeze() + ) + get_attn_backend().forward_metadata.actual_seq_lengths_kv = ( + total_kv_len_prev_tensor, + total_kv_len_next_tensor, + ) + else: + get_attn_backend().forward_metadata.actual_seq_lengths_kv = ( + forward_batch.attn_cp_metadata.kv_len_prev_tensor, + forward_batch.attn_cp_metadata.kv_len_next_tensor, + ) + actual_seq_lengths_q = ( + get_attn_backend().forward_metadata.actual_seq_lengths_q + ) + actual_seq_lengths_kv = ( + get_attn_backend().forward_metadata.actual_seq_lengths_kv + ) + else: + actual_seq_lengths_kv = forward_batch.seq_lens + actual_seq_lengths_q = forward_batch.extend_seq_lens.cumsum(dim=0) + else: + if get_attn_backend().forward_metadata.actual_seq_lengths_q is None: + if ( + forward_batch.forward_mode.is_draft_extend_v2() + or forward_batch.forward_mode.is_target_verify() + ): + num_draft_tokens = get_attn_backend().speculative_num_draft_tokens + actual_seq_lengths_q = torch.arange( + num_draft_tokens, + num_draft_tokens + bs, + num_draft_tokens, + dtype=torch.int32, + device=k.device, + ) + else: + actual_seq_lengths_q = torch.tensor( + [1 + i * 1 for i in range(bs)], + dtype=torch.int32, + device=k.device, + ) + else: + actual_seq_lengths_q = ( + get_attn_backend().forward_metadata.actual_seq_lengths_q + ) + + past_key_states = get_token_to_kv_pool().get_index_k_buffer(layer_id) + + if self.rotary_emb.is_neox_style and self.alt_stream is not None: + torch.npu.current_stream().wait_event(q_rope_event) + if envs.SGLANG_NPU_USE_MULTI_STREAM.get(): + torch.npu.current_stream().wait_event(weights_event) + if ( + _use_ag_after_qlora + and layer_scatter_modes.layer_input_mode == ScatterMode.SCATTERED + and layer_scatter_modes.attn_mode == ScatterMode.TP_ATTN_FULL + ): + weights = scattered_to_tp_attn_full(weights, forward_batch) + block_table = get_attn_backend().forward_metadata.block_tables + if ( + is_prefill + and self.dsa_enable_prefill_cp + and forward_batch.attn_cp_metadata is not None + ): + block_table = block_table[: actual_seq_lengths_q[0].numel()] + topk_indices = self.do_npu_cp_balance_indexer( + q.view(-1, self.n_heads, self.head_dim), + past_key_states, + weights, + actual_seq_lengths_q, + actual_seq_lengths_kv, + block_table, + ) + return topk_indices + else: + block_table = ( + block_table[: actual_seq_lengths_q.size()[0]] + if is_prefill + else block_table + ) + + topk_indices = torch_npu.npu_lightning_indexer( + query=q.view(-1, self.n_heads, self.head_dim), + key=past_key_states, + weights=weights, + actual_seq_lengths_query=actual_seq_lengths_q.to(torch.int32), + actual_seq_lengths_key=actual_seq_lengths_kv.to(k.device).to( + torch.int32 + ), + block_table=block_table, + layout_query="TND", + layout_key="PA_BSND", + sparse_count=self.index_topk, + sparse_mode=3, + ) + # Keep DSA top-k as [T, K]; NPU attention expands it when needed. + return topk_indices[0].squeeze(1) + + def do_npu_cp_balance_indexer( + self, + q, + past_key_states, + indexer_weights, + actual_seq_lengths_q, + actual_seq_lengths_kv, + block_table, + ): + q_prev, q_next = torch.split(q, (q.size(0) + 1) // 2, dim=0) + weights_prev, weights_next = None, None + if indexer_weights is not None: + weights_prev, weights_next = torch.split( + indexer_weights, (indexer_weights.size(0) + 1) // 2, dim=0 + ) + weights_prev = weights_prev.contiguous().view(-1, weights_prev.shape[-1]) + weights_next = weights_next.contiguous().view(-1, weights_next.shape[-1]) + + actual_seq_lengths_q_prev, actual_seq_lengths_q_next = actual_seq_lengths_q + actual_seq_lengths_kv_prev, actual_seq_lengths_kv_next = actual_seq_lengths_kv + + topk_indices_prev = torch_npu.npu_lightning_indexer( + query=q_prev, + key=past_key_states, + weights=weights_prev, + actual_seq_lengths_query=actual_seq_lengths_q_prev.to( + device=q.device, dtype=torch.int32 + ), + actual_seq_lengths_key=actual_seq_lengths_kv_prev.to( + device=q.device, dtype=torch.int32 + ), + block_table=block_table, + layout_query="TND", + layout_key="PA_BSND", + sparse_count=self.index_topk, + sparse_mode=3, + ) + topk_indices_next = torch_npu.npu_lightning_indexer( + query=q_next, + key=past_key_states, + weights=weights_next, + actual_seq_lengths_query=actual_seq_lengths_q_next.to( + device=q.device, dtype=torch.int32 + ), + actual_seq_lengths_key=actual_seq_lengths_kv_next.to( + device=q.device, dtype=torch.int32 + ), + block_table=block_table, + layout_query="TND", + layout_key="PA_BSND", + sparse_count=self.index_topk, + sparse_mode=3, + ) + return topk_indices_prev[0], topk_indices_next[0] + + +@register_custom_op(mutates_args=["topk_result"]) +@register_split_op() +def pcg_dsa_indexer_prefill_split( + layer_id: int, + x: torch.Tensor, + q_lora: torch.Tensor, + positions: torch.Tensor, + topk_result: torch.Tensor, +) -> None: + # Default in-graph indexer path for non-CP prefill: runs the whole indexer + # (q/k proj, head gate, k-cache store, topk) as one eager split op. PCG calls + # this as a split op; BCG uses the explicit eager wrapper below. + # + # Output contract (differs from the eager `forward` path): a split op returns + # None, so results are delivered only by mutating `topk_result` in place. The + # call site pre-allocates it at a static, padded shape and a downstream + # captured graph reads it at a fixed address; eager code instead allocates + # and returns a fresh, naturally-sized tensor each call. + assert _is_cuda, "Internal error: DSA graph dispatch is only supported on CUDA" + from sglang.kernels.ops.attention.dsa.triton_kernel import act_quant + + forward_context = get_tc_piecewise_forward_context() + forward_batch = forward_context.forward_batch + indexer = forward_context.dsa_indexers[layer_id] + metadata = get_attn_backend().get_indexer_metadata(layer_id, forward_batch) + + extend_num_tokens = forward_batch.extend_num_tokens + # Empty buffer encodes return_indices=False for graph dispatch. + return_indices = topk_result.numel() != 0 + k_only = not return_indices or ( + indexer._should_skip_logits_computation(forward_batch) + and not indexer.dsa_enable_prefill_cp + ) + if k_only: + indexer._forward_cuda_k_only( + x, + positions, + forward_batch, + layer_id, + act_quant, + enable_dual_stream=False, + metadata=metadata, + return_indices=return_indices, + num_tokens=extend_num_tokens, + topk_result=topk_result, + ) + return + + # Fused path stores K (no-Hadamard) and computes q_fp8 + head gate in the + # fused kernels, sliced to the unpadded count. Single stream: the split op is + # captured, so the dual-stream overlap is disabled. + if indexer.use_dsa_indexer_fusion: + q_fp8, weights = indexer._fused_q_prepare_and_store( + x, + q_lora, + positions, + forward_batch, + layer_id, + act_quant, + num_tokens=extend_num_tokens, + enable_dual_stream=False, + ) + indexer._get_topk_ragged( + False, + forward_batch, + layer_id, + q_fp8, + weights, + metadata, + topk_result, + ) + return + + query, key, _ = indexer._get_q_k_bf16( + q_lora, + x, + positions, + enable_dual_stream=False, + forward_batch=forward_batch, + ) + q_fp8, q_scale = act_quant(query, indexer.block_size, indexer.scale_fmt) + # Reuse the compiled head-gate util shared with the eager path. + weights = indexer._get_logits_head_gate(x, q_scale) + # Store K cache + ragged top-k, sliced to the unpadded count and writing into + # the static padded topk_result buffer (the graph contract). Mirrors the eager + # path's store + _get_topk_ragged. + indexer._store_index_k_cache( + forward_batch=forward_batch, + layer_id=layer_id, + key=key[:extend_num_tokens], + act_quant=act_quant, + out_cache_loc=forward_batch.out_cache_loc[:extend_num_tokens], + ) + indexer._get_topk_ragged( + False, + forward_batch, + layer_id, + q_fp8[:extend_num_tokens], + weights, + metadata, + topk_result, + ) + + +bcg_dsa_indexer_prefill_split = eager_on_graph(True)(pcg_dsa_indexer_prefill_split) + + +def scattered_to_tp_attn_full( + hidden_states: torch.Tensor, + forward_batch, +) -> torch.Tensor: + hidden_states, local_hidden_states = ( + torch.empty( + (forward_batch.input_ids.shape[0], hidden_states.shape[1]), + dtype=hidden_states.dtype, + device=hidden_states.device, + ), + hidden_states, + ) + attn_tp_all_gather_into_tensor(hidden_states, local_hidden_states.contiguous()) + return hidden_states diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/generate_scenarios.py b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/generate_scenarios.py new file mode 100755 index 0000000..f722479 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/generate_scenarios.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +"""Generate the scenario list for the TP×DP matrix experiment. + +Reads matrix.json and prints TSV lines: + + mark input_len output_len concurrency num_prompts + +mark is one of Y/P/N. The caller (run_bench.sh) decides how to treat each. +""" +import argparse +import json +import math +import os +from pathlib import Path + + +def sample_concurrency(low: int, high: int, target: int) -> list[int]: + """Return only the low and high concurrency values in [low, high]. + + For the TP×DP matrix we only need the two endpoints of the concurrency + range (e.g. 1 and 128 for ISL=1024). The `target` argument is kept for + API compatibility but is ignored. + """ + assert 1 <= low <= high, f"invalid concurrency range: {low}-{high}" + if low == high: + return [low] + return [low, high] + + +def generate_scenarios(matrix_path: Path, mode: str, target_samples: int) -> list[dict]: + with open(matrix_path, "r", encoding="utf-8") as f: + data = json.load(f) + + matrix = data["matrix"] + concurrency_cfg = data["concurrency"] + + scenarios = [] + for isl_str in sorted(matrix.keys(), key=int): + dsl_map = matrix[isl_str] + low = concurrency_cfg[isl_str]["low"] + high = concurrency_cfg[isl_str]["high"] + concurrencies = sample_concurrency(low, high, target_samples) + + for dsl_str in sorted(dsl_map.keys(), key=int): + mark = dsl_map[dsl_str] + if mode == "Y" and mark != "Y": + continue + if mode == "Y+P" and mark not in ("Y", "P"): + continue + # mode == "all" keeps everything, including N. + + for conc in concurrencies: + scenarios.append( + { + "mark": mark, + "input_len": int(isl_str), + "output_len": int(dsl_str), + "concurrency": conc, + "num_prompts": conc * 5, + } + ) + + return scenarios + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--matrix", type=Path, default=Path("matrix.json")) + parser.add_argument("--mode", choices=["Y", "Y+P", "all"], default=None, + help="Scenario selection mode. Defaults to matrix.mode.") + parser.add_argument("--target-samples", type=int, default=0, + help="Target number of concurrency samples. 0 = heuristic (6-8).") + args = parser.parse_args() + + with open(args.matrix, "r", encoding="utf-8") as f: + data = json.load(f) + + mode = args.mode if args.mode else data.get("mode", "Y+P") + + target_samples = args.target_samples + if target_samples <= 0: + env_samples = os.getenv("CONCURRENCY_SAMPLES", "0") + try: + target_samples = int(env_samples) + except ValueError: + target_samples = 0 + if target_samples <= 0: + target_samples = 7 + + scenarios = generate_scenarios(args.matrix, mode, target_samples) + + print("mark\tinput_len\toutput_len\tconcurrency\tnum_prompts") + for s in scenarios: + print(f"{s['mark']}\t{s['input_len']}\t{s['output_len']}\t{s['concurrency']}\t{s['num_prompts']}") + + +if __name__ == "__main__": + main() diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/matrix.json b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/matrix.json new file mode 100644 index 0000000..dfc7572 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/matrix.json @@ -0,0 +1,128 @@ +{ + "comment": "Only mandatory (Y) combinations are tested; 1M ISL is excluded per user request.", + "mode": "Y", + "matrix": { + "1024": { + "128": "Y", + "256": "Y", + "512": "Y", + "1024": "Y", + "2048": "Y", + "4096": "Y" + }, + "4096": { + "128": "Y", + "256": "Y", + "512": "Y", + "1024": "Y", + "2048": "Y", + "4096": "Y" + }, + "8192": { + "128": "Y", + "256": "Y", + "512": "Y", + "1024": "Y", + "2048": "Y", + "4096": "Y" + }, + "16384": { + "128": "Y", + "256": "Y", + "512": "Y", + "1024": "Y", + "2048": "Y", + "4096": "P" + }, + "32768": { + "128": "Y", + "256": "Y", + "512": "Y", + "1024": "Y", + "2048": "Y", + "4096": "P" + }, + "65536": { + "128": "Y", + "256": "Y", + "512": "Y", + "1024": "Y", + "2048": "P", + "4096": "N" + }, + "131072": { + "128": "Y", + "256": "Y", + "512": "Y", + "1024": "P", + "2048": "N", + "4096": "N" + }, + "262144": { + "128": "Y", + "256": "Y", + "512": "P", + "1024": "N", + "2048": "N", + "4096": "N" + }, + "524288": { + "128": "Y", + "256": "P", + "512": "N", + "1024": "N", + "2048": "N", + "4096": "N" + }, + "1048576": { + "128": "Y", + "256": "P", + "512": "N", + "1024": "N", + "2048": "N", + "4096": "N" + } + }, + "concurrency": { + "1024": { + "low": 1, + "high": 128 + }, + "4096": { + "low": 1, + "high": 64 + }, + "8192": { + "low": 1, + "high": 64 + }, + "16384": { + "low": 1, + "high": 32 + }, + "32768": { + "low": 1, + "high": 16 + }, + "65536": { + "low": 1, + "high": 8 + }, + "131072": { + "low": 1, + "high": 4 + }, + "262144": { + "low": 1, + "high": 2 + }, + "524288": { + "low": 1, + "high": 2 + }, + "1048576": { + "low": 1, + "high": 2 + } + } +} diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-134240/comparison.md b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-134240/comparison.md new file mode 100644 index 0000000..bdd9237 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-134240/comparison.md @@ -0,0 +1,20 @@ +# SGLang TP×DP matrix comparison (8x NVIDIA RTX 6000D 84GB) + +## Summary + +- Model: `/data/hf_models/GLM-5.2-FP8` +- Hardware: 8x NVIDIA RTX 6000D 84GB +- Backend: SGLang (Docker) +- Benchmark client: `sglang.benchmark.serving` +- SLO reference: TTFT P95 < 3000.0ms, TPOT mean < 50.0ms + +### Configurations + +| Config | TP | DP | GPUs/replica | Notes | +|---|---:|---:|---:|---| +| tp16_dp1 | ? | ? | ? | server args recorded per ISL in results.json | + +## Side-by-side results + +| Scenario | ISL | DSL | Config | Conc | Req/s | OutTok/s | TTFT P95(ms) | TTFT P99(ms) | TPOT Mean(ms) | TPOT P95(ms) | TPOT P99(ms) | E2E P99(ms) | Peak GPU mem | SLO | +|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---| diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-134240/tp16_dp1/results.json b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-134240/tp16_dp1/results.json new file mode 100644 index 0000000..1bc8da7 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-134240/tp16_dp1/results.json @@ -0,0 +1,489 @@ +{ + "metadata": { + "experiment": "glm52_pro6000_sglang_multinode_tp16_tp16_dp1", + "run_id": "20260724-134240", + "timestamp": "2026-07-24T13:42:41+08:00", + "model": "/data/hf_models/GLM-5.2-FP8", + "backend": "sglang", + "engine": "sglang", + "hardware": "8x NVIDIA RTX 6000D 84GB", + "accelerator": "NVIDIA RTX 6000D", + "chip": "nvidia_rtx6000d", + "script": "experiments/glm52_pro6000_sglang_multinode_tp16/run_bench.sh", + "env": "lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45", + "git_commit": "d13f61f", + "git_dirty": "dirty", + "description": "H200 SGLang TP×DP matrix for DeepSeek-V4-Flash on RTX 6000D" + }, + "config": {}, + "scenarios": [ + { + "name": "c1_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + } + ] +} \ No newline at end of file diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-134240/tp16_dp1/scenarios.tsv b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-134240/tp16_dp1/scenarios.tsv new file mode 100644 index 0000000..70b2ee8 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-134240/tp16_dp1/scenarios.tsv @@ -0,0 +1,79 @@ +mark input_len output_len concurrency num_prompts +Y 1024 128 1 5 +Y 1024 128 128 640 +Y 1024 256 1 5 +Y 1024 256 128 640 +Y 1024 512 1 5 +Y 1024 512 128 640 +Y 1024 1024 1 5 +Y 1024 1024 128 640 +Y 1024 2048 1 5 +Y 1024 2048 128 640 +Y 1024 4096 1 5 +Y 1024 4096 128 640 +Y 4096 128 1 5 +Y 4096 128 64 320 +Y 4096 256 1 5 +Y 4096 256 64 320 +Y 4096 512 1 5 +Y 4096 512 64 320 +Y 4096 1024 1 5 +Y 4096 1024 64 320 +Y 4096 2048 1 5 +Y 4096 2048 64 320 +Y 4096 4096 1 5 +Y 4096 4096 64 320 +Y 8192 128 1 5 +Y 8192 128 64 320 +Y 8192 256 1 5 +Y 8192 256 64 320 +Y 8192 512 1 5 +Y 8192 512 64 320 +Y 8192 1024 1 5 +Y 8192 1024 64 320 +Y 8192 2048 1 5 +Y 8192 2048 64 320 +Y 8192 4096 1 5 +Y 8192 4096 64 320 +Y 16384 128 1 5 +Y 16384 128 32 160 +Y 16384 256 1 5 +Y 16384 256 32 160 +Y 16384 512 1 5 +Y 16384 512 32 160 +Y 16384 1024 1 5 +Y 16384 1024 32 160 +Y 16384 2048 1 5 +Y 16384 2048 32 160 +Y 32768 128 1 5 +Y 32768 128 16 80 +Y 32768 256 1 5 +Y 32768 256 16 80 +Y 32768 512 1 5 +Y 32768 512 16 80 +Y 32768 1024 1 5 +Y 32768 1024 16 80 +Y 32768 2048 1 5 +Y 32768 2048 16 80 +Y 65536 128 1 5 +Y 65536 128 8 40 +Y 65536 256 1 5 +Y 65536 256 8 40 +Y 65536 512 1 5 +Y 65536 512 8 40 +Y 65536 1024 1 5 +Y 65536 1024 8 40 +Y 131072 128 1 5 +Y 131072 128 4 20 +Y 131072 256 1 5 +Y 131072 256 4 20 +Y 131072 512 1 5 +Y 131072 512 4 20 +Y 262144 128 1 5 +Y 262144 128 2 10 +Y 262144 256 1 5 +Y 262144 256 2 10 +Y 524288 128 1 5 +Y 524288 128 2 10 +Y 1048576 128 1 5 +Y 1048576 128 2 10 diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-143107/comparison.md b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-143107/comparison.md new file mode 100644 index 0000000..bdd9237 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-143107/comparison.md @@ -0,0 +1,20 @@ +# SGLang TP×DP matrix comparison (8x NVIDIA RTX 6000D 84GB) + +## Summary + +- Model: `/data/hf_models/GLM-5.2-FP8` +- Hardware: 8x NVIDIA RTX 6000D 84GB +- Backend: SGLang (Docker) +- Benchmark client: `sglang.benchmark.serving` +- SLO reference: TTFT P95 < 3000.0ms, TPOT mean < 50.0ms + +### Configurations + +| Config | TP | DP | GPUs/replica | Notes | +|---|---:|---:|---:|---| +| tp16_dp1 | ? | ? | ? | server args recorded per ISL in results.json | + +## Side-by-side results + +| Scenario | ISL | DSL | Config | Conc | Req/s | OutTok/s | TTFT P95(ms) | TTFT P99(ms) | TPOT Mean(ms) | TPOT P95(ms) | TPOT P99(ms) | E2E P99(ms) | Peak GPU mem | SLO | +|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---| diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-143107/tp16_dp1/results.json b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-143107/tp16_dp1/results.json new file mode 100644 index 0000000..15d00f2 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-143107/tp16_dp1/results.json @@ -0,0 +1,489 @@ +{ + "metadata": { + "experiment": "glm52_pro6000_sglang_multinode_tp16_tp16_dp1", + "run_id": "20260724-143107", + "timestamp": "2026-07-24T14:31:22+08:00", + "model": "/data/hf_models/GLM-5.2-FP8", + "backend": "sglang", + "engine": "sglang", + "hardware": "8x NVIDIA RTX 6000D 84GB", + "accelerator": "NVIDIA RTX 6000D", + "chip": "nvidia_rtx6000d", + "script": "experiments/glm52_pro6000_sglang_multinode_tp16/run_bench.sh", + "env": "lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45", + "git_commit": "d13f61f", + "git_dirty": "dirty", + "description": "H200 SGLang TP×DP matrix for DeepSeek-V4-Flash on RTX 6000D" + }, + "config": {}, + "scenarios": [ + { + "name": "c1_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + } + ] +} \ No newline at end of file diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-143107/tp16_dp1/scenarios.tsv b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-143107/tp16_dp1/scenarios.tsv new file mode 100644 index 0000000..70b2ee8 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-143107/tp16_dp1/scenarios.tsv @@ -0,0 +1,79 @@ +mark input_len output_len concurrency num_prompts +Y 1024 128 1 5 +Y 1024 128 128 640 +Y 1024 256 1 5 +Y 1024 256 128 640 +Y 1024 512 1 5 +Y 1024 512 128 640 +Y 1024 1024 1 5 +Y 1024 1024 128 640 +Y 1024 2048 1 5 +Y 1024 2048 128 640 +Y 1024 4096 1 5 +Y 1024 4096 128 640 +Y 4096 128 1 5 +Y 4096 128 64 320 +Y 4096 256 1 5 +Y 4096 256 64 320 +Y 4096 512 1 5 +Y 4096 512 64 320 +Y 4096 1024 1 5 +Y 4096 1024 64 320 +Y 4096 2048 1 5 +Y 4096 2048 64 320 +Y 4096 4096 1 5 +Y 4096 4096 64 320 +Y 8192 128 1 5 +Y 8192 128 64 320 +Y 8192 256 1 5 +Y 8192 256 64 320 +Y 8192 512 1 5 +Y 8192 512 64 320 +Y 8192 1024 1 5 +Y 8192 1024 64 320 +Y 8192 2048 1 5 +Y 8192 2048 64 320 +Y 8192 4096 1 5 +Y 8192 4096 64 320 +Y 16384 128 1 5 +Y 16384 128 32 160 +Y 16384 256 1 5 +Y 16384 256 32 160 +Y 16384 512 1 5 +Y 16384 512 32 160 +Y 16384 1024 1 5 +Y 16384 1024 32 160 +Y 16384 2048 1 5 +Y 16384 2048 32 160 +Y 32768 128 1 5 +Y 32768 128 16 80 +Y 32768 256 1 5 +Y 32768 256 16 80 +Y 32768 512 1 5 +Y 32768 512 16 80 +Y 32768 1024 1 5 +Y 32768 1024 16 80 +Y 32768 2048 1 5 +Y 32768 2048 16 80 +Y 65536 128 1 5 +Y 65536 128 8 40 +Y 65536 256 1 5 +Y 65536 256 8 40 +Y 65536 512 1 5 +Y 65536 512 8 40 +Y 65536 1024 1 5 +Y 65536 1024 8 40 +Y 131072 128 1 5 +Y 131072 128 4 20 +Y 131072 256 1 5 +Y 131072 256 4 20 +Y 131072 512 1 5 +Y 131072 512 4 20 +Y 262144 128 1 5 +Y 262144 128 2 10 +Y 262144 256 1 5 +Y 262144 256 2 10 +Y 524288 128 1 5 +Y 524288 128 2 10 +Y 1048576 128 1 5 +Y 1048576 128 2 10 diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-150539/tp16_dp1/results.json b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-150539/tp16_dp1/results.json new file mode 100644 index 0000000..e6d058a --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-150539/tp16_dp1/results.json @@ -0,0 +1,489 @@ +{ + "metadata": { + "experiment": "glm52_pro6000_sglang_multinode_tp16_tp16_dp1", + "run_id": "20260724-150539", + "timestamp": "2026-07-24T15:05:40+08:00", + "model": "/data/hf_models/GLM-5.2-FP8", + "backend": "sglang", + "engine": "sglang", + "hardware": "8x NVIDIA RTX 6000D 84GB", + "accelerator": "NVIDIA RTX 6000D", + "chip": "nvidia_rtx6000d", + "script": "experiments/glm52_pro6000_sglang_multinode_tp16/run_bench.sh", + "env": "lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45", + "git_commit": "d13f61f", + "git_dirty": "dirty", + "description": "H200 SGLang TP×DP matrix for DeepSeek-V4-Flash on RTX 6000D" + }, + "config": {}, + "scenarios": [ + { + "name": "c1_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + } + ] +} \ No newline at end of file diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-150539/tp16_dp1/scenarios.tsv b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-150539/tp16_dp1/scenarios.tsv new file mode 100644 index 0000000..70b2ee8 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-150539/tp16_dp1/scenarios.tsv @@ -0,0 +1,79 @@ +mark input_len output_len concurrency num_prompts +Y 1024 128 1 5 +Y 1024 128 128 640 +Y 1024 256 1 5 +Y 1024 256 128 640 +Y 1024 512 1 5 +Y 1024 512 128 640 +Y 1024 1024 1 5 +Y 1024 1024 128 640 +Y 1024 2048 1 5 +Y 1024 2048 128 640 +Y 1024 4096 1 5 +Y 1024 4096 128 640 +Y 4096 128 1 5 +Y 4096 128 64 320 +Y 4096 256 1 5 +Y 4096 256 64 320 +Y 4096 512 1 5 +Y 4096 512 64 320 +Y 4096 1024 1 5 +Y 4096 1024 64 320 +Y 4096 2048 1 5 +Y 4096 2048 64 320 +Y 4096 4096 1 5 +Y 4096 4096 64 320 +Y 8192 128 1 5 +Y 8192 128 64 320 +Y 8192 256 1 5 +Y 8192 256 64 320 +Y 8192 512 1 5 +Y 8192 512 64 320 +Y 8192 1024 1 5 +Y 8192 1024 64 320 +Y 8192 2048 1 5 +Y 8192 2048 64 320 +Y 8192 4096 1 5 +Y 8192 4096 64 320 +Y 16384 128 1 5 +Y 16384 128 32 160 +Y 16384 256 1 5 +Y 16384 256 32 160 +Y 16384 512 1 5 +Y 16384 512 32 160 +Y 16384 1024 1 5 +Y 16384 1024 32 160 +Y 16384 2048 1 5 +Y 16384 2048 32 160 +Y 32768 128 1 5 +Y 32768 128 16 80 +Y 32768 256 1 5 +Y 32768 256 16 80 +Y 32768 512 1 5 +Y 32768 512 16 80 +Y 32768 1024 1 5 +Y 32768 1024 16 80 +Y 32768 2048 1 5 +Y 32768 2048 16 80 +Y 65536 128 1 5 +Y 65536 128 8 40 +Y 65536 256 1 5 +Y 65536 256 8 40 +Y 65536 512 1 5 +Y 65536 512 8 40 +Y 65536 1024 1 5 +Y 65536 1024 8 40 +Y 131072 128 1 5 +Y 131072 128 4 20 +Y 131072 256 1 5 +Y 131072 256 4 20 +Y 131072 512 1 5 +Y 131072 512 4 20 +Y 262144 128 1 5 +Y 262144 128 2 10 +Y 262144 256 1 5 +Y 262144 256 2 10 +Y 524288 128 1 5 +Y 524288 128 2 10 +Y 1048576 128 1 5 +Y 1048576 128 2 10 diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-151233/comparison.md b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-151233/comparison.md new file mode 100644 index 0000000..bdd9237 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-151233/comparison.md @@ -0,0 +1,20 @@ +# SGLang TP×DP matrix comparison (8x NVIDIA RTX 6000D 84GB) + +## Summary + +- Model: `/data/hf_models/GLM-5.2-FP8` +- Hardware: 8x NVIDIA RTX 6000D 84GB +- Backend: SGLang (Docker) +- Benchmark client: `sglang.benchmark.serving` +- SLO reference: TTFT P95 < 3000.0ms, TPOT mean < 50.0ms + +### Configurations + +| Config | TP | DP | GPUs/replica | Notes | +|---|---:|---:|---:|---| +| tp16_dp1 | ? | ? | ? | server args recorded per ISL in results.json | + +## Side-by-side results + +| Scenario | ISL | DSL | Config | Conc | Req/s | OutTok/s | TTFT P95(ms) | TTFT P99(ms) | TPOT Mean(ms) | TPOT P95(ms) | TPOT P99(ms) | E2E P99(ms) | Peak GPU mem | SLO | +|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---| diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-151233/tp16_dp1/results.json b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-151233/tp16_dp1/results.json new file mode 100644 index 0000000..725fb15 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-151233/tp16_dp1/results.json @@ -0,0 +1,489 @@ +{ + "metadata": { + "experiment": "glm52_pro6000_sglang_multinode_tp16_tp16_dp1", + "run_id": "20260724-151233", + "timestamp": "2026-07-24T15:12:34+08:00", + "model": "/data/hf_models/GLM-5.2-FP8", + "backend": "sglang", + "engine": "sglang", + "hardware": "8x NVIDIA RTX 6000D 84GB", + "accelerator": "NVIDIA RTX 6000D", + "chip": "nvidia_rtx6000d", + "script": "experiments/glm52_pro6000_sglang_multinode_tp16/run_bench.sh", + "env": "lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45", + "git_commit": "d13f61f", + "git_dirty": "dirty", + "description": "H200 SGLang TP×DP matrix for DeepSeek-V4-Flash on RTX 6000D" + }, + "config": {}, + "scenarios": [ + { + "name": "c1_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + } + ] +} \ No newline at end of file diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-151233/tp16_dp1/scenarios.tsv b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-151233/tp16_dp1/scenarios.tsv new file mode 100644 index 0000000..70b2ee8 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-151233/tp16_dp1/scenarios.tsv @@ -0,0 +1,79 @@ +mark input_len output_len concurrency num_prompts +Y 1024 128 1 5 +Y 1024 128 128 640 +Y 1024 256 1 5 +Y 1024 256 128 640 +Y 1024 512 1 5 +Y 1024 512 128 640 +Y 1024 1024 1 5 +Y 1024 1024 128 640 +Y 1024 2048 1 5 +Y 1024 2048 128 640 +Y 1024 4096 1 5 +Y 1024 4096 128 640 +Y 4096 128 1 5 +Y 4096 128 64 320 +Y 4096 256 1 5 +Y 4096 256 64 320 +Y 4096 512 1 5 +Y 4096 512 64 320 +Y 4096 1024 1 5 +Y 4096 1024 64 320 +Y 4096 2048 1 5 +Y 4096 2048 64 320 +Y 4096 4096 1 5 +Y 4096 4096 64 320 +Y 8192 128 1 5 +Y 8192 128 64 320 +Y 8192 256 1 5 +Y 8192 256 64 320 +Y 8192 512 1 5 +Y 8192 512 64 320 +Y 8192 1024 1 5 +Y 8192 1024 64 320 +Y 8192 2048 1 5 +Y 8192 2048 64 320 +Y 8192 4096 1 5 +Y 8192 4096 64 320 +Y 16384 128 1 5 +Y 16384 128 32 160 +Y 16384 256 1 5 +Y 16384 256 32 160 +Y 16384 512 1 5 +Y 16384 512 32 160 +Y 16384 1024 1 5 +Y 16384 1024 32 160 +Y 16384 2048 1 5 +Y 16384 2048 32 160 +Y 32768 128 1 5 +Y 32768 128 16 80 +Y 32768 256 1 5 +Y 32768 256 16 80 +Y 32768 512 1 5 +Y 32768 512 16 80 +Y 32768 1024 1 5 +Y 32768 1024 16 80 +Y 32768 2048 1 5 +Y 32768 2048 16 80 +Y 65536 128 1 5 +Y 65536 128 8 40 +Y 65536 256 1 5 +Y 65536 256 8 40 +Y 65536 512 1 5 +Y 65536 512 8 40 +Y 65536 1024 1 5 +Y 65536 1024 8 40 +Y 131072 128 1 5 +Y 131072 128 4 20 +Y 131072 256 1 5 +Y 131072 256 4 20 +Y 131072 512 1 5 +Y 131072 512 4 20 +Y 262144 128 1 5 +Y 262144 128 2 10 +Y 262144 256 1 5 +Y 262144 256 2 10 +Y 524288 128 1 5 +Y 524288 128 2 10 +Y 1048576 128 1 5 +Y 1048576 128 2 10 diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-152532/comparison.md b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-152532/comparison.md new file mode 100644 index 0000000..bdd9237 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-152532/comparison.md @@ -0,0 +1,20 @@ +# SGLang TP×DP matrix comparison (8x NVIDIA RTX 6000D 84GB) + +## Summary + +- Model: `/data/hf_models/GLM-5.2-FP8` +- Hardware: 8x NVIDIA RTX 6000D 84GB +- Backend: SGLang (Docker) +- Benchmark client: `sglang.benchmark.serving` +- SLO reference: TTFT P95 < 3000.0ms, TPOT mean < 50.0ms + +### Configurations + +| Config | TP | DP | GPUs/replica | Notes | +|---|---:|---:|---:|---| +| tp16_dp1 | ? | ? | ? | server args recorded per ISL in results.json | + +## Side-by-side results + +| Scenario | ISL | DSL | Config | Conc | Req/s | OutTok/s | TTFT P95(ms) | TTFT P99(ms) | TPOT Mean(ms) | TPOT P95(ms) | TPOT P99(ms) | E2E P99(ms) | Peak GPU mem | SLO | +|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---| diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-152532/tp16_dp1/results.json b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-152532/tp16_dp1/results.json new file mode 100644 index 0000000..209f0d4 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-152532/tp16_dp1/results.json @@ -0,0 +1,489 @@ +{ + "metadata": { + "experiment": "glm52_pro6000_sglang_multinode_tp16_tp16_dp1", + "run_id": "20260724-152532", + "timestamp": "2026-07-24T15:25:33+08:00", + "model": "/data/hf_models/GLM-5.2-FP8", + "backend": "sglang", + "engine": "sglang", + "hardware": "8x NVIDIA RTX 6000D 84GB", + "accelerator": "NVIDIA RTX 6000D", + "chip": "nvidia_rtx6000d", + "script": "experiments/glm52_pro6000_sglang_multinode_tp16/run_bench.sh", + "env": "lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45", + "git_commit": "d13f61f", + "git_dirty": "dirty", + "description": "H200 SGLang TP×DP matrix for DeepSeek-V4-Flash on RTX 6000D" + }, + "config": {}, + "scenarios": [ + { + "name": "c1_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + } + ] +} \ No newline at end of file diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-152532/tp16_dp1/scenarios.tsv b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-152532/tp16_dp1/scenarios.tsv new file mode 100644 index 0000000..70b2ee8 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-152532/tp16_dp1/scenarios.tsv @@ -0,0 +1,79 @@ +mark input_len output_len concurrency num_prompts +Y 1024 128 1 5 +Y 1024 128 128 640 +Y 1024 256 1 5 +Y 1024 256 128 640 +Y 1024 512 1 5 +Y 1024 512 128 640 +Y 1024 1024 1 5 +Y 1024 1024 128 640 +Y 1024 2048 1 5 +Y 1024 2048 128 640 +Y 1024 4096 1 5 +Y 1024 4096 128 640 +Y 4096 128 1 5 +Y 4096 128 64 320 +Y 4096 256 1 5 +Y 4096 256 64 320 +Y 4096 512 1 5 +Y 4096 512 64 320 +Y 4096 1024 1 5 +Y 4096 1024 64 320 +Y 4096 2048 1 5 +Y 4096 2048 64 320 +Y 4096 4096 1 5 +Y 4096 4096 64 320 +Y 8192 128 1 5 +Y 8192 128 64 320 +Y 8192 256 1 5 +Y 8192 256 64 320 +Y 8192 512 1 5 +Y 8192 512 64 320 +Y 8192 1024 1 5 +Y 8192 1024 64 320 +Y 8192 2048 1 5 +Y 8192 2048 64 320 +Y 8192 4096 1 5 +Y 8192 4096 64 320 +Y 16384 128 1 5 +Y 16384 128 32 160 +Y 16384 256 1 5 +Y 16384 256 32 160 +Y 16384 512 1 5 +Y 16384 512 32 160 +Y 16384 1024 1 5 +Y 16384 1024 32 160 +Y 16384 2048 1 5 +Y 16384 2048 32 160 +Y 32768 128 1 5 +Y 32768 128 16 80 +Y 32768 256 1 5 +Y 32768 256 16 80 +Y 32768 512 1 5 +Y 32768 512 16 80 +Y 32768 1024 1 5 +Y 32768 1024 16 80 +Y 32768 2048 1 5 +Y 32768 2048 16 80 +Y 65536 128 1 5 +Y 65536 128 8 40 +Y 65536 256 1 5 +Y 65536 256 8 40 +Y 65536 512 1 5 +Y 65536 512 8 40 +Y 65536 1024 1 5 +Y 65536 1024 8 40 +Y 131072 128 1 5 +Y 131072 128 4 20 +Y 131072 256 1 5 +Y 131072 256 4 20 +Y 131072 512 1 5 +Y 131072 512 4 20 +Y 262144 128 1 5 +Y 262144 128 2 10 +Y 262144 256 1 5 +Y 262144 256 2 10 +Y 524288 128 1 5 +Y 524288 128 2 10 +Y 1048576 128 1 5 +Y 1048576 128 2 10 diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-153034/comparison.md b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-153034/comparison.md new file mode 100644 index 0000000..bdd9237 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-153034/comparison.md @@ -0,0 +1,20 @@ +# SGLang TP×DP matrix comparison (8x NVIDIA RTX 6000D 84GB) + +## Summary + +- Model: `/data/hf_models/GLM-5.2-FP8` +- Hardware: 8x NVIDIA RTX 6000D 84GB +- Backend: SGLang (Docker) +- Benchmark client: `sglang.benchmark.serving` +- SLO reference: TTFT P95 < 3000.0ms, TPOT mean < 50.0ms + +### Configurations + +| Config | TP | DP | GPUs/replica | Notes | +|---|---:|---:|---:|---| +| tp16_dp1 | ? | ? | ? | server args recorded per ISL in results.json | + +## Side-by-side results + +| Scenario | ISL | DSL | Config | Conc | Req/s | OutTok/s | TTFT P95(ms) | TTFT P99(ms) | TPOT Mean(ms) | TPOT P95(ms) | TPOT P99(ms) | E2E P99(ms) | Peak GPU mem | SLO | +|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---| diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-153034/tp16_dp1/results.json b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-153034/tp16_dp1/results.json new file mode 100644 index 0000000..df653ee --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-153034/tp16_dp1/results.json @@ -0,0 +1,489 @@ +{ + "metadata": { + "experiment": "glm52_pro6000_sglang_multinode_tp16_tp16_dp1", + "run_id": "20260724-153034", + "timestamp": "2026-07-24T15:30:34+08:00", + "model": "/data/hf_models/GLM-5.2-FP8", + "backend": "sglang", + "engine": "sglang", + "hardware": "8x NVIDIA RTX 6000D 84GB", + "accelerator": "NVIDIA RTX 6000D", + "chip": "nvidia_rtx6000d", + "script": "experiments/glm52_pro6000_sglang_multinode_tp16/run_bench.sh", + "env": "lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45", + "git_commit": "d13f61f", + "git_dirty": "dirty", + "description": "H200 SGLang TP×DP matrix for DeepSeek-V4-Flash on RTX 6000D" + }, + "config": {}, + "scenarios": [ + { + "name": "c1_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + } + ] +} \ No newline at end of file diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-153034/tp16_dp1/scenarios.tsv b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-153034/tp16_dp1/scenarios.tsv new file mode 100644 index 0000000..70b2ee8 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-153034/tp16_dp1/scenarios.tsv @@ -0,0 +1,79 @@ +mark input_len output_len concurrency num_prompts +Y 1024 128 1 5 +Y 1024 128 128 640 +Y 1024 256 1 5 +Y 1024 256 128 640 +Y 1024 512 1 5 +Y 1024 512 128 640 +Y 1024 1024 1 5 +Y 1024 1024 128 640 +Y 1024 2048 1 5 +Y 1024 2048 128 640 +Y 1024 4096 1 5 +Y 1024 4096 128 640 +Y 4096 128 1 5 +Y 4096 128 64 320 +Y 4096 256 1 5 +Y 4096 256 64 320 +Y 4096 512 1 5 +Y 4096 512 64 320 +Y 4096 1024 1 5 +Y 4096 1024 64 320 +Y 4096 2048 1 5 +Y 4096 2048 64 320 +Y 4096 4096 1 5 +Y 4096 4096 64 320 +Y 8192 128 1 5 +Y 8192 128 64 320 +Y 8192 256 1 5 +Y 8192 256 64 320 +Y 8192 512 1 5 +Y 8192 512 64 320 +Y 8192 1024 1 5 +Y 8192 1024 64 320 +Y 8192 2048 1 5 +Y 8192 2048 64 320 +Y 8192 4096 1 5 +Y 8192 4096 64 320 +Y 16384 128 1 5 +Y 16384 128 32 160 +Y 16384 256 1 5 +Y 16384 256 32 160 +Y 16384 512 1 5 +Y 16384 512 32 160 +Y 16384 1024 1 5 +Y 16384 1024 32 160 +Y 16384 2048 1 5 +Y 16384 2048 32 160 +Y 32768 128 1 5 +Y 32768 128 16 80 +Y 32768 256 1 5 +Y 32768 256 16 80 +Y 32768 512 1 5 +Y 32768 512 16 80 +Y 32768 1024 1 5 +Y 32768 1024 16 80 +Y 32768 2048 1 5 +Y 32768 2048 16 80 +Y 65536 128 1 5 +Y 65536 128 8 40 +Y 65536 256 1 5 +Y 65536 256 8 40 +Y 65536 512 1 5 +Y 65536 512 8 40 +Y 65536 1024 1 5 +Y 65536 1024 8 40 +Y 131072 128 1 5 +Y 131072 128 4 20 +Y 131072 256 1 5 +Y 131072 256 4 20 +Y 131072 512 1 5 +Y 131072 512 4 20 +Y 262144 128 1 5 +Y 262144 128 2 10 +Y 262144 256 1 5 +Y 262144 256 2 10 +Y 524288 128 1 5 +Y 524288 128 2 10 +Y 1048576 128 1 5 +Y 1048576 128 2 10 diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-154223/comparison.md b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-154223/comparison.md new file mode 100644 index 0000000..bdd9237 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-154223/comparison.md @@ -0,0 +1,20 @@ +# SGLang TP×DP matrix comparison (8x NVIDIA RTX 6000D 84GB) + +## Summary + +- Model: `/data/hf_models/GLM-5.2-FP8` +- Hardware: 8x NVIDIA RTX 6000D 84GB +- Backend: SGLang (Docker) +- Benchmark client: `sglang.benchmark.serving` +- SLO reference: TTFT P95 < 3000.0ms, TPOT mean < 50.0ms + +### Configurations + +| Config | TP | DP | GPUs/replica | Notes | +|---|---:|---:|---:|---| +| tp16_dp1 | ? | ? | ? | server args recorded per ISL in results.json | + +## Side-by-side results + +| Scenario | ISL | DSL | Config | Conc | Req/s | OutTok/s | TTFT P95(ms) | TTFT P99(ms) | TPOT Mean(ms) | TPOT P95(ms) | TPOT P99(ms) | E2E P99(ms) | Peak GPU mem | SLO | +|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---| diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-154223/tp16_dp1/results.json b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-154223/tp16_dp1/results.json new file mode 100644 index 0000000..07a6eba --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-154223/tp16_dp1/results.json @@ -0,0 +1,489 @@ +{ + "metadata": { + "experiment": "glm52_pro6000_sglang_multinode_tp16_tp16_dp1", + "run_id": "20260724-154223", + "timestamp": "2026-07-24T15:42:24+08:00", + "model": "/data/hf_models/GLM-5.2-FP8", + "backend": "sglang", + "engine": "sglang", + "hardware": "8x NVIDIA RTX 6000D 84GB", + "accelerator": "NVIDIA RTX 6000D", + "chip": "nvidia_rtx6000d", + "script": "experiments/glm52_pro6000_sglang_multinode_tp16/run_bench.sh", + "env": "lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45", + "git_commit": "d13f61f", + "git_dirty": "dirty", + "description": "H200 SGLang TP×DP matrix for DeepSeek-V4-Flash on RTX 6000D" + }, + "config": {}, + "scenarios": [ + { + "name": "c1_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + } + ] +} \ No newline at end of file diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-154223/tp16_dp1/scenarios.tsv b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-154223/tp16_dp1/scenarios.tsv new file mode 100644 index 0000000..70b2ee8 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-154223/tp16_dp1/scenarios.tsv @@ -0,0 +1,79 @@ +mark input_len output_len concurrency num_prompts +Y 1024 128 1 5 +Y 1024 128 128 640 +Y 1024 256 1 5 +Y 1024 256 128 640 +Y 1024 512 1 5 +Y 1024 512 128 640 +Y 1024 1024 1 5 +Y 1024 1024 128 640 +Y 1024 2048 1 5 +Y 1024 2048 128 640 +Y 1024 4096 1 5 +Y 1024 4096 128 640 +Y 4096 128 1 5 +Y 4096 128 64 320 +Y 4096 256 1 5 +Y 4096 256 64 320 +Y 4096 512 1 5 +Y 4096 512 64 320 +Y 4096 1024 1 5 +Y 4096 1024 64 320 +Y 4096 2048 1 5 +Y 4096 2048 64 320 +Y 4096 4096 1 5 +Y 4096 4096 64 320 +Y 8192 128 1 5 +Y 8192 128 64 320 +Y 8192 256 1 5 +Y 8192 256 64 320 +Y 8192 512 1 5 +Y 8192 512 64 320 +Y 8192 1024 1 5 +Y 8192 1024 64 320 +Y 8192 2048 1 5 +Y 8192 2048 64 320 +Y 8192 4096 1 5 +Y 8192 4096 64 320 +Y 16384 128 1 5 +Y 16384 128 32 160 +Y 16384 256 1 5 +Y 16384 256 32 160 +Y 16384 512 1 5 +Y 16384 512 32 160 +Y 16384 1024 1 5 +Y 16384 1024 32 160 +Y 16384 2048 1 5 +Y 16384 2048 32 160 +Y 32768 128 1 5 +Y 32768 128 16 80 +Y 32768 256 1 5 +Y 32768 256 16 80 +Y 32768 512 1 5 +Y 32768 512 16 80 +Y 32768 1024 1 5 +Y 32768 1024 16 80 +Y 32768 2048 1 5 +Y 32768 2048 16 80 +Y 65536 128 1 5 +Y 65536 128 8 40 +Y 65536 256 1 5 +Y 65536 256 8 40 +Y 65536 512 1 5 +Y 65536 512 8 40 +Y 65536 1024 1 5 +Y 65536 1024 8 40 +Y 131072 128 1 5 +Y 131072 128 4 20 +Y 131072 256 1 5 +Y 131072 256 4 20 +Y 131072 512 1 5 +Y 131072 512 4 20 +Y 262144 128 1 5 +Y 262144 128 2 10 +Y 262144 256 1 5 +Y 262144 256 2 10 +Y 524288 128 1 5 +Y 524288 128 2 10 +Y 1048576 128 1 5 +Y 1048576 128 2 10 diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-154918/tp16_dp1/results.json b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-154918/tp16_dp1/results.json new file mode 100644 index 0000000..50519c7 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-154918/tp16_dp1/results.json @@ -0,0 +1,20 @@ +{ + "metadata": { + "experiment": "glm52_pro6000_sglang_multinode_tp16_tp16_dp1", + "run_id": "20260724-154918", + "timestamp": "2026-07-24T15:49:19+08:00", + "model": "/data/hf_models/GLM-5.2-FP8", + "backend": "sglang", + "engine": "sglang", + "hardware": "8x NVIDIA RTX 6000D 84GB", + "accelerator": "NVIDIA RTX 6000D", + "chip": "nvidia_rtx6000d", + "script": "experiments/glm52_pro6000_sglang_multinode_tp16/run_bench.sh", + "env": "lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45", + "git_commit": "d13f61f", + "git_dirty": "dirty", + "description": "H200 SGLang TP×DP matrix for DeepSeek-V4-Flash on RTX 6000D" + }, + "config": {}, + "scenarios": [] +} \ No newline at end of file diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-154918/tp16_dp1/scenarios.tsv b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-154918/tp16_dp1/scenarios.tsv new file mode 100644 index 0000000..70b2ee8 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-154918/tp16_dp1/scenarios.tsv @@ -0,0 +1,79 @@ +mark input_len output_len concurrency num_prompts +Y 1024 128 1 5 +Y 1024 128 128 640 +Y 1024 256 1 5 +Y 1024 256 128 640 +Y 1024 512 1 5 +Y 1024 512 128 640 +Y 1024 1024 1 5 +Y 1024 1024 128 640 +Y 1024 2048 1 5 +Y 1024 2048 128 640 +Y 1024 4096 1 5 +Y 1024 4096 128 640 +Y 4096 128 1 5 +Y 4096 128 64 320 +Y 4096 256 1 5 +Y 4096 256 64 320 +Y 4096 512 1 5 +Y 4096 512 64 320 +Y 4096 1024 1 5 +Y 4096 1024 64 320 +Y 4096 2048 1 5 +Y 4096 2048 64 320 +Y 4096 4096 1 5 +Y 4096 4096 64 320 +Y 8192 128 1 5 +Y 8192 128 64 320 +Y 8192 256 1 5 +Y 8192 256 64 320 +Y 8192 512 1 5 +Y 8192 512 64 320 +Y 8192 1024 1 5 +Y 8192 1024 64 320 +Y 8192 2048 1 5 +Y 8192 2048 64 320 +Y 8192 4096 1 5 +Y 8192 4096 64 320 +Y 16384 128 1 5 +Y 16384 128 32 160 +Y 16384 256 1 5 +Y 16384 256 32 160 +Y 16384 512 1 5 +Y 16384 512 32 160 +Y 16384 1024 1 5 +Y 16384 1024 32 160 +Y 16384 2048 1 5 +Y 16384 2048 32 160 +Y 32768 128 1 5 +Y 32768 128 16 80 +Y 32768 256 1 5 +Y 32768 256 16 80 +Y 32768 512 1 5 +Y 32768 512 16 80 +Y 32768 1024 1 5 +Y 32768 1024 16 80 +Y 32768 2048 1 5 +Y 32768 2048 16 80 +Y 65536 128 1 5 +Y 65536 128 8 40 +Y 65536 256 1 5 +Y 65536 256 8 40 +Y 65536 512 1 5 +Y 65536 512 8 40 +Y 65536 1024 1 5 +Y 65536 1024 8 40 +Y 131072 128 1 5 +Y 131072 128 4 20 +Y 131072 256 1 5 +Y 131072 256 4 20 +Y 131072 512 1 5 +Y 131072 512 4 20 +Y 262144 128 1 5 +Y 262144 128 2 10 +Y 262144 256 1 5 +Y 262144 256 2 10 +Y 524288 128 1 5 +Y 524288 128 2 10 +Y 1048576 128 1 5 +Y 1048576 128 2 10 diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-160131/tp16_dp1/results.json b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-160131/tp16_dp1/results.json new file mode 100644 index 0000000..a6c4dae --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-160131/tp16_dp1/results.json @@ -0,0 +1,20 @@ +{ + "metadata": { + "experiment": "glm52_pro6000_sglang_multinode_tp16_tp16_dp1", + "run_id": "20260724-160131", + "timestamp": "2026-07-24T16:01:33+08:00", + "model": "/data/hf_models/GLM-5.2-FP8", + "backend": "sglang", + "engine": "sglang", + "hardware": "8x NVIDIA RTX 6000D 84GB", + "accelerator": "NVIDIA RTX 6000D", + "chip": "nvidia_rtx6000d", + "script": "experiments/glm52_pro6000_sglang_multinode_tp16/run_bench.sh", + "env": "lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45", + "git_commit": "d13f61f", + "git_dirty": "dirty", + "description": "H200 SGLang TP×DP matrix for DeepSeek-V4-Flash on RTX 6000D" + }, + "config": {}, + "scenarios": [] +} \ No newline at end of file diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-160131/tp16_dp1/scenarios.tsv b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-160131/tp16_dp1/scenarios.tsv new file mode 100644 index 0000000..70b2ee8 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-160131/tp16_dp1/scenarios.tsv @@ -0,0 +1,79 @@ +mark input_len output_len concurrency num_prompts +Y 1024 128 1 5 +Y 1024 128 128 640 +Y 1024 256 1 5 +Y 1024 256 128 640 +Y 1024 512 1 5 +Y 1024 512 128 640 +Y 1024 1024 1 5 +Y 1024 1024 128 640 +Y 1024 2048 1 5 +Y 1024 2048 128 640 +Y 1024 4096 1 5 +Y 1024 4096 128 640 +Y 4096 128 1 5 +Y 4096 128 64 320 +Y 4096 256 1 5 +Y 4096 256 64 320 +Y 4096 512 1 5 +Y 4096 512 64 320 +Y 4096 1024 1 5 +Y 4096 1024 64 320 +Y 4096 2048 1 5 +Y 4096 2048 64 320 +Y 4096 4096 1 5 +Y 4096 4096 64 320 +Y 8192 128 1 5 +Y 8192 128 64 320 +Y 8192 256 1 5 +Y 8192 256 64 320 +Y 8192 512 1 5 +Y 8192 512 64 320 +Y 8192 1024 1 5 +Y 8192 1024 64 320 +Y 8192 2048 1 5 +Y 8192 2048 64 320 +Y 8192 4096 1 5 +Y 8192 4096 64 320 +Y 16384 128 1 5 +Y 16384 128 32 160 +Y 16384 256 1 5 +Y 16384 256 32 160 +Y 16384 512 1 5 +Y 16384 512 32 160 +Y 16384 1024 1 5 +Y 16384 1024 32 160 +Y 16384 2048 1 5 +Y 16384 2048 32 160 +Y 32768 128 1 5 +Y 32768 128 16 80 +Y 32768 256 1 5 +Y 32768 256 16 80 +Y 32768 512 1 5 +Y 32768 512 16 80 +Y 32768 1024 1 5 +Y 32768 1024 16 80 +Y 32768 2048 1 5 +Y 32768 2048 16 80 +Y 65536 128 1 5 +Y 65536 128 8 40 +Y 65536 256 1 5 +Y 65536 256 8 40 +Y 65536 512 1 5 +Y 65536 512 8 40 +Y 65536 1024 1 5 +Y 65536 1024 8 40 +Y 131072 128 1 5 +Y 131072 128 4 20 +Y 131072 256 1 5 +Y 131072 256 4 20 +Y 131072 512 1 5 +Y 131072 512 4 20 +Y 262144 128 1 5 +Y 262144 128 2 10 +Y 262144 256 1 5 +Y 262144 256 2 10 +Y 524288 128 1 5 +Y 524288 128 2 10 +Y 1048576 128 1 5 +Y 1048576 128 2 10 diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-161611/comparison.md b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-161611/comparison.md new file mode 100644 index 0000000..bdd9237 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-161611/comparison.md @@ -0,0 +1,20 @@ +# SGLang TP×DP matrix comparison (8x NVIDIA RTX 6000D 84GB) + +## Summary + +- Model: `/data/hf_models/GLM-5.2-FP8` +- Hardware: 8x NVIDIA RTX 6000D 84GB +- Backend: SGLang (Docker) +- Benchmark client: `sglang.benchmark.serving` +- SLO reference: TTFT P95 < 3000.0ms, TPOT mean < 50.0ms + +### Configurations + +| Config | TP | DP | GPUs/replica | Notes | +|---|---:|---:|---:|---| +| tp16_dp1 | ? | ? | ? | server args recorded per ISL in results.json | + +## Side-by-side results + +| Scenario | ISL | DSL | Config | Conc | Req/s | OutTok/s | TTFT P95(ms) | TTFT P99(ms) | TPOT Mean(ms) | TPOT P95(ms) | TPOT P99(ms) | E2E P99(ms) | Peak GPU mem | SLO | +|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---| diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-161611/tp16_dp1/results.json b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-161611/tp16_dp1/results.json new file mode 100644 index 0000000..f3eab9a --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-161611/tp16_dp1/results.json @@ -0,0 +1,489 @@ +{ + "metadata": { + "experiment": "glm52_pro6000_sglang_multinode_tp16_tp16_dp1", + "run_id": "20260724-161611", + "timestamp": "2026-07-24T16:16:12+08:00", + "model": "/data/hf_models/GLM-5.2-FP8", + "backend": "sglang", + "engine": "sglang", + "hardware": "8x NVIDIA RTX 6000D 84GB", + "accelerator": "NVIDIA RTX 6000D", + "chip": "nvidia_rtx6000d", + "script": "experiments/glm52_pro6000_sglang_multinode_tp16/run_bench.sh", + "env": "lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45", + "git_commit": "d13f61f", + "git_dirty": "dirty", + "description": "H200 SGLang TP×DP matrix for DeepSeek-V4-Flash on RTX 6000D" + }, + "config": {}, + "scenarios": [ + { + "name": "c1_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + } + ] +} \ No newline at end of file diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-161611/tp16_dp1/scenarios.tsv b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-161611/tp16_dp1/scenarios.tsv new file mode 100644 index 0000000..70b2ee8 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-161611/tp16_dp1/scenarios.tsv @@ -0,0 +1,79 @@ +mark input_len output_len concurrency num_prompts +Y 1024 128 1 5 +Y 1024 128 128 640 +Y 1024 256 1 5 +Y 1024 256 128 640 +Y 1024 512 1 5 +Y 1024 512 128 640 +Y 1024 1024 1 5 +Y 1024 1024 128 640 +Y 1024 2048 1 5 +Y 1024 2048 128 640 +Y 1024 4096 1 5 +Y 1024 4096 128 640 +Y 4096 128 1 5 +Y 4096 128 64 320 +Y 4096 256 1 5 +Y 4096 256 64 320 +Y 4096 512 1 5 +Y 4096 512 64 320 +Y 4096 1024 1 5 +Y 4096 1024 64 320 +Y 4096 2048 1 5 +Y 4096 2048 64 320 +Y 4096 4096 1 5 +Y 4096 4096 64 320 +Y 8192 128 1 5 +Y 8192 128 64 320 +Y 8192 256 1 5 +Y 8192 256 64 320 +Y 8192 512 1 5 +Y 8192 512 64 320 +Y 8192 1024 1 5 +Y 8192 1024 64 320 +Y 8192 2048 1 5 +Y 8192 2048 64 320 +Y 8192 4096 1 5 +Y 8192 4096 64 320 +Y 16384 128 1 5 +Y 16384 128 32 160 +Y 16384 256 1 5 +Y 16384 256 32 160 +Y 16384 512 1 5 +Y 16384 512 32 160 +Y 16384 1024 1 5 +Y 16384 1024 32 160 +Y 16384 2048 1 5 +Y 16384 2048 32 160 +Y 32768 128 1 5 +Y 32768 128 16 80 +Y 32768 256 1 5 +Y 32768 256 16 80 +Y 32768 512 1 5 +Y 32768 512 16 80 +Y 32768 1024 1 5 +Y 32768 1024 16 80 +Y 32768 2048 1 5 +Y 32768 2048 16 80 +Y 65536 128 1 5 +Y 65536 128 8 40 +Y 65536 256 1 5 +Y 65536 256 8 40 +Y 65536 512 1 5 +Y 65536 512 8 40 +Y 65536 1024 1 5 +Y 65536 1024 8 40 +Y 131072 128 1 5 +Y 131072 128 4 20 +Y 131072 256 1 5 +Y 131072 256 4 20 +Y 131072 512 1 5 +Y 131072 512 4 20 +Y 262144 128 1 5 +Y 262144 128 2 10 +Y 262144 256 1 5 +Y 262144 256 2 10 +Y 524288 128 1 5 +Y 524288 128 2 10 +Y 1048576 128 1 5 +Y 1048576 128 2 10 diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-161700/comparison.md b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-161700/comparison.md new file mode 100644 index 0000000..bdd9237 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-161700/comparison.md @@ -0,0 +1,20 @@ +# SGLang TP×DP matrix comparison (8x NVIDIA RTX 6000D 84GB) + +## Summary + +- Model: `/data/hf_models/GLM-5.2-FP8` +- Hardware: 8x NVIDIA RTX 6000D 84GB +- Backend: SGLang (Docker) +- Benchmark client: `sglang.benchmark.serving` +- SLO reference: TTFT P95 < 3000.0ms, TPOT mean < 50.0ms + +### Configurations + +| Config | TP | DP | GPUs/replica | Notes | +|---|---:|---:|---:|---| +| tp16_dp1 | ? | ? | ? | server args recorded per ISL in results.json | + +## Side-by-side results + +| Scenario | ISL | DSL | Config | Conc | Req/s | OutTok/s | TTFT P95(ms) | TTFT P99(ms) | TPOT Mean(ms) | TPOT P95(ms) | TPOT P99(ms) | E2E P99(ms) | Peak GPU mem | SLO | +|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---| diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-161700/tp16_dp1/results.json b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-161700/tp16_dp1/results.json new file mode 100644 index 0000000..90297eb --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-161700/tp16_dp1/results.json @@ -0,0 +1,489 @@ +{ + "metadata": { + "experiment": "glm52_pro6000_sglang_multinode_tp16_tp16_dp1", + "run_id": "20260724-161700", + "timestamp": "2026-07-24T16:17:00+08:00", + "model": "/data/hf_models/GLM-5.2-FP8", + "backend": "sglang", + "engine": "sglang", + "hardware": "8x NVIDIA RTX 6000D 84GB", + "accelerator": "NVIDIA RTX 6000D", + "chip": "nvidia_rtx6000d", + "script": "experiments/glm52_pro6000_sglang_multinode_tp16/run_bench.sh", + "env": "lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45", + "git_commit": "d13f61f", + "git_dirty": "dirty", + "description": "H200 SGLang TP×DP matrix for DeepSeek-V4-Flash on RTX 6000D" + }, + "config": {}, + "scenarios": [ + { + "name": "c1_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + } + ] +} \ No newline at end of file diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-161700/tp16_dp1/scenarios.tsv b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-161700/tp16_dp1/scenarios.tsv new file mode 100644 index 0000000..70b2ee8 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-161700/tp16_dp1/scenarios.tsv @@ -0,0 +1,79 @@ +mark input_len output_len concurrency num_prompts +Y 1024 128 1 5 +Y 1024 128 128 640 +Y 1024 256 1 5 +Y 1024 256 128 640 +Y 1024 512 1 5 +Y 1024 512 128 640 +Y 1024 1024 1 5 +Y 1024 1024 128 640 +Y 1024 2048 1 5 +Y 1024 2048 128 640 +Y 1024 4096 1 5 +Y 1024 4096 128 640 +Y 4096 128 1 5 +Y 4096 128 64 320 +Y 4096 256 1 5 +Y 4096 256 64 320 +Y 4096 512 1 5 +Y 4096 512 64 320 +Y 4096 1024 1 5 +Y 4096 1024 64 320 +Y 4096 2048 1 5 +Y 4096 2048 64 320 +Y 4096 4096 1 5 +Y 4096 4096 64 320 +Y 8192 128 1 5 +Y 8192 128 64 320 +Y 8192 256 1 5 +Y 8192 256 64 320 +Y 8192 512 1 5 +Y 8192 512 64 320 +Y 8192 1024 1 5 +Y 8192 1024 64 320 +Y 8192 2048 1 5 +Y 8192 2048 64 320 +Y 8192 4096 1 5 +Y 8192 4096 64 320 +Y 16384 128 1 5 +Y 16384 128 32 160 +Y 16384 256 1 5 +Y 16384 256 32 160 +Y 16384 512 1 5 +Y 16384 512 32 160 +Y 16384 1024 1 5 +Y 16384 1024 32 160 +Y 16384 2048 1 5 +Y 16384 2048 32 160 +Y 32768 128 1 5 +Y 32768 128 16 80 +Y 32768 256 1 5 +Y 32768 256 16 80 +Y 32768 512 1 5 +Y 32768 512 16 80 +Y 32768 1024 1 5 +Y 32768 1024 16 80 +Y 32768 2048 1 5 +Y 32768 2048 16 80 +Y 65536 128 1 5 +Y 65536 128 8 40 +Y 65536 256 1 5 +Y 65536 256 8 40 +Y 65536 512 1 5 +Y 65536 512 8 40 +Y 65536 1024 1 5 +Y 65536 1024 8 40 +Y 131072 128 1 5 +Y 131072 128 4 20 +Y 131072 256 1 5 +Y 131072 256 4 20 +Y 131072 512 1 5 +Y 131072 512 4 20 +Y 262144 128 1 5 +Y 262144 128 2 10 +Y 262144 256 1 5 +Y 262144 256 2 10 +Y 524288 128 1 5 +Y 524288 128 2 10 +Y 1048576 128 1 5 +Y 1048576 128 2 10 diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-164734/comparison.md b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-164734/comparison.md new file mode 100644 index 0000000..bdd9237 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-164734/comparison.md @@ -0,0 +1,20 @@ +# SGLang TP×DP matrix comparison (8x NVIDIA RTX 6000D 84GB) + +## Summary + +- Model: `/data/hf_models/GLM-5.2-FP8` +- Hardware: 8x NVIDIA RTX 6000D 84GB +- Backend: SGLang (Docker) +- Benchmark client: `sglang.benchmark.serving` +- SLO reference: TTFT P95 < 3000.0ms, TPOT mean < 50.0ms + +### Configurations + +| Config | TP | DP | GPUs/replica | Notes | +|---|---:|---:|---:|---| +| tp16_dp1 | ? | ? | ? | server args recorded per ISL in results.json | + +## Side-by-side results + +| Scenario | ISL | DSL | Config | Conc | Req/s | OutTok/s | TTFT P95(ms) | TTFT P99(ms) | TPOT Mean(ms) | TPOT P95(ms) | TPOT P99(ms) | E2E P99(ms) | Peak GPU mem | SLO | +|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---| diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-164734/tp16_dp1/results.json b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-164734/tp16_dp1/results.json new file mode 100644 index 0000000..7e4588f --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-164734/tp16_dp1/results.json @@ -0,0 +1,489 @@ +{ + "metadata": { + "experiment": "glm52_pro6000_sglang_multinode_tp16_tp16_dp1", + "run_id": "20260724-164734", + "timestamp": "2026-07-24T16:47:35+08:00", + "model": "/data/hf_models/GLM-5.2-FP8", + "backend": "sglang", + "engine": "sglang", + "hardware": "8x NVIDIA RTX 6000D 84GB", + "accelerator": "NVIDIA RTX 6000D", + "chip": "nvidia_rtx6000d", + "script": "experiments/glm52_pro6000_sglang_multinode_tp16/run_bench.sh", + "env": "lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45", + "git_commit": "d13f61f", + "git_dirty": "dirty", + "description": "H200 SGLang TP×DP matrix for DeepSeek-V4-Flash on RTX 6000D" + }, + "config": {}, + "scenarios": [ + { + "name": "c1_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c128_i1024_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i4096_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c64_i8192_o4096", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c32_i16384_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c16_i32768_o2048", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c8_i65536_o1024", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c4_i131072_o512", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i262144_o256", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i524288_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c1_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + }, + { + "name": "c2_i1048576_o128", + "config": "", + "status": "SKIPPED_SERVICE_START_FAILED", + "note": "service failed to start" + } + ] +} \ No newline at end of file diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-164734/tp16_dp1/scenarios.tsv b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-164734/tp16_dp1/scenarios.tsv new file mode 100644 index 0000000..70b2ee8 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/results/20260724-164734/tp16_dp1/scenarios.tsv @@ -0,0 +1,79 @@ +mark input_len output_len concurrency num_prompts +Y 1024 128 1 5 +Y 1024 128 128 640 +Y 1024 256 1 5 +Y 1024 256 128 640 +Y 1024 512 1 5 +Y 1024 512 128 640 +Y 1024 1024 1 5 +Y 1024 1024 128 640 +Y 1024 2048 1 5 +Y 1024 2048 128 640 +Y 1024 4096 1 5 +Y 1024 4096 128 640 +Y 4096 128 1 5 +Y 4096 128 64 320 +Y 4096 256 1 5 +Y 4096 256 64 320 +Y 4096 512 1 5 +Y 4096 512 64 320 +Y 4096 1024 1 5 +Y 4096 1024 64 320 +Y 4096 2048 1 5 +Y 4096 2048 64 320 +Y 4096 4096 1 5 +Y 4096 4096 64 320 +Y 8192 128 1 5 +Y 8192 128 64 320 +Y 8192 256 1 5 +Y 8192 256 64 320 +Y 8192 512 1 5 +Y 8192 512 64 320 +Y 8192 1024 1 5 +Y 8192 1024 64 320 +Y 8192 2048 1 5 +Y 8192 2048 64 320 +Y 8192 4096 1 5 +Y 8192 4096 64 320 +Y 16384 128 1 5 +Y 16384 128 32 160 +Y 16384 256 1 5 +Y 16384 256 32 160 +Y 16384 512 1 5 +Y 16384 512 32 160 +Y 16384 1024 1 5 +Y 16384 1024 32 160 +Y 16384 2048 1 5 +Y 16384 2048 32 160 +Y 32768 128 1 5 +Y 32768 128 16 80 +Y 32768 256 1 5 +Y 32768 256 16 80 +Y 32768 512 1 5 +Y 32768 512 16 80 +Y 32768 1024 1 5 +Y 32768 1024 16 80 +Y 32768 2048 1 5 +Y 32768 2048 16 80 +Y 65536 128 1 5 +Y 65536 128 8 40 +Y 65536 256 1 5 +Y 65536 256 8 40 +Y 65536 512 1 5 +Y 65536 512 8 40 +Y 65536 1024 1 5 +Y 65536 1024 8 40 +Y 131072 128 1 5 +Y 131072 128 4 20 +Y 131072 256 1 5 +Y 131072 256 4 20 +Y 131072 512 1 5 +Y 131072 512 4 20 +Y 262144 128 1 5 +Y 262144 128 2 10 +Y 262144 256 1 5 +Y 262144 256 2 10 +Y 524288 128 1 5 +Y 524288 128 2 10 +Y 1048576 128 1 5 +Y 1048576 128 2 10 diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/run_adaptive_concurrency.sh b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/run_adaptive_concurrency.sh new file mode 100755 index 0000000..8b26367 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/run_adaptive_concurrency.sh @@ -0,0 +1,154 @@ +#!/usr/bin/env bash +# Find the Total-TPS saturation concurrency for each SGLang TP/DP/ISL/OSL shape. +set -Eeuo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +EXPERIMENT_NAME="$(basename "$SCRIPT_DIR")" + +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/platform.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/config.env" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/adaptive_config.env" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/adaptive_bench_lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" + +DEPLOY_PROFILE="${DEPLOY_PROFILE:-pro6000/glm52_pro6000_sglang_tp16}" + +ENGINE="sglang" +ENGINE_PORT="$SGLANG_PORT" +RESULT_BASE="${RESULT_BASE:-${SCRIPT_DIR}/adaptive_results}" +ACTIVE_ENGINE_SERVER_LOG="" + +if [[ -x "${VENV_CLIENT}/bin/python" ]]; then + PYTHON="${VENV_CLIENT}/bin/python" +else + PYTHON="$(command -v python3)" +fi + +DOCKER_IMAGE="${DOCKER_IMAGE:-lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45}" + +engine_is_healthy() { + curl --fail --silent --show-error --max-time 5 \ + "http://127.0.0.1:${ENGINE_PORT}/health" >/dev/null 2>&1 +} + +engine_stop_server() { + local tp="$1" + local dp="$2" + + log "stopping sglang server tp=${tp} dp=${dp} via deploy profile" + bash "${SCRIPT_DIR}/stop_sglang_multinode.sh" "$tp" "$dp" >> "${ADAPTIVE_LOG_DIR}/sglang_tp${tp}_dp${dp}.server.outer.log" 2>&1 || true + ACTIVE_ENGINE_SERVER_LOG="" + sleep 2 +} + +engine_build_server_args() { + local tp="$1" + local dp="$2" + deploy_render_args "$DEPLOY_PROFILE" "$tp" "$dp" "$ENGINE_PORT" "$MODEL_PATH" +} + +engine_start_server() { + local tp="$1" + local dp="$2" + local outer_log="${ADAPTIVE_LOG_DIR}/sglang_tp${tp}_dp${dp}.server.outer.log" + log "starting sglang server tp=${tp} dp=${dp}" + if [[ -n "${CONTAINER_NAME:-}" ]]; then + bash "${SCRIPT_DIR}/run_sglang_in_container.sh" "$tp" "$dp" >> "$outer_log" 2>&1 + else + bash "${SCRIPT_DIR}/start_sglang_dp.sh" "$tp" "$dp" >> "$outer_log" 2>&1 + fi + if ! engine_is_healthy; then + log "ERROR: sglang health check failed tp=${tp} dp=${dp}" + return 1 + fi + if [[ -z "${CONTAINER_NAME:-}" ]]; then + ACTIVE_ENGINE_SERVER_LOG="$( + find "${RUNTIME_BASE}/logs" -maxdepth 1 -type f \ + -name "${EXPERIMENT}_sglang*tp${tp}_dp${dp}_*.log" \ + -printf '%T@ %p\n' 2>/dev/null | sort -nr | head -n 1 | cut -d' ' -f2- + )" + fi + log "sglang server healthy tp=${tp} dp=${dp} log=${ACTIVE_ENGINE_SERVER_LOG:-container:/tmp/sglang_server.log}" +} + +engine_detect_oom() { + local detail_log="$1" + local tp="$2" + local dp="$3" + local pattern='CUDA out of memory|torch\.OutOfMemoryError|OutOfMemory|out of memory|OOM|RESOURCE_EXHAUSTED|Failed to allocate memory' + local outer_log="${ADAPTIVE_LOG_DIR}/sglang_tp${tp}_dp${dp}.server.outer.log" + local -a logs=("$detail_log" "$outer_log") + if [[ -n "$ACTIVE_ENGINE_SERVER_LOG" ]]; then + logs+=("$ACTIVE_ENGINE_SERVER_LOG") + fi + if grep -Eiq "$pattern" "${logs[@]}" 2>/dev/null; then + return 0 + fi + if [[ -n "${CONTAINER_NAME:-}" ]]; then + docker exec "$CONTAINER_NAME" grep -Eiq "$pattern" /tmp/sglang_server.log 2>/dev/null + return $? + fi + return 1 +} + +engine_run_bench() { + local isl="$1" + local osl="$2" + local concurrency="$3" + local num_prompts="$4" + local output_file="$5" + local warmup_requests + warmup_requests="$(adaptive_warmup_request_count "$concurrency")" + local -a bench_args=( + --backend sglang + --host 127.0.0.1 + --port "$ENGINE_PORT" + --dataset-name "$BENCH_DATASET_NAME" + --random-input-len "$isl" + --random-output-len "$osl" + --random-range-ratio "$RANDOM_RANGE_RATIO" + --num-prompts "$num_prompts" + --max-concurrency "$concurrency" + --request-rate 10000 + --warmup-requests "$warmup_requests" + --output-file "$output_file" + --output-details + --disable-tqdm + ) + if [[ "$BENCH_DATASET_NAME" == "random" ]]; then + bench_args+=(--dataset-path "$DATASET_PATH") + elif [[ "$BENCH_DATASET_NAME" == "random-ids" ]]; then + : # random-ids does not need --tokenize-prompt + else + bench_args+=(--tokenize-prompt) + fi + + if [[ "$USE_DOCKER_CLIENT" == "1" ]]; then + local -a volume_args=(-v "${MODEL_PATH}:${MODEL_PATH}:ro" -v "${RESULT_BASE}:${RESULT_BASE}") + if [[ "$BENCH_DATASET_NAME" == "random" ]]; then + volume_args+=(-v "${DATASET_PATH}:${DATASET_PATH}:ro") + fi + docker run --rm \ + --network host \ + "${volume_args[@]}" \ + -e PYTHONUNBUFFERED=1 \ + --entrypoint python3 \ + "$DOCKER_IMAGE" \ + -m "$SGLANG_BENCH_MODULE" "${bench_args[@]}" + else + "$PYTHON" -m "$SGLANG_BENCH_MODULE" "${bench_args[@]}" + fi +} + +export -f engine_run_bench +export ENGINE_PORT MODEL_PATH RESULT_BASE DOCKER_IMAGE USE_DOCKER_CLIENT +export BENCH_DATASET_NAME DATASET_PATH RANDOM_RANGE_RATIO BENCH_WARMUP_MAX_REQUESTS PYTHON SGLANG_BENCH_MODULE + +adaptive_main "$@" diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/run_adaptive_concurrency_add16.sh b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/run_adaptive_concurrency_add16.sh new file mode 100755 index 0000000..75a19b4 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/run_adaptive_concurrency_add16.sh @@ -0,0 +1,163 @@ +#!/usr/bin/env bash +# Find the Total-TPS saturation concurrency for each SGLang TP/DP/ISL/OSL shape. +set -Eeuo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +EXPERIMENT_NAME="$(basename "$SCRIPT_DIR")" + +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/platform.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/config.env" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/adaptive_config.env" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/adaptive_bench_lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" + +DEPLOY_PROFILE="${DEPLOY_PROFILE:-pro6000/glm52_pro6000_sglang_tp16}" + +ENGINE="sglang" +ENGINE_PORT="$SGLANG_PORT" +RESULT_BASE="${RESULT_BASE:-${SCRIPT_DIR}/adaptive_results}" +ACTIVE_ENGINE_SERVER_LOG="" + +if [[ -x "${VENV_CLIENT}/bin/python" ]]; then + PYTHON="${VENV_CLIENT}/bin/python" +else + PYTHON="$(command -v python3)" +fi + +DOCKER_IMAGE="${DOCKER_IMAGE:-lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45}" + +engine_is_healthy() { + curl --fail --silent --show-error --max-time 5 \ + "http://127.0.0.1:${ENGINE_PORT}/health" >/dev/null 2>&1 +} + +engine_stop_server() { + local tp="$1" + local dp="$2" + + log "stopping sglang server tp=${tp} dp=${dp} via deploy profile" + bash "${SCRIPT_DIR}/stop_sglang_multinode.sh" "$tp" "$dp" >> "${ADAPTIVE_LOG_DIR}/sglang_tp${tp}_dp${dp}.server.outer.log" 2>&1 || true + ACTIVE_ENGINE_SERVER_LOG="" + sleep 2 +} + +engine_build_server_args() { + local tp="$1" + local dp="$2" + deploy_render_args "$DEPLOY_PROFILE" "$tp" "$dp" "$ENGINE_PORT" "$MODEL_PATH" +} + +engine_start_server() { + local tp="$1" + local dp="$2" + local outer_log="${ADAPTIVE_LOG_DIR}/sglang_tp${tp}_dp${dp}.server.outer.log" + log "starting sglang server tp=${tp} dp=${dp}" + if [[ -n "${CONTAINER_NAME:-}" ]]; then + bash "${SCRIPT_DIR}/run_sglang_in_container.sh" "$tp" "$dp" >> "$outer_log" 2>&1 + else + bash "${SCRIPT_DIR}/start_sglang_dp.sh" "$tp" "$dp" >> "$outer_log" 2>&1 + fi + if ! engine_is_healthy; then + log "ERROR: sglang health check failed tp=${tp} dp=${dp}" + return 1 + fi + if [[ -z "${CONTAINER_NAME:-}" ]]; then + ACTIVE_ENGINE_SERVER_LOG="$( + find "${RUNTIME_BASE}/logs" -maxdepth 1 -type f \ + -name "${EXPERIMENT}_sglang*tp${tp}_dp${dp}_*.log" \ + -printf '%T@ %p\n' 2>/dev/null | sort -nr | head -n 1 | cut -d' ' -f2- + )" + fi + log "sglang server healthy tp=${tp} dp=${dp} log=${ACTIVE_ENGINE_SERVER_LOG:-container:/tmp/sglang_server.log}" +} + +engine_detect_oom() { + local detail_log="$1" + local tp="$2" + local dp="$3" + local pattern='CUDA out of memory|torch\.OutOfMemoryError|OutOfMemory|out of memory|OOM|RESOURCE_EXHAUSTED|Failed to allocate memory' + local outer_log="${ADAPTIVE_LOG_DIR}/sglang_tp${tp}_dp${dp}.server.outer.log" + local -a logs=("$detail_log" "$outer_log") + if [[ -n "$ACTIVE_ENGINE_SERVER_LOG" ]]; then + logs+=("$ACTIVE_ENGINE_SERVER_LOG") + fi + if grep -Eiq "$pattern" "${logs[@]}" 2>/dev/null; then + return 0 + fi + if [[ -n "${CONTAINER_NAME:-}" ]]; then + docker exec "$CONTAINER_NAME" grep -Eiq "$pattern" /tmp/sglang_server.log 2>/dev/null + return $? + fi + return 1 +} + +engine_run_bench() { + local isl="$1" + local osl="$2" + local concurrency="$3" + local num_prompts="$4" + local output_file="$5" + local warmup_requests + warmup_requests="$(adaptive_warmup_request_count "$concurrency")" + local -a bench_args=( + --backend sglang + --host 127.0.0.1 + --port "$ENGINE_PORT" + --dataset-name "$BENCH_DATASET_NAME" + --random-input-len "$isl" + --random-output-len "$osl" + --random-range-ratio "$RANDOM_RANGE_RATIO" + --num-prompts "$num_prompts" + --max-concurrency "$concurrency" + --request-rate 10000 + --warmup-requests "$warmup_requests" + --output-file "$output_file" + --output-details + --disable-tqdm + ) + if [[ "$BENCH_DATASET_NAME" == "random" ]]; then + bench_args+=(--dataset-path "$DATASET_PATH") + elif [[ "$BENCH_DATASET_NAME" == "random-ids" ]]; then + : # random-ids does not need --tokenize-prompt + else + bench_args+=(--tokenize-prompt) + fi + + if [[ "$USE_DOCKER_CLIENT" == "1" ]]; then + local -a volume_args=(-v "${MODEL_PATH}:${MODEL_PATH}:ro" -v "${RESULT_BASE}:${RESULT_BASE}") + if [[ "$BENCH_DATASET_NAME" == "random" ]]; then + volume_args+=(-v "${DATASET_PATH}:${DATASET_PATH}:ro") + fi + docker run --rm \ + --network host \ + "${volume_args[@]}" \ + -e PYTHONUNBUFFERED=1 \ + --entrypoint python3 \ + "$DOCKER_IMAGE" \ + -m "$SGLANG_BENCH_MODULE" "${bench_args[@]}" + else + "$PYTHON" -m "$SGLANG_BENCH_MODULE" "${bench_args[@]}" + fi +} + +export -f engine_run_bench +export ENGINE_PORT MODEL_PATH RESULT_BASE DOCKER_IMAGE USE_DOCKER_CLIENT +export BENCH_DATASET_NAME DATASET_PATH RANDOM_RANGE_RATIO BENCH_WARMUP_MAX_REQUESTS PYTHON SGLANG_BENCH_MODULE + +adaptive_main "$@" + +export SEARCH_START_CONCURRENCY=16 +export SEARCH_ADDEND=16 +# If the initial concurrency violates the TTFT SLO, search downward. Stop at +# the first acceptable value (16 -> 8; only try 1 when 8 still violates it). +export SEARCH_INITIAL_BACKOFF_CONCURRENCIES="8 1" +# When concurrency 1 still has a severely excessive TTFT, stop the remaining +# shapes in this TP/DP group. Zero disables this rule. +export TTFT_GROUP_SKIP_MS="${TTFT_GROUP_SKIP_MS:-8000}" diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/run_bench.sh b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/run_bench.sh new file mode 100755 index 0000000..f505055 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/run_bench.sh @@ -0,0 +1,536 @@ +#!/usr/bin/env bash +# TP×DP matrix benchmark for DeepSeek-V4-Flash on SGLang (Docker). +set -Eeuo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +EXPERIMENT_NAME="$(basename "$SCRIPT_DIR")" + +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/platform.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/config.env" + +RUN_ID="${RUN_ID:-$(date '+%Y%m%d-%H%M%S')}" +RESULT_BASE="${SCRIPT_DIR}/results" +MATRIX_FILE="${MATRIX_FILE:-${SCRIPT_DIR}/matrix.json}" +MATRIX_MODE="${MATRIX_MODE:-Y}" +SCENARIO_TIMEOUT_S="${SCENARIO_TIMEOUT_S:-1800}" +GPU_MEM_SAMPLE_INTERVAL_S="${GPU_MEM_SAMPLE_INTERVAL_S:-1}" +DRY_RUN="${DRY_RUN:-0}" +GRID_LIMIT="${GRID_LIMIT:-0}" + +# Export variables used inside functions that are called via bash -c subshells. +export BENCH_DATASET_NAME DATASET_PATH MODEL_PATH RESULT_BASE DOCKER_IMAGE USE_DOCKER_CLIENT SGLANG_BENCH_MODULE + +if [[ -x "${VENV_CLIENT}/bin/python" ]]; then + PYTHON="${VENV_CLIENT}/bin/python" +else + PYTHON="$(command -v python3)" +fi + +DOCKER_IMAGE="${DOCKER_IMAGE:-lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45}" + +log_dir_global="${RESULT_BASE}/${RUN_ID}/logs" +mkdir -p "$log_dir_global" +log_init "${log_dir_global}/orchestrator.log" + +log "experiment=${EXPERIMENT_NAME} run_id=${RUN_ID} platform=${PLATFORM} hardware=${HARDWARE}" +log "matrix_mode=${MATRIX_MODE} matrix_file=${MATRIX_FILE} dry_run=${DRY_RUN} grid_limit=${GRID_LIMIT}" + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +is_server_healthy() { + curl --fail --silent --show-error --max-time 10 "http://127.0.0.1:${SGLANG_PORT}/v1/models" >/dev/null 2>&1 +} + +stop_server() { + local tp="$1" + local dp="$2" + + log "stopping multi-node sglang (tp=${tp}, dp=${dp}) on both nodes" + # Delegate to the multi-node stop orchestrator, which tears down node0 + # (local) and node1 (remote via SSH) together: pid file, docker rm, pkill. + bash "${SCRIPT_DIR}/stop_sglang_multinode.sh" "$tp" "$dp" \ + >> "${log_dir_global}/sglang_tp${tp}_dp${dp}.stop.log" 2>&1 || true +} + +build_server_args() { + local tp="$1" + local dp="$2" + + local args=( + "python3 -m sglang.launch_server" --model-path "$MODEL_PATH" + --trust-remote-code + --tp-size "$tp" + --dp-size "$dp" + --quantization "$QUANTIZATION" + --moe-runner-backend "$MOE_RUNNER_BACKEND" + --mem-fraction-static "$MEM_FRACTION_STATIC" + --context-length "$CONTEXT_LENGTH" + --max-running-requests "$MAX_RUNNING_REQUESTS" + --host 0.0.0.0 + --port "$SGLANG_PORT" + --dist-init-addr "${NODE0_IP}:${DIST_INIT_PORT}" + --nnodes "$NNODES" + --node-rank 0 + ) + printf '%s ' "${args[@]}" +} + +start_server() { + local tp="$1" + local dp="$2" + + log "starting sglang server tp=${tp} dp=${dp}" + if [[ -n "${CONTAINER_NAME:-}" ]]; then + bash "${SCRIPT_DIR}/run_sglang_in_container.sh" "$tp" "$dp" \ + >> "${log_dir_global}/sglang_tp${tp}_dp${dp}.server.outer.log" 2>&1 + else + bash "${SCRIPT_DIR}/start_sglang_dp.sh" "$tp" "$dp" \ + >> "${log_dir_global}/sglang_tp${tp}_dp${dp}.server.outer.log" 2>&1 + fi + + if ! is_server_healthy; then + log "error: sglang server tp=${tp} dp=${dp} failed health check on port ${SGLANG_PORT}" + return 1 + fi + log "sglang server tp=${tp} dp=${dp} is healthy on port ${SGLANG_PORT}" +} + +restart_server() { + local tp="$1" + local dp="$2" + log "restarting sglang server tp=${tp} dp=${dp} after non-OOM failure" + stop_server "$tp" "$dp" + sleep 10 + start_server "$tp" "$dp" +} + +run_bench_serving() { + # Inject the offline workload choice consistently for warmup and main runs. + local -a dataset_args=(--dataset-name "$BENCH_DATASET_NAME") + if [[ "$BENCH_DATASET_NAME" == "random" ]]; then + dataset_args+=(--dataset-path "$DATASET_PATH") + elif [[ "$BENCH_DATASET_NAME" == "random-ids" ]]; then + : # random-ids does not need --tokenize-prompt + else + dataset_args+=(--tokenize-prompt) + fi + if [[ "${USE_DOCKER_CLIENT:-1}" == "1" ]]; then + local vol_args=() + vol_args+=("-v" "${MODEL_PATH}:${MODEL_PATH}:ro") + if [[ "$BENCH_DATASET_NAME" == "random" ]]; then + vol_args+=("-v" "${DATASET_PATH}:${DATASET_PATH}:ro") + fi + vol_args+=("-v" "${RESULT_BASE}:${RESULT_BASE}") + docker run --rm \ + --network host \ + "${vol_args[@]}" \ + -e PYTHONUNBUFFERED=1 \ + --entrypoint python3 \ + "${DOCKER_IMAGE}" \ + -m "$SGLANG_BENCH_MODULE" "${dataset_args[@]}" "$@" + else + "$PYTHON" -m "$SGLANG_BENCH_MODULE" "${dataset_args[@]}" "$@" + fi +} +export -f run_bench_serving + +run_warmup() { + local input_len="$1" + local output_len="$2" + + log "warming up (input=${input_len}, output=${output_len}, num=1)" + bash -c ' + run_bench_serving \ + --backend sglang \ + --host 127.0.0.1 \ + --port "'"$SGLANG_PORT"'" \ + --random-input-len "'"$input_len"'" \ + --random-output-len "'"$output_len"'" \ + --num-prompts 1 \ + --max-concurrency 1 \ + --request-rate 10000 \ + --output-file /dev/null \ + --output-details \ + >> "'"${log_dir_global}/warmup.log"'" 2>&1 + ' + log "warmup completed" +} + +scenario_already_completed() { + local output_file="$1" + local expected="$2" + [[ -s "$output_file" ]] || return 1 + local completed + completed="$("$PYTHON" -c " +import json, sys +path = sys.argv[1] +try: + with open(path, 'r', encoding='utf-8') as f: + for line in f: + line = line.strip() + if line: + data = json.loads(line) + print(data.get('completed', 0)) + break +except Exception: + print(0) +" "$output_file")" + [[ "${completed:-0}" -ge "$expected" ]] +} + +scenario_already_processed() { + local result_root="$1" + local scenario_name="$2" + local json_path="${result_root}/results.json" + [[ -f "$json_path" ]] || return 1 + "$PYTHON" -c " +import json, sys +path, name = sys.argv[1], sys.argv[2] +try: + with open(path, 'r', encoding='utf-8') as f: + data = json.load(f) + for s in data.get('scenarios', []): + if s.get('name') == name: + if s.get('status') or s.get('metrics', {}).get('success', 0) > 0: + sys.exit(0) +except Exception: + pass +sys.exit(1) +" "$json_path" "$scenario_name" +} + +detect_oom() { + local detail_log="$1" + local server_outer_log="$2" + local pattern='CUDA out of memory|torch\.OutOfMemoryError|OutOfMemory|out of memory|OOM|RESOURCE_EXHAUSTED|Failed to allocate memory' + if grep -Eiq "$pattern" "$detail_log" "$server_outer_log" 2>/dev/null; then + return 0 + fi + return 1 +} + +start_gpu_monitor() { + local csv_path="$1" + mkdir -p "$(dirname "$csv_path")" + nvidia-smi \ + --query-gpu=timestamp,index,memory.used,memory.total,utilization.gpu \ + --format=csv \ + -l "$GPU_MEM_SAMPLE_INTERVAL_S" \ + > "$csv_path" 2>/dev/null & + echo $! +} + +stop_gpu_monitor() { + local pid="$1" + if kill -0 "$pid" 2>/dev/null; then + kill "$pid" 2>/dev/null || true + sleep 1 + kill -9 "$pid" 2>/dev/null || true + fi +} + +append_scenario_record() { + local result_root="$1" + local json_path="$result_root/results.json" + shift + local scenario_json + scenario_json="$("$PYTHON" -c " +import json, sys +pairs = [a.split('=', 1) for a in sys.argv[1:]] +d = {} +for k, v in pairs: + try: + d[k] = json.loads(v) + except json.JSONDecodeError: + d[k] = v +print(json.dumps(d, ensure_ascii=False)) +" "$@")" + PYTHON="$PYTHON" append_scenario_to_json "$json_path" "$scenario_json" +} + +record_skipped_csv() { + local csv_path="$1" + shift + # Args: key=value + local row + row="$("$PYTHON" -c " +import csv, json, sys, io +pairs = [a.split('=', 1) for a in sys.argv[1:]] +d = {} +for k, v in pairs: + try: + d[k] = json.loads(v) + except json.JSONDecodeError: + d[k] = v +buf = io.StringIO() +writer = csv.DictWriter(buf, fieldnames=['engine','tp','dp','mark','isl','dsl','concurrency','status','reason','detail_log'], extrasaction='ignore') +writer.writerow(d) +print(buf.getvalue().strip()) +" "$@")" + echo "$row" >> "$csv_path" +} + +skip_remaining_scenarios() { + local result_root="$1" + local scenario_tsv="$2" + local start_index="$3" + local status="$4" + local reason="$5" + local tp="$6" + local dp="$7" + local skipped_csv="${RESULT_BASE}/${RUN_ID}/skipped_after_oom.csv" + + local i=0 + tail -n +2 "$scenario_tsv" | while IFS=$'\t' read -r mark isl dsl conc num; do + if (( i < start_index )); then + i=$((i + 1)) + continue + fi + i=$((i + 1)) + local sname="c${conc}_i${isl}_o${dsl}" + if scenario_already_processed "$result_root" "$sname"; then + continue + fi + append_scenario_record "$result_root" \ + "name=${sname}" \ + "config=$(jq -n --arg phase main --argjson c "$conc" --argjson i "$isl" --argjson o "$dsl" --arg dataset "$BENCH_DATASET_NAME" --argjson n "$num" '{phase: $phase, concurrency: $c, input_len: $i, output_len: $o, dataset: $dataset, num_prompts: $n}')" \ + "status=\"${status}\"" \ + "note=\"${reason}\"" + record_skipped_csv "$skipped_csv" \ + "engine=sglang" "tp=${tp}" "dp=${dp}" "mark=${mark}" "isl=${isl}" "dsl=${dsl}" "concurrency=${conc}" "status=${status}" "reason=${reason}" + done +} + +# --------------------------------------------------------------------------- +# Per-configuration runner +# --------------------------------------------------------------------------- + +run_parallel_config() { + local tp="$1" + local dp="$2" + local config_label="tp${tp}_dp${dp}" + local result_root="${RESULT_BASE}/${RUN_ID}/${config_label}" + local raw_dir="${result_root}/raw_outputs" + local gpu_log_dir="${result_root}/gpu_logs" + local phase_log_dir="${result_root}/logs" + mkdir -p "$raw_dir" "$gpu_log_dir" "$phase_log_dir" + + log "===== ${config_label} START =====" + + # Generate scenario list for this config. + local scenario_tsv="${result_root}/scenarios.tsv" + "$PYTHON" "${SCRIPT_DIR}/generate_scenarios.py" \ + --matrix "$MATRIX_FILE" \ + --mode "$MATRIX_MODE" \ + > "$scenario_tsv" + local total_scenarios + total_scenarios="$(tail -n +2 "$scenario_tsv" | wc -l)" + log "generated ${total_scenarios} scenarios for ${config_label}" + + # Write metadata. + ensure_result_root "$result_root" + write_metadata_json \ + "${result_root}/results.json" \ + "${EXPERIMENT_NAME}_${config_label}" \ + "$RUN_ID" \ + "$MODEL_PATH" \ + "sglang" \ + "sglang" \ + "$HARDWARE" \ + "$ACCELERATOR" \ + "$CHIP" \ + "experiments/${EXPERIMENT_NAME}/run_bench.sh" \ + "$DOCKER_IMAGE" \ + "H200 SGLang TP×DP matrix for DeepSeek-V4-Flash on RTX 6000D" + + local server_args_str + server_args_str="$(build_server_args "$tp" "$dp")" + jq --arg tp "$tp" --arg dp "$dp" --arg cuda "$CUDA_VISIBLE_DEVICES" --arg args "$server_args_str" \ + '.config = { + "tp": ($tp | tonumber), + "dp": ($dp | tonumber), + "cuda_visible_devices": $cuda, + "backend": "sglang", + "server_start_script": "experiments/'${EXPERIMENT_NAME}'/start_sglang_dp.sh", + "server_args": $args + }' "${result_root}/results.json" > "${result_root}/results.json.tmp" && \ + mv "${result_root}/results.json.tmp" "${result_root}/results.json" + + if [[ "$DRY_RUN" == "1" ]]; then + log "DRY_RUN: would start server with args: ${server_args_str}" + local line + tail -n +2 "$scenario_tsv" | while IFS=$'\t' read -r mark isl dsl conc num; do + log "DRY_RUN: ${config_label} scenario mark=${mark} c=${conc} i=${isl} o=${dsl} n=${num}" + done + log "===== ${config_label} DONE (dry run) =====" + return 0 + fi + + # Initialize skipped_after_oom.csv for this run. + local skipped_csv="${RESULT_BASE}/${RUN_ID}/skipped_after_oom.csv" + if [[ ! -f "$skipped_csv" ]]; then + echo "engine,tp,dp,mark,isl,dsl,concurrency,status,reason,detail_log" > "$skipped_csv" + fi + + # Start server once for this TP×DP config. + if ! start_server "$tp" "$dp"; then + log "ERROR: ${config_label} failed to start; skipping all scenarios" + skip_remaining_scenarios "$result_root" "$scenario_tsv" 0 "SKIPPED_SERVICE_START_FAILED" "service failed to start" "$tp" "$dp" + log "===== ${config_label} DONE =====" + return 0 + fi + + # Warmup with a small prompt before the first scenario. + run_warmup 1024 128 || true + + # Read scenarios into an array so we can skip remaining entries on failure. + local -a scenarios=() + while IFS= read -r line; do + scenarios+=("$line") + done < <(tail -n +2 "$scenario_tsv") + + local i mark isl dsl conc num + local output_file detail_log gpu_csv sname bench_rc + for (( i = 0; i < ${#scenarios[@]}; i++ )); do + IFS=$'\t' read -r mark isl dsl conc num <<< "${scenarios[$i]}" + + if [[ "$GRID_LIMIT" -gt 0 && "$i" -ge "$GRID_LIMIT" ]]; then + log "GRID_LIMIT=${GRID_LIMIT} reached; skipping remaining scenarios" + skip_remaining_scenarios "$result_root" "$scenario_tsv" "$i" "SKIPPED_GRID_LIMIT" "GRID_LIMIT reached" "$tp" "$dp" + break + fi + + sname="c${conc}_i${isl}_o${dsl}" + output_file="${raw_dir}/sglang_main_${conc}_${isl}_${dsl}.jsonl" + detail_log="${phase_log_dir}/sglang_${config_label}_${sname}.log" + gpu_csv="${gpu_log_dir}/gpu_mem_${conc}_${isl}_${dsl}.csv" + + if scenario_already_completed "$output_file" "$num" || scenario_already_processed "$result_root" "$sname"; then + log "skipping already-processed ${config_label} scenario: ${sname}" + continue + fi + + log "running ${config_label} scenario: mark=${mark} c=${conc} i=${isl} o=${dsl} n=${num}" + + local gpu_pid + gpu_pid="$(start_gpu_monitor "$gpu_csv")" + + bench_rc=0 + timeout "$SCENARIO_TIMEOUT_S" bash -c ' + run_bench_serving \ + --backend sglang \ + --host 127.0.0.1 \ + --port "'"$SGLANG_PORT"'" \ + --random-input-len "'"$isl"'" \ + --random-output-len "'"$dsl"'" \ + --num-prompts "'"$num"'" \ + --max-concurrency "'"$conc"'" \ + --request-rate 10000 \ + --output-file "'"$output_file"'" \ + --output-details \ + > "'"$detail_log"'" 2>&1 + ' || bench_rc=$? + + stop_gpu_monitor "$gpu_pid" + + if [[ "$bench_rc" -eq 0 ]]; then + log "finished ${config_label} scenario: output=${output_file}" + append_scenario_record "$result_root" \ + "name=${sname}" \ + "config=$(jq -n --arg phase main --argjson c "$conc" --argjson i "$isl" --argjson o "$dsl" --arg dataset "$BENCH_DATASET_NAME" --argjson n "$num" '{phase: $phase, concurrency: $c, input_len: $i, output_len: $o, dataset: $dataset, num_prompts: $n}')" \ + "status=\"completed\"" \ + "note=\"benchmark finished successfully\"" + continue + fi + + # Failure handling. + if detect_oom "$detail_log" "${log_dir_global}/sglang_tp${tp}_dp${dp}.server.outer.log"; then + log "ERROR: ${config_label} scenario ${sname} triggered OOM; stopping config" + append_scenario_record "$result_root" \ + "name=${sname}" \ + "config=$(jq -n --arg phase main --argjson c "$conc" --argjson i "$isl" --argjson o "$dsl" --arg dataset "$BENCH_DATASET_NAME" --argjson n "$num" '{phase: $phase, concurrency: $c, input_len: $i, output_len: $o, dataset: $dataset, num_prompts: $n}')" \ + "status=\"OOM\"" \ + "note=\"detected CUDA out-of-memory\"" + record_skipped_csv "$skipped_csv" \ + "engine=sglang" "tp=${tp}" "dp=${dp}" "mark=${mark}" "isl=${isl}" "dsl=${dsl}" "concurrency=${conc}" "status=OOM" "reason=detected CUDA out-of-memory" "detail_log=${detail_log}" + stop_server "$tp" "$dp" + skip_remaining_scenarios "$result_root" "$scenario_tsv" "$((i + 1))" "SKIPPED_AFTER_OOM" "previous case OOM" "$tp" "$dp" + break + fi + + log "ERROR: ${config_label} scenario ${sname} failed (rc=${bench_rc}); see ${detail_log}" + if [[ "$mark" == "P" ]]; then + log "optional (P) scenario failed; recording as skipped and continuing" + append_scenario_record "$result_root" \ + "name=${sname}" \ + "config=$(jq -n --arg phase main --argjson c "$conc" --argjson i "$isl" --argjson o "$dsl" --arg dataset "$BENCH_DATASET_NAME" --argjson n "$num" '{phase: $phase, concurrency: $c, input_len: $i, output_len: $o, dataset: $dataset, num_prompts: $n}')" \ + "status=\"skipped_optional\"" \ + "note=\"optional scenario failed (rc=${bench_rc})\"" + record_skipped_csv "$skipped_csv" \ + "engine=sglang" "tp=${tp}" "dp=${dp}" "mark=${mark}" "isl=${isl}" "dsl=${dsl}" "concurrency=${conc}" "status=skipped_optional" "reason=optional scenario failed (rc=${bench_rc})" "detail_log=${detail_log}" + continue + fi + + # Mandatory scenario failed but not OOM: try to restart the server. + if restart_server "$tp" "$dp"; then + run_warmup 1024 128 || true + log "resuming ${config_label} after server restart" + continue + fi + + log "ERROR: ${config_label} server restart failed; skipping remaining scenarios" + append_scenario_record "$result_root" \ + "name=${sname}" \ + "config=$(jq -n --arg phase main --argjson c "$conc" --argjson i "$isl" --argjson o "$dsl" --arg dataset "$BENCH_DATASET_NAME" --argjson n "$num" '{phase: $phase, concurrency: $c, input_len: $i, output_len: $o, dataset: $dataset, num_prompts: $n}')" \ + "status=\"FAILED\"" \ + "note=\"scenario failed and server restart failed (rc=${bench_rc})\"" + record_skipped_csv "$skipped_csv" \ + "engine=sglang" "tp=${tp}" "dp=${dp}" "mark=${mark}" "isl=${isl}" "dsl=${dsl}" "concurrency=${conc}" "status=FAILED" "reason=scenario failed and server restart failed" "detail_log=${detail_log}" + skip_remaining_scenarios "$result_root" "$scenario_tsv" "$((i + 1))" "SKIPPED_RESTART_FAILED" "server restart failed" "$tp" "$dp" + break + done + + stop_server "$tp" "$dp" + + # Parse results. + log "parsing ${config_label} results" + "$PYTHON" "${SCRIPT_DIR}/../../../scripts/common/parse_backend.py" "$result_root" --backend sglang \ + >> "${phase_log_dir}/parse.log" 2>&1 || { + log "WARNING: parser failed for ${config_label}; see ${phase_log_dir}/parse.log" + } + + log "===== ${config_label} DONE =====" +} + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +# Cleanup any leftovers. +for cfg in "${PARALLEL_CONFIGS[@]}"; do + read -r tp dp <<< "$cfg" + stop_server "$tp" "$dp" +done + +# Run each parallel configuration. +for cfg in "${PARALLEL_CONFIGS[@]}"; do + read -r tp dp <<< "$cfg" + run_parallel_config "$tp" "$dp" +done + +# Generate cross-configuration comparison. +log "generating comparison report" +"$PYTHON" "${SCRIPT_DIR}/compare.py" \ + --run-root "${RESULT_BASE}/${RUN_ID}" \ + --output "${RESULT_BASE}/${RUN_ID}/comparison.md" \ + >> "${log_dir_global}/compare.log" 2>&1 || { + log "WARNING: comparison script failed; see ${log_dir_global}/compare.log" + } + +log "all results saved to ${RESULT_BASE}/${RUN_ID}" diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/start_sglang_dp.sh b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/start_sglang_dp.sh new file mode 100755 index 0000000..32470a9 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/start_sglang_dp.sh @@ -0,0 +1,16 @@ +#!/usr/bin/env bash +# Entry point for the matrix runner (run_bench.sh calls this as +# start_sglang_dp.sh ). For this multi-node experiment it always +# delegates to the multi-node orchestrator, which brings up both nodes. +# +# Usage: start_sglang_dp.sh +set -e + +TP="${1:?usage: start_sglang_dp.sh }" +DP="${2:?usage: start_sglang_dp.sh }" + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# This experiment is multi-node only; there is no single-node venv path. +# USE_DOCKER is kept for config.env compatibility but ignored here. +exec "${SCRIPT_DIR}/start_sglang_multinode.sh" "$TP" "$DP" diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/start_sglang_multinode.sh b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/start_sglang_multinode.sh new file mode 100755 index 0000000..acc2834 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/start_sglang_multinode.sh @@ -0,0 +1,45 @@ +#!/usr/bin/env bash +# Multi-node orchestrator for GLM-5.2-FP8 TP=16 across pro6000D.1 + pro6000D.3, +# through the shared deployment layer (`python -m sskj.deploy start` with +# NODE_HOSTS orchestration). Rank 0 (local) is started without ssh. +# +# Usage: start_sglang_multinode.sh +set -e + +TP="${1:-${TP_SIZE:-16}}" +DP="${2:-${DP_SIZE:-1}}" + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/config.env" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/platform.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" + +RUNTIME_BASE="${RUNTIME_BASE:-${SCRIPT_DIR}/runtime}" +mkdir -p "${RUNTIME_BASE}/logs" "${RUNTIME_BASE}/tmp" + +DRY_RUN="${DRY_RUN:-0}" +DEPLOY_PROFILE="${DEPLOY_PROFILE:-pro6000/glm52_pro6000_sglang_tp16}" + +echo "=== Multi-node start: TP=${TP} DP=${DP} nnodes=${NNODES} via deploy profile ${DEPLOY_PROFILE} ===" + +if [[ "$DRY_RUN" == "1" ]]; then + PYTHONPATH="${ROOT_DIR}/src" "${DEPLOY_PYTHON:-python3}" -m sskj.deploy start \ + --profile "$(deploy_profile_abs "$DEPLOY_PROFILE")" \ + --tp "$TP" --dp "$DP" \ + --port "${SGLANG_PORT:-30031}" \ + --model-path "$MODEL_PATH" \ + --log-dir "${RUNTIME_BASE}/logs" \ + --dry-run +else + deploy_start_multinode \ + "$DEPLOY_PROFILE" "$TP" "$DP" \ + "${RUNTIME_BASE}/logs" \ + "${SGLANG_PORT:-30031}" \ + "$MODEL_PATH" \ + "$EXPERIMENT" +fi diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/start_sglang_node.sh b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/start_sglang_node.sh new file mode 100755 index 0000000..d8e4acd --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/start_sglang_node.sh @@ -0,0 +1,46 @@ +#!/usr/bin/env bash +# Start one GLM-5.2-FP8 SGLang node through the shared deployment layer. +# Convenience for per-node debugging; the multinode orchestrator +# (start_sglang_multinode.sh) does not use this anymore. +# +# Usage: start_sglang_node.sh +set -e + +NODE_RANK="${1:?usage: start_sglang_node.sh }" + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/config.env" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/platform.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" + +RUNTIME_BASE="${RUNTIME_BASE:-${SCRIPT_DIR}/runtime}" +mkdir -p "${RUNTIME_BASE}/logs" "${RUNTIME_BASE}/tmp" + +DRY_RUN="${DRY_RUN:-0}" +DEPLOY_PROFILE="${DEPLOY_PROFILE:-pro6000/glm52_pro6000_sglang_tp16}" + +echo "=== Start node${NODE_RANK} (TP=${TP_SIZE}, DP=${DP_SIZE}, nnodes=${NNODES}) via deploy profile ===" + +if [[ "$DRY_RUN" == "1" ]]; then + PYTHONPATH="${ROOT_DIR}/src" "${DEPLOY_PYTHON:-python3}" -m sskj.deploy start \ + --profile "$(deploy_profile_abs "$DEPLOY_PROFILE")" \ + --tp "${TP_SIZE:-16}" --dp "${DP_SIZE:-1}" \ + --port "${SGLANG_PORT:-30031}" \ + --model-path "$MODEL_PATH" \ + --node-rank "$NODE_RANK" \ + --log-dir "${RUNTIME_BASE}/logs" \ + --dry-run +else + deploy_start \ + "$DEPLOY_PROFILE" "${TP_SIZE:-16}" "${DP_SIZE:-1}" \ + "${RUNTIME_BASE}/logs" \ + "${SGLANG_PORT:-30031}" \ + "$MODEL_PATH" \ + "$EXPERIMENT" \ + "$NODE_RANK" +fi diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/stop_sglang_multinode.sh b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/stop_sglang_multinode.sh new file mode 100755 index 0000000..345678a --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/stop_sglang_multinode.sh @@ -0,0 +1,45 @@ +#!/usr/bin/env bash +# Tear down the multi-node GLM-5.2-FP8 deployment (both nodes) through the +# shared deployment layer. +# +# Usage: stop_sglang_multinode.sh +set -e + +TP="${1:-${TP_SIZE:-16}}" +DP="${2:-${DP_SIZE:-1}}" + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/config.env" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/platform.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" + +DRY_RUN="${DRY_RUN:-0}" +DEPLOY_PROFILE="${DEPLOY_PROFILE:-pro6000/glm52_pro6000_sglang_tp16}" + +echo "=== Multi-node stop: TP=${TP} DP=${DP} via deploy profile ${DEPLOY_PROFILE} ===" + +if [[ "$DRY_RUN" == "1" ]]; then + PYTHONPATH="${ROOT_DIR}/src" "${DEPLOY_PYTHON:-python3}" -m sskj.deploy stop \ + --profile "$(deploy_profile_abs "$DEPLOY_PROFILE")" \ + --tp "$TP" --dp "$DP" \ + --port "${SGLANG_PORT:-30031}" \ + --model-path "$MODEL_PATH" \ + --container-name "$EXPERIMENT" \ + --dry-run +else + deploy_stop \ + "$DEPLOY_PROFILE" "$TP" "$DP" \ + "${SGLANG_PORT:-30031}" \ + "$MODEL_PATH" \ + "$EXPERIMENT" +fi + +# Local fallback: kill any stray launch_server for this model. +pkill -9 -f "sglang.launch_server.*${MODEL_NAME}" 2>/dev/null || true +pkill -9 -f "sglang.launch_server.*${MODEL_PATH}" 2>/dev/null || true +echo "=== both nodes stopped ===" diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/test_glm52_tilelang_partial.py b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/test_glm52_tilelang_partial.py new file mode 100644 index 0000000..9f93fb6 --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/test_glm52_tilelang_partial.py @@ -0,0 +1,72 @@ +import math +import os + +import torch + +from sglang.kernels.ops.attention.dsa.tilelang_kernel import ( + sparse_mla_fwd_decode_combine, + sparse_mla_fwd_decode_partial, + tilelang_sparse_fwd, +) + + +def main() -> None: + torch.manual_seed(0) + heads = int(os.environ.get("HEADS", "4")) + d_v = 256 + d_tail = 64 + topk = 2048 + block_i = 32 + threads = 128 + scale = 1.0 / math.sqrt(d_v + d_tail) + + q = torch.randn((1, heads, d_v + d_tail), device="cuda", dtype=torch.bfloat16) + kv = torch.randn( + (topk, 1, d_v + d_tail), device="cuda", dtype=torch.bfloat16 + ) + indices = torch.arange(topk, device="cuda", dtype=torch.int32).view(1, 1, topk) + + if os.environ.get("USE_WRAPPER") == "1": + output = tilelang_sparse_fwd(q, kv, indices, scale, d_v) + else: + partial_kernel = sparse_mla_fwd_decode_partial( + heads, + d_v, + d_tail, + topk, + sm_scale=scale, + block_I=block_i, + inner_iter=1, + num_stages=1, + threads=threads, + ) + partial_o, partial_lse = partial_kernel( + q.unsqueeze(0), kv.unsqueeze(0), indices.unsqueeze(0) + ) + groups = topk // block_i + combine_kernel = sparse_mla_fwd_decode_combine( + heads, + d_v, + groups * block_i, + head_per_block=4, + block_I=block_i, + threads=threads, + ) + output = combine_kernel(partial_o, partial_lse) + + q_float = q[0].float() + k_float = kv[:, 0].float() + scores = q_float @ k_float.transpose(0, 1) * scale + probs = torch.softmax(scores, dim=-1) + reference = probs @ k_float[:, :d_v] + actual = output.reshape(heads, d_v).float() + error = (actual - reference).abs() + + print(f"output_shape={tuple(output.shape)}") + print(f"max_abs_error={error.max().item():.6f}") + print(f"mean_abs_error={error.mean().item():.6f}") + print(f"all_finite={torch.isfinite(actual).all().item()}") + + +if __name__ == "__main__": + main() diff --git a/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/tilelang_kernel_sm120.py b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/tilelang_kernel_sm120.py new file mode 100644 index 0000000..b5752bb --- /dev/null +++ b/experiments/pro6000/glm52_pro6000_sglang_multinode_tp16/tilelang_kernel_sm120.py @@ -0,0 +1,2610 @@ +import functools +from functools import lru_cache +from typing import Any, Optional, Tuple + +import tilelang +import tilelang.language as T +import torch + +from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz +from sglang.srt.utils import is_gfx95_supported, is_hip + +tilelang.set_log_level("WARNING") + +# Workaround a tilelang bug: BaseKernelAdapter._legalize_result_idx mutates the +# `out_idx` list in place when normalising negative indices to positive ones. +# That breaks any @tilelang.jit factory that compiles two prim_funcs with +# different param counts (e.g. our unified single/dual partial kernel) — the +# second compile sees indices already-converted for the first's len(params) +# and silently builds the wrong adapter, leading to IndexError at call time. +# Patch once on import to copy the list before mutation. +from tilelang.jit.adapter.base import ( # noqa: E402 + BaseKernelAdapter as _BaseKernelAdapter, +) + +if not getattr(_BaseKernelAdapter, "_legalize_result_idx_patched", False): + _orig_legalize = _BaseKernelAdapter._legalize_result_idx + + def _legalize_result_idx_safe(self, result_idx): + if isinstance(result_idx, list): + result_idx = list(result_idx) + return _orig_legalize(self, result_idx) + + _BaseKernelAdapter._legalize_result_idx = _legalize_result_idx_safe + _BaseKernelAdapter._legalize_result_idx_patched = True + +pass_configs = { + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, +} +# TL_DISABLE_FAST_MATH has deprecated in v0.1.7.post1 tilelang +if hasattr(tilelang.PassConfigKey, "TL_DISABLE_FAST_MATH"): + pass_configs[tilelang.PassConfigKey.TL_DISABLE_FAST_MATH] = True +elif hasattr(tilelang.PassConfigKey, "TL_ENABLE_FAST_MATH"): + pass_configs[tilelang.PassConfigKey.TL_ENABLE_FAST_MATH] = False + +_is_hip = is_hip() +_is_gfx95_supported = is_gfx95_supported() +_is_fp8_fnuz = is_fp8_fnuz() + +BF16 = "bfloat16" +FP8 = "float8_e4m3fnuz" if _is_fp8_fnuz else "float8_e4m3fn" +FP8_DTYPE = torch.float8_e4m3fnuz if _is_fp8_fnuz else torch.float8_e4m3fn +FP32 = "float32" +INT32 = "int32" +UINT8 = "uint8" + + +def fast_log2_ceil(x): + bits_x = T.reinterpret("uint32", x) + exp_x = (bits_x >> 23) & 0xFF + man_bits = bits_x & ((1 << 23) - 1) + return T.Cast("int32", exp_x - 127 + T.if_then_else(man_bits != 0, 1, 0)) + + +def fast_pow2(x): + bits_x = (x + 127) << 23 + return T.reinterpret("float32", bits_x) + + +def fast_round_scale(amax, fp8_max_inv): + return fast_pow2(fast_log2_ceil(amax * fp8_max_inv)) + + +@lru_cache(maxsize=8) +def _pick_inner_iter(seq: int, ni: int, cu: int, block_per_cu: int) -> int: + """ + Pick the largest valid inner_iter (power-of-two divisor of ni) that keeps + enough work per CU (seq * ni / inner_iter / cu >= block_per_cu), so we avoid + under-utilization while minimizing the number of partial groups. + """ + + max_it = int(seq * ni / (cu * block_per_cu)) + it = ni + while it >= 2: + if it <= max_it and ni % it == 0: + return it + it //= 2 + return 1 + + +@tilelang.jit(pass_configs=pass_configs) +def act_quant_kernel( + N, in_dtype=BF16, out_dtype=FP8, scale_dtype=FP32, round_scale=False +): + M = T.symbolic("M") + fp8_min = -224.0 if _is_fp8_fnuz else -448.0 + fp8_max = 224.0 if _is_fp8_fnuz else 448.0 + fp8_max_inv = 1 / fp8_max + num_stages = 0 if round_scale else 2 + blk_m = 32 + group_size = 128 + + @T.prim_func + def act_quant_kernel_( + X: T.Tensor[(M, N), in_dtype], + Y: T.Tensor[(M, N), out_dtype], + S: T.Tensor[(M, T.ceildiv(N, group_size)), scale_dtype], + ): + with T.Kernel(T.ceildiv(M, blk_m), T.ceildiv(N, group_size), threads=128) as ( + pid_m, + pid_n, + ): + x_shared = T.alloc_shared((blk_m, group_size), in_dtype) + x_local = T.alloc_fragment((blk_m, group_size), in_dtype) + amax_local = T.alloc_fragment((blk_m,), scale_dtype) + s_local = T.alloc_fragment((blk_m,), scale_dtype) + y_local = T.alloc_fragment((blk_m, group_size), out_dtype) + y_shared = T.alloc_shared((blk_m, group_size), out_dtype) + + for _ in T.Pipelined(1, num_stages=num_stages): + T.copy(X[pid_m * blk_m, pid_n * group_size], x_shared) + T.copy(x_shared, x_local) + T.reduce_absmax(x_local, amax_local, dim=1) + for i in T.Parallel(blk_m): + amax_local[i] = T.max(amax_local[i], 1e-4) + if round_scale: + s_local[i] = fast_round_scale(amax_local[i], fp8_max_inv) + else: + s_local[i] = amax_local[i] * fp8_max_inv + for i, j in T.Parallel(blk_m, group_size): + y_local[i, j] = T.clamp( + x_local[i, j] / s_local[i], fp8_min, fp8_max + ) + for i in T.Parallel(blk_m): + S[pid_m * blk_m + i, pid_n] = s_local[i] + T.copy(y_local, y_shared) + T.copy(y_shared, Y[pid_m * blk_m, pid_n * group_size]) + + return act_quant_kernel_ + + +def act_quant( + x: torch.Tensor, block_size: int = 128, scale_fmt: Optional[str] = None +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Quantizes the input tensor `x` using block-wise quantization. + + Args: + x (torch.Tensor): The input tensor to be quantized. Must be contiguous and its last dimension size must be divisible by `block_size`. + block_size (int, optional): The size of the blocks to be used for quantization. Default is 128. + scale_fmt (Optional[str], optional): The format of the scale. Default is None. + Returns: + Tuple[torch.Tensor, torch.Tensor]: A tuple containing: + - The quantized tensor with dtype `torch.float8_e4m3fn`. + - A tensor of scaling factors with dtype `torch.float32`. + """ + assert x.is_contiguous(), "Input tensor must be contiguous" + assert ( + x.size(-1) % block_size == 0 + ), f"Last dimension size must be divisible by block_size (block_size={block_size})" + N = x.size(-1) + if _is_fp8_fnuz: + y = torch.empty_like(x, dtype=torch.float8_e4m3fnuz) + else: + y = torch.empty_like(x, dtype=torch.float8_e4m3fn) + s = x.new_empty(*x.size()[:-1], N // block_size, dtype=torch.float32) + kernel = act_quant_kernel(N, round_scale=scale_fmt is not None) + kernel(x.view(-1, N), y.view(-1, N), s.view(-1, N // block_size)) + return y, s + + +@tilelang.jit(out_idx=[4], pass_configs=pass_configs) +def fp8_index_kernel(h: int, d: int, clear_accum=True): + b = T.symbolic("b") + m = T.symbolic("m") + n = T.symbolic("n") + + blk_n1 = 512 + blk_n2 = 128 + + @T.prim_func + def fp8_index_kernel_( + q: T.Tensor[(b, m, h, d), FP8], + q_s: T.Tensor[(b, m, h), FP32], + k: T.Tensor[(b, n, d), FP8], + k_s: T.Tensor[(b, n), FP32], + o: T.Tensor[(b, m, n), FP32], + ) -> None: + with T.Kernel(b, m, T.ceildiv(n, blk_n1)) as (i_b, i_m, i1_n): + q_smem = T.alloc_shared((h, d), FP8) + T.copy(q[i_b, i_m, 0, 0], q_smem) + + q_s_frag = T.alloc_fragment(h, FP32) + T.copy(q_s[i_b, i_m, 0], q_s_frag) + + for i2_n in T.Pipelined(blk_n1 // blk_n2, num_stages=2): + k_smem = T.alloc_shared((blk_n2, d), FP8) + T.copy(k[i_b, i1_n * blk_n1 + i2_n * blk_n2, 0], k_smem) + + k_s_frag = T.alloc_fragment(blk_n2, FP32) + T.copy(k_s[i_b, i1_n * blk_n1 + i2_n * blk_n2], k_s_frag) + + logits = T.alloc_fragment((blk_n2, h), FP32) + if not clear_accum: + T.fill(logits, 0) + T.gemm( + k_smem, + q_smem, + logits, + transpose_A=False, + transpose_B=True, + clear_accum=clear_accum, + ) + + for i_h, i3_n in T.Parallel(h, blk_n2): + logits[i3_n, i_h] = T.max(logits[i3_n, i_h], 0) * q_s_frag[i_h] + + logits_sum = T.alloc_fragment(blk_n2, FP32) + T.reduce_sum(logits, logits_sum, dim=1) + + for i3_n in T.Parallel(blk_n2): + logits_sum[i3_n] *= k_s_frag[i3_n] + + T.copy(logits_sum, o[i_b, i_m, i1_n * blk_n1 + i2_n * blk_n2]) + + return fp8_index_kernel_ + + +def fp8_index( + q: torch.Tensor, + q_s: torch.Tensor, + k: torch.Tensor, + k_s: torch.Tensor, +) -> torch.Tensor: + """ + Perform index score using FP8 precision. + + Args: + q (torch.Tensor): The Q tensor, must be contiguous. + q_s (torch.Tensor): The scaling factor for Q (float), must be contiguous. + k (torch.Tensor): The K tensor, must be contiguous. + k_s (torch.Tensor): The scaling factor for K (e8m0 here), must be contiguous. + + fp8 q @ fp8 k -> fp32 logits + relu(fp32 logits) * q_s (weights) -> fp32 logits + fp32 logits -> fp32 logits_sum + fp32 logits_sum * k_s (e8m0) -> fp32 index_score + """ + if _is_hip: + return fp8_index_kernel(q.shape[2], q.shape[3], False)(q, q_s, k, k_s) + else: + return fp8_index_kernel(q.shape[2], q.shape[3])(q, q_s, k, k_s) + + +@tilelang.jit( + out_idx=[-1], + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + }, +) +def sparse_attention_fwd_kernel_v1( + num_heads, + dim, + tail_dim, + topk, + *, + kv_group=1, + sm_scale=None, + is_causal=True, + block_I=64, + num_stages=2, + threads=256, +): + assert dim == tilelang.math.next_power_of_2( + dim + ), f"haven't check padding correctness yet, dim={dim}" + assert tail_dim == tilelang.math.next_power_of_2( + tail_dim + ), f"haven't check padding correctness yet, dim={tail_dim}" + assert is_causal == True, "non-casual is not supported" + assert ( + topk % block_I == 0 + ), "otherwise will load some index=0 thus causing wrong kv to be loaded" + if sm_scale is None: + sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e) + else: + sm_scale = sm_scale * 1.44269504 # log2(e) + + batch = T.symbolic("batch") + seq_len = T.symbolic("seq_len") + seq_len_kv = T.symbolic("seq_len_kv") + + head_kv = num_heads // kv_group + q_shape = [batch, seq_len, num_heads, dim + tail_dim] + kv_shape = [batch, seq_len_kv, kv_group, dim + tail_dim] + o_shape = [batch, seq_len, num_heads, dim] + indices_shape = [batch, seq_len, kv_group, topk] + indices_dtype = "int32" + dtype = "bfloat16" + accum_dtype = "float" + + H = head_kv + padded_H = max(tilelang.math.next_power_of_2(head_kv), 16) + if padded_H != H: + assert kv_group == 1 + BI = block_I + NI = tilelang.cdiv(topk, block_I) + D = dim + D_tail = tail_dim + + if head_kv > 64: + assert head_kv % 64 == 0, "head_kv should be a multiple of 64" + REPLICATE_H = head_kv // 64 + else: + REPLICATE_H = 1 + + H_per_block = padded_H if REPLICATE_H == 1 else 64 + + @T.prim_func + def main( + Q: T.Tensor(q_shape, dtype), # type: ignore + KV: T.Tensor(kv_shape, dtype), # type: ignore + Indices: T.Tensor(indices_shape, indices_dtype), # type: ignore + Output: T.Tensor(o_shape, dtype), # type: ignore + ): + with T.Kernel(seq_len * REPLICATE_H, batch, kv_group, threads=threads) as ( + bx, + by, + bz, + ): + Q_shared = T.alloc_shared([H_per_block, D], dtype) + Q_tail_shared = T.alloc_shared([H_per_block, D_tail], dtype) + KV_shared = T.alloc_shared([BI, D], dtype) + K_tail_shared = T.alloc_shared([BI, D_tail], dtype) + O_shared = T.alloc_shared([H_per_block, D], dtype) + mask = T.alloc_fragment([BI], "bool") + + acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) + acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) + S_shared = T.alloc_shared([H_per_block, BI], dtype) + sumexp = T.alloc_fragment([H_per_block], accum_dtype) + sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) + alpha = T.alloc_fragment([H_per_block], accum_dtype) + m_i = T.alloc_fragment([H_per_block], accum_dtype) + m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) + + T.fill(acc_o, 0) + T.fill(sumexp, 0) + T.fill(m_i, -(2**30)) # avoid -inf - inf to cause nan + + b_i, g_i = by, bz + s_i = bx if REPLICATE_H == 1 else (bx // REPLICATE_H) + q_i = s_i + max_kv_i = q_i + + H0 = g_i * padded_H + (0 if REPLICATE_H == 1 else (bx % REPLICATE_H) * 64) + H1 = H0 + H_per_block + + T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared) + T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared) + + for i_i in T.Pipelined(NI, num_stages=num_stages): + + for bi_i in T.Parallel(BI): + mask[bi_i] = Indices[b_i, s_i, g_i, i_i * BI + bi_i] >= 0 + + for bi_i, d_i in T.Parallel(BI, D): + KV_shared[bi_i, d_i] = KV[ + b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, d_i + ] + for bi_i, d_i in T.Parallel(BI, D_tail): + K_tail_shared[bi_i, d_i] = KV[ + b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, D + d_i + ] + + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else( + mask[bi_i], 0, -T.infinity(acc_s.dtype) + ) + T.gemm( + Q_shared, + KV_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullCol, + ) + T.gemm( + Q_tail_shared, + K_tail_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullCol, + ) + T.copy(m_i, m_i_prev) + T.reduce_max(acc_s, m_i, dim=1, clear=False) + for h_i in T.Parallel(H_per_block): + alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.exp2( + acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale + ) + T.reduce_sum(acc_s, sumexp_i, dim=1) # is this a accumulate operator? + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] = acc_o[h_i, d_i] * alpha[h_i] + + T.copy(acc_s, S_shared) + T.gemm(S_shared, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullCol) + + # Rescale + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] /= sumexp[h_i] + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale + + T.copy(acc_o, O_shared) + T.copy(acc_o, Output[b_i, s_i, H0:H1, :]) + + return main + + +@tilelang.jit( + out_idx=[-1], + compile_flags=[ + "-O3", + "-Wno-deprecated-declarations", + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-U__CUDA_NO_HALF2_OPERATORS__", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + "--expt-relaxed-constexpr", + "--expt-extended-lambda", + "--ptxas-options=-v,--register-usage-level=10", + "-DNDEBUG", + ], +) # type: ignore +def sparse_attention_fwd_kernel_v2( + num_heads: int, + dim: int, + tail_dim: int, + topk: int, + *, + kv_group: int = 1, + sm_scale: Optional[float] = None, + block_I: int = 64, +): + assert dim == tilelang.math.next_power_of_2( + dim + ), f"haven't check padding correctness yet, dim={dim}" + assert tail_dim == tilelang.math.next_power_of_2( + tail_dim + ), f"haven't check padding correctness yet, dim={tail_dim}" + assert ( + topk % block_I == 0 + ), "otherwise will load some index=0 thus causing wrong kv to be loaded" + if sm_scale is None: + sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e) + else: + sm_scale = sm_scale * 1.44269504 # log2(e) + threads = 384 + + batch = T.symbolic("batch") + qo_len = T.symbolic("seq_len") + num_pages = T.symbolic("num_pages") + + q_shape = [batch, qo_len, num_heads, dim + tail_dim] + kv_shape = [batch, num_pages, kv_group, dim + tail_dim] + o_shape = [batch, qo_len, num_heads, dim] + indices_shape = [batch, qo_len, kv_group, topk] + + indices_dtype = "int32" + dtype = "bfloat16" + accum_dtype = "float" + + H = num_heads + padded_H = max(tilelang.math.next_power_of_2(num_heads), 16) + if padded_H != H: + assert kv_group == 1 + BI = block_I + NI = tilelang.cdiv(topk, block_I) + assert NI % 2 == 0, "NI should be a multiple of 2" + D = dim + D_tail = tail_dim + if num_heads > 64: + assert num_heads % 64 == 0, "head_kv should be a multiple of 64" + REPLICATE_H = num_heads // 64 + else: + REPLICATE_H = 1 + + H_per_block = padded_H if REPLICATE_H == 1 else 64 + + @T.prim_func + def main( + Q: T.Tensor(q_shape, dtype), # type: ignore + KV: T.Tensor(kv_shape, dtype), # type: ignore + Indices: T.Tensor(indices_shape, indices_dtype), # type: ignore + Output: T.Tensor(o_shape, dtype), # type: ignore + ): + """ + Q: [b, qo_len, H, D + D_tail] (bfloat16) + KV: [b, num_pages, kv_group, D + D_tail] (bfloat16) + Indices: [b, qo_len, kv_group, topk] (int32) + """ + + with T.Kernel(qo_len * REPLICATE_H, batch, 1, threads=threads) as (bx, by, bz): # type: ignore + Q_shared_l = T.alloc_shared([H_per_block, D // 2], dtype) + Q_shared_r = T.alloc_shared([H_per_block, D // 2], dtype) + Q_tail_shared = T.alloc_shared([H_per_block, D_tail], dtype) + KV_shared_0_l = T.alloc_shared([BI, D // 2], dtype) + KV_shared_0_r = T.alloc_shared([BI, D // 2], dtype) + KV_shared_1_l = T.alloc_shared([BI, D // 2], dtype) + KV_shared_1_r = T.alloc_shared([BI, D // 2], dtype) + K_tail_shared_0 = T.alloc_shared([BI, D_tail], dtype) + K_tail_shared_1 = T.alloc_shared([BI, D_tail], dtype) + O_shared_l = Q_shared_l + O_shared_r = Q_shared_r + is_kv_valid_0 = T.alloc_shared([BI], "bool", scope="shared") + is_kv_valid_1 = T.alloc_shared([BI], "bool", scope="shared") + + acc_o_l = T.alloc_fragment([H_per_block, D // 2], accum_dtype) + acc_o_r = T.alloc_fragment([H_per_block, D // 2], accum_dtype) + acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) + S_shared = T.alloc_shared([H_per_block, BI], dtype) + sumexp = T.alloc_fragment([H_per_block], accum_dtype) + sum_exp_shared = T.alloc_shared([H_per_block], accum_dtype) + sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) + alpha_shared = T.alloc_shared([H_per_block], accum_dtype, scope="shared") + alpha_local = T.alloc_fragment([H_per_block], accum_dtype) + m_i = T.alloc_fragment([H_per_block], accum_dtype) + m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) + indices_local = T.alloc_local([1], indices_dtype) + indices_tmp = T.alloc_local([1], indices_dtype) + + bar_q = T.alloc_barrier(arrive_count=384) + bar_k_0_ready = T.alloc_barrier(arrive_count=128) + bar_k_1_ready = T.alloc_barrier(arrive_count=128) + bar_k_0_free = T.alloc_barrier(arrive_count=256) + bar_k_1_free = T.alloc_barrier(arrive_count=256) + bar_sScale_and_sS_ready = T.alloc_barrier(arrive_count=256) + bar_sScale_and_sS_free = T.alloc_barrier(arrive_count=256) + + bar_0_128 = T.alloc_barrier(arrive_count=128) + bar_1_128 = T.alloc_barrier(arrive_count=128) + bar_2_128 = T.alloc_barrier(arrive_count=128) + bar_final = T.alloc_barrier(arrive_count=128) + + b_i, g_i = by, bz + s_i = bx if REPLICATE_H == 1 else bx // REPLICATE_H + + H0 = g_i * padded_H + (0 if REPLICATE_H == 1 else (bx % REPLICATE_H) * 64) + H1 = H0 + H_per_block + + tx = T.get_thread_binding() + + T.copy(Q[b_i, s_i, H0:H1, 0 : D // 2], Q_shared_l) + T.copy(Q[b_i, s_i, H0:H1, D // 2 : D], Q_shared_r) + T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared) + T.barrier_arrive(bar_q) + + if tx < 128: + T.set_max_nreg(240, 1) + T.fill(sumexp, 0) + T.fill(m_i, -(2**30)) # avoid -inf - inf to cause nan + T.fill(acc_o_l, 0) + T.barrier_wait(bar_q, 0) + + for i_i in T.serial(T.ceildiv(NI, 2)): + # Buffer 0 + # with sync_at(bar_0_128, 0): + T.barrier_wait(bar_k_0_ready[0], (i_i & 1)) + T.barrier_arrive(bar_0_128) + T.barrier_wait(bar_0_128, 0) + + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else( + is_kv_valid_0[bi_i], 0, -T.infinity(acc_s.dtype) + ) + T.gemm(Q_shared_l, KV_shared_0_l, acc_s, transpose_B=True) + T.gemm(Q_shared_r, KV_shared_0_r, acc_s, transpose_B=True) + T.gemm( + Q_tail_shared, + K_tail_shared_0, + acc_s, + transpose_B=True, + ) + if i_i != 0: + T.barrier_arrive(bar_sScale_and_sS_free) + T.barrier_wait(bar_sScale_and_sS_free, ((i_i * 2) & 1) ^ 1) + + T.copy(m_i, m_i_prev) + T.reduce_max(acc_s, m_i, dim=1, clear=False) + for h_i in T.Parallel(H_per_block): + alpha_local[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.exp2( + acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale + ) + T.reduce_sum( + acc_s, sumexp_i, dim=1 + ) # is this a accumulate operator? + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = sumexp[h_i] * alpha_local[h_i] + sumexp_i[h_i] + for h_i, d_i in T.Parallel(H_per_block, D // 2): + acc_o_l[h_i, d_i] *= alpha_local[h_i] + T.copy(alpha_local, alpha_shared) + + T.copy(acc_s, S_shared) + T.gemm(S_shared, KV_shared_0_l, acc_o_l) + + T.barrier_arrive(bar_sScale_and_sS_ready) + T.barrier_arrive(bar_k_0_free[0]) + + # Buffer 1 + T.barrier_wait(bar_k_1_ready[0], (i_i & 1)) + T.barrier_arrive(bar_0_128) + T.barrier_wait(bar_0_128, 1) + + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else( + is_kv_valid_1[bi_i], 0, -T.infinity(acc_s.dtype) + ) + T.gemm(Q_shared_l, KV_shared_1_l, acc_s, transpose_B=True) + T.gemm(Q_shared_r, KV_shared_1_r, acc_s, transpose_B=True) + T.gemm( + Q_tail_shared, + K_tail_shared_1, + acc_s, + transpose_B=True, + ) + T.barrier_arrive(bar_sScale_and_sS_free) + T.barrier_wait(bar_sScale_and_sS_free, ((i_i * 2 + 1) & 1) ^ 1) + + T.copy(m_i, m_i_prev) + T.reduce_max(acc_s, m_i, dim=1, clear=False) + for h_i in T.Parallel(H_per_block): + alpha_local[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.exp2( + acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale + ) + T.reduce_sum( + acc_s, sumexp_i, dim=1 + ) # is this a accumulate operator? + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = sumexp[h_i] * alpha_local[h_i] + sumexp_i[h_i] + for h_i, d_i in T.Parallel(H_per_block, D // 2): + acc_o_l[h_i, d_i] *= alpha_local[h_i] + T.copy(alpha_local, alpha_shared) + + T.copy(acc_s, S_shared) + T.gemm(S_shared, KV_shared_1_l, acc_o_l) + + T.barrier_arrive(bar_sScale_and_sS_ready) + T.barrier_arrive(bar_k_1_free[0]) + + # Rescale + for h_i in T.Parallel(H_per_block): + sum_exp_shared[h_i] = sumexp[h_i] + T.barrier_arrive(bar_final) + for h_i, d_i in T.Parallel(H_per_block, D // 2): + acc_o_l[h_i, d_i] /= sumexp[h_i] + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale + T.copy(acc_o_l, O_shared_l) + T.copy(O_shared_l, Output[b_i, s_i, H0:H1, 0 : D // 2]) + elif tx >= 128 and tx < 256: + # T.set_max_nreg(168, 1) + T.fill(acc_o_r, 0) + for i_i in T.serial(T.ceildiv(NI, 2)): + # Buffer 0 + T.barrier_arrive(bar_sScale_and_sS_ready) + T.barrier_wait(bar_sScale_and_sS_ready, ((i_i * 2) & 1)) + T.barrier_arrive(bar_1_128) + T.barrier_wait(bar_1_128, 0) + for h_i, d_i in T.Parallel(H_per_block, D // 2): + acc_o_r[h_i, d_i] *= alpha_shared[h_i] + T.gemm(S_shared, KV_shared_0_r, acc_o_r) + T.barrier_arrive(bar_k_0_free[0]) + T.barrier_arrive(bar_sScale_and_sS_free) + + # Buffer 1 + T.barrier_arrive(bar_sScale_and_sS_ready) + T.barrier_wait(bar_sScale_and_sS_ready, ((i_i * 2 + 1) & 1)) + T.barrier_arrive(bar_1_128) + T.barrier_wait(bar_1_128, 1) + for h_i, d_i in T.Parallel(H_per_block, D // 2): + acc_o_r[h_i, d_i] *= alpha_shared[h_i] + T.gemm(S_shared, KV_shared_1_r, acc_o_r) + T.barrier_arrive(bar_k_1_free[0]) + if i_i != T.ceildiv(NI, 2) - 1: + T.barrier_arrive(bar_sScale_and_sS_free) + + # Rescale + T.barrier_wait(bar_final, 0) + for h_i, d_i in T.Parallel(H_per_block, D // 2): + acc_o_r[h_i, d_i] /= sum_exp_shared[h_i] + + T.copy(acc_o_r, O_shared_r) + T.copy(O_shared_r, Output[b_i, s_i, H0:H1, D // 2 : D]) + elif tx >= 256: + # producer + T.set_max_nreg(80, 0) + indices_local[0] = 0 + for i_i in T.serial(T.ceildiv(NI, 2)): + # Buffer 0 + T.barrier_wait(bar_k_0_free[0], ((i_i & 1) ^ 1)) + T.barrier_arrive(bar_2_128) + T.barrier_wait(bar_2_128, 0) + + for r in T.serial(4): + indices_tmp[0] = Indices[ + b_i, s_i, g_i, (i_i * 2) * BI + r * 16 + (tx - 256) // 8 + ] + is_kv_valid_0[r * 16 + (tx - 256) // 8] = indices_tmp[0] >= 0 + if is_kv_valid_0[r * 16 + (tx - 256) // 8]: + indices_local[0] = indices_tmp[0] + + with T.attr("default", "async_scope", 1): # type: ignore + for u in T.serial(4): + for v in T.vectorized(8): + KV_shared_0_l[ + r * 16 + (tx - 256) // 8, + 64 * u + (tx - 256) % 8 * 8 + v, + ] = KV[ + b_i, + indices_local[0], + g_i, + 64 * u + (tx - 256) % 8 * 8 + v, + ] + KV_shared_0_r[ + r * 16 + (tx - 256) // 8, + 64 * u + (tx - 256) % 8 * 8 + v, + ] = KV[ + b_i, + indices_local[0], + g_i, + D // 2 + 64 * u + (tx - 256) % 8 * 8 + v, + ] + with T.attr("default", "async_scope", 1): # type: ignore + for v in T.vectorized(8): + K_tail_shared_0[ + r * 16 + (tx - 256) // 8, (tx - 256) % 8 * 8 + v + ] = KV[ + b_i, + indices_local[0], + g_i, + D + (tx - 256) % 8 * 8 + v, + ] + + T.cp_async_barrier_noinc(bar_k_0_ready[0]) + + # Buffer 1 + T.barrier_wait(bar_k_1_free[0], ((i_i & 1) ^ 1)) + T.barrier_arrive(bar_2_128) + T.barrier_wait(bar_2_128, 1) + + for r in T.serial(4): + indices_tmp[0] = Indices[ + b_i, s_i, g_i, (i_i * 2 + 1) * BI + r * 16 + (tx - 256) // 8 + ] + is_kv_valid_1[r * 16 + (tx - 256) // 8] = indices_tmp[0] >= 0 + if is_kv_valid_1[r * 16 + (tx - 256) // 8]: + indices_local[0] = indices_tmp[0] + + with T.attr("default", "async_scope", 1): # type: ignore + for u in T.serial(4): + for v in T.vectorized(8): + KV_shared_1_l[ + r * 16 + (tx - 256) // 8, + 64 * u + (tx - 256) % 8 * 8 + v, + ] = KV[ + b_i, + indices_local[0], + g_i, + 64 * u + (tx - 256) % 8 * 8 + v, + ] + KV_shared_1_r[ + r * 16 + (tx - 256) // 8, + 64 * u + (tx - 256) % 8 * 8 + v, + ] = KV[ + b_i, + indices_local[0], + g_i, + D // 2 + 64 * u + (tx - 256) % 8 * 8 + v, + ] + with T.attr("default", "async_scope", 1): # type: ignore + for v in T.vectorized(8): + K_tail_shared_1[ + r * 16 + (tx - 256) // 8, (tx - 256) % 8 * 8 + v + ] = KV[ + b_i, + indices_local[0], + g_i, + D + (tx - 256) % 8 * 8 + v, + ] + + T.cp_async_barrier_noinc(bar_k_1_ready[0]) + + return main + + +@tilelang.jit( + out_idx=[-2, -1], + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + }, +) +def sparse_mla_fwd_decode_partial( + heads, + dim, + tail_dim, + topk, + *, + kv_group=1, + sm_scale=None, + is_causal=True, + block_I=64, + inner_iter=1, + num_stages=1, + threads=256, +): + """ + grid: (seq_len * REPLICATE_H, top_k / block_I / inner_iter) + Each GPU block processes `inner_iter` consecutive KV tiles and writes one (partial_o, partial_lse) entry. + """ + + assert is_causal == True, "non-causal is not supported" + assert kv_group == 1 + assert topk % block_I == 0 + assert topk % (block_I * inner_iter) == 0, ( + f"topk ({topk}) must be divisible by block_I * inner_iter = " + f"{block_I} * {inner_iter}" + ) + + # log2(e) = 1.44269504 + if sm_scale is None: + sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 + else: + sm_scale = sm_scale * 1.44269504 + + batch = 1 + seq_len = T.dynamic("seq_len") + seq_len_kv = T.dynamic("seq_len_kv") + + head_kv = heads // kv_group + padded_H = max(tilelang.math.next_power_of_2(head_kv), 16) + REPLICATE_H = (head_kv // 64) if head_kv > 64 else 1 + H_per_block = padded_H if REPLICATE_H == 1 else 64 + N_GROUPS = topk // (block_I * inner_iter) + BI = block_I + D = dim + D_tail = tail_dim + + q_shape = [batch, seq_len, heads, dim + tail_dim] + kv_shape = [batch, seq_len_kv, kv_group, dim + tail_dim] + indices_shape = [batch, seq_len, kv_group, topk] + partial_o_shape = [batch, seq_len, N_GROUPS, heads, dim] + partial_lse_shape = [batch, seq_len, N_GROUPS, heads] + indices_dtype = T.int32 + dtype = T.bfloat16 + accum_dtype = T.float32 + + _q_in_shared = inner_iter == 1 + + @T.prim_func + def main( + Q: T.Tensor(q_shape, dtype), + KV: T.Tensor(kv_shape, dtype), + Indices: T.Tensor(indices_shape, indices_dtype), + Partial_O: T.Tensor(partial_o_shape, dtype), + Partial_Lse: T.Tensor(partial_lse_shape, accum_dtype), + ): + with T.Kernel(seq_len * REPLICATE_H, N_GROUPS, threads=threads) as (bx, by): + if _q_in_shared: + Q_buf = T.alloc_shared([H_per_block, D], dtype) + Q_tail_buf = T.alloc_shared([H_per_block, D_tail], dtype) + else: + Q_buf = T.alloc_fragment([H_per_block, D], dtype) + Q_tail_buf = T.alloc_fragment([H_per_block, D_tail], dtype) + + KV_shared = T.alloc_shared([BI, D], dtype) + K_tail_shared = T.alloc_shared([BI, D_tail], dtype) + S_shared = T.alloc_shared([H_per_block, BI], dtype) + mask = T.alloc_fragment([BI], T.bool) + + acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) + acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) + sumexp = T.alloc_fragment([H_per_block], accum_dtype) + sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) + alpha = T.alloc_fragment([H_per_block], accum_dtype) + m_i = T.alloc_fragment([H_per_block], accum_dtype) + m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) + + T.fill(acc_o, 0) + T.fill(sumexp, 0) + T.fill(m_i, -(2**30)) + + b_i, g_i = 0, 0 + s_i = bx if REPLICATE_H == 1 else (bx // REPLICATE_H) + group_i = by + H0 = 0 if REPLICATE_H == 1 else (bx % REPLICATE_H) * 64 + H1 = H0 + H_per_block + + T.copy(Q[b_i, s_i, H0:H1, :D], Q_buf) + T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_buf) + + for k_i in T.Pipelined(inner_iter, num_stages=num_stages): + topk_block_i = group_i * inner_iter + k_i + + for bi_i in T.Parallel(BI): + mask[bi_i] = Indices[b_i, s_i, g_i, topk_block_i * BI + bi_i] >= 0 + for bi_i, d_i in T.Parallel(BI, D): + idx = Indices[b_i, s_i, g_i, topk_block_i * BI + bi_i] + KV_shared[bi_i, d_i] = KV[ + b_i, T.if_then_else(idx >= 0, idx, 0), g_i, d_i + ] + for bi_i, d_i in T.Parallel(BI, D_tail): + idx = Indices[b_i, s_i, g_i, topk_block_i * BI + bi_i] + K_tail_shared[bi_i, d_i] = KV[ + b_i, T.if_then_else(idx >= 0, idx, 0), g_i, D + d_i + ] + + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else( + mask[bi_i], 0, -T.infinity(acc_s.dtype) + ) + + T.gemm( + Q_buf, + KV_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullCol, + ) + T.gemm( + Q_tail_buf, + K_tail_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullCol, + ) + + T.copy(m_i, m_i_prev) + T.reduce_max(acc_s, m_i, dim=1, clear=False) + for h_i in T.Parallel(H_per_block): + alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.exp2( + acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale + ) + T.reduce_sum(acc_s, sumexp_i, dim=1) + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] *= alpha[h_i] + + T.copy(acc_s, S_shared) + T.gemm(S_shared, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullCol) + + # sumexp==0 (all masked), divide by 1 to get 0 and avoid nan + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] = acc_o[h_i, d_i] / T.if_then_else( + sumexp[h_i] == 0.0, 1.0, sumexp[h_i] + ) + # sumexp==0 (all masked), use large negative so combine ignores this split + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = T.if_then_else( + sumexp[h_i] == 0.0, + -(2**30), + T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale, + ) + + T.copy(acc_o, Partial_O[b_i, s_i, group_i, H0:H1, :]) + T.copy(sumexp, Partial_Lse[b_i, s_i, group_i, H0:H1]) + + return main + + +@tilelang.jit( + out_idx=[-1], + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + }, +) +def sparse_mla_fwd_decode_combine( + heads, + dim, + topk, + head_per_block, + *, + block_I=64, + threads=256, +): + """ + grid: (seq_len * REPLICATE_H). batch=1, kv_group=1. + Each block does one tile of heads (e.g. 4 or 8 for decode). + """ + + assert heads % head_per_block == 0, f"head_per_block must divide heads" + + batch = 1 + seq_len = T.dynamic("seq_len") + + NI = topk // block_I + H_per_block = head_per_block + REPLICATE_H = heads // H_per_block + + partial_o_shape = [batch, seq_len, NI, heads, dim] + partial_lse_shape = [batch, seq_len, NI, heads] + o_shape = [batch, seq_len, heads, dim] + dtype = T.bfloat16 + accum_dtype = T.float32 + + @T.prim_func + def main( + Partial_O: T.Tensor(partial_o_shape, dtype), + Partial_Lse: T.Tensor(partial_lse_shape, accum_dtype), + Output: T.Tensor(o_shape, dtype), + ): + with T.Kernel(seq_len * REPLICATE_H, threads=threads) as (bx,): + shared_lse = T.alloc_shared([NI, H_per_block], accum_dtype) + + lse_max = T.alloc_fragment([H_per_block], accum_dtype) + lse_sum = T.alloc_fragment([H_per_block], accum_dtype) + scale = T.alloc_fragment([H_per_block, NI], accum_dtype) + acc_o = T.alloc_fragment([H_per_block, dim], accum_dtype) + + b_i = 0 + s_i = bx if REPLICATE_H == 1 else (bx // REPLICATE_H) + H0 = 0 if REPLICATE_H == 1 else (bx % REPLICATE_H) * H_per_block + H1 = H0 + H_per_block + + for k in T.serial(NI): + T.copy(Partial_Lse[b_i, s_i, k, H0:H1], shared_lse[k, :]) + + T.fill(lse_max, -(2**30)) + for k in T.serial(NI): + for h_i in T.Parallel(H_per_block): + lse_max[h_i] = T.max(lse_max[h_i], shared_lse[k, h_i]) + T.fill(lse_sum, 0) + for k in T.serial(NI): + for h_i in T.Parallel(H_per_block): + lse_sum[h_i] = lse_sum[h_i] + T.exp2( + shared_lse[k, h_i] - lse_max[h_i] + ) + for k in T.serial(NI): + for h_i in T.Parallel(H_per_block): + scale[h_i, k] = T.exp2( + shared_lse[k, h_i] - lse_max[h_i] - T.log2(lse_sum[h_i]) + ) + + T.fill(acc_o, 0) + for k in T.serial(NI): + for h_i, d_i in T.Parallel(H_per_block, dim): + acc_o[h_i, d_i] = acc_o[h_i, d_i] + scale[h_i, k] * Partial_O[ + b_i, s_i, k, H0 + h_i, d_i + ].astype(accum_dtype) + + T.copy(acc_o, Output[b_i, s_i, H0:H1, :]) + + return main + + +@tilelang.jit(out_idx=[-2, -1], pass_configs=pass_configs) +def sparse_mla_fwd_decode_partial_fp8( + num_heads: int, + d_v: int, + d_tail: int, + topk: int, + *, + sm_scale=None, + block_I=64, + inner_iter=1, + threads=256, +): + assert d_v == 512, f"only support d_v=512" + assert ( + topk % block_I == 0 + ), "otherwise will load some index=0 thus causing wrong kv to be loaded" + + # Softmax scores are in [0, 1]. We scale by fp8_max_val before FP8 cast + # to better utilize FP8 dynamic range, then apply the inverse scale after GEMM. + # This is numerically safe because softmax output is bounded by 1. + fp8_dtype = "float8_e4m3fnuz" if _is_fp8_fnuz else "float8_e4m3fn" + fp8_max_val = 240.0 if _is_fp8_fnuz else 448.0 + s_inv_scale_const = fp8_max_val + s_scale_const = 1.0 / fp8_max_val + + BI = block_I + group_size = 128 + dim_quant_fp8 = d_v + d_tail + rope_offset_fp8 = d_v + n_groups = topk // (BI * inner_iter) + + if sm_scale is None: + sm_scale = (1.0 / (d_v + d_tail)) ** 0.5 * 1.44269504 + else: + sm_scale = sm_scale * 1.44269504 + + h_per_block = 16 + # Match bf16 partial behavior: keep fixed 16-head tiles and use + # sliced T.copy on H0:H1 for tail handling. + assert ( + num_heads <= h_per_block or num_heads % h_per_block == 0 + ), "num_heads must be <=16 or divisible by 16" + head_blocks_per_seq = (num_heads + h_per_block - 1) // h_per_block + + batch = 1 + kv_group = 1 + seq_len = T.symbolic("seq_len") + num_pages = T.symbolic("num_pages") + + q_fp8_shape = [batch, seq_len, num_heads, d_v + d_tail] + kv_fp8_shape = [batch, num_pages, kv_group, dim_quant_fp8] + idx_shape = [batch, seq_len, kv_group, topk] + partial_o_shape = [batch, seq_len, n_groups, num_heads, d_v] + partial_lse_shape = [batch, seq_len, n_groups, num_heads] + + accum_dtype = T.float32 + dtype_bf16 = T.bfloat16 + + @T.prim_func + def main( + q_fp8: T.Tensor(q_fp8_shape, fp8_dtype), + kv_fp8: T.Tensor(kv_fp8_shape, fp8_dtype), + indices: T.Tensor(idx_shape, T.int32), + partial_o: T.Tensor(partial_o_shape, dtype_bf16), + partial_lse: T.Tensor(partial_lse_shape, accum_dtype), + ): + with T.Kernel(seq_len * head_blocks_per_seq, n_groups, threads=threads) as ( + bx, + by, + ): + b_i, g_i = 0, 0 + s_i = bx // head_blocks_per_seq + group_i = by + H0 = (bx % head_blocks_per_seq) * h_per_block + H1 = H0 + h_per_block + + # We intentionally split the K=512 GEMM into 4x128 tiles. + # Although this adds extra intermediate memory traffic, + # it shortens the MFMA accumulation dependency chain and improves performance. + q_tile0 = T.alloc_shared([h_per_block, group_size], fp8_dtype) + q_tile1 = T.alloc_shared([h_per_block, group_size], fp8_dtype) + q_tile2 = T.alloc_shared([h_per_block, group_size], fp8_dtype) + q_tile3 = T.alloc_shared([h_per_block, group_size], fp8_dtype) + kv_tile0 = T.alloc_shared([BI, group_size], fp8_dtype) + kv_tile1 = T.alloc_shared([BI, group_size], fp8_dtype) + kv_tile2 = T.alloc_shared([BI, group_size], fp8_dtype) + kv_tile3 = T.alloc_shared([BI, group_size], fp8_dtype) + q_tail_buf = T.alloc_shared([h_per_block, d_tail], fp8_dtype) + k_tail_shared = T.alloc_shared([BI, d_tail], fp8_dtype) + s_fp8_shared = T.alloc_shared([h_per_block, BI], fp8_dtype) + page_idx_shared = T.alloc_shared([BI], T.int32) + + mask = T.alloc_fragment([BI], T.bool) + acc_s = T.alloc_fragment([h_per_block, BI], accum_dtype) + acc_tile = T.alloc_fragment([h_per_block, BI], accum_dtype) + sv_tile = T.alloc_fragment([h_per_block, group_size], accum_dtype) + sumexp = T.alloc_fragment([h_per_block], accum_dtype) + sumexp_i = T.alloc_fragment([h_per_block], accum_dtype) + alpha = T.alloc_fragment([h_per_block], accum_dtype) + m_i = T.alloc_fragment([h_per_block], accum_dtype) + m_i_prev = T.alloc_fragment([h_per_block], accum_dtype) + inv_denom = T.alloc_fragment([h_per_block], accum_dtype) + + acc_o_tile0 = T.alloc_fragment([h_per_block, group_size], accum_dtype) + acc_o_tile1 = T.alloc_fragment([h_per_block, group_size], accum_dtype) + acc_o_tile2 = T.alloc_fragment([h_per_block, group_size], accum_dtype) + acc_o_tile3 = T.alloc_fragment([h_per_block, group_size], accum_dtype) + + T.fill(acc_o_tile0, 0) + T.fill(acc_o_tile1, 0) + T.fill(acc_o_tile2, 0) + T.fill(acc_o_tile3, 0) + T.fill(sumexp, 0) + T.fill(m_i, -(2**30)) + + T.copy(q_fp8[b_i, s_i, H0:H1, d_v:], q_tail_buf) + T.copy(q_fp8[b_i, s_i, H0:H1, 0 * group_size : 1 * group_size], q_tile0) + T.copy(q_fp8[b_i, s_i, H0:H1, 1 * group_size : 2 * group_size], q_tile1) + T.copy(q_fp8[b_i, s_i, H0:H1, 2 * group_size : 3 * group_size], q_tile2) + T.copy(q_fp8[b_i, s_i, H0:H1, 3 * group_size : 4 * group_size], q_tile3) + + for k_i in T.serial(inner_iter): + topk_block_i = group_i * inner_iter + k_i + + for bi_i in T.Parallel(BI): + idx = indices[b_i, s_i, g_i, topk_block_i * BI + bi_i] + valid = idx >= 0 + page_idx_shared[bi_i] = T.if_then_else(valid, idx, 0) + mask[bi_i] = valid + + for bi_i, j in T.Parallel(BI, group_size): + page = page_idx_shared[bi_i] + kv_tile0[bi_i, j] = kv_fp8[b_i, page, g_i, 0 * group_size + j] + kv_tile1[bi_i, j] = kv_fp8[b_i, page, g_i, 1 * group_size + j] + kv_tile2[bi_i, j] = kv_fp8[b_i, page, g_i, 2 * group_size + j] + kv_tile3[bi_i, j] = kv_fp8[b_i, page, g_i, 3 * group_size + j] + + for bi_i, j in T.Parallel(BI, d_tail): + page = page_idx_shared[bi_i] + k_tail_shared[bi_i, j] = kv_fp8[b_i, page, g_i, rope_offset_fp8 + j] + + for h_i, bi_i in T.Parallel(h_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else( + mask[bi_i], 0, -T.infinity(acc_s.dtype) + ) + + T.gemm(q_tile0, kv_tile0, acc_s, transpose_B=True, clear_accum=False) + T.gemm(q_tile1, kv_tile1, acc_tile, transpose_B=True, clear_accum=True) + for h_i, bi_i in T.Parallel(h_per_block, BI): + acc_s[h_i, bi_i] += acc_tile[h_i, bi_i] + T.gemm(q_tile2, kv_tile2, acc_tile, transpose_B=True, clear_accum=True) + for h_i, bi_i in T.Parallel(h_per_block, BI): + acc_s[h_i, bi_i] += acc_tile[h_i, bi_i] + T.gemm(q_tile3, kv_tile3, acc_tile, transpose_B=True, clear_accum=True) + for h_i, bi_i in T.Parallel(h_per_block, BI): + acc_s[h_i, bi_i] += acc_tile[h_i, bi_i] + T.gemm( + q_tail_buf, + k_tail_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullCol, + ) + + T.copy(m_i, m_i_prev) + T.reduce_max(acc_s, m_i, dim=1, clear=False) + for h_i in T.Parallel(h_per_block): + alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) + for h_i, bi_i in T.Parallel(h_per_block, BI): + acc_s[h_i, bi_i] = T.exp2( + acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale + ) + T.reduce_sum(acc_s, sumexp_i, dim=1) + for h_i in T.Parallel(h_per_block): + sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] + for h_i, j in T.Parallel(h_per_block, group_size): + acc_o_tile0[h_i, j] = acc_o_tile0[h_i, j] * alpha[h_i] + acc_o_tile1[h_i, j] = acc_o_tile1[h_i, j] * alpha[h_i] + acc_o_tile2[h_i, j] = acc_o_tile2[h_i, j] * alpha[h_i] + acc_o_tile3[h_i, j] = acc_o_tile3[h_i, j] * alpha[h_i] + + for h_i, bi_i in T.Parallel(h_per_block, BI): + s_fp8_shared[h_i, bi_i] = T.clamp( + acc_s[h_i, bi_i] * s_inv_scale_const, + -fp8_max_val, + fp8_max_val, + ) + T.gemm(s_fp8_shared, kv_tile0, sv_tile, clear_accum=True) + for h_i, j in T.Parallel(h_per_block, group_size): + acc_o_tile0[h_i, j] = ( + acc_o_tile0[h_i, j] + sv_tile[h_i, j] * s_scale_const + ) + + T.gemm(s_fp8_shared, kv_tile1, sv_tile, clear_accum=True) + for h_i, j in T.Parallel(h_per_block, group_size): + acc_o_tile1[h_i, j] = ( + acc_o_tile1[h_i, j] + sv_tile[h_i, j] * s_scale_const + ) + + T.gemm(s_fp8_shared, kv_tile2, sv_tile, clear_accum=True) + for h_i, j in T.Parallel(h_per_block, group_size): + acc_o_tile2[h_i, j] = ( + acc_o_tile2[h_i, j] + sv_tile[h_i, j] * s_scale_const + ) + + T.gemm(s_fp8_shared, kv_tile3, sv_tile, clear_accum=True) + for h_i, j in T.Parallel(h_per_block, group_size): + acc_o_tile3[h_i, j] = ( + acc_o_tile3[h_i, j] + sv_tile[h_i, j] * s_scale_const + ) + + for h_i in T.Parallel(h_per_block): + denom = T.if_then_else(sumexp[h_i] == 0.0, 1.0, sumexp[h_i]) + inv_denom[h_i] = 1.0 / denom + for h_i, j in T.Parallel(h_per_block, group_size): + acc_o_tile0[h_i, j] = acc_o_tile0[h_i, j] * inv_denom[h_i] + acc_o_tile1[h_i, j] = acc_o_tile1[h_i, j] * inv_denom[h_i] + acc_o_tile2[h_i, j] = acc_o_tile2[h_i, j] * inv_denom[h_i] + acc_o_tile3[h_i, j] = acc_o_tile3[h_i, j] * inv_denom[h_i] + + for h_i in T.Parallel(h_per_block): + sumexp[h_i] = T.if_then_else( + sumexp[h_i] == 0.0, + -(2**30), + T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale, + ) + + T.copy( + acc_o_tile0, + partial_o[b_i, s_i, group_i, H0:H1, 0 * group_size : 1 * group_size], + ) + T.copy( + acc_o_tile1, + partial_o[b_i, s_i, group_i, H0:H1, 1 * group_size : 2 * group_size], + ) + T.copy( + acc_o_tile2, + partial_o[b_i, s_i, group_i, H0:H1, 2 * group_size : 3 * group_size], + ) + T.copy( + acc_o_tile3, + partial_o[b_i, s_i, group_i, H0:H1, 3 * group_size : 4 * group_size], + ) + + T.copy(sumexp, partial_lse[b_i, s_i, group_i, H0:H1]) + + return main + + +def tilelang_sparse_fwd( + q: torch.Tensor, + kv: torch.Tensor, + indices: torch.Tensor, + sm_scale: float, + d_v: int = 512, +) -> torch.Tensor: + assert q.dim() == 3 and kv.dim() == 3 and indices.dim() == 3 + num_heads = q.shape[1] + dim = q.shape[2] + tail_dim = dim - d_v + topk = indices.shape[-1] + assert topk == 2048 + + if _is_hip: + is_fp8_kv = kv.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz) + if is_fp8_kv: + if q.dtype != kv.dtype: + q = q.to(kv.dtype) + if _is_gfx95_supported: + block_I, threads, block_per_cu, cu = 64, 256, 2, 256 + else: + block_I, threads, block_per_cu, cu = 64, 256, 1, 304 + ni = topk // block_I + inner_iter = _pick_inner_iter(q.shape[0], ni, cu, block_per_cu) + kernel_partial = sparse_mla_fwd_decode_partial_fp8( + num_heads, + d_v, + tail_dim, + topk, + sm_scale=sm_scale, + block_I=block_I, + inner_iter=inner_iter, + threads=threads, + ) + else: + if _is_gfx95_supported: + block_I, threads, block_per_cu, cu = 64, 256, 2, 256 + else: + block_I, threads, block_per_cu, cu = 32, 128, 1, 304 + ni = topk // block_I + inner_iter = _pick_inner_iter(q.shape[0], ni, cu, block_per_cu) + kernel_partial = sparse_mla_fwd_decode_partial( + num_heads, + d_v, + tail_dim, + topk, + sm_scale=sm_scale, + block_I=block_I, + inner_iter=inner_iter, + threads=threads, + ) + partial_o_batched, partial_lse_batched = kernel_partial( + q.unsqueeze(0), kv.unsqueeze(0), indices.unsqueeze(0) + ) + n_groups = ni // inner_iter + kernel_combine = sparse_mla_fwd_decode_combine( + num_heads, + d_v, + n_groups * block_I, + head_per_block=4, + block_I=block_I, + threads=threads, + ) + out = kernel_combine(partial_o_batched, partial_lse_batched) + elif torch.cuda.get_device_capability(q.device)[0] == 12: + # The default CUDA kernel double-buffers a 64-token KV tile and asks + # for about 170 KiB of dynamic shared memory for GLM-5 DSA. SM120 + # client Blackwell GPUs expose only 99 KiB per block. Reuse the + # numerically equivalent partial/combine implementation with a + # 32-token tile so the kernel fits the device limit. + block_I = 32 + threads = 128 + inner_iter = 1 + ni = topk // block_I + kernel_partial = sparse_mla_fwd_decode_partial( + num_heads, + d_v, + tail_dim, + topk, + sm_scale=sm_scale, + block_I=block_I, + inner_iter=inner_iter, + num_stages=1, + threads=threads, + ) + partial_o_batched, partial_lse_batched = kernel_partial( + q.unsqueeze(0), kv.unsqueeze(0), indices.unsqueeze(0) + ) + n_groups = ni // inner_iter + kernel_combine = sparse_mla_fwd_decode_combine( + num_heads, + d_v, + n_groups * block_I, + head_per_block=4, + block_I=block_I, + threads=threads, + ) + out = kernel_combine(partial_o_batched, partial_lse_batched) + else: + kernel = sparse_attention_fwd_kernel_v2( + num_heads, d_v, tail_dim, topk, sm_scale=sm_scale + ) + out = kernel(q.unsqueeze(0), kv.unsqueeze(0), indices.unsqueeze(0)) # type: ignore + return out + + +@functools.cache +def fp8_paged_mqa_logits_kernel( + head_dim: int = 128, + num_heads: int = 64, + block_size: int = 64, + clear_accum: bool = True, + split_kv: int = 1, +) -> Any: + N = T.symbolic("batch_size") + L = T.symbolic("max_table_length") + S = T.symbolic("max_seq_len") + C = T.symbolic("num_blocks") + B = block_size + D = head_dim + H = num_heads + SK = int(split_kv) + BLOCK_BYTES = B * (D + 4) + SCALE_OFFSET = B * D + + assert D % 4 == 0 + assert H % 4 == 0 + assert D == 128 + assert SK >= 1 + + @tilelang.jit( + pass_configs={ + **pass_configs, + tilelang.PassConfigKey.TL_DISABLE_SAFE_MEMORY_ACCESS: True, + } + ) + def fp8_paged_mqa_logits( + q: T.Tensor[(N, H, D), FP8], + kvcache_u8: T.Tensor[(C, BLOCK_BYTES), UINT8], + weight: T.Tensor[(N, H), FP32], + seq_lens: T.Tensor[(N,), INT32], + page_table: T.Tensor[(N, L), INT32], + o: T.Tensor[(N, S), FP32], + ) -> None: + _ = N, L, S, C, D, H, B + with T.Kernel(N * SK) as bxs: + bx = bxs % N + pid_split = bxs // N + seq_len = seq_lens[bx] + np_total = T.ceildiv(seq_len, B) + stride = T.ceildiv(np_total, SK) + i_start = pid_split * stride + n_iters = T.max(0, T.min(stride, np_total - i_start)) + + q_smem = T.alloc_shared((H, D), FP8) + q_s_frag = T.alloc_fragment((H,), FP32) + T.copy(q[bx, 0, 0], q_smem) + T.copy(weight[bx, 0], q_s_frag) + + for j in T.Pipelined(n_iters, num_stages=2): + i = i_start + j + page = page_table[bx, i] + k_smem_u8 = T.alloc_shared((1, B * D), UINT8) + T.copy(kvcache_u8[page : page + 1, 0:SCALE_OFFSET], k_smem_u8) + k_smem = T.view(k_smem_u8, (B, D), FP8) + k_s_smem_u8 = T.alloc_shared((1, B * 4), UINT8) + T.copy( + kvcache_u8[page : page + 1, SCALE_OFFSET:BLOCK_BYTES], + k_s_smem_u8, + ) + k_s_smem = T.view(k_s_smem_u8, (B,), FP32) + k_s_frag = T.alloc_fragment((B,), FP32) + T.copy(k_s_smem, k_s_frag) + + logits = T.alloc_fragment((B, H), FP32) + if not clear_accum: + T.fill(logits, 0.0) + T.gemm( + k_smem, + q_smem, + logits, + transpose_A=False, + transpose_B=True, + clear_accum=clear_accum, + ) + + # post processing + for h, j2 in T.Parallel(H, B): + logits[j2, h] = T.max(logits[j2, h], 0.0) * q_s_frag[h] + logits_sum = T.alloc_fragment((B,), FP32) + T.reduce_sum(logits, logits_sum, dim=1) + for j2 in T.Parallel(B): + logits_sum[j2] *= k_s_frag[j2] + T.copy(logits_sum, o[bx, i * B]) + + return fp8_paged_mqa_logits + + +def tilelang_fp8_paged_mqa_logits( + q_fp8: torch.Tensor, + kvcache_fp8: torch.Tensor, + weight: torch.Tensor, + seq_lens: torch.Tensor, + page_table: torch.Tensor, + deep_gemm_metadata: Any, + max_seq_len: int, + clean_logits: bool = True, +) -> torch.Tensor: + _ = deep_gemm_metadata + batch_size, _, num_heads, head_dim = q_fp8.shape + block_size = kvcache_fp8.shape[1] + assert head_dim == 128, "TODO" + assert block_size == 64, "TODO" + assert q_fp8.shape == (batch_size, 1, num_heads, head_dim) + assert kvcache_fp8.shape[1:] == (block_size, 1, head_dim + 4) + assert weight.shape == (batch_size, num_heads) + assert seq_lens.shape == (batch_size,) + assert page_table.shape[0] == batch_size + assert clean_logits == False + + logits = page_table.new_empty((batch_size, max_seq_len), dtype=torch.float32) + + NUM_CU = 256 + split_kv = split_kv = max(1, min(max_seq_len // block_size, NUM_CU // batch_size)) + kernel = fp8_paged_mqa_logits_kernel( + head_dim=head_dim, + num_heads=num_heads, + block_size=block_size, + clear_accum=clean_logits, + split_kv=split_kv, + ) + q_fp8 = q_fp8.view(batch_size, num_heads, head_dim) + kvcache_u8 = kvcache_fp8.view(-1, block_size * (head_dim + 4)) + kernel(q_fp8, kvcache_u8, weight, seq_lens, page_table, logits) + return logits + + +def _build_fp8_combined_view(k_cache: torch.Tensor) -> Tuple[torch.Tensor, int, int]: + """ + Reinterpret a MODEL1_FP8Sparse KV cache as a contiguous uint32 view. + Input: k_cache (num_blocks, block_size, 1, d_qk) fp8/uint8 + — per-block storage also holds scales + padding past d_qk. + Output: (num_blocks, block_pad_u32) uint32 covering the full block + stride. Same storage ashe input, no copy. + """ + k_u8 = k_cache.view(torch.uint8) if k_cache.dtype != torch.uint8 else k_cache + num_blocks = k_u8.shape[0] + block_size = k_u8.shape[1] + block_pad_u32 = k_u8.stride(0) // 4 + storage = k_u8.untyped_storage() + flat_u32 = torch.empty(0, dtype=torch.uint32, device=k_u8.device).set_( + storage, 0, (storage.nbytes() // 4,), (1,) + ) + k_combined = torch.as_strided( + flat_u32, + size=(num_blocks, block_pad_u32), + stride=(block_pad_u32, 1), + storage_offset=k_u8.storage_offset() // 4, + ) + return k_combined, num_blocks, block_size + + +_TOPK_LEN_SENTINEL_CACHE: dict = {} +_INT32_MAX = 2**30 + + +def _topk_length_sentinel(device: torch.device, batch: int) -> torch.Tensor: + """Cached `(batch,) int32 INT_MAX` tensor used when `topk_length` is None.""" + cur = _TOPK_LEN_SENTINEL_CACHE.get(device) + if cur is None or cur.numel() < batch: + cur = torch.full( + (max(batch, 256),), _INT32_MAX, dtype=torch.int32, device=device + ) + _TOPK_LEN_SENTINEL_CACHE[device] = cur + return cur[:batch] + + +@tilelang.jit( + out_idx=[-2, -1], + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, +) +def dpsk_v4_fp8_partial_kernel( + num_heads: int, + topk_1: int, + block_size_kv_1: int, + topk_2: int = 0, + block_size_kv_2: int = 0, + *, + dim: int = 448, + tail_dim: int = 64, + sm_scale: float = 0.0, + block_I: int = 64, + inner_iter_1: int = 1, + inner_iter_2: int = 0, + num_stages: int = 0, + threads: int = 512, +) -> Any: + """ + Read FP8 K cache directly, dequantise to BF16 in-kernel, do flash-attn + online softmax with split-K. Supports a second cache (`topk_2>0`) and + `attn_sink` is folded later by the combine kernel. + """ + log2e: float = 1.44269504 + if sm_scale <= 0.0: + sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * log2e + else: + sm_scale = sm_scale * log2e + assert dim == 448 and tail_dim == 64 + assert topk_1 % block_I == 0 + assert ( + topk_1 // block_I + ) % inner_iter_1 == 0, ( + f"NI_1={topk_1 // block_I} must be divisible by inner_iter_1={inner_iter_1}" + ) + assert block_size_kv_1 > 0 and (block_size_kv_1 & (block_size_kv_1 - 1)) == 0 + + is_dual = topk_2 > 0 + if is_dual: + assert inner_iter_2 > 0, "dual-cache call requires inner_iter_2 > 0" + assert topk_2 % block_I == 0 + assert ( + topk_2 // block_I + ) % inner_iter_2 == 0, ( + f"NI_2={topk_2 // block_I} must be divisible by inner_iter_2={inner_iter_2}" + ) + assert block_size_kv_2 > 0 and (block_size_kv_2 & (block_size_kv_2 - 1)) == 0 + + PACKED_W = dim + 2 * tail_dim + NOPE_TILE = 64 + NUM_TILES = dim // NOPE_TILE + SCALE_W = 8 + PACKED_W4 = PACKED_W // 4 + SCALE_W4 = SCALE_W // 4 + + kv_group = 1 + batch = T.symbolic("batch") + seq_len = T.symbolic("seq_len") + num_blocks_kv_1 = T.symbolic("num_blocks_kv_1") + block_pad_u32_1 = T.symbolic("block_pad_u32_1") + if is_dual: + num_blocks_kv_2 = T.symbolic("num_blocks_kv_2") + block_pad_u32_2 = T.symbolic("block_pad_u32_2") + + head_kv = num_heads // kv_group + D = dim + D_tail = tail_dim + BI = block_I + padded_H = max(tilelang.math.next_power_of_2(head_kv), 16) + if head_kv > 64: + assert head_kv % 64 == 0 + REPLICATE_H = (head_kv + 63) // 64 if head_kv > 64 else 1 + H_per_block = 64 if REPLICATE_H > 1 else padded_H + + NI_1 = topk_1 // BI + n_groups_1 = NI_1 // inner_iter_1 + NI_2 = (topk_2 // BI) if is_dual else 0 + n_groups_2 = (NI_2 // inner_iter_2) if is_dual else 0 + n_groups = n_groups_1 + n_groups_2 + + BS_KV_1 = block_size_kv_1 + NOPE_ROPE_U32_PER_BLOCK_1 = BS_KV_1 * PACKED_W4 + if is_dual: + BS_KV_2 = block_size_kv_2 + NOPE_ROPE_U32_PER_BLOCK_2 = BS_KV_2 * PACKED_W4 + + q_shape = [batch, seq_len, num_heads, D + D_tail] + k1_shape = [num_blocks_kv_1, block_pad_u32_1] + indices1_shape = [batch, seq_len, topk_1] + topk_length_shape = [batch] + partial_o_shape = [batch, seq_len, n_groups, num_heads, D + D_tail] + partial_lse_shape = [batch, seq_len, n_groups, num_heads] + if is_dual: + k2_shape = [num_blocks_kv_2, block_pad_u32_2] + indices2_shape = [batch, seq_len, topk_2] + + accum_dtype = "float" + indices_dtype = INT32 + + if is_dual: + + @T.prim_func + def main( + Q: T.Tensor(q_shape, BF16), # type: ignore + K_combined_1: T.Tensor(k1_shape, "uint32"), # type: ignore + Indices_1: T.Tensor(indices1_shape, indices_dtype), # type: ignore + Topk_length_1: T.Tensor(topk_length_shape, indices_dtype), # type: ignore + K_combined_2: T.Tensor(k2_shape, "uint32"), # type: ignore + Indices_2: T.Tensor(indices2_shape, indices_dtype), # type: ignore + Topk_length_2: T.Tensor(topk_length_shape, indices_dtype), # type: ignore + Partial_O: T.Tensor(partial_o_shape, BF16), # type: ignore + Partial_LSE: T.Tensor(partial_lse_shape, accum_dtype), # type: ignore + ) -> None: + """ + grid: (seq_len * REPLICATE_H * n_groups, batch, 1) + Each block processes `inner_iter_1` (or `inner_iter_2`) consecutive + KV tiles of one phase and writes one (partial_o, partial_lse) entry. + """ + with T.Kernel( + seq_len * REPLICATE_H * n_groups, batch, kv_group, threads=threads + ) as (bx, by, bz): + Q_shared = T.alloc_fragment([H_per_block, D], BF16) + Q_tail_shared = T.alloc_fragment([H_per_block, D_tail], BF16) + K_packed_shared = T.alloc_shared([BI, PACKED_W4], "uint32") + K_scale_shared = T.alloc_shared([BI, SCALE_W4], "uint32") + KV_shared = T.alloc_shared([BI, D], BF16) + K_tail_shared = T.alloc_shared([BI, D_tail], BF16) + S_shared = T.alloc_shared([H_per_block, BI], BF16) + page_idx_shared = T.alloc_shared([BI], INT32) + + mask = T.alloc_fragment([BI], "bool") + scale_byte_local = T.alloc_fragment([BI, NUM_TILES], "uint32") + + acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) + acc_o_tail = T.alloc_fragment([H_per_block, D_tail], accum_dtype) + acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) + sumexp = T.alloc_fragment([H_per_block], accum_dtype) + sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) + alpha = T.alloc_fragment([H_per_block], accum_dtype) + m_i = T.alloc_fragment([H_per_block], accum_dtype) + m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) + + T.fill(acc_o, 0) + T.fill(acc_o_tail, 0) + T.fill(sumexp, 0) + T.fill(m_i, -(2**30)) + + b_i, g_i = by, bz + # bx encodes (s_i, h_replicate, group_i). + spans_per_seq = REPLICATE_H * n_groups + s_i = bx // spans_per_seq + rest = bx % spans_per_seq + group_i = rest // REPLICATE_H + h_rep = rest % REPLICATE_H + H0 = g_i * padded_H + (0 if REPLICATE_H == 1 else h_rep * 64) + H1 = H0 + H_per_block + + tk_len_1 = Topk_length_1[b_i] + tk_len_2 = Topk_length_2[b_i] + actual_n_groups_1 = T.ceildiv(tk_len_1, BI * inner_iter_1) + actual_n_groups_2 = T.ceildiv(tk_len_2, BI * inner_iter_2) + + if (group_i < n_groups_1) & (group_i < actual_n_groups_1): + # Phase 1 active: SWA cache work + Partial_O write. + T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared) + T.copy(Q[b_i, s_i, H0:H1, D : D + D_tail], Q_tail_shared) + for k_i in T.Pipelined(inner_iter_1, num_stages=num_stages): + iter_i = group_i * inner_iter_1 + k_i + for bi_i in T.Parallel(BI): + pos = iter_i * BI + bi_i + idx = Indices_1[b_i, s_i, pos] + valid = (idx >= 0) & (pos < tk_len_1) + page_idx_shared[bi_i] = T.if_then_else(valid, idx, 0) + mask[bi_i] = valid + + for bi_i, w_i in T.Parallel(BI, PACKED_W4): + page = page_idx_shared[bi_i] + block_id = page // BS_KV_1 + t_in_block = page % BS_KV_1 + K_packed_shared[bi_i, w_i] = K_combined_1[ + block_id, t_in_block * PACKED_W4 + w_i + ] + + for bi_i, w_i in T.Parallel(BI, SCALE_W4): + page = page_idx_shared[bi_i] + block_id = page // BS_KV_1 + t_in_block = page % BS_KV_1 + K_scale_shared[bi_i, w_i] = K_combined_1[ + block_id, + NOPE_ROPE_U32_PER_BLOCK_1 + t_in_block * SCALE_W4 + w_i, + ] + + for bi_i, ti in T.Parallel(BI, NUM_TILES): + word_idx = ti // 4 + byte_in_word = ti % 4 + word = K_scale_shared[bi_i, word_idx] + scale_byte_local[bi_i, ti] = ( + word >> T.Cast("uint32", byte_in_word * 8) + ) & T.uint32(0xFF) + + for bi_i, d_i in T.Parallel(BI, D): + word_idx = d_i // 4 + byte_in_word = d_i % 4 + word = K_packed_shared[bi_i, word_idx] + b_u32 = ( + word >> T.Cast("uint32", byte_in_word * 8) + ) & T.uint32(0xFF) + sign_bf = (b_u32 & T.uint32(0x80)) * T.uint32(0x100) + exp_e4 = (b_u32 & T.uint32(0x78)) >> T.uint32(3) + mant_bf = (b_u32 & T.uint32(0x7)) * T.uint32(0x10) + scale_byte = scale_byte_local[bi_i, d_i // NOPE_TILE] + exp_combined = exp_e4 + scale_byte - T.uint32(7) + bf16_bits = ( + sign_bf | (exp_combined << T.uint32(7)) | mant_bf + ) + KV_shared[bi_i, d_i] = T.reinterpret( + BF16, T.Cast("uint16", bf16_bits) + ) + + for bi_i, j in T.Parallel(BI, D_tail): + abs_off = D + 2 * j + word_idx = abs_off // 4 + word_off = abs_off % 4 + word = K_packed_shared[bi_i, word_idx] + half_u32 = T.if_then_else( + word_off == 0, + word & T.uint32(0xFFFF), + (word >> T.uint32(16)) & T.uint32(0xFFFF), + ) + K_tail_shared[bi_i, j] = T.reinterpret( + BF16, T.Cast("uint16", half_u32) + ) + + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else( + mask[bi_i], 0, -T.infinity(acc_s.dtype) + ) + T.gemm( + Q_shared, + KV_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.gemm( + Q_tail_shared, + K_tail_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.copy(m_i, m_i_prev) + T.reduce_max(acc_s, m_i, dim=1, clear=False) + for h_i in T.Parallel(H_per_block): + m_i[h_i] = T.max(m_i[h_i], m_i_prev[h_i]) + for h_i in T.Parallel(H_per_block): + alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.exp2( + acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale + ) + T.reduce_sum(acc_s, sumexp_i, dim=1) + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] *= alpha[h_i] + for h_i, d_i in T.Parallel(H_per_block, D_tail): + acc_o_tail[h_i, d_i] *= alpha[h_i] + T.copy(acc_s, S_shared) + T.gemm( + S_shared, + KV_shared, + acc_o, + policy=T.GemmWarpPolicy.FullRow, + ) + T.gemm( + S_shared, + K_tail_shared, + acc_o_tail, + policy=T.GemmWarpPolicy.FullRow, + ) + # ---- finalize phase 1 (active) ---- + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] = acc_o[h_i, d_i] / T.if_then_else( + sumexp[h_i] == 0.0, 1.0, sumexp[h_i] + ) + for h_i, d_i in T.Parallel(H_per_block, D_tail): + acc_o_tail[h_i, d_i] = acc_o_tail[h_i, d_i] / T.if_then_else( + sumexp[h_i] == 0.0, 1.0, sumexp[h_i] + ) + for h_i in T.Parallel(H_per_block): + m_i[h_i] = T.if_then_else( + sumexp[h_i] == 0.0, + -(2.0**30), + T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale, + ) + T.copy(acc_o, Partial_O[b_i, s_i, group_i, H0:H1, :D]) + T.copy( + acc_o_tail, + Partial_O[b_i, s_i, group_i, H0:H1, D : D + D_tail], + ) + T.copy(m_i, Partial_LSE[b_i, s_i, group_i, H0:H1]) + elif group_i < n_groups_1: + # Phase 1 skipped: m_i is still the -2^30 + T.copy(m_i, Partial_LSE[b_i, s_i, group_i, H0:H1]) + elif (group_i - n_groups_1) < actual_n_groups_2: + # Phase 2 active: c128 cache work + Partial_O write. + T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared) + T.copy(Q[b_i, s_i, H0:H1, D : D + D_tail], Q_tail_shared) + for k_i in T.Pipelined(inner_iter_2, num_stages=num_stages): + iter_i = (group_i - n_groups_1) * inner_iter_2 + k_i + for bi_i in T.Parallel(BI): + pos = iter_i * BI + bi_i + idx = Indices_2[b_i, s_i, pos] + valid = (idx >= 0) & (pos < tk_len_2) + page_idx_shared[bi_i] = T.if_then_else(valid, idx, 0) + mask[bi_i] = valid + + for bi_i, w_i in T.Parallel(BI, PACKED_W4): + page = page_idx_shared[bi_i] + block_id = page // BS_KV_2 + t_in_block = page % BS_KV_2 + K_packed_shared[bi_i, w_i] = K_combined_2[ + block_id, t_in_block * PACKED_W4 + w_i + ] + + for bi_i, w_i in T.Parallel(BI, SCALE_W4): + page = page_idx_shared[bi_i] + block_id = page // BS_KV_2 + t_in_block = page % BS_KV_2 + K_scale_shared[bi_i, w_i] = K_combined_2[ + block_id, + NOPE_ROPE_U32_PER_BLOCK_2 + t_in_block * SCALE_W4 + w_i, + ] + + for bi_i, ti in T.Parallel(BI, NUM_TILES): + word_idx = ti // 4 + byte_in_word = ti % 4 + word = K_scale_shared[bi_i, word_idx] + scale_byte_local[bi_i, ti] = ( + word >> T.Cast("uint32", byte_in_word * 8) + ) & T.uint32(0xFF) + + for bi_i, d_i in T.Parallel(BI, D): + word_idx = d_i // 4 + byte_in_word = d_i % 4 + word = K_packed_shared[bi_i, word_idx] + b_u32 = ( + word >> T.Cast("uint32", byte_in_word * 8) + ) & T.uint32(0xFF) + sign_bf = (b_u32 & T.uint32(0x80)) * T.uint32(0x100) + exp_e4 = (b_u32 & T.uint32(0x78)) >> T.uint32(3) + mant_bf = (b_u32 & T.uint32(0x7)) * T.uint32(0x10) + scale_byte = scale_byte_local[bi_i, d_i // NOPE_TILE] + exp_combined = exp_e4 + scale_byte - T.uint32(7) + bf16_bits = ( + sign_bf | (exp_combined << T.uint32(7)) | mant_bf + ) + KV_shared[bi_i, d_i] = T.reinterpret( + BF16, T.Cast("uint16", bf16_bits) + ) + + for bi_i, j in T.Parallel(BI, D_tail): + abs_off = D + 2 * j + word_idx = abs_off // 4 + word_off = abs_off % 4 + word = K_packed_shared[bi_i, word_idx] + half_u32 = T.if_then_else( + word_off == 0, + word & T.uint32(0xFFFF), + (word >> T.uint32(16)) & T.uint32(0xFFFF), + ) + K_tail_shared[bi_i, j] = T.reinterpret( + BF16, T.Cast("uint16", half_u32) + ) + + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else( + mask[bi_i], 0, -T.infinity(acc_s.dtype) + ) + T.gemm( + Q_shared, + KV_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.gemm( + Q_tail_shared, + K_tail_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.copy(m_i, m_i_prev) + T.reduce_max(acc_s, m_i, dim=1, clear=False) + for h_i in T.Parallel(H_per_block): + m_i[h_i] = T.max(m_i[h_i], m_i_prev[h_i]) + for h_i in T.Parallel(H_per_block): + alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.exp2( + acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale + ) + T.reduce_sum(acc_s, sumexp_i, dim=1) + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] *= alpha[h_i] + for h_i, d_i in T.Parallel(H_per_block, D_tail): + acc_o_tail[h_i, d_i] *= alpha[h_i] + T.copy(acc_s, S_shared) + T.gemm( + S_shared, + KV_shared, + acc_o, + policy=T.GemmWarpPolicy.FullRow, + ) + T.gemm( + S_shared, + K_tail_shared, + acc_o_tail, + policy=T.GemmWarpPolicy.FullRow, + ) + # ---- finalize phase 2 (active) ---- + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] = acc_o[h_i, d_i] / T.if_then_else( + sumexp[h_i] == 0.0, 1.0, sumexp[h_i] + ) + for h_i, d_i in T.Parallel(H_per_block, D_tail): + acc_o_tail[h_i, d_i] = acc_o_tail[h_i, d_i] / T.if_then_else( + sumexp[h_i] == 0.0, 1.0, sumexp[h_i] + ) + for h_i in T.Parallel(H_per_block): + m_i[h_i] = T.if_then_else( + sumexp[h_i] == 0.0, + -(2.0**30), + T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale, + ) + T.copy(acc_o, Partial_O[b_i, s_i, group_i, H0:H1, :D]) + T.copy( + acc_o_tail, + Partial_O[b_i, s_i, group_i, H0:H1, D : D + D_tail], + ) + T.copy(m_i, Partial_LSE[b_i, s_i, group_i, H0:H1]) + else: + # Phase 2 skipped: m_i is still the -2^30 + T.copy(m_i, Partial_LSE[b_i, s_i, group_i, H0:H1]) + + return main + + @T.prim_func + def main( + Q: T.Tensor(q_shape, BF16), # type: ignore + K_combined_1: T.Tensor(k1_shape, "uint32"), # type: ignore + Indices_1: T.Tensor(indices1_shape, indices_dtype), # type: ignore + Topk_length_1: T.Tensor(topk_length_shape, indices_dtype), # type: ignore + Partial_O: T.Tensor(partial_o_shape, BF16), # type: ignore + Partial_LSE: T.Tensor(partial_lse_shape, accum_dtype), # type: ignore + ) -> None: + """ + grid: (seq_len * REPLICATE_H * n_groups, batch, 1) + Each block processes `inner_iter_1` consecutive KV tiles and writes + one (partial_o, partial_lse) entry. + """ + with T.Kernel( + seq_len * REPLICATE_H * n_groups, batch, kv_group, threads=threads + ) as (bx, by, bz): + Q_shared = T.alloc_fragment([H_per_block, D], BF16) + Q_tail_shared = T.alloc_fragment([H_per_block, D_tail], BF16) + K_packed_shared = T.alloc_shared([BI, PACKED_W4], "uint32") + K_scale_shared = T.alloc_shared([BI, SCALE_W4], "uint32") + KV_shared = T.alloc_shared([BI, D], BF16) + K_tail_shared = T.alloc_shared([BI, D_tail], BF16) + S_shared = T.alloc_shared([H_per_block, BI], BF16) + page_idx_shared = T.alloc_shared([BI], INT32) + + mask = T.alloc_fragment([BI], "bool") + scale_byte_local = T.alloc_fragment([BI, NUM_TILES], "uint32") + + acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) + acc_o_tail = T.alloc_fragment([H_per_block, D_tail], accum_dtype) + acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) + sumexp = T.alloc_fragment([H_per_block], accum_dtype) + sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) + alpha = T.alloc_fragment([H_per_block], accum_dtype) + m_i = T.alloc_fragment([H_per_block], accum_dtype) + m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) + + T.fill(acc_o, 0) + T.fill(acc_o_tail, 0) + T.fill(sumexp, 0) + T.fill(m_i, -(2**30)) + + b_i, g_i = by, bz + spans_per_seq = REPLICATE_H * n_groups + s_i = bx // spans_per_seq + rest = bx % spans_per_seq + group_i = rest // REPLICATE_H + h_rep = rest % REPLICATE_H + H0 = g_i * padded_H + (0 if REPLICATE_H == 1 else h_rep * 64) + H1 = H0 + H_per_block + + T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared) + T.copy(Q[b_i, s_i, H0:H1, D : D + D_tail], Q_tail_shared) + + tk_len_1 = Topk_length_1[b_i] + + for k_i in T.Pipelined(inner_iter_1, num_stages=num_stages): + iter_i = group_i * inner_iter_1 + k_i + for bi_i in T.Parallel(BI): + pos = iter_i * BI + bi_i + idx = Indices_1[b_i, s_i, pos] + valid = (idx >= 0) & (pos < tk_len_1) + page_idx_shared[bi_i] = T.if_then_else(valid, idx, 0) + mask[bi_i] = valid + + for bi_i, w_i in T.Parallel(BI, PACKED_W4): + page = page_idx_shared[bi_i] + block_id = page // BS_KV_1 + t_in_block = page % BS_KV_1 + K_packed_shared[bi_i, w_i] = K_combined_1[ + block_id, t_in_block * PACKED_W4 + w_i + ] + + for bi_i, w_i in T.Parallel(BI, SCALE_W4): + page = page_idx_shared[bi_i] + block_id = page // BS_KV_1 + t_in_block = page % BS_KV_1 + K_scale_shared[bi_i, w_i] = K_combined_1[ + block_id, + NOPE_ROPE_U32_PER_BLOCK_1 + t_in_block * SCALE_W4 + w_i, + ] + + for bi_i, ti in T.Parallel(BI, NUM_TILES): + word_idx = ti // 4 + byte_in_word = ti % 4 + word = K_scale_shared[bi_i, word_idx] + scale_byte_local[bi_i, ti] = ( + word >> T.Cast("uint32", byte_in_word * 8) + ) & T.uint32(0xFF) + + for bi_i, d_i in T.Parallel(BI, D): + word_idx = d_i // 4 + byte_in_word = d_i % 4 + word = K_packed_shared[bi_i, word_idx] + b_u32 = (word >> T.Cast("uint32", byte_in_word * 8)) & T.uint32( + 0xFF + ) + sign_bf = (b_u32 & T.uint32(0x80)) * T.uint32(0x100) + exp_e4 = (b_u32 & T.uint32(0x78)) >> T.uint32(3) + mant_bf = (b_u32 & T.uint32(0x7)) * T.uint32(0x10) + scale_byte = scale_byte_local[bi_i, d_i // NOPE_TILE] + exp_combined = exp_e4 + scale_byte - T.uint32(7) + bf16_bits = sign_bf | (exp_combined << T.uint32(7)) | mant_bf + KV_shared[bi_i, d_i] = T.reinterpret( + BF16, T.Cast("uint16", bf16_bits) + ) + + for bi_i, j in T.Parallel(BI, D_tail): + abs_off = D + 2 * j + word_idx = abs_off // 4 + word_off = abs_off % 4 + word = K_packed_shared[bi_i, word_idx] + half_u32 = T.if_then_else( + word_off == 0, + word & T.uint32(0xFFFF), + (word >> T.uint32(16)) & T.uint32(0xFFFF), + ) + K_tail_shared[bi_i, j] = T.reinterpret( + BF16, T.Cast("uint16", half_u32) + ) + + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else( + mask[bi_i], 0, -T.infinity(acc_s.dtype) + ) + T.gemm( + Q_shared, + KV_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.gemm( + Q_tail_shared, + K_tail_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.copy(m_i, m_i_prev) + T.reduce_max(acc_s, m_i, dim=1, clear=False) + for h_i in T.Parallel(H_per_block): + m_i[h_i] = T.max(m_i[h_i], m_i_prev[h_i]) + for h_i in T.Parallel(H_per_block): + alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.exp2( + acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale + ) + T.reduce_sum(acc_s, sumexp_i, dim=1) + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] *= alpha[h_i] + for h_i, d_i in T.Parallel(H_per_block, D_tail): + acc_o_tail[h_i, d_i] *= alpha[h_i] + T.copy(acc_s, S_shared) + T.gemm(S_shared, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) + T.gemm( + S_shared, K_tail_shared, acc_o_tail, policy=T.GemmWarpPolicy.FullRow + ) + + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] = acc_o[h_i, d_i] / T.if_then_else( + sumexp[h_i] == 0.0, 1.0, sumexp[h_i] + ) + for h_i, d_i in T.Parallel(H_per_block, D_tail): + acc_o_tail[h_i, d_i] = acc_o_tail[h_i, d_i] / T.if_then_else( + sumexp[h_i] == 0.0, 1.0, sumexp[h_i] + ) + for h_i in T.Parallel(H_per_block): + m_i[h_i] = T.if_then_else( + sumexp[h_i] == 0.0, + -(2.0**30), + T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale, + ) + T.copy(acc_o, Partial_O[b_i, s_i, group_i, H0:H1, :D]) + T.copy(acc_o_tail, Partial_O[b_i, s_i, group_i, H0:H1, D : D + D_tail]) + T.copy(m_i, Partial_LSE[b_i, s_i, group_i, H0:H1]) + + return main + + +@tilelang.jit( + out_idx=[-2, -1], + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, +) +def dpsk_v4_combine_kernel( + num_heads: int, + n_groups_1: int, + n_groups_2: int = 0, + *, + block_I: int = 64, + inner_iter_1: int = 1, + inner_iter_2: int = 1, + dim: int = 448, + tail_dim: int = 64, + head_per_block: int = 16, + threads: int = 256, + use_attn_sink: bool = False, +) -> Any: + """ + Combine `n_groups` flash-attention partials into the final output. + + Inputs: + Partial_O : (batch, seq_len, n_groups, num_heads, dim+tail_dim) bf16 + Partial_LSE : (batch, seq_len, n_groups, num_heads) fp32, log2 form + Topk_length_1: (batch,) int32, actual phase-1 length + Topk_length_2: (batch,) int32, actual phase-2 length (dual only) + Attn_sink : (num_heads,) fp32 + Outputs: + Output : (batch, seq_len, num_heads, dim+tail_dim) bf16 + LSE : (batch, seq_len, num_heads) fp32, natural log + + Each grid block handles `head_per_block` heads of one (batch, seq) row. + """ + log2e: float = 1.44269504 + ln2: float = 0.69314718 + assert num_heads % head_per_block == 0 + + is_dual = n_groups_2 > 0 + n_groups = n_groups_1 + n_groups_2 + + H_per_block = head_per_block + HEAD_BLOCKS = num_heads // H_per_block + DT = dim + tail_dim + + batch = T.symbolic("batch") + seq_len = T.symbolic("seq_len") + + accum_dtype = "float" + + if is_dual: + + @T.prim_func + def main( + Partial_O: T.Tensor( + [batch, seq_len, n_groups, num_heads, DT], BF16 + ), # type: ignore + Partial_LSE: T.Tensor( + [batch, seq_len, n_groups, num_heads], accum_dtype + ), # type: ignore + Topk_length_1: T.Tensor([batch], INT32), # type: ignore + Topk_length_2: T.Tensor([batch], INT32), # type: ignore + Attn_sink: T.Tensor([num_heads], FP32), # type: ignore + Output: T.Tensor([batch, seq_len, num_heads, DT], BF16), # type: ignore + LSE: T.Tensor([batch, seq_len, num_heads], accum_dtype), # type: ignore + ) -> None: + with T.Kernel(seq_len * HEAD_BLOCKS, batch, threads=threads) as ( + bx, + by, + ): + shared_lse = T.alloc_shared([n_groups, H_per_block], accum_dtype) + lse_max = T.alloc_fragment([H_per_block], accum_dtype) + lse_sum = T.alloc_fragment([H_per_block], accum_dtype) + scale = T.alloc_fragment([H_per_block, n_groups], accum_dtype) + acc_o = T.alloc_fragment([H_per_block, DT], accum_dtype) + attn_sink_frag = T.alloc_fragment([H_per_block], accum_dtype) + o_scale_frag = T.alloc_fragment([H_per_block], accum_dtype) + final_lse = T.alloc_fragment([H_per_block], accum_dtype) + + b_i = by + s_i = bx // HEAD_BLOCKS + head_block = bx % HEAD_BLOCKS + H0 = head_block * H_per_block + H1 = H0 + H_per_block + + # Clamp to the captured-shape upper bounds so callers passing + # the INT32_MAX sentinel (= "all valid") still iterate exactly + # n_groups groups, not 33M. + actual_n_groups_1 = T.min( + T.ceildiv(Topk_length_1[b_i], block_I * inner_iter_1), + n_groups_1, + ) + actual_n_groups_2 = T.min( + T.ceildiv(Topk_length_2[b_i], block_I * inner_iter_2), + n_groups - n_groups_1, + ) + actual_n_groups = actual_n_groups_1 + actual_n_groups_2 + + # Pass 1: load only active groups' LSE into compact slots. + for k_c in T.serial(actual_n_groups): + k = T.if_then_else( + k_c < actual_n_groups_1, + k_c, + n_groups_1 + (k_c - actual_n_groups_1), + ) + T.copy(Partial_LSE[b_i, s_i, k, H0:H1], shared_lse[k_c, :]) + + T.fill(lse_max, -(2**30)) + for k_c in T.serial(actual_n_groups): + for h_i in T.Parallel(H_per_block): + lse_max[h_i] = T.max(lse_max[h_i], shared_lse[k_c, h_i]) + T.fill(lse_sum, 0) + for k_c in T.serial(actual_n_groups): + for h_i in T.Parallel(H_per_block): + lse_sum[h_i] = lse_sum[h_i] + T.exp2( + shared_lse[k_c, h_i] - lse_max[h_i] + ) + for k_c in T.serial(actual_n_groups): + for h_i in T.Parallel(H_per_block): + scale[h_i, k_c] = T.exp2( + shared_lse[k_c, h_i] - lse_max[h_i] - T.log2(lse_sum[h_i]) + ) + + T.fill(acc_o, 0) + for k_c in T.serial(actual_n_groups): + k = T.if_then_else( + k_c < actual_n_groups_1, + k_c, + n_groups_1 + (k_c - actual_n_groups_1), + ) + for h_i, d_i in T.Parallel(H_per_block, DT): + acc_o[h_i, d_i] = acc_o[h_i, d_i] + scale[h_i, k_c] * Partial_O[ + b_i, s_i, k, H0 + h_i, d_i + ].astype(accum_dtype) + + for h_i in T.Parallel(H_per_block): + empty = lse_max[h_i] <= -(2**29) + final_lse[h_i] = T.if_then_else( + empty, + T.infinity(accum_dtype), + (lse_max[h_i] + T.log2(lse_sum[h_i])) * ln2, + ) + + if use_attn_sink: + for h_i in T.Parallel(H_per_block): + attn_sink_frag[h_i] = Attn_sink[H0 + h_i] + for h_i in T.Parallel(H_per_block): + empty = lse_max[h_i] <= -(2**29) + o_scale_frag[h_i] = T.if_then_else( + empty, + 0.0, + 1.0 + / ( + 1.0 + + T.exp2((attn_sink_frag[h_i] - final_lse[h_i]) * log2e) + ), + ) + for h_i, d_i in T.Parallel(H_per_block, DT): + acc_o[h_i, d_i] = acc_o[h_i, d_i] * o_scale_frag[h_i] + + T.copy(acc_o, Output[b_i, s_i, H0:H1, :]) + T.copy(final_lse, LSE[b_i, s_i, H0:H1]) + + return main + + @T.prim_func + def main( + Partial_O: T.Tensor( + [batch, seq_len, n_groups, num_heads, DT], BF16 + ), # type: ignore + Partial_LSE: T.Tensor( + [batch, seq_len, n_groups, num_heads], accum_dtype + ), # type: ignore + Attn_sink: T.Tensor([num_heads], FP32), # type: ignore + Output: T.Tensor([batch, seq_len, num_heads, DT], BF16), # type: ignore + LSE: T.Tensor([batch, seq_len, num_heads], accum_dtype), # type: ignore + ) -> None: + with T.Kernel(seq_len * HEAD_BLOCKS, batch, threads=threads) as (bx, by): + shared_lse = T.alloc_shared([n_groups, H_per_block], accum_dtype) + + lse_max = T.alloc_fragment([H_per_block], accum_dtype) + lse_sum = T.alloc_fragment([H_per_block], accum_dtype) + scale = T.alloc_fragment([H_per_block, n_groups], accum_dtype) + acc_o = T.alloc_fragment([H_per_block, DT], accum_dtype) + attn_sink_frag = T.alloc_fragment([H_per_block], accum_dtype) + o_scale_frag = T.alloc_fragment([H_per_block], accum_dtype) + final_lse = T.alloc_fragment([H_per_block], accum_dtype) + + b_i = by + s_i = bx // HEAD_BLOCKS + head_block = bx % HEAD_BLOCKS + H0 = head_block * H_per_block + H1 = H0 + H_per_block + + for k in T.serial(n_groups): + T.copy(Partial_LSE[b_i, s_i, k, H0:H1], shared_lse[k, :]) + + T.fill(lse_max, -(2**30)) + for k in T.serial(n_groups): + for h_i in T.Parallel(H_per_block): + lse_max[h_i] = T.max(lse_max[h_i], shared_lse[k, h_i]) + T.fill(lse_sum, 0) + for k in T.serial(n_groups): + for h_i in T.Parallel(H_per_block): + lse_sum[h_i] = lse_sum[h_i] + T.exp2( + shared_lse[k, h_i] - lse_max[h_i] + ) + for k in T.serial(n_groups): + for h_i in T.Parallel(H_per_block): + scale[h_i, k] = T.exp2( + shared_lse[k, h_i] - lse_max[h_i] - T.log2(lse_sum[h_i]) + ) + + T.fill(acc_o, 0) + for k in T.serial(n_groups): + for h_i, d_i in T.Parallel(H_per_block, DT): + acc_o[h_i, d_i] = acc_o[h_i, d_i] + scale[h_i, k] * Partial_O[ + b_i, s_i, k, H0 + h_i, d_i + ].astype(accum_dtype) + + for h_i in T.Parallel(H_per_block): + empty = lse_max[h_i] <= -(2**29) + final_lse[h_i] = T.if_then_else( + empty, + T.infinity(accum_dtype), + (lse_max[h_i] + T.log2(lse_sum[h_i])) * ln2, + ) + + if use_attn_sink: + for h_i in T.Parallel(H_per_block): + attn_sink_frag[h_i] = Attn_sink[H0 + h_i] + for h_i in T.Parallel(H_per_block): + empty = lse_max[h_i] <= -(2**29) + o_scale_frag[h_i] = T.if_then_else( + empty, + 0.0, + 1.0 + / ( + 1.0 + T.exp2((attn_sink_frag[h_i] - final_lse[h_i]) * log2e) + ), + ) + for h_i, d_i in T.Parallel(H_per_block, DT): + acc_o[h_i, d_i] = acc_o[h_i, d_i] * o_scale_frag[h_i] + + T.copy(acc_o, Output[b_i, s_i, H0:H1, :]) + T.copy(final_lse, LSE[b_i, s_i, H0:H1]) + + return main + + +""" +2-stage attention kernel (partial + combine) over an FP8 KV cache, +with optional second cache (`extra_k_cache`). +""" + + +def dpsk_v4_fp8_attention_fwd( + q: torch.Tensor, + k_cache: torch.Tensor, + block_table: Optional[torch.Tensor], + cache_seqlens: Optional[torch.Tensor], + head_dim_v: int, + tile_scheduler_metadata: Any, + num_splits: None = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + is_fp8_kvcache: bool = False, + indices: Optional[torch.Tensor] = None, + attn_sink: Optional[torch.Tensor] = None, + extra_k_cache: Optional[torch.Tensor] = None, + extra_indices_in_kvcache: Optional[torch.Tensor] = None, + topk_length: Optional[torch.Tensor] = None, + extra_topk_length: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Follows the original `flash_mla.flash_mla_with_kvcache` signature. + """ + if _is_gfx95_supported: + block_I, threads, num_stages, block_per_cu, cu = 64, 512, 0, 2, 256 + else: + block_I, threads, num_stages, block_per_cu, cu = 32, 128, 1, 1, 304 + + batch, seq_len, num_heads, _ = q.shape + # Partial grid is (seq_len * REPLICATE_H * n_groups, batch, kv_group); the + # heuristic in _pick_inner_iter assumes `total_blocks = seq * ni / inner_iter`, + # so `seq` must include REPLICATE_H or n_groups doubles for medium batches. + replicate_h = max((num_heads + 63) // 64, 1) + seq = batch * seq_len * replicate_h + + k1, _, bs_kv_1 = _build_fp8_combined_view(k_cache) + topk_1 = indices.shape[-1] + ni_1 = topk_1 // block_I + tk_len_1 = ( + topk_length + if topk_length is not None + else _topk_length_sentinel(q.device, batch) + ) + if attn_sink is None: + attn_sink = torch.full( + (num_heads,), float("-inf"), dtype=torch.float32, device=q.device + ) + + has_extra = extra_k_cache is not None + if not has_extra: + inner_iter_1 = _pick_inner_iter(seq, ni_1, cu, block_per_cu) + inner_iter_2 = 1 + n_groups_1 = ni_1 // inner_iter_1 + n_groups_2 = 0 + partial = dpsk_v4_fp8_partial_kernel( + num_heads, + topk_1, + bs_kv_1, + sm_scale=softmax_scale, + block_I=block_I, + inner_iter_1=inner_iter_1, + num_stages=num_stages, + threads=threads, + ) + partial_o, partial_lse = partial(q, k1, indices, tk_len_1) + else: + k2, _, bs_kv_2 = _build_fp8_combined_view(extra_k_cache) + topk_2 = extra_indices_in_kvcache.shape[-1] + ni_2 = topk_2 // block_I + # Each phase picks its own optimal split-K independently — kernel + # body uses two T.Pipelined loops with separate compile-time iter + # counts, no shared-divisor constraint. + inner_iter_1 = _pick_inner_iter(seq, ni_1, cu, block_per_cu) + inner_iter_2 = _pick_inner_iter(seq, ni_2, cu, block_per_cu) + n_groups_1 = ni_1 // inner_iter_1 + n_groups_2 = ni_2 // inner_iter_2 + tk_len_2 = ( + extra_topk_length + if extra_topk_length is not None + else _topk_length_sentinel(q.device, batch) + ) + partial = dpsk_v4_fp8_partial_kernel( + num_heads, + topk_1, + bs_kv_1, + topk_2, + bs_kv_2, + sm_scale=softmax_scale, + block_I=block_I, + inner_iter_1=inner_iter_1, + inner_iter_2=inner_iter_2, + num_stages=num_stages, + threads=threads, + ) + partial_o, partial_lse = partial( + q, + k1, + indices, + tk_len_1, + k2, + extra_indices_in_kvcache, + tk_len_2, + ) + + combine = dpsk_v4_combine_kernel( + num_heads, + n_groups_1, + n_groups_2, + block_I=block_I, + inner_iter_1=inner_iter_1, + inner_iter_2=inner_iter_2, + head_per_block=4, + threads=256, + use_attn_sink=True, + ) + if has_extra: + return combine(partial_o, partial_lse, tk_len_1, tk_len_2, attn_sink) + return combine(partial_o, partial_lse, attn_sink) diff --git a/experiments/pro6000/qwen3_235b_pro6000_sglang_tp8/config.env b/experiments/pro6000/qwen3_235b_pro6000_sglang_tp8/config.env index d3e094d..813cc17 100644 --- a/experiments/pro6000/qwen3_235b_pro6000_sglang_tp8/config.env +++ b/experiments/pro6000/qwen3_235b_pro6000_sglang_tp8/config.env @@ -17,6 +17,8 @@ DATASET_PATH="${DATASET_PATH:-/data/yy/sskj/dataset/ShareGPT_V3_unfiltered_clean # Serving port + container. PORT="${PORT:-30010}" +# Alias for the unified bench layer (`python -m sskj.bench`). +SGLANG_PORT="${SGLANG_PORT:-${PORT:-30010}}" CONTAINER_NAME="${CONTAINER_NAME:-qwen3_235b_pro6000_sglang_tp8}" # Docker image: sm120-capable SGLang (built/verified on this RTX 6000D box), diff --git a/experiments/pro6000/qwen3_235b_pro6000_sglang_tp8/run_bench.sh b/experiments/pro6000/qwen3_235b_pro6000_sglang_tp8/run_bench.sh index 7524e6e..44bda11 100755 --- a/experiments/pro6000/qwen3_235b_pro6000_sglang_tp8/run_bench.sh +++ b/experiments/pro6000/qwen3_235b_pro6000_sglang_tp8/run_bench.sh @@ -85,14 +85,8 @@ fi stop_server() { if [[ -z "${SKIP_MANAGE_SERVER:-}" ]]; then - log "stopping server (kill docker run PID -> --rm removes container)" - if [[ -f "${RUNTIME_BASE}/${EXPERIMENT_NAME}.pid" ]]; then - local pid - pid="$(cat "${RUNTIME_BASE}/${EXPERIMENT_NAME}.pid")" - kill "$pid" 2>/dev/null || true - sleep 3 - kill -9 "$pid" 2>/dev/null || true - fi + log "stopping server via deploy profile" + bash "${SCRIPT_DIR}/stop_server.sh" >> "${LOG_DIR}/stop_server.outer.log" 2>&1 || true docker rm -f "$CONTAINER_NAME" >/dev/null 2>&1 || true fi } diff --git a/experiments/pro6000/qwen3_235b_pro6000_sglang_tp8/start_server.sh b/experiments/pro6000/qwen3_235b_pro6000_sglang_tp8/start_server.sh index 613d8bd..5b1bef0 100755 --- a/experiments/pro6000/qwen3_235b_pro6000_sglang_tp8/start_server.sh +++ b/experiments/pro6000/qwen3_235b_pro6000_sglang_tp8/start_server.sh @@ -1,92 +1,31 @@ #!/usr/bin/env bash -# Start the SGLang server for Qwen3-235B-A22B on NVIDIA RTX 6000D (TP=8, 8 GPUs). -# -# NVIDIA counterpart of experiments/p800/qwen3_235b_p800_sglang_tp8/start_server.sh. -# Uses the standard sglang qwen3_moe path (no Kunlun patch, no XPU env). -# Launch pattern mirrors pro6000/dsv4_pro6000_sglang_tiny_1k_output/start_sglang_docker.sh. +# Start the Qwen3-235B-A22B SGLang TP=8 server through the shared deployment layer. set -Eeuo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" EXPERIMENT_NAME="$(basename "$SCRIPT_DIR")" -# Source config before platform so our CONTAINER_NAME/PORT defaults win. # shellcheck source=/dev/null source "${SCRIPT_DIR}/../../../scripts/common/lib.sh" # shellcheck source=/dev/null source "${SCRIPT_DIR}/config.env" # shellcheck source=/dev/null source "${SCRIPT_DIR}/../../../scripts/common/platform.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" -RUN_ID="${RUN_ID:-$(date '+%Y%m%d-%H%M%S')}" -RESULT_ROOT="${RESULT_ROOT:-${SCRIPT_DIR}/results/${RUN_ID}}" RUNTIME_BASE="${RUNTIME_BASE:-${SCRIPT_DIR}/runtime}" -LOG_DIR="${RESULT_ROOT}/logs" -mkdir -p "$LOG_DIR" "${RUNTIME_BASE}/tmp" "$CACHE_DIR" +mkdir -p "${RUNTIME_BASE}/logs" "${RUNTIME_BASE}/tmp" -log_init "${LOG_DIR}/start_server.log" -log "starting server for ${EXPERIMENT_NAME} (NVIDIA, TP=${TP})" +log "starting server for ${EXPERIMENT_NAME} (NVIDIA, TP=${TP}) via deploy profile" log "model: ${MODEL_PATH}" log "image: ${DOCKER_IMAGE}" log "port: ${PORT} container: ${CONTAINER_NAME}" -SERVER_LOG="${LOG_DIR}/server.log" -PID_FILE="${RUNTIME_BASE}/${EXPERIMENT_NAME}.pid" -rm -f "$PID_FILE" - -# Remove any stale container with the same name. -docker rm -f "$CONTAINER_NAME" >/dev/null 2>&1 || true - -# Standard sglang launch for qwen3_moe on NVIDIA (bf16 auto, flashinfer attn). -# Mount /data so model + experiment results dir + dataset are all at the same -# path inside the container (bench client writes results there directly). -DTYPE_ARGS=() -if [[ "${DTYPE}" != "auto" ]]; then - DTYPE_ARGS=(--dtype "${DTYPE}") -fi - -nohup docker run --rm \ - --name "$CONTAINER_NAME" \ - --gpus all \ - --privileged \ - --ipc=host \ - --network host \ - --ulimit memlock=-1 \ - --ulimit stack=67108864 \ - --entrypoint python3 \ - -v /data:/data \ - -v "${CACHE_DIR}:/root/.cache" \ - -v "${RUNTIME_BASE}/tmp:/tmp" \ - -e CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" \ - -e PYTHONUNBUFFERED=1 \ - -e HF_HUB_OFFLINE=1 \ - -e TRANSFORMERS_OFFLINE=1 \ - -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \ - "$DOCKER_IMAGE" \ - -m sglang.launch_server \ - --model-path "${MODEL_PATH}" \ - --trust-remote-code \ - --tp-size "${TP}" \ - --mem-fraction-static "${MEM_FRACTION_STATIC}" \ - --context-length "${CONTEXT_LENGTH}" \ - --max-running-requests "${MAX_RUNNING_REQUESTS}" \ - --host 0.0.0.0 \ - --port "${PORT}" \ - "${DTYPE_ARGS[@]}" \ - > "$SERVER_LOG" 2>&1 & - -PID=$! -echo "$PID" > "$PID_FILE" -log "PID: ${PID}" -HEALTH_TIMEOUT="${HEALTH_TIMEOUT:-1200}" -log "waiting for health on http://127.0.0.1:${PORT}/health (timeout=${HEALTH_TIMEOUT}s)" - -if health_check 127.0.0.1 "$PORT" "$HEALTH_TIMEOUT"; then - log "sglang server is READY at http://127.0.0.1:${PORT}" - log "server log: ${SERVER_LOG} (container: ${CONTAINER_NAME})" - exit 0 -else - log "ERROR: server not healthy after ${HEALTH_TIMEOUT}s" - log "----- last 200 lines of server log -----" - tail -200 "$SERVER_LOG" 2>/dev/null || docker logs --tail 200 "$CONTAINER_NAME" 2>&1 | tail -200 - exit 1 -fi +deploy_start \ + "${DEPLOY_PROFILE:-pro6000/qwen3_235b_pro6000_sglang_tp8}" \ + "${TP:-8}" "1" \ + "${RUNTIME_BASE}/logs" \ + "${PORT:-30010}" \ + "$MODEL_PATH" \ + "$CONTAINER_NAME" diff --git a/experiments/pro6000/qwen3_235b_pro6000_sglang_tp8/stop_server.sh b/experiments/pro6000/qwen3_235b_pro6000_sglang_tp8/stop_server.sh new file mode 100755 index 0000000..5448832 --- /dev/null +++ b/experiments/pro6000/qwen3_235b_pro6000_sglang_tp8/stop_server.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +# Stop the Qwen3-235B-A22B SGLang server through the shared deployment layer. +set -Eeuo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +EXPERIMENT_NAME="$(basename "$SCRIPT_DIR")" + +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/lib.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/config.env" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/platform.sh" +# shellcheck source=/dev/null +source "${SCRIPT_DIR}/../../../scripts/common/deploy_cli.sh" + +log "stopping server for ${EXPERIMENT_NAME} via deploy profile" +deploy_stop \ + "${DEPLOY_PROFILE:-pro6000/qwen3_235b_pro6000_sglang_tp8}" \ + "${TP:-8}" "1" \ + "${PORT:-30010}" \ + "$MODEL_PATH" \ + "$CONTAINER_NAME" diff --git a/ops/README.md b/ops/README.md index 65a32ab..9d48cfd 100644 --- a/ops/README.md +++ b/ops/README.md @@ -80,7 +80,48 @@ experiments///results// 跨机测试时可用 `--client docker --client-image lmsysorg/sglang:latest` 指定外部 client。 -## 7. 常见问题 +## 7. pro6000（RTX 6000D）平台说明 + +实验目录 `experiments/pro6000/`，部署 profile 见 `deploy/profiles/pro6000/`。 + +可测实验： + +```bash +PYTHONPATH=src python3 -m sskj.bench list --platform pro6000 +# dsv4_pro6000_sglang_tp16 / dsv4_pro6000_sglang_tp16_eagle +# dsv4_pro6000_sglang_tp_dp_matrix / dsv4_pro6000_vllm_tp_dp_matrix +# glm52_pro6000_sglang_multinode_tp16 / qwen3_235b_pro6000_sglang_tp8 +``` + +单节点实验（dsv4/vllm tp_dp_matrix、qwen3，均 8x RTX 6000D 单机）： + +```bash +PYTHONPATH=src python3 -m sskj.bench run \ + --url http://:30031 \ + --platform pro6000 \ + --experiment dsv4_pro6000_sglang_tp_dp_matrix \ + --profile smoke +``` + +多节点实验（tp16 / tp16_eagle / glm52，2 台机器 16x RTX 6000D）： +- head 节点 pro6000D.1（`10.101.0.11` / `174.1.51.5`，HTTP API 所在） +- worker 节点 pro6000D.3（`10.101.0.13` / `174.1.51.7`，纯计算） +- `--url` 固定指向 head 节点端口（tp16/eagle: 30000，glm52: 30031）， + bench client 在 head 节点跑，worker 由部署层通过 ssh 管理： + +```bash +PYTHONPATH=src python3 -m sskj.bench run \ + --url http://10.101.0.11:30000 \ + --platform pro6000 \ + --experiment dsv4_pro6000_sglang_tp16 \ + --profile smoke +``` + +vLLM 实验（`dsv4_pro6000_vllm_tp_dp_matrix`）的 bench client 使用 SGLang +镜像（vLLM 镜像不含 bench_serving），由 `DOCKER_CLIENT_IMAGE` 自动选择， +无需额外参数。 + +## 8. 常见问题 - 服务不健康：确认 `--url` 的端口可从测试机访问，服务已启动且 `/health` 可通。 - 数据集缺失：仓库 `datasets/` 下没有 ShareGPT 文件时，统一 CLI 会自动退回 diff --git a/scripts/common/deploy_cli.sh b/scripts/common/deploy_cli.sh index 4766cec..0d83f76 100644 --- a/scripts/common/deploy_cli.sh +++ b/scripts/common/deploy_cli.sh @@ -19,6 +19,11 @@ deploy_start() { local port="$5" local model_path="$6" local container="$7" + local node_rank="${8:-}" + local -a extra=() + if [[ -n "$node_rank" ]]; then + extra+=(--node-rank "$node_rank") + fi PYTHONPATH="${ROOT_DIR}/src" "${DEPLOY_PYTHON}" -m sskj.deploy start \ --profile "$(deploy_profile_abs "$profile")" \ --tp "$tp" \ @@ -26,7 +31,8 @@ deploy_start() { --log-dir "$log_dir" \ --port "$port" \ --model-path "$model_path" \ - --container-name "$container" + --container-name "$container" \ + "${extra[@]}" } deploy_render_args() { @@ -42,3 +48,63 @@ deploy_render_args() { --port "$port" \ --model-path "$model_path" } + +# Multi-node convenience helpers: orchestrate all NODE_HOSTS from one host. + +deploy_start_multinode() { + local profile="$1" + local tp="$2" + local dp="$3" + local log_dir="$4" + local port="$5" + local model_path="$6" + local container="$7" + PYTHONPATH="${ROOT_DIR}/src" "${DEPLOY_PYTHON}" -m sskj.deploy start \ + --profile "$(deploy_profile_abs "$profile")" \ + --tp "$tp" \ + --dp "$dp" \ + --log-dir "$log_dir" \ + --port "$port" \ + --model-path "$model_path" \ + --container-name "$container" +} + +deploy_stop() { + local profile="$1" + local tp="$2" + local dp="$3" + local port="$4" + local model_path="$5" + local container="$6" + local node_rank="${7:-}" + local -a extra=() + if [[ -n "$node_rank" ]]; then + extra+=(--node-rank "$node_rank") + fi + PYTHONPATH="${ROOT_DIR}/src" "${DEPLOY_PYTHON}" -m sskj.deploy stop \ + --profile "$(deploy_profile_abs "$profile")" \ + --tp "$tp" \ + --dp "$dp" \ + --port "$port" \ + --model-path "$model_path" \ + --container-name "$container" \ + "${extra[@]}" +} + +deploy_status() { + local profile="$1" + local tp="${2:-}" + local dp="${3:-}" + local port="${4:-}" + local model_path="${5:-}" + local container="${6:-}" + local -a cmd=( + --profile "$(deploy_profile_abs "$profile")" + ) + [[ -n "$tp" ]] && cmd+=(--tp "$tp") + [[ -n "$dp" ]] && cmd+=(--dp "$dp") + [[ -n "$port" ]] && cmd+=(--port "$port") + [[ -n "$model_path" ]] && cmd+=(--model-path "$model_path") + [[ -n "$container" ]] && cmd+=(--container-name "$container") + PYTHONPATH="${ROOT_DIR}/src" "${DEPLOY_PYTHON}" -m sskj.deploy status "${cmd[@]}" +} diff --git a/src/sskj/common/env.py b/src/sskj/common/env.py index 3bdc497..42a134e 100644 --- a/src/sskj/common/env.py +++ b/src/sskj/common/env.py @@ -3,7 +3,8 @@ The repository keeps deployment and platform configuration in shell `.env` files so existing bash scripts can still source them. This parser understands the scalar subset those files use: `KEY=value`, `export KEY=value`, quoted -values, comments, and `${VAR:-default}` substitution. +values, comments, and `${VAR:-default}` substitution. Nested defaults such as +`${A:-${B}/path}` are supported via balanced-brace scanning. """ from __future__ import annotations @@ -13,7 +14,7 @@ import re from pathlib import Path _ASSIGN_RE = re.compile(r"^(?:export\s+)?([A-Za-z_][A-Za-z0-9_]*)\s*=\s*(.*)$") -_VAR_RE = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)(?::-([^}]*))?\}") +_NAME_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") def _strip_quotes(value: str) -> str: @@ -42,20 +43,61 @@ def _strip_inline_comment(value: str) -> str: return value +def _find_brace_end(value: str, start: int) -> int | None: + """Index of the `}` closing the `${` at `start`, handling nested `${...}`.""" + depth = 0 + i = start + n = len(value) + while i < n: + if value.startswith("${", i): + depth += 1 + i += 2 + continue + if value[i] == "}": + depth -= 1 + if depth == 0: + return i + i += 1 + return None + + +def _split_default(inner: str) -> tuple[str, str | None]: + """Split `NAME` or `NAME:-default` (default may itself contain `${...}`).""" + for sep in (":-", "-"): + idx = inner.find(sep) + if idx > 0 and _NAME_RE.match(inner[:idx]): + return inner[:idx], inner[idx + len(sep):] + if _NAME_RE.match(inner): + return inner, None + return inner, None + + def _expand(value: str, env: dict[str, str], depth: int = 0) -> str: if depth > 6: return value - def repl(match: re.Match[str]) -> str: - name = match.group(1) - default = match.group(2) - if name in env and env[name] != "": - return env[name] - if default is not None: - return _expand(default, env, depth + 1) - return match.group(0) - - return _VAR_RE.sub(repl, value) + out: list[str] = [] + i = 0 + n = len(value) + while i < n: + if value.startswith("${", i): + end = _find_brace_end(value, i) + if end is None: + out.append(value[i:]) + break + inner = value[i + 2:end] + name, default = _split_default(inner) + if _NAME_RE.match(name) and name in env and env[name] != "": + out.append(env[name]) + elif _NAME_RE.match(name) and default is not None: + out.append(_expand(default, env, depth + 1)) + else: + out.append(value[i:end + 1]) + i = end + 1 + continue + out.append(value[i]) + i += 1 + return "".join(out) def parse_env_file( diff --git a/src/sskj/deploy/runtime.py b/src/sskj/deploy/runtime.py index 2e04488..21dccd5 100644 --- a/src/sskj/deploy/runtime.py +++ b/src/sskj/deploy/runtime.py @@ -1,14 +1,33 @@ -"""Docker/native server lifecycle for deployment profiles.""" +"""Docker/native server lifecycle for deployment profiles. + +Supports single-node and multi-node (NODE_HOSTS) docker deployments. For +multi-node profiles, start/stop/status distribute docker commands to each +node over ssh; health checks only target the head (rank 0) node. +""" from __future__ import annotations import base64 +import re import shlex import subprocess +import time from pathlib import Path from sskj.bench.runner import split_env_assignments, wait_health +_NODE_RANK_RE = re.compile(r"\$\{NODE_RANK(?::-([^}]*))?\}") +_TEMPLATED_KEYS = ( + "CONTAINER_NAME", + "LAUNCH_ARGS", + "DEVICE_VARS", + "ENGINE_ENV", + "DOCKER_FLAGS", + "VOLUMES", + "PATCH_MOUNTS", + "BOOTSTRAP", +) + def _log_dir(root: Path, log_dir: str | None) -> Path: path = Path(log_dir) if log_dir else root / "deploy" / "logs" @@ -21,62 +40,218 @@ def _health_url(profile: dict[str, str]) -> str: return f"http://{host}:{profile['PORT']}" +def _node_hosts(profile: dict[str, str]) -> list[str]: + return [h for h in shlex.split(profile.get("NODE_HOSTS", "")) if h] + + +def _ssh_host(profile: dict[str, str], host: str) -> str: + user = profile.get("NODE_SSH_USER", "root") + if not user or "@" in host: + return host + return f"{user}@{host}" + + +def _sub_node_rank(profile: dict[str, str], rank: int) -> dict[str, str]: + """Return a profile copy with NODE_RANK rendered for the given rank.""" + rendered = dict(profile) + container_name = str(rendered.get("CONTAINER_NAME", "server")) + has_template = _NODE_RANK_RE.search(container_name) is not None + + def _sub(value: str | None) -> str: + if not value: + return value or "" + return _NODE_RANK_RE.sub(lambda m: m.group(1) or str(rank), str(value)) + + for key in _TEMPLATED_KEYS: + if key in rendered: + rendered[key] = _sub(rendered[key]) + if not has_template: + rendered["CONTAINER_NAME"] = f"{container_name}_node{rank}" + rendered["NODE_RANK"] = str(rank) + return rendered + + +def _server_cmd(profile: dict[str, str]) -> str: + """Launch command prefix (without LAUNCH_ARGS).""" + if profile.get("SERVER_CMD"): + return profile["SERVER_CMD"] + engine = (profile.get("ENGINE") or "sglang").lower() + python_bin = profile.get("CONTAINER_PYTHON", "python") + if engine == "vllm": + return f"{python_bin} -m vllm.entrypoints.openai.api_server" + return f"{python_bin} -m sglang.launch_server" + + +def _docker_run_cmd(profile: dict[str, str]) -> list[str]: + container = profile["CONTAINER_NAME"] + cmd = ["docker", "run", "-d", "--name", container] + if profile.get("DOCKER_FLAGS"): + cmd += shlex.split(profile["DOCKER_FLAGS"]) + if profile.get("NETWORK_MODE") == "bridge" and profile.get("PORT_MAP") == "1": + cmd += ["-p", f"{profile['PORT']}:{profile['PORT']}"] + for volume in split_env_assignments(profile.get("VOLUMES", "") + " " + profile.get("PATCH_MOUNTS", "")): + cmd += ["-v", volume] + for assignment in split_env_assignments( + profile.get("DEVICE_VARS", "") + " " + profile.get("ENGINE_ENV", "") + ): + cmd += ["-e", assignment] + cmd.append(profile["DOCKER_IMAGE"]) + + bootstrap = profile.get("BOOTSTRAP", "") + if bootstrap: + encoded = base64.b64encode(bootstrap.encode("utf-8")).decode("ascii") + cmd += ["bash", "-c", f"echo {encoded} | base64 -d | bash"] + else: + cmd += ["bash", "-c", f"{_server_cmd(profile)} {profile['LAUNCH_ARGS']}"] + return cmd + + +def _is_multinode(profile: dict[str, str], node_rank: str | None) -> bool: + return len(_node_hosts(profile)) > 1 and node_rank is None + + +def _local_rank(profile: dict[str, str]) -> int: + """Rank that runs on this host (no ssh); defaults to 0 (head).""" + try: + return int(profile.get("LOCAL_NODE_RANK", "0")) + except ValueError: + return 0 + + +def _node_target(profile: dict[str, str], rank: int) -> str | None: + """ssh host for the rank, or None when the rank runs locally.""" + if rank == _local_rank(profile): + return None + return _ssh_host(profile, _node_hosts(profile)[rank]) + + +def _run_docker_cmd( + profile: dict[str, str], + cmd: list[str], + host: str | None, + dry_run: bool, + log_path: Path | None = None, +) -> int: + joined = shlex.join(cmd) + if log_path is not None: + log_path.write_text(joined + "\n", encoding="utf-8") + if dry_run: + print(joined) + return 0 + if host: + remote = f"bash -c {shlex.quote(joined)}" + result = subprocess.run( + ["ssh", "-o", "StrictHostKeyChecking=no", host, remote], + capture_output=True, + text=True, + ) + if result.returncode != 0: + print(result.stdout, result.stderr) + return result.returncode + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0: + print(result.stdout, result.stderr) + return result.returncode + + +def _start_docker( + profile: dict[str, str], + root: Path, + dry_run: bool, + logs: Path, +) -> int: + if not profile.get("DOCKER_IMAGE"): + raise SystemExit("docker runtime requires DOCKER_IMAGE") + if not profile.get("CONTAINER_NAME"): + raise SystemExit("docker runtime requires CONTAINER_NAME") + + hosts = _node_hosts(profile) + node_rank = profile.get("NODE_RANK") + health_wait = int(profile.get("HEALTH_WAIT_S", "600") or 600) + + if _is_multinode(profile, node_rank): + # Workers first (rank > 0), then head (rank 0), mirroring the legacy + # multi-node scripts so NCCL bootstrap finds the head up. The + # LOCAL_NODE_RANK rank runs locally (no ssh). + order = [r for r in range(len(hosts)) if r != 0] + [0] + for rank in order: + host = hosts[rank] + target = _node_target(profile, rank) + node_profile = _sub_node_rank(profile, rank) + cmd = _docker_run_cmd(node_profile) + if not dry_run: + if target: + subprocess.run( + ["ssh", "-o", "StrictHostKeyChecking=no", target, f"docker rm -f {shlex.quote(node_profile['CONTAINER_NAME'])}"], + capture_output=True, + check=False, + ) + else: + subprocess.run( + ["docker", "rm", "-f", node_profile["CONTAINER_NAME"]], + capture_output=True, + check=False, + ) + print(f"starting container {node_profile['CONTAINER_NAME']} on {host} (rank {rank}) ...") + rc = _run_docker_cmd( + profile, + cmd, + target, + dry_run, + logs / f"{node_profile['CONTAINER_NAME']}.cmd.txt", + ) + if rc != 0: + return rc + if rank != 0: + print(" sleeping 5s before head ...") + if not dry_run: + time.sleep(5) + + head_url = f"http://{profile.get('HEALTH_HOST', hosts[0])}:{profile['PORT']}" + if dry_run: + print(f"[dry-run] health: {head_url}{profile.get('HEALTH_PATH', '/health')}") + return 0 + if wait_health(head_url, profile.get("HEALTH_PATH", "/health"), health_wait): + print(f"server on {hosts[0]} is healthy") + return 0 + print("ERROR: head failed health check; inspect per-node logs") + return 1 + + # Single-node or explicit --node-rank: run locally on this host. + if not dry_run: + subprocess.run(["docker", "rm", "-f", profile["CONTAINER_NAME"]], capture_output=True, check=False) + + cmd = _docker_run_cmd(profile) + container = profile["CONTAINER_NAME"] + print(f"starting container {container} ...") + rc = _run_docker_cmd(profile, cmd, None, dry_run, logs / f"{container}.cmd.txt") + if rc != 0: + return rc + if node_rank is not None and node_rank != "0": + print(f"node rank {node_rank} started; skipping health check") + return 0 + if wait_health(_health_url(profile), profile.get("HEALTH_PATH", "/health"), health_wait): + print(f"container {container} is healthy") + return 0 + subprocess.run(["docker", "logs", "--tail", "100", container], check=False) + print(f"ERROR: container {container} failed health check") + return 1 + + def start(profile: dict[str, str], root: Path, dry_run: bool = False, log_dir: str | None = None) -> int: validate_runtime(profile) logs = _log_dir(root, log_dir) - runtime = profile["RUNTIME"] - container = profile.get("CONTAINER_NAME", "") - health_wait = int(profile.get("HEALTH_WAIT_S", "600") or 600) - - if runtime == "docker": - if not profile.get("DOCKER_IMAGE"): - raise SystemExit("docker runtime requires DOCKER_IMAGE") - if not container: - raise SystemExit("docker runtime requires CONTAINER_NAME") - if not dry_run: - subprocess.run(["docker", "rm", "-f", container], capture_output=True, check=False) - - cmd = ["docker", "run", "-d", "--name", container] - if profile.get("DOCKER_FLAGS"): - cmd += shlex.split(profile["DOCKER_FLAGS"]) - if profile.get("NETWORK_MODE") == "bridge" and profile.get("PORT_MAP") == "1": - cmd += ["-p", f"{profile['PORT']}:{profile['PORT']}"] - for volume in split_env_assignments(profile.get("VOLUMES", "") + " " + profile.get("PATCH_MOUNTS", "")): - cmd += ["-v", volume] - for assignment in split_env_assignments( - profile.get("DEVICE_VARS", "") + " " + profile.get("ENGINE_ENV", "") - ): - cmd += ["-e", assignment] - cmd.append(profile["DOCKER_IMAGE"]) - - bootstrap = profile.get("BOOTSTRAP", "") - if bootstrap: - encoded = base64.b64encode(bootstrap.encode("utf-8")).decode("ascii") - cmd += ["bash", "-c", f"echo {encoded} | base64 -d | bash"] - else: - python_bin = profile.get("CONTAINER_PYTHON", "python") - cmd += ["bash", "-c", f"{python_bin} -m sglang.launch_server {profile['LAUNCH_ARGS']}"] - - (logs / f"{container}.cmd.txt").write_text(shlex.join(cmd) + "\n", encoding="utf-8") - if dry_run: - print(shlex.join(cmd)) - return 0 - print(f"starting container {container} ...") - result = subprocess.run(cmd, capture_output=True, text=True) - if result.returncode != 0: - print(result.stdout, result.stderr) - return result.returncode - if wait_health(_health_url(profile), profile.get("HEALTH_PATH", "/health"), health_wait): - print(f"container {container} is healthy") - return 0 - subprocess.run(["docker", "logs", "--tail", "100", container], check=False) - print(f"ERROR: container {container} failed health check") - return 1 + if profile["RUNTIME"] == "docker": + return _start_docker(profile, root, dry_run, logs) # native runtime python_bin = profile.get("PYTHON_BIN") or profile.get("CONTAINER_PYTHON") or "python3" - launch_args = shlex.split(profile.get("LAUNCH_ARGS", "")) - cmd = [python_bin, "-m", "sglang.launch_server", *launch_args] + if profile.get("SERVER_CMD"): + cmd = shlex.split(profile["SERVER_CMD"]) + shlex.split(profile.get("LAUNCH_ARGS", "")) + else: + engine = (profile.get("ENGINE") or "sglang").lower() + module = "vllm.entrypoints.openai.api_server" if engine == "vllm" else "sglang.launch_server" + cmd = [python_bin, "-m", module, *shlex.split(profile.get("LAUNCH_ARGS", ""))] log_path = logs / f"{profile.get('MODEL_NAME', 'server')}.log" pid_file = logs / f"{profile.get('MODEL_NAME', 'server')}.pid" if dry_run: @@ -86,6 +261,7 @@ def start(profile: dict[str, str], root: Path, dry_run: bool = False, log_dir: s proc = subprocess.Popen(cmd, stdout=f, stderr=subprocess.STDOUT) pid_file.write_text(str(proc.pid), encoding="utf-8") print(f"started native server pid={proc.pid} log={log_path}") + health_wait = int(profile.get("HEALTH_WAIT_S", "600") or 600) if wait_health(_health_url(profile), profile.get("HEALTH_PATH", "/health"), health_wait): print("native server is healthy") return 0 @@ -93,15 +269,45 @@ def start(profile: dict[str, str], root: Path, dry_run: bool = False, log_dir: s return 1 +def _docker_stop(profile: dict[str, str], dry_run: bool = False) -> int: + node_rank = profile.get("NODE_RANK") + if _is_multinode(profile, node_rank): + for rank, host in enumerate(_node_hosts(profile)): + node_profile = _sub_node_rank(profile, rank) + name = node_profile["CONTAINER_NAME"] + target = _node_target(profile, rank) + if dry_run: + if target: + print(f"[dry-run] ssh {target} docker rm -f {name}") + else: + print(f"[dry-run] docker rm -f {name} (local)") + continue + if target: + subprocess.run( + ["ssh", "-o", "StrictHostKeyChecking=no", target, f"docker rm -f {shlex.quote(name)}"], + capture_output=True, + check=False, + ) + print(f"container {name} removed on {host}") + else: + subprocess.run(["docker", "rm", "-f", name], capture_output=True, check=False) + print(f"container {name} removed (local)") + return 0 + container = profile.get("CONTAINER_NAME") + if not container: + raise SystemExit("docker runtime requires CONTAINER_NAME") + if dry_run: + print(f"[dry-run] docker rm -f {container}") + return 0 + subprocess.run(["docker", "rm", "-f", container], capture_output=True, check=False) + print(f"container {container} removed") + return 0 + + def stop(profile: dict[str, str], root: Path) -> int: validate_runtime(profile) if profile["RUNTIME"] == "docker": - container = profile.get("CONTAINER_NAME") - if not container: - raise SystemExit("docker runtime requires CONTAINER_NAME") - subprocess.run(["docker", "rm", "-f", container], capture_output=True, check=False) - print(f"container {container} removed") - return 0 + return _docker_stop(profile) logs = _log_dir(root, None) pid_file = logs / f"{profile.get('MODEL_NAME', 'server')}.pid" if pid_file.exists(): @@ -113,18 +319,52 @@ def stop(profile: dict[str, str], root: Path) -> int: return 0 +def _docker_status(profile: dict[str, str]) -> int: + node_rank = profile.get("NODE_RANK") + if _is_multinode(profile, node_rank): + all_ok = 0 + for rank, host in enumerate(_node_hosts(profile)): + node_profile = _sub_node_rank(profile, rank) + name = node_profile["CONTAINER_NAME"] + target = _node_target(profile, rank) + if target: + result = subprocess.run( + ["ssh", "-o", "StrictHostKeyChecking=no", target, f"docker inspect -f '{{{{.State.Status}}}}' {shlex.quote(name)}"], + capture_output=True, + text=True, + check=False, + ) + state = result.stdout.strip() or "missing" + print(f"{host}: {name} -> {state}") + if result.returncode != 0 or state != "running": + all_ok = 1 + else: + result = subprocess.run( + ["docker", "inspect", "-f", "{{.State.Status}}", name], + capture_output=True, + text=True, + check=False, + ) + state = result.stdout.strip() or "missing" + print(f"{host}: {name} -> {state}") + if result.returncode != 0 or state != "running": + all_ok = 1 + return all_ok + container = profile.get("CONTAINER_NAME", "") + result = subprocess.run( + ["docker", "inspect", "-f", "{{.State.Status}}", container], + capture_output=True, + text=True, + check=False, + ) + print(result.stdout.strip() or "missing") + return 0 if result.returncode == 0 else 1 + + def status(profile: dict[str, str], root: Path) -> int: validate_runtime(profile) if profile["RUNTIME"] == "docker": - container = profile.get("CONTAINER_NAME", "") - result = subprocess.run( - ["docker", "inspect", "-f", "{{.State.Status}}", container], - capture_output=True, - text=True, - check=False, - ) - print(result.stdout.strip() or "missing") - return 0 if result.returncode == 0 else 1 + return _docker_status(profile) logs = _log_dir(root, None) pid_file = logs / f"{profile.get('MODEL_NAME', 'server')}.pid" if not pid_file.exists():