# Adaptive speculation length: the scheduler picks how many of the 7 DFlash2
# drafts each step verifies (adaptive_k.py). The per-batch-size schedule only
# makes vLLM capture CUDA graphs for every level; the scheduler ignores it.
# draft-trunc/: after the drafts exist, each request drops its drafts whose
# predicted survival (vLLM's acceptance estimator) is below
# VLLM_DRAFT_TRUNC_TAU. Dead drafts keep the verify shape, skip their MoE
# experts and are rejected, so outputs are unchanged. Every
# VLLM_DRAFT_TRUNC_EXPLORE-th step skips the cut and trains the estimator.
# disable_eagle_block_drop: vLLM drops the last matching prefix-cache block for
# EAGLE-style drafters, 2,304 tokens here at TP=4, and every cache hit then
# recomputes it. With the drop off, an agent turn's time to first token falls by
# about a third and the output is identical (#54).
services:
  glm53:
    volumes:
      - ../experimental/adaptive-k/adaptive_k.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/adaptive_k.py:ro
      - ../experimental/draft-trunc/draft_trunc.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu/spec_decode/draft_trunc.py:ro
      - ../experimental/draft-trunc/model_runner.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu/model_runner.py:ro
      - ../experimental/draft-trunc/rejection_sampler.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py:ro
      - ../experimental/draft-trunc/grouped_topk_router.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py:ro
      - ../experimental/draft-trunc/speculator.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu/spec_decode/speculator.py:ro
      - ../experimental/draft-trunc/dflash2_speculator.py:/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py:ro
    environment:
      - EXTRA_ARGS=--scheduler-cls vllm.v1.core.sched.adaptive_k.AdaptiveKScheduler
      # Merged into the entrypoint's --speculative-config (method, drafter and
      # width stay SPEC_METHOD, DFLASH_MODEL and SPEC_TOKENS).
      - SPEC_EXTRA={"disable_eagle_block_drop":true,"attention_backend":"TRITON_ATTN","num_speculative_tokens_per_batch_size":[[1,1,7],[2,2,5],[3,3,4],[4,4,3],[5,5,2],[6,32,7]]}
      - VLLM_ADAPTIVE_K_CONTROL=/root/.cache/adaptive-k
      - VLLM_DRAFT_TRUNC_TAU=${VLLM_DRAFT_TRUNC_TAU:-0.3}
      - VLLM_DRAFT_TRUNC_EXPLORE=${VLLM_DRAFT_TRUNC_EXPLORE:-4}
      - VLLM_DRAFT_TRUNC_CONTROL=/root/.cache/draft-trunc
      - VLLM_DRAFT_TRUNC_LOG=${VLLM_DRAFT_TRUNC_LOG:-}
      - VLLM_ADAPTIVE_K_PER_REQUEST=${VLLM_ADAPTIVE_K_PER_REQUEST:-0}
      - VLLM_ADAPTIVE_K_CODE_MODE=${VLLM_ADAPTIVE_K_CODE_MODE:-gated}
      - VLLM_ADAPTIVE_K_ONLINE=${VLLM_ADAPTIVE_K_ONLINE:-1}
      - VLLM_ADAPTIVE_K_MODEL=${VLLM_ADAPTIVE_K_MODEL:-}
      - VLLM_ADAPTIVE_K_COST_MS=${VLLM_ADAPTIVE_K_COST_MS:-2:29.7,3:33.1,4:36.5,5:39.8,6:43.1,7:46.0,8:48.9}