sglang dsv4 MTP

benchmarks/single_node/dsv4_fp4_b300_sglang_mtp.sh

@@ -0,0 +1,170 @@
+#!/usr/bin/env bash
+
+source "$(dirname "$0")/../benchmark_lib.sh"
+
+check_env_vars \
+    MODEL \
+    TP \
+    CONC \
+    ISL \
+    OSL \
+    RANDOM_RANGE_RATIO \
+    RESULT_FILENAME
+
+if [[ -n "$SLURM_JOB_ID" ]]; then
+  echo "JOB $SLURM_JOB_ID running on $SLURMD_NODENAME"
+fi
+
+# The B300 runner overrides MODEL to a pre-staged /data/models path, so skip
+# `hf download`. Only fetch when MODEL looks like a HF repo ID.
+if [[ "$MODEL" != /* ]]; then
+    hf download "$MODEL"
+fi
+
+nvidia-smi
+
+export SGLANG_JIT_DEEPGEMM_PRECOMPILE=0
+
+# TODO(Cam): the deepseek-v4 sglang images install sglang editable at
+# /workspace/sglang/python; prior sglang tags used /sgl-workspace/sglang.
+# The runner mounts our repo at a non-/workspace path for these images so the
+# editable install stays visible. Paths in this script are $PWD-relative for
+# that reason. Drop the runner conditional once lmsys moves sglang back out of
+# /workspace.
+
+SERVER_LOG="$PWD/server.log"
+PORT=${PORT:-8888}
+
+echo "TP: $TP, CONC: $CONC, ISL: $ISL, OSL: $OSL"
+
+EVAL_CONTEXT_ARGS=""
+if [ "${EVAL_ONLY}" = "true" ]; then
+    setup_eval_context
+    EVAL_CONTEXT_ARGS="--context-length $EVAL_MAX_MODEL_LEN"
+fi
+
+start_gpu_monitor --output "$PWD/gpu_metrics.csv"
+
+# Three recipes from https://docs.sglang.io/cookbook/autoregressive/DeepSeek/DeepSeek-V4
+# with EAGLE / MTP enabled:
+#   - low-latency    (CONC <= 32):        TP-only, flashinfer_mxfp4 MoE
+#   - balanced       (32 < CONC <= 128):  + DP-attn, mega-moe EP
+#   - max-throughput (CONC > 128):        + DP-attn, mega-moe EP, max-running-requests=512
+# Speculative-decoding flags follow the cookbook EAGLE config.
+DEEPEP_CONFIG='{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}'
+
+# MTP (EAGLE) speculative-decoding flags applied to every recipe.
+SPEC_FLAGS=(
+    --speculative-algorithm EAGLE
+    --speculative-num-steps 3
+    --speculative-eagle-topk 1
+    --speculative-num-draft-tokens 4
+)
+
+if [[ $CONC -le 32 ]]; then
+    RECIPE=low-latency
+    export SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT=1
+    # common optimizations
+    export SGLANG_OPT_USE_JIT_NORM=1
+    export SGLANG_OPT_USE_JIT_INDEXER_METADATA=1
+    export SGLANG_OPT_USE_TOPK_V2=1
+    export SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2=1
+    RECIPE_FLAGS=(
+        --moe-runner-backend flashinfer_mxfp4
+        --chunked-prefill-size 32768
+        --disable-flashinfer-autotune
+        --mem-fraction-static 0.90
+        --max-running-requests 32
+        --swa-full-tokens-ratio 0.1
+    )
+elif [[ $CONC -le 128 ]]; then
+    RECIPE=balanced
+    export SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT=1
+    # common optimizations
+    export SGLANG_OPT_USE_JIT_NORM=1
+    export SGLANG_OPT_USE_JIT_INDEXER_METADATA=1
+    export SGLANG_OPT_USE_TOPK_V2=1
+    export SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2=1
+    # MoE EP related flags
+    export SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE=1
+    export SGLANG_OPT_FIX_HASH_MEGA_MOE=1
+    export SGLANG_OPT_USE_FAST_MASK_EP=1
+    export SGLANG_OPT_FIX_MEGA_MOE_MEMORY=1
+    export SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=4096
+    export SGLANG_OPT_FIX_NEXTN_MEGA_MOE=1
+    export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=0
+    RECIPE_FLAGS=(
+        --dp-size "$TP"
+        --enable-dp-attention
+        --moe-a2a-backend deepep
+        --deepep-config "$DEEPEP_CONFIG"
+        --mem-fraction-static 0.83
+        --max-running-requests 128
+        --chunked-prefill-size 32768
+        --swa-full-tokens-ratio 0.1
+    )
+else
+    RECIPE=max-throughput
+    export SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT=1
+    # common optimizations
+    export SGLANG_OPT_USE_JIT_NORM=1
+    export SGLANG_OPT_USE_JIT_INDEXER_METADATA=1
+    export SGLANG_OPT_USE_TOPK_V2=1
+    export SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2=1
+    # MoE EP related flags
+    export SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE=1
+    export SGLANG_OPT_FIX_HASH_MEGA_MOE=1
+    export SGLANG_OPT_USE_FAST_MASK_EP=1
+    export SGLANG_OPT_FIX_MEGA_MOE_MEMORY=1
+    export SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=4096
+    export SGLANG_OPT_FIX_NEXTN_MEGA_MOE=1
+    export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=0
+    RECIPE_FLAGS=(
+        --dp-size "$TP"
+        --enable-dp-attention
+        --moe-a2a-backend deepep
+        --deepep-config "$DEEPEP_CONFIG"
+        --mem-fraction-static 0.90
+        --max-running-requests 512
+        --chunked-prefill-size 32768
+        --swa-full-tokens-ratio 0.1
+    )
+fi
+echo "Recipe: $RECIPE (CONC=$CONC)"
+
+set -x
+PYTHONNOUSERSITE=1 sglang serve \
+    --model-path $MODEL \
+    --host 0.0.0.0 \
+    --port $PORT \
+    --trust-remote-code \
+    --tp $TP \
+    "${SPEC_FLAGS[@]}" \
+    "${RECIPE_FLAGS[@]}" $EVAL_CONTEXT_ARGS > $SERVER_LOG 2>&1 &
+
+SERVER_PID=$!
+
+wait_for_server_ready --port "$PORT" --server-log "$SERVER_LOG" --server-pid "$SERVER_PID"
+
+pip install -q datasets pandas
+
+run_benchmark_serving \
+    --model "$MODEL" \
+    --port "$PORT" \
+    --backend vllm \
+    --input-len "$ISL" \
+    --output-len "$OSL" \
+    --random-range-ratio "$RANDOM_RANGE_RATIO" \
+    --num-prompts $((CONC * 10)) \
+    --max-concurrency "$CONC" \
+    --result-filename "$RESULT_FILENAME" \
+    --result-dir "$PWD/" \
+    --use-chat-template
+
+if [ "${RUN_EVAL}" = "true" ]; then
+    run_eval --framework lm-eval --port "$PORT"
+    append_lm_eval_summary
+fi
+
+stop_gpu_monitor
+set +x