feat(LightSpec): Support LightSpec Dynamic Speculative Decoding with EAGLE3, DSpark, and DFlash

deploy_dspark_1p1d.sh

@@ -0,0 +1,175 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+PORT="${PORT:-16666}"
+MODEL_DIR="${MODEL_DIR:-/mnt/afs/models/Qwen3.5-27B}"
+DRAFT_MODEL_DIR="${DRAFT_MODEL_DIR:-${draft_model:-}}"
+MODEL_NAME="${MODEL_NAME:-sensenova-flash-lite-v41-20260816-fp8-step3k-dspark}"
+
+CPU_CACHE_SIZE="${CPU_CACHE_SIZE:-600}"
+CHAT_TEMPLATE="${CHAT_TEMPLATE:-${MODEL_DIR}/chat_template.jinja}"
+PD_MASTER_IP="${PD_MASTER_IP:-127.0.0.1}"
+
+# DSpark checkpoint block_size. Prefill and decode must use the same value.
+MTP_STEP="${MTP_STEP:-5}"
+
+# DSpark expands each logical decode request into speculative verification rows.
+# Start conservatively; raise both values together after checking GPU headroom.
+D_MAX_REQ_SIZE="${D_MAX_REQ_SIZE:-16}"
+D_GRAPH_MAX_BATCH_SIZE="${D_GRAPH_MAX_BATCH_SIZE:-16}"
+
+if [[ -z "${DRAFT_MODEL_DIR}" ]]; then
+  echo "DRAFT_MODEL_DIR is required (the DSpark draft checkpoint directory)." >&2
+  exit 2
+fi
+
+export LOADWORKER="${LOADWORKER:-8}"
+export LIGHTLLM_TRITON_AUTOTUNE_LEVEL="${LIGHTLLM_TRITON_AUTOTUNE_LEVEL:-1}"
+export LIGHTLLM_ANTHROPIC_ENABLE_PDF_PARSING="${LIGHTLLM_ANTHROPIC_ENABLE_PDF_PARSING:-1}"
+export LIGHTLLM_LOG_LEVEL="${LIGHTLLM_LOG_LEVEL:-debug}"
+export PYTHONUNBUFFERED=1
+export LIGHTLLM_PD_SPLIT_MAX_NEW_TOKENS="${LIGHTLLM_PD_SPLIT_MAX_NEW_TOKENS:-4096}"
+
+# Keep same-host PD WebSocket traffic away from HTTP/WebSocket proxies.
+export NO_PROXY="${NO_PROXY:+${NO_PROXY},}${PD_MASTER_IP},localhost"
+export no_proxy="${NO_PROXY}"
+
+RUN_NAME="${RUN_NAME:-qwen35-dspark-1p1d}"
+LOG_DIR="${LOG_ROOT:-/mnt/afs/lightllm-runs}/${RUN_NAME}"
+mkdir -p "${LOG_DIR}"
+
+P_COMMON_ARGS=(
+  --model_dir "${MODEL_DIR}"
+  --model_name "${MODEL_NAME}"
+  --graph_max_batch_size 32
+  --running_max_req_size 32
+  --mem_fraction 0.75
+  --max_image_token_count 4096
+  --max_image_pixels 3686400
+  --batch_max_tokens 32768
+  --chunked_prefill_size 16384
+  --linear_att_cache_size 500
+  --linear_att_hash_page_size 2048
+  --linear_att_page_block_num 8
+  --quant_type fp8w8a8-pt-sgl
+  --mtp_mode dspark
+  --mtp_draft_model_dir "${DRAFT_MODEL_DIR}"
+  --mtp_step "${MTP_STEP}"
+  --chat_template "${CHAT_TEMPLATE}"
+  --pd_trans_mode nixl
+  --pd_kv_page_size 4096
+  --pd_master_ip "${PD_MASTER_IP}"
+  --pd_master_port "${PORT}"
+  --llm_prefill_att_backend fa3 flashqla
+)
+
+D_COMMON_ARGS=(
+  --model_dir "${MODEL_DIR}"
+  --model_name "${MODEL_NAME}"
+  --graph_max_batch_size "${D_GRAPH_MAX_BATCH_SIZE}"
+  --running_max_req_size "${D_MAX_REQ_SIZE}"
+  --mem_fraction 0.75
+  --max_image_token_count 4096
+  --max_image_pixels 3686400
+  --batch_max_tokens 256
+  --linear_att_cache_size 500
+  --quant_type fp8w8a8-pt-sgl
+  --mtp_mode dspark
+  --mtp_draft_model_dir "${DRAFT_MODEL_DIR}"
+  --mtp_step "${MTP_STEP}"
+  --chat_template "${CHAT_TEMPLATE}"
+  --pd_trans_mode nixl
+  --pd_kv_page_size 4096
+  --pd_master_ip "${PD_MASTER_IP}"
+  --pd_master_port "${PORT}"
+)
+
+PIDS=()
+
+cleanup() {
+  if ((${#PIDS[@]} > 0)); then
+    kill -TERM "${PIDS[@]}" 2>/dev/null || true
+    wait "${PIDS[@]}" 2>/dev/null || true
+  fi
+}
+
+trap cleanup EXIT
+trap 'exit 130' INT
+trap 'exit 143' TERM
+
+start_prefill() {
+  local name="$1"
+  local devices="$2"
+  local http_port="$3"
+  local nccl_port="$4"
+
+  echo "Starting ${name}: GPUs=${devices}, port=${http_port}"
+
+  CUDA_VISIBLE_DEVICES="${devices}" \
+    python -m lightllm.server.api_server \
+      "${P_COMMON_ARGS[@]}" \
+      --run_mode prefill \
+      --enable_cpu_cache \
+      --cpu_cache_storage_size "${CPU_CACHE_SIZE}" \
+      --tp 4 \
+      --dp 1 \
+      --visual_dp 4 \
+      --host 0.0.0.0 \
+      --port "${http_port}" \
+      --nccl_port "${nccl_port}" \
+      >>"${LOG_DIR}/${name}.log" 2>&1 &
+
+  PIDS+=("$!")
+}
+
+start_decode() {
+  local name="$1"
+  local devices="$2"
+  local http_port="$3"
+  local nccl_port="$4"
+
+  echo "Starting ${name}: GPUs=${devices}, port=${http_port}"
+
+  CUDA_VISIBLE_DEVICES="${devices}" \
+    python -m lightllm.server.api_server \
+      "${D_COMMON_ARGS[@]}" \
+      --run_mode decode \
+      --tp 4 \
+      --dp 1 \
+      --visual_dp 4 \
+      --host 0.0.0.0 \
+      --port "${http_port}" \
+      --nccl_port "${nccl_port}" \
+      >>"${LOG_DIR}/${name}.log" 2>&1 &
+
+  PIDS+=("$!")
+}
+
+# Single node: Prefill on GPUs 0-3, Decode on GPUs 4-7.
+start_prefill p1 "0,1,2,3" 28761 29761
+start_decode d1 "4,5,6,7" 28762 29762
+
+# CPU-only PD Master. This is the public service endpoint.
+python -m lightllm.server.api_server \
+  --model_dir "${MODEL_DIR}" \
+  --model_name "${MODEL_NAME}" \
+  --run_mode pd_master \
+  --pd_master_mode 1p1d \
+  --host 0.0.0.0 \
+  --port "${PORT}" \
+  --max_image_token_count 4096 \
+  --max_image_pixels 3686400 \
+  --chat_template "${CHAT_TEMPLATE}" \
+  >>"${LOG_DIR}/pd-master.log" 2>&1 &
+PIDS+=("$!")
+
+echo "PD Master is starting at http://${PD_MASTER_IP}:${PORT}"
+echo "logs: ${LOG_DIR}"
+
+# A fixed 1P1D deployment is incomplete as soon as any component exits.
+if wait -n "${PIDS[@]}"; then
+  echo "A LightLLM component exited unexpectedly" >&2
+else
+  echo "A LightLLM component failed" >&2
+fi
+exit 1

lightllm/common/basemodel/attention/base_att.py

@@ -1,42 +1,79 @@
+import threading
+
 import torch
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from typing import Optional, TYPE_CHECKING, Tuple, Union, Dict
 
+from lightllm.utils.dist_utils import get_current_device_id
+from lightllm.utils.envs_utils import get_env_start_args
+
 if TYPE_CHECKING:
     from lightllm.common.basemodel.basemodel import TpPartBaseModel
     from lightllm.common.basemodel.infer_struct import InferStateInfo
 
 
 class BaseAttBackend:
     """
-    用于创建支持各种不同的AttBackend, 如 fa3, flashinfer, triton 实现等，
-    这个是单列模式, 每种backend只有一个实例
+    用于创建支持各种不同的AttBackend, 如 fa3, flashinfer, triton 实现等。
+    每个 model 复用一个 backend 实例。
     """
 
     _instances = {}
+    _workspace_buffers = {}
+    _workspace_buffer_lock = threading.Lock()
 
     def __new__(cls, *args, **kwargs):
         """
-        重写__new__方法实现单例模式
+        Main 和 speculative draft model 可能使用不同的 CUDA graph 上限
+        和缓存布局，不能只按 backend class 共享实例。
         """
-        # 检查是否已经有该类的实例
-        if cls not in cls._instances:
-            # 创建新实例并存储
+        model = kwargs.get("model", args[0] if args else None)
+        instance_key = (cls, model)
+        if instance_key not in cls._instances:
             instance = super().__new__(cls)
-            cls._instances[cls] = instance
-        # 返回已有的实例
-        return cls._instances[cls]
+            cls._instances[instance_key] = instance
+        return cls._instances[instance_key]
 
     def __init__(self, model: "TpPartBaseModel"):
         self.model = model
 
+    @staticmethod
+    def get_gpu_workspace_buffer(key_name: str, workspace_size: int, dtype: torch.dtype = torch.int8) -> torch.Tensor:
+        """Return a process-local workspace shared by key name and CUDA device."""
+        if not key_name:
+            raise ValueError("workspace key_name must not be empty")
+        if workspace_size <= 0:
+            raise ValueError(f"workspace_size must be positive, got {workspace_size}")
+
+        device_id = get_current_device_id()
+        buffer_key = (device_id, key_name, workspace_size, dtype)
+        with BaseAttBackend._workspace_buffer_lock:
+            workspace_buffer = BaseAttBackend._workspace_buffers.get(buffer_key)
+            if workspace_buffer is None:
+                workspace_buffer = torch.empty(workspace_size, dtype=dtype, device=device_id)
+                BaseAttBackend._workspace_buffers[buffer_key] = workspace_buffer
+            return workspace_buffer
+
     def create_att_prefill_state(self) -> "BasePrefillAttState":
         raise NotImplementedError("not impl")
 
     def create_att_decode_state(self) -> "BaseDecodeAttState":
         raise NotImplementedError("not impl")
 
+    def uses_dynamic_spec_verify_layout(self) -> bool:
+        args = get_env_start_args()
+        draft_step = self.model.mtp_manager.get_decode_draft_step(self.model.is_mtp_draft_model)
+        is_main_model = not self.model.is_mtp_draft_model
+        has_decode_draft_step = draft_step > 0
+        dynamic_verify_enabled = args.mtp_dynamic_verify
+        return is_main_model and has_decode_draft_step and dynamic_verify_enabled
+
+    def uses_causal_attention(self) -> bool:
+        args = get_env_start_args()
+        is_parallel_block_draft = self.model.is_mtp_draft_model and args.mtp_mode in ("dspark", "dflash")
+        return not is_parallel_block_draft
+
     def _find_layer_index(
         self, k: torch.Tensor, v: torch.Tensor, att_state: Union["BasePrefillAttState", "BaseDecodeAttState"]
     ) -> int: