test: benchmark DP2 TP32 with DFlash

experiments/pro6000/kimi3_pro6000_dp2_nospec_ab/README.md

@@ -4,8 +4,9 @@
 
 - `tp32`：每副本 `TP32/PP1/EP4`
 - `pp8`：每副本 `TP4/PP8/EP4`
+- `tp32-dflash`：每副本 `TP32/PP1/EP4`，开启 DFlash，draft tokens 16；`mem-fraction-static=0.83` 为 16K Prefill 的 FlashInfer MoE workspace 留出余量
 
-测试口径固定为 SPEED-Bench、`16K -> 512`、`C=8`、40 请求、chunk 8K、BF16 KV、FlashInfer MXFP4，不启用投机解码和 Prefix Cache。
+测试口径固定为 SPEED-Bench、`16K -> 512`、`C=8`、40 请求、chunk 8K、BF16 KV、FlashInfer MXFP4，并关闭 Prefix Cache。`tp32`、`pp8` 不启用投机解码；`tp32-dflash` 只增加 DFlash 及其必要参数。
 
 ```bash
 cd /data/hzy/sskj/experiments/pro6000/kimi3_pro6000_dp2_nospec_ab
@@ -13,10 +14,12 @@ cd /data/hzy/sskj/experiments/pro6000/kimi3_pro6000_dp2_nospec_ab
 # 只展开命令
 bash run_dp2_nospec_ab.sh dry-run tp32
 bash run_dp2_nospec_ab.sh dry-run pp8
+bash run_dp2_nospec_ab.sh dry-run tp32-dflash
 
 # 完整实验
 RUN_ID=dp2-tp32-$(date +%Y%m%d-%H%M%S) bash run_dp2_nospec_ab.sh all tp32
 RUN_ID=dp2-pp8-$(date +%Y%m%d-%H%M%S) bash run_dp2_nospec_ab.sh all pp8
+RUN_ID=dp2-tp32-dflash-$(date +%Y%m%d-%H%M%S) bash run_dp2_nospec_ab.sh all tp32-dflash
 ```
 
 每轮结果保存在 `results/<RUN_ID>/`。正式计时前，脚本分别给两个副本发送一个 `16K -> 1` 请求，避免 Router 只预热到一个副本。
@@ -27,10 +30,29 @@ RUN_ID=dp2-pp8-$(date +%Y%m%d-%H%M%S) bash run_dp2_nospec_ab.sh all pp8
 |---|---:|---:|---:|---:|---:|
 | `2 x TP32/PP1/EP4` | 2797.64 | 85.99 | 8.35s / 20.79s | 75.93ms / 88.06ms | 47.15s |
 | `2 x TP4/PP8/EP4` | 1799.46 | 55.31 | 7.84s / 11.68s | 126.86ms / 145.13ms | 72.66s |
+| `2 x TP32/PP1/EP4 + DFlash` | 2218.94 | 68.20 | 8.12s / 21.75s | 98.44ms / 154.58ms | 58.42s |
+
+### TP32 DFlash 对照
+
+`tp32-dflash` 的投机参数与 PD+DFlash 保持一致：draft model 为 `Kimi-K3-DFlash`，draft KV 为 BF16，`speculative-num-draft-tokens=16`。Router 后的 `/server_info` 不汇总接受长度；两副本 Decode 日志共 37 个周期采样的平均接受长度为 3.76，平均接受率为 0.184。
+
+| 指标 | TP32 无投机 | TP32 + DFlash | 变化 |
+|---|---:|---:|---:|
+| Benchmark duration | 238.18s | 300.30s | +26.08% |
+| Input / Output TPS | 2797.64 / 85.99 | 2218.94 / 68.20 | -20.69% |
+| TTFT mean / p95 | 8.35s / 20.79s | 8.12s / 21.75s | -2.73% / +4.62% |
+| TPOT mean / p95 | 75.93ms / 88.06ms | 98.44ms / 154.58ms | +29.64% / +75.54% |
+| E2E mean / p95 | 47.15s / 59.25s | 58.42s / 88.76s | +23.90% / +49.83% |
+
+在当前 `16K -> 512, C=8` 口径下，DFlash 没有带来端到端收益：平均 TTFT 仅下降 2.73%，P95 TTFT 反而增加 4.62%，吞吐下降 20.69%，TPOT 和 E2E 明显恶化。当前参数不建议用于 DP2 正式部署。
+
+第一次 DFlash 正式轮使用 `mem-fraction-static=0.86`，在首个 16K Prefill 的 FlashInfer MXFP4 MoE workspace 分配阶段 OOM；DP2 单实例同时承担 Prefill 和 Decode，不能直接复用 PD 的 D-only 显存比例。仅将该模式降到 0.83 后完整通过，投机参数未改变。
 
 证据目录：
 
 - `results/dp2-tp32-final-20260902-123051/`
 - `results/dp2-pp8-final-20260902-130711/`
+- `results/dp2-tp32-dflash-mem083-20260902-143710/`
+- 首轮 OOM 证据：`results/dp2-tp32-dflash-final-20260902-141001/`
 
-两轮均为 40/40 请求成功。Router 使用固定 `round_robin`，两副本负载接近均分。正式运行后补充了 Router 的模型目录挂载，避免后续 tokenizer 注册产生与本轮 round-robin 转发无关的错误日志。
+三轮最终结果均为 40/40 请求成功。Router 使用固定 `round_robin`，两副本均实际处理请求。正式运行后补充了 Router 的模型目录挂载，避免后续 tokenizer 注册产生与 round-robin 转发无关的错误日志。

experiments/pro6000/kimi3_pro6000_dp2_nospec_ab/config.env

@@ -1,6 +1,14 @@
 EXPERIMENT_BASE=kimi3_pro6000_dp2_nospec_ab
 IMAGE=local/sglang:kimi-k3-pp-dflash-33863-fi0618-situ4460-kvbounds1
 MODEL_PATH=/data/hf_models/Kimi-K3
+DRAFT_MODEL_PATH=/data/hf_models/Kimi-K3-DFlash
+DFLASH_FIX_SOURCE=/data/hzy/sskj/experiments/pro6000/cdflash_fix
+DFLASH_FIX_DIR=/data/hzy/kimi3_dp2_dflash_fix
+DFLASH_FIX_MAP=(
+  arg_groups/speculative_hook.py
+  arg_groups/validation_hook.py
+  speculative/dflash_worker_v2.py
+)
 
 REPLICA_A_NODES=(174.1.60.1 174.1.60.2 174.1.60.3 174.1.60.4)
 REPLICA_B_NODES=(174.1.60.5 174.1.60.6 174.1.60.7 174.1.60.8)
@@ -16,6 +24,7 @@ PAGE_SIZE=64
 KV_DTYPE=bfloat16
 MAX_RUNNING_REQUESTS=8
 CUDA_GRAPH_MAX_BS=8
+DRAFT_TOKENS=16
 IB_DEVICES=mlx5_0,mlx5_1,mlx5_2,mlx5_3
 NCCL_DEBUG_LEVEL=${NCCL_DEBUG_LEVEL:-WARN}
 JIT_CACHE=/data/hzy/cache/kimi-dflash-fi0618-situ4460

experiments/pro6000/kimi3_pro6000_dp2_nospec_ab/run_dp2_nospec_ab.sh

@@ -7,15 +7,20 @@ source "${SCRIPT_DIR}/config.env"
 ACTION="${1:-status}"
 MODE="${2:-tp32}"
 DRY_RUN="${DRY_RUN:-0}"
-[[ "$MODE" == tp32 || "$MODE" == pp8 ]] || {
+[[ "$MODE" == tp32 || "$MODE" == pp8 || "$MODE" == tp32-dflash ]] || {
-  echo "MODE must be tp32 or pp8" >&2
+  echo "MODE must be tp32, pp8, or tp32-dflash" >&2
   exit 2
 }
 
-if [[ "$MODE" == tp32 ]]; then
+SPEC_ENABLED=0
+if [[ "$MODE" == tp32 || "$MODE" == tp32-dflash ]]; then
   TP_SIZE=32
   PP_SIZE=1
-  MEM_FRACTION=0.86
+  if [[ "$MODE" == tp32-dflash ]]; then
+    MEM_FRACTION=0.83
+  else
+    MEM_FRACTION=0.86
+  fi
   MAMBA_RATIO=0.21
 else
   TP_SIZE=4
@@ -23,6 +28,7 @@ else
   MEM_FRACTION=0.84
   MAMBA_RATIO=0.36
 fi
+[[ "$MODE" == tp32-dflash ]] && SPEC_ENABLED=1
 
 RUN_ID="${RUN_ID:-${EXPERIMENT_BASE}-${MODE}-$(date +%Y%m%d-%H%M%S)}"
 [[ "$RUN_ID" =~ ^[a-zA-Z0-9_.-]+$ ]] || exit 2
@@ -80,7 +86,8 @@ emit() {
 }
 
 build_server_command() {
-  local replica="$1" rank="$2" host head
+  local replica="$1" rank="$2" host head rel
+  local -a dflash_mounts=() dflash_args=()
   if [[ "$replica" == a ]]; then
     host="${REPLICA_A_NODES[$rank]}"
     head="${REPLICA_A_NODES[0]}"
@@ -88,12 +95,25 @@ build_server_command() {
     host="${REPLICA_B_NODES[$rank]}"
     head="${REPLICA_B_NODES[0]}"
   fi
+  if [[ "$SPEC_ENABLED" == 1 ]]; then
+    dflash_mounts+=(-v "${DRAFT_MODEL_PATH}:${DRAFT_MODEL_PATH}:ro")
+    for rel in "${DFLASH_FIX_MAP[@]}"; do
+      dflash_mounts+=(-v "${DFLASH_FIX_DIR}/$(basename "$rel"):/opt/kimi-dflash/python/sglang/srt/${rel}:ro")
+    done
+    dflash_args=(
+      --speculative-draft-kv-cache-dtype "$KV_DTYPE"
+      --speculative-algorithm DFLASH
+      --speculative-draft-model-path "$DRAFT_MODEL_PATH"
+      --speculative-num-draft-tokens "$DRAFT_TOKENS"
+    )
+  fi
   HOST="$host"
   CMD=(docker run -d --name "$(container_name "$replica" "$rank")"
     --gpus all --network host --ipc=host --shm-size 32g
     --ulimit memlock=-1 --device /dev/infiniband
     -v "${MODEL_PATH}:${MODEL_PATH}:ro"
     -v "${JIT_CACHE}:/cache"
+    "${dflash_mounts[@]}"
     -e "SGLANG_HOST_IP=${host}" -e PYTHONUNBUFFERED=1
     -e HF_HUB_OFFLINE=1 -e TRANSFORMERS_OFFLINE=1
     -e GLOO_SOCKET_IFNAME=bond0 -e NCCL_SOCKET_IFNAME=bond1
@@ -113,6 +133,7 @@ build_server_command() {
     --nnodes 4 --node-rank "$rank" --dist-init-addr "${head}:${DIST_PORT}"
     --moe-runner-backend flashinfer_mxfp4 --moe-a2a-backend none
     --kv-cache-dtype "$KV_DTYPE"
+    "${dflash_args[@]}"
     --chunked-prefill-size "$CHUNK_SIZE" --page-size "$PAGE_SIZE"
     --mem-fraction-static "$MEM_FRACTION"
     --mamba-full-memory-ratio "$MAMBA_RATIO"
@@ -123,6 +144,23 @@ build_server_command() {
     --host 0.0.0.0 --port "$SERVER_PORT")
 }
 
+stage_dflash_fixes() {
+  local host rel
+  [[ "$SPEC_ENABLED" == 1 ]] || return 0
+  for host in "${ALL_NODES[@]}"; do
+    remote "$host" mkdir -p "$DFLASH_FIX_DIR"
+    for rel in "${DFLASH_FIX_MAP[@]}"; do
+      if [[ "$host" == "${REPLICA_A_NODES[0]}" ]]; then
+        cp "${DFLASH_FIX_SOURCE}/$(basename "$rel")" \
+          "${DFLASH_FIX_DIR}/$(basename "$rel")"
+      else
+        scp "${SSH_OPTIONS[@]}" -q "${DFLASH_FIX_SOURCE}/$(basename "$rel")" \
+          "${SSH_USER}@${host}:${DFLASH_FIX_DIR}/$(basename "$rel")"
+      fi
+    done
+  done
+}
+
 preflight() {
   local host image reference="" gpu ports
   for host in "${ALL_NODES[@]}"; do
@@ -133,6 +171,9 @@ preflight() {
       return 1
     }
     remote "$host" test -r "${MODEL_PATH}/config.json"
+    if [[ "$SPEC_ENABLED" == 1 ]]; then
+      remote "$host" test -r "${DRAFT_MODEL_PATH}/model.safetensors"
+    fi
     remote "$host" test -e /dev/infiniband/uverbs0
     root "$host" mkdir -p "${JIT_CACHE}/tmp"
     gpu="$(remote "$host" nvidia-smi --query-compute-apps=pid --format=csv,noheader)"
@@ -159,6 +200,7 @@ for p in map(int,sys.argv[1:]):
     printf 'RUN_ID=%s\nMODE=%s\n' "$RUN_ID" "$MODE"
     printf 'TP_SIZE=%s\nPP_SIZE=%s\nEP_SIZE=%s\n' "$TP_SIZE" "$PP_SIZE" "$EP_SIZE"
     printf 'IMAGE=%s\nCHUNK_SIZE=%s\nKV_DTYPE=%s\n' "$IMAGE" "$CHUNK_SIZE" "$KV_DTYPE"
+    printf 'SPEC_ENABLED=%s\nDRAFT_TOKENS=%s\n' "$SPEC_ENABLED" "$DRAFT_TOKENS"
     printf 'NUM_PROMPTS=%s\nOUTPUT_LEN=%s\nCONCURRENCY=%s\n' \
       "$NUM_PROMPTS" "$OUTPUT_LEN" "$CONCURRENCY"
   } >"${RESULT_ROOT}/metadata/manifest.env"
@@ -186,6 +228,7 @@ wait_health() {
 start() {
   local replica rank
   [[ "$DRY_RUN" == 1 ]] || preflight
+  [[ "$DRY_RUN" == 1 ]] || stage_dflash_fixes
   for replica in a b; do
     for rank in 1 2 3 0; do
       build_server_command "$replica" "$rank"
@@ -311,7 +354,7 @@ case "$ACTION" in
     stop
     ;;
   *)
-    echo 'Usage: run_dp2_nospec_ab.sh {dry-run|preflight|start|bench|logs|stop|status|all} {tp32|pp8}' >&2
+    echo 'Usage: run_dp2_nospec_ab.sh {dry-run|preflight|start|bench|logs|stop|status|all} {tp32|pp8|tp32-dflash}' >&2
     exit 2
     ;;
 esac