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

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:

lightllm/common/basemodel/attention/fa3/fp.py

@@ -4,9 +4,13 @@
 from typing import Optional, TYPE_CHECKING
 from lightllm.utils.dist_utils import get_current_device_id
 from lightllm.utils.sgl_utils import flash_attn_with_kvcache, flash_attn_with_kvcache_autotune
-from lightllm.utils.envs_utils import get_env_start_args
-from lightllm.common.basemodel.triton_kernel.fa3_utils import page_table_copy
+from lightllm.common.basemodel.triton_kernel.fa3_utils import (
+    build_dynamic_spec_fa3_decode_params,
+    page_table_copy,
+)
 from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_cumsum_pad0_tensor
+from lightllm.common.basemodel.triton_kernel.mtp_utils import build_mtp_shared_group_markers
+from lightllm.utils.envs_utils import get_env_start_args
 
 
 class Fa3AttBackend(BaseAttBackend):
@@ -20,13 +24,19 @@ def get_page_table_buffer(self):
         """
         model = self.model
         if not hasattr(self, "_shared_page_table_buffer"):
+            max_att_batch_size = model.graph_max_batch_size
+            if not get_env_start_args().mtp_dynamic_verify:
+                # FA3 merges each fixed speculative block into one attention sequence.
+                max_att_batch_size //= model.mtp_manager.get_decode_batch_multiplier(model.is_mtp_draft_model)
+
+            buffer_count = 2 if model.args.enable_decode_microbatch_overlap else 1
             self._shared_page_table_buffer = [
-                torch.empty(model.graph_max_batch_size * model.graph_max_len_in_batch, dtype=torch.int32).to(
-                    get_current_device_id()
-                ),
-                torch.empty(model.graph_max_batch_size * model.graph_max_len_in_batch, dtype=torch.int32).to(
-                    get_current_device_id()
-                ),
+                torch.empty(
+                    max_att_batch_size * model.graph_max_len_in_batch,
+                    dtype=torch.int32,
+                    device=get_current_device_id(),
+                )
+                for _ in range(buffer_count)
             ]
         return self._shared_page_table_buffer
 
@@ -42,8 +52,10 @@ class Fa3PrefillAttState(BasePrefillAttState):
     cu_seqlens_q: torch.Tensor = None
     cu_seqlens_k: torch.Tensor = None
     page_table: torch.Tensor = None
+    causal: bool = None
 
     def init_state(self):
+        self.causal = self.backend.uses_causal_attention()
         self.cu_seqlens_q = self.infer_state.b1_cu_q_seq_len.int()
         self.cu_seqlens_k = self.infer_state.b1_cu_kv_seq_len.int()
         self.page_table = torch.empty(
@@ -102,7 +114,7 @@ def _nomarl_prefill_att(
             cu_seqlens_k_new=self.cu_seqlens_k,
             max_seqlen_q=self.infer_state.max_q_seq_len,
             softmax_scale=sm_scale,
-            causal=True,
+            causal=self.causal,
             window_size=window_size,
             softcap=0.0,
             k_descale=k_descale,
@@ -121,31 +133,71 @@ class Fa3DecodeAttState(BaseDecodeAttState):
     b_att_seq_len: torch.Tensor = None
     # 在是否开启mtp 的不同模式下，其设置不同的值，可以加速算子的运行。
     decode_max_q_seq_len: int = None
+    causal: bool = None
 
     def init_state(self):
         self.backend: Fa3AttBackend = self.backend
-        args_mtp_step = get_env_start_args().mtp_step
-
-        if args_mtp_step > 0:
-            # 修正 mtp 在 fa3 下的输入。
-            mtp_size = args_mtp_step + 1
-            b_q_seq_len = torch.full(
-                (self.infer_state.b_seq_len.shape[0] // mtp_size,),
-                fill_value=mtp_size,
-                dtype=torch.int32,
-                device=self.infer_state.b_seq_len.device,
-            )
-            b_kv_seq_len = self.infer_state.b_seq_len[mtp_size - 1 :: mtp_size]
-            b1_cu_q_seq_len, b1_cu_kv_seq_len = gen_cumsum_pad0_tensor(b_q_seq_len, b_kv_seq_len)
-            self.cu_seqlens_q = b1_cu_q_seq_len.int()
-            self.cu_seqlens_k = b1_cu_kv_seq_len.int()
+        self.causal = self.backend.uses_causal_attention()
+        draft_step = self.backend.model.mtp_manager.get_decode_draft_step(self.backend.model.is_mtp_draft_model)
+        if self.backend.uses_dynamic_spec_verify_layout():
+            b_att_req_idx = self._init_dynamic_spec_verify_state(draft_step)
+        elif draft_step > 0:
+            b_att_req_idx = self._init_fixed_spec_decode_state(draft_step)
         else:
-            self.cu_seqlens_q = self.infer_state.b1_cu_q_seq_len.int()
-            self.cu_seqlens_k = self.infer_state.b1_cu_kv_seq_len.int()
+            b_att_req_idx = self._init_normal_decode_state()
 
-        att_batch_size = self.infer_state.batch_size // (args_mtp_step + 1)
-        assert self.infer_state.batch_size % (args_mtp_step + 1) == 0
+        self._init_page_table(b_att_req_idx)
 
+    def _init_dynamic_spec_verify_state(self, draft_step: int) -> torch.Tensor:
+        b_mark_mtp_shared_group = build_mtp_shared_group_markers(
+            self.infer_state.b_req_idx,
+            hold_req_id=self.backend.model.req_manager.HOLD_REQUEST_ID,
+        )
+        b_q_seq_len, b_kv_seq_len, b_att_req_idx, self.b_att_seq_len = build_dynamic_spec_fa3_decode_params(
+            b_req_idx=self.infer_state.b_req_idx,
+            b_seq_len=self.infer_state.b_seq_len,
+            b_mark_mtp_shared_group=b_mark_mtp_shared_group,
+            att_batch_size=self.infer_state.batch_size,
+            hold_req_id=self.backend.model.req_manager.HOLD_REQUEST_ID,
+        )
+        self._init_spec_decode_cu_seqlens(b_q_seq_len, b_kv_seq_len)
+        self.decode_max_q_seq_len = draft_step + 1
+        return b_att_req_idx
+
+    def _init_fixed_spec_decode_state(self, draft_step: int) -> torch.Tensor:
+        mtp_size = draft_step + 1
+        assert self.infer_state.batch_size % mtp_size == 0, (
+            "FA3 fixed-layout decode requires batch_size to be divisible by draft_step + 1, "
+            f"got batch_size={self.infer_state.batch_size}, draft_step={draft_step}."
+        )
+
+        b_q_seq_len = torch.full(
+            (self.infer_state.b_seq_len.shape[0] // mtp_size,),
+            fill_value=mtp_size,
+            dtype=torch.int32,
+            device=self.infer_state.b_seq_len.device,
+        )
+        b_kv_seq_len = self.infer_state.b_seq_len[draft_step::mtp_size]
+        b_att_req_idx = self.infer_state.b_req_idx[draft_step::mtp_size]
+        self.b_att_seq_len = b_kv_seq_len.contiguous()
+        self._init_spec_decode_cu_seqlens(b_q_seq_len, b_kv_seq_len)
+        self.decode_max_q_seq_len = mtp_size
+        return b_att_req_idx
+
+    def _init_normal_decode_state(self) -> torch.Tensor:
+        self.cu_seqlens_q = self.infer_state.b1_cu_q_seq_len.int()
+        self.cu_seqlens_k = self.infer_state.b1_cu_kv_seq_len.int()
+        self.b_att_seq_len = self.infer_state.b_seq_len
+        self.decode_max_q_seq_len = 1
+        return self.infer_state.b_req_idx
+
+    def _init_spec_decode_cu_seqlens(self, b_q_seq_len: torch.Tensor, b_kv_seq_len: torch.Tensor):
+        b1_cu_q_seq_len, b1_cu_kv_seq_len = gen_cumsum_pad0_tensor(b_q_seq_len, b_kv_seq_len)
+        self.cu_seqlens_q = b1_cu_q_seq_len.int()
+        self.cu_seqlens_k = b1_cu_kv_seq_len.int()
+
+    def _init_page_table(self, b_att_req_idx: torch.Tensor):
+        att_batch_size = b_att_req_idx.shape[0]
         model = self.backend.model
         # 可以使用 cuda graph的时候从 buffer中申请
         if (
@@ -163,23 +215,11 @@ def init_state(self):
                 device=self.infer_state.input_ids.device,
             )
 
-        if args_mtp_step > 0:
-            page_table_copy(
-                page_table=self.page_table[:, : self.infer_state.max_kv_seq_len],
-                req_to_token_indexs=model.req_manager.req_to_token_indexs,
-                b_req_idx=self.infer_state.b_req_idx[args_mtp_step :: (args_mtp_step + 1)],
-            )
-            self.b_att_seq_len = self.infer_state.b_seq_len[args_mtp_step :: (args_mtp_step + 1)].contiguous()
-            self.decode_max_q_seq_len = args_mtp_step + 1
-        else:
-            page_table_copy(
-                page_table=self.page_table[:, : self.infer_state.max_kv_seq_len],
-                req_to_token_indexs=model.req_manager.req_to_token_indexs,
-                b_req_idx=self.infer_state.b_req_idx,
-            )
-            self.b_att_seq_len = self.infer_state.b_seq_len
-            self.decode_max_q_seq_len = 1
-        return
+        page_table_copy(
+            page_table=self.page_table[:, : self.infer_state.max_kv_seq_len],
+            req_to_token_indexs=model.req_manager.req_to_token_indexs,
+            b_req_idx=b_att_req_idx,
+        )
 
     def copy_for_decode_cuda_graph(self, new_state: "Fa3DecodeAttState"):
         super().copy_for_decode_cuda_graph(new_state)
@@ -232,7 +272,7 @@ def _normal_decode_att(
             cu_seqlens_k_new=self.cu_seqlens_k,
             max_seqlen_q=self.decode_max_q_seq_len,
             softmax_scale=sm_scale,
-            causal=True,
+            causal=self.causal,
             window_size=window_size,
             softcap=0.0,
             k_descale=k_descale,

lightllm/common/basemodel/attention/fa3/fp8.py

@@ -3,7 +3,6 @@
 from ..base_att import AttControl
 from typing import Optional, TYPE_CHECKING
 from lightllm.utils.sgl_utils import flash_attn_with_kvcache
-from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.common.basemodel.triton_kernel.quantization.q_per_head_fp8_quant import q_per_head_fp8_quant
 from lightllm.utils.vllm_utils import HAS_VLLM, vllm_ops
 from typing import Union
@@ -99,7 +98,7 @@ def _fp8_prefill_att(
             cu_seqlens_q=self.cu_seqlens_q,
             cu_seqlens_k_new=self.cu_seqlens_k,
             max_seqlen_q=self.infer_state.max_q_seq_len,
-            causal=True,
+            causal=self.causal,
             window_size=(-1, -1),
             softcap=0.0,
             q_descale=q_scale,
@@ -119,22 +118,18 @@ def init_state(self):
         super().init_state()
         self.backend: Fp8Fa3AttBackend = self.backend
 
-        args_mtp_step = get_env_start_args().mtp_step
-        att_batch_size = self.infer_state.batch_size // (args_mtp_step + 1)
-        assert self.infer_state.batch_size % (args_mtp_step + 1) == 0
-
-        batch_size = att_batch_size
+        att_batch_size = self.b_att_seq_len.shape[0]
         mem_manager = self.backend.model.mem_manager
 
         offline_scales: torch.Tensor = mem_manager.scales
         head_num = mem_manager.head_num
 
         # 为了减少推理计算量，在推理外部初始化k_descale和v_descale
         self.k_descale = (
-            offline_scales[:, :head_num].view(-1, 1, head_num).expand(offline_scales.shape[0], batch_size, head_num)
+            offline_scales[:, :head_num].view(-1, 1, head_num).expand(offline_scales.shape[0], att_batch_size, head_num)
         )
         self.v_descale = (
-            offline_scales[:, head_num:].view(-1, 1, head_num).expand(offline_scales.shape[0], batch_size, head_num)
+            offline_scales[:, head_num:].view(-1, 1, head_num).expand(offline_scales.shape[0], att_batch_size, head_num)
         )
 
         return
@@ -190,7 +185,7 @@ def _fp8_decode_att(
             cu_seqlens_q=self.cu_seqlens_q,
             cu_seqlens_k_new=self.cu_seqlens_k,
             max_seqlen_q=self.decode_max_q_seq_len,
-            causal=True,
+            causal=self.causal,
             window_size=(-1, -1),
             softcap=0.0,
             q_descale=q_scale.view(self.infer_state.batch_size, k_head_num),