chore: import upstream snapshot with attribution

python/sglang/srt/speculative/ragged_verify.py

@@ -0,0 +1,304 @@
+from __future__ import annotations
+
+import bisect
+from enum import Enum
+from typing import List, Optional, Sequence, Tuple
+
+import msgspec
+import torch
+
+from sglang.srt.environ import envs
+
+
+class RaggedVerifyMode(str, Enum):
+    STATIC = "static"
+    CAP_ACCEPT = "cap-accept"
+    COMPACT = "compact"
+
+
+def read_ragged_verify_mode() -> RaggedVerifyMode:
+    value = envs.SGLANG_RAGGED_VERIFY_MODE.get()
+    for mode in RaggedVerifyMode:
+        if value == mode.value:
+            return mode
+    raise ValueError(
+        f"invalid SGLANG_RAGGED_VERIFY_MODE={value!r}; expected one of "
+        f"{', '.join(repr(m.value) for m in RaggedVerifyMode)}"
+    )
+
+
+def ragged_verify_compact_enabled() -> bool:
+    return read_ragged_verify_mode() == RaggedVerifyMode.COMPACT
+
+
+def round_up_grid(total: int, grid: Sequence[int]) -> int:
+    if not grid:
+        raise ValueError("round_up_grid requires a non-empty grid")
+    if total > grid[-1]:
+        raise ValueError(
+            f"total {total} exceeds max grid tier {grid[-1]}; "
+            "the caller must reject this batch before selecting a graph tier"
+        )
+    index = bisect.bisect_left(grid, total)
+    return grid[index]
+
+
+class RaggedVerifyLayout(msgspec.Struct, frozen=True):
+    verify_lens: torch.Tensor
+    graph_num_tokens: int
+    extend_start_loc: torch.Tensor
+    qo_indptr_device: torch.Tensor
+    verify_lens_cpu: Optional[list[int]] = None
+    total_verify_tokens: Optional[int] = None
+    qo_indptr_host: Optional[torch.Tensor] = None
+    kv_indptr_host: Optional[torch.Tensor] = None
+    kv_lens_host: Optional[torch.Tensor] = None
+    max_q_len: Optional[int] = None
+    max_kv_len: Optional[int] = None
+
+    def __post_init__(self) -> None:
+        if self.verify_lens_cpu is None:
+            return
+        if not self.verify_lens_cpu:
+            raise ValueError("RaggedVerifyLayout requires at least one request")
+        if min(self.verify_lens_cpu) < 1:
+            raise ValueError(
+                f"every request must verify the anchor (verify_len >= 1), got "
+                f"{self.verify_lens_cpu}"
+            )
+        if self.total_verify_tokens != sum(self.verify_lens_cpu):
+            raise ValueError(
+                f"total_verify_tokens {self.total_verify_tokens} != "
+                f"sum(verify_lens_cpu) {sum(self.verify_lens_cpu)}"
+            )
+        if not (self.total_verify_tokens <= self.graph_num_tokens):
+            raise ValueError(
+                f"total_verify_tokens {self.total_verify_tokens} exceeds "
+                f"graph_num_tokens {self.graph_num_tokens}"
+            )
+
+    @property
+    def bs(self) -> int:
+        return int(self.verify_lens.shape[0])
+
+    @classmethod
+    def _assemble_device(
+        cls,
+        *,
+        verify_lens: torch.Tensor,
+        graph_num_tokens: int,
+        verify_lens_cpu: Optional[list[int]] = None,
+        total_verify_tokens: Optional[int] = None,
+    ) -> RaggedVerifyLayout:
+        from sglang.srt.speculative.ragged_verify_kernels import (
+            BuildQoIndptr,
+        )
+
+        verify_lens = verify_lens.to(torch.int32)
+        indptr = BuildQoIndptr.execute(verify_lens=verify_lens)
+        return cls(
+            verify_lens=verify_lens,
+            graph_num_tokens=graph_num_tokens,
+            extend_start_loc=indptr.extend_start_loc,
+            qo_indptr_device=indptr.qo_indptr,
+            verify_lens_cpu=verify_lens_cpu,
+            total_verify_tokens=total_verify_tokens,
+        )
+
+    @classmethod
+    def _assemble(
+        cls,
+        *,
+        verify_lens_cpu: list[int],
+        total_verify_tokens: int,
+        graph_num_tokens: int,
+        device: torch.device,
+    ) -> RaggedVerifyLayout:
+        verify_lens = torch.tensor(verify_lens_cpu, dtype=torch.int32, device=device)
+        return cls._assemble_device(
+            verify_lens=verify_lens,
+            graph_num_tokens=graph_num_tokens,
+            verify_lens_cpu=verify_lens_cpu,
+            total_verify_tokens=total_verify_tokens,
+        )
+
+    @classmethod
+    def from_verify_lens_device(
+        cls,
+        *,
+        verify_lens: torch.Tensor,
+        graph_num_tokens: int,
+    ) -> RaggedVerifyLayout:
+        return cls._assemble_device(
+            verify_lens=verify_lens, graph_num_tokens=graph_num_tokens
+        )
+
+    @classmethod
+    def from_verify_lens(
+        cls,
+        *,
+        verify_lens_cpu: Sequence[int],
+        device: torch.device,
+        grid: Sequence[int],
+        graph_num_tokens_floor: int = 0,
+    ) -> RaggedVerifyLayout:
+        verify_lens_list = [int(v) for v in verify_lens_cpu]
+        total_verify_tokens = sum(verify_lens_list)
+        bucket_input = max(total_verify_tokens, graph_num_tokens_floor)
+        graph_num_tokens = round_up_grid(total=bucket_input, grid=grid)
+
+        return cls._assemble(
+            verify_lens_cpu=verify_lens_list,
+            total_verify_tokens=total_verify_tokens,
+            graph_num_tokens=graph_num_tokens,
+            device=device,
+        )
+
+    def padded_to_bucket(self, *, padded_bs: int) -> RaggedVerifyLayout:
+        from sglang.srt.speculative.ragged_verify_kernels import (
+            PaddedToBucket,
+        )
+
+        padded = PaddedToBucket.execute(
+            verify_lens=self.verify_lens,
+            graph_num_tokens=self.graph_num_tokens,
+            bs=self.bs,
+            padded_bs=padded_bs,
+        )
+
+        return RaggedVerifyLayout._assemble_device(
+            verify_lens=padded,
+            graph_num_tokens=self.graph_num_tokens,
+            total_verify_tokens=self.graph_num_tokens,
+        )
+
+
+def build_capture_verify_lens(
+    *,
+    num_tokens: int,
+    num_slots: int,
+    num_draft_tokens: int,
+) -> list[int]:
+    if num_slots < 1 or num_tokens < num_slots:
+        raise ValueError(
+            f"capture layout needs 1 <= num_slots <= num_tokens, got "
+            f"num_slots={num_slots}, num_tokens={num_tokens}"
+        )
+    if num_tokens > num_slots * num_draft_tokens:
+        raise ValueError(
+            f"capture layout cannot pack num_tokens={num_tokens} into "
+            f"{num_slots} rows of at most {num_draft_tokens} tokens"
+        )
+    base = num_tokens // num_slots
+    rem = num_tokens - base * num_slots
+    return [base + 1] * rem + [base] * (num_slots - rem)
+
+
+def resolve_ragged_verify_layout(forward_batch) -> Optional[RaggedVerifyLayout]:
+    """Layout riding the batch's spec input, or None. Tolerates the runner's
+    ad-hoc replay batch views, which may not carry spec_info at all."""
+    spec_info = getattr(forward_batch, "spec_info", None)
+    if spec_info is None:
+        return None
+    return spec_info.ragged_verify_layout
+
+
+class RaggedTargetVerifyGeometry(msgspec.Struct):
+    cache_seqlens_int32: torch.Tensor
+    cu_seqlens_q: torch.Tensor
+    cu_seqlens_k: torch.Tensor
+    max_seq_len_q: Optional[int]
+
+
+def build_ragged_target_verify_geometry(
+    *,
+    seq_lens: torch.Tensor,
+    layout: RaggedVerifyLayout,
+) -> RaggedTargetVerifyGeometry:
+    cache_seqlens_int32 = (seq_lens + layout.verify_lens).to(torch.int32)
+    cu_seqlens_q = layout.qo_indptr_device.to(torch.int32)
+    cu_seqlens_k = torch.nn.functional.pad(
+        torch.cumsum(cache_seqlens_int32, dim=0, dtype=torch.int32), (1, 0)
+    )
+    max_seq_len_q = (
+        max(layout.verify_lens_cpu) if layout.verify_lens_cpu is not None else None
+    )
+    return RaggedTargetVerifyGeometry(
+        cache_seqlens_int32=cache_seqlens_int32,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seq_len_q=max_seq_len_q,
+    )
+
+
+def compute_target_verify_graph_key(
+    *,
+    bs: int,
+    num_draft_tokens: int,
+    ragged_layout: Optional[RaggedVerifyLayout],
+) -> Tuple[int, int]:
+    num_tokens_full_block = num_draft_tokens * bs
+    if ragged_layout is None:
+        return bs, num_tokens_full_block
+    graph_num_tokens = ragged_layout.graph_num_tokens
+    assert graph_num_tokens <= num_tokens_full_block, (
+        f"ragged verify graph_num_tokens={graph_num_tokens} exceeds full block "
+        f"num_draft*bs={num_tokens_full_block}"
+    )
+    total_verify_tokens = ragged_layout.total_verify_tokens
+    if total_verify_tokens is not None:
+        assert total_verify_tokens <= graph_num_tokens, (
+            f"ragged verify total_verify_tokens={total_verify_tokens} exceeds the "
+            f"round-up bucket graph_num_tokens={graph_num_tokens}"
+        )
+    return graph_num_tokens, graph_num_tokens
+
+
+class VerifyExtendLengths(msgspec.Struct, frozen=True):
+    seq_lens_extended: torch.Tensor
+    seq_lens_cpu_extended: List[int]
+    extend_seq_lens_cpu: List[int]
+    num_tokens: int
+    extend_start_loc: Optional[torch.Tensor]
+
+
+def compute_uniform_extend_lengths(
+    *,
+    seq_lens: torch.Tensor,
+    seq_lens_cpu: List[int],
+    extend_len: int,
+) -> VerifyExtendLengths:
+    batch_size = len(seq_lens_cpu)
+    seq_lens_extended = seq_lens + extend_len
+    seq_lens_cpu_extended = [x + extend_len for x in seq_lens_cpu]
+    extend_seq_lens_cpu = [extend_len] * batch_size
+    num_tokens = extend_len * batch_size
+    return VerifyExtendLengths(
+        seq_lens_extended=seq_lens_extended,
+        seq_lens_cpu_extended=seq_lens_cpu_extended,
+        extend_seq_lens_cpu=extend_seq_lens_cpu,
+        num_tokens=num_tokens,
+        extend_start_loc=None,
+    )
+
+
+def compute_ragged_extend_lengths(
+    *,
+    seq_lens: torch.Tensor,
+    seq_lens_cpu: List[int],
+    ragged_layout: RaggedVerifyLayout,
+) -> VerifyExtendLengths:
+    extend_seq_lens_cpu = list(ragged_layout.verify_lens_cpu)
+    seq_lens_extended = seq_lens + ragged_layout.verify_lens
+    seq_lens_cpu_extended = [
+        raw + length for raw, length in zip(seq_lens_cpu, extend_seq_lens_cpu)
+    ]
+    num_tokens = ragged_layout.total_verify_tokens
+    extend_start_loc = ragged_layout.extend_start_loc
+    return VerifyExtendLengths(
+        seq_lens_extended=seq_lens_extended,
+        seq_lens_cpu_extended=seq_lens_cpu_extended,
+        extend_seq_lens_cpu=extend_seq_lens_cpu,
+        num_tokens=num_tokens,
+        extend_start_loc=extend_start_loc,
+    )