Organize spec-related data structures (#10735)

python/sglang/srt/speculative/spec_info.py

@@ -1,4 +1,8 @@
+from abc import ABC, abstractmethod
 from enum import IntEnum, auto
+from typing import List, Tuple
+
+from sglang.srt.managers.schedule_batch import ModelWorkerBatch
 
 
 class SpeculativeAlgorithm(IntEnum):
@@ -35,3 +39,41 @@ class SpeculativeAlgorithm(IntEnum):
         if name is not None:
             name = name.upper()
         return name_map[name]
+
+
+class SpecInputType(IntEnum):
+    # NOTE: introduce this to distinguish the SpecInput types of multiple algorithms when asserting in attention backends.
+    # If all algorithms can share the same datastrucutre of draft_input and verify_input, consider simplify it
+    EAGLE_DRAFT = auto()
+    EAGLE_VERIFY = auto()
+    NGRAM_VERIFY = auto()
+
+
+class SpecInput(ABC):
+    def __init__(self, spec_input_type: SpecInputType):
+        self.spec_input_type = spec_input_type
+
+    def is_draft_input(self) -> bool:
+        # FIXME: remove this function which is only used for assertion
+        # or use another variable name like `draft_input` to substitute `spec_info`
+        return self.spec_input_type == SpecInputType.EAGLE_DRAFT
+
+    def is_verify_input(self) -> bool:
+        return self.spec_input_type in {
+            SpecInputType.EAGLE_VERIFY,
+            SpecInputType.NGRAM_VERIFY,
+        }
+
+    @abstractmethod
+    def get_spec_adjust_token_coefficient(self) -> Tuple[int, int]:
+        pass
+
+    def get_spec_adjusted_global_num_tokens(
+        self, forward_batch: ModelWorkerBatch
+    ) -> Tuple[List[int], List[int]]:
+        c1, c2 = self.get_spec_adjust_token_coefficient()
+        global_num_tokens = [x * c1 for x in forward_batch.global_num_tokens]
+        global_num_tokens_for_logprob = [
+            x * c2 for x in forward_batch.global_num_tokens_for_logprob
+        ]
+        return global_num_tokens, global_num_tokens_for_logprob