minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Speculative Decoding] Add native UNO serving support (#37667)

Browse Source 
Co-authored-by: drproduck <drproduck@MacBook-Air-2.local>
Co-authored-by: BBuf <1182563586@qq.com>
This commit is contained in:
Yash Akhauri
2026-09-03 20:08:41 +08:00
committed by GitHub
co-authored by drproduck BBuf
parent 354ed6d66b
commit 2bb25dc18b
51 changed files with 6300 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/__init__.py
+1
View File
@@ -16,6 +16,7 @@ _TRITON_KERNELS = [
("extend_attention", "build_unified_kv_indices"),
("prefill_attention", "context_attention_fwd"),
("merge_state", "merge_state_triton"),
("suffix_attention_merge", "merge_suffix_attention_in_place"),
("metadata", "get_num_kv_splits_triton"),
("metadata", "prepare_swa_spec_page_table_triton"),
("metadata", "normal_decode_set_metadata"),
python/sglang/kernels/ops/attention/suffix_attention_merge.py
+203
View File
@@ -0,0 +1,203 @@
"""Fused attention over a short sparse suffix and merge with a prefix state."""
from __future__ import annotations
import torch
import triton
import triton.language as tl
def can_use_fused_suffix_attention_merge(
*,
layer,
q: torch.Tensor,
key_cache: torch.Tensor,
value_cache: torch.Tensor,
extra_kwargs: dict,
) -> bool:
"""Whether attention can use the specialized suffix merge."""
return bool(
q.dtype in (torch.float16, torch.bfloat16)
and key_cache.dtype == q.dtype
and value_cache.dtype == q.dtype
and layer.head_dim == layer.v_head_dim
and not layer.is_cross_attention
and not layer.logit_cap
and not extra_kwargs
)
@triton.jit
def _fused_suffix_attention_merge_kernel(
q_ptr,
k_cache_ptr,
v_cache_ptr,
page_table_ptr,
suffix_seqlens_ptr,
prefix_ptr,
prefix_lse_ptr,
scale,
q_stride_t,
q_stride_h,
q_stride_d,
k_stride_t,
k_stride_h,
k_stride_d,
v_stride_t,
v_stride_h,
v_stride_d,
page_stride_t,
page_stride_s,
prefix_stride_t,
prefix_stride_h,
prefix_stride_d,
prefix_lse_stride_h,
prefix_lse_stride_t,
NUM_Q_HEADS: tl.constexpr,
NUM_KV_HEADS: tl.constexpr,
HEAD_DIM: tl.constexpr,
BLOCK_SUFFIX: tl.constexpr,
BLOCK_D: tl.constexpr,
):
token = tl.program_id(0)
q_head = tl.program_id(1)
kv_head = q_head // (NUM_Q_HEADS // NUM_KV_HEADS)
suffix_offsets = tl.arange(0, BLOCK_SUFFIX)
suffix_length = tl.load(suffix_seqlens_ptr + token)
suffix_valid = suffix_offsets < suffix_length
slots = tl.load(
page_table_ptr + token * page_stride_t + suffix_offsets * page_stride_s,
mask=suffix_valid,
other=0,
).to(tl.int64)
dims = tl.arange(0, BLOCK_D)
dim_valid = dims < HEAD_DIM
q = tl.load(
q_ptr + token * q_stride_t + q_head * q_stride_h + dims * q_stride_d,
mask=dim_valid,
other=0.0,
).to(tl.float32)
k = tl.load(
k_cache_ptr
+ slots[:, None] * k_stride_t
+ kv_head * k_stride_h
+ dims[None, :] * k_stride_d,
mask=suffix_valid[:, None] & dim_valid[None, :],
other=0.0,
).to(tl.float32)
scores = tl.sum(k * q[None, :], axis=1) * scale
scores = tl.where(suffix_valid, scores, -float("inf"))
suffix_max = tl.max(scores, axis=0)
prefix_lse = tl.load(
prefix_lse_ptr + q_head * prefix_lse_stride_h + token * prefix_lse_stride_t
).to(tl.float32)
global_max = tl.maximum(prefix_lse, suffix_max)
prefix_weight = tl.exp(prefix_lse - global_max)
suffix_weights = tl.exp(scores - global_max)
denominator = prefix_weight + tl.sum(suffix_weights, axis=0)
v = tl.load(
v_cache_ptr
+ slots[:, None] * v_stride_t
+ kv_head * v_stride_h
+ dims[None, :] * v_stride_d,
mask=suffix_valid[:, None] & dim_valid[None, :],
other=0.0,
).to(tl.float32)
suffix_numerator = tl.sum(suffix_weights[:, None] * v, axis=0)
prefix = tl.load(
prefix_ptr
+ token * prefix_stride_t
+ q_head * prefix_stride_h
+ dims * prefix_stride_d,
mask=dim_valid,
other=0.0,
).to(tl.float32)
output = (prefix * prefix_weight + suffix_numerator) / denominator
tl.store(
prefix_ptr
+ token * prefix_stride_t
+ q_head * prefix_stride_h
+ dims * prefix_stride_d,
output,
mask=dim_valid,
)
def merge_suffix_attention_in_place(
q: torch.Tensor,
k_cache: torch.Tensor,
v_cache: torch.Tensor,
suffix_page_table: torch.Tensor,
suffix_cache_seqlens: torch.Tensor,
prefix: torch.Tensor,
prefix_lse: torch.Tensor,
softmax_scale: float,
) -> torch.Tensor:
"""Compute a short sparse suffix and merge it into ``prefix`` in place.
``prefix_lse`` uses FlashAttention's varlen layout ``[num_q_heads,
num_queries]``. The suffix page table contains physical token slots, one
row per query, and only its first ``suffix_cache_seqlens[row]`` entries are
visible.
"""
if q.ndim != 3:
raise ValueError("q must have shape [num_queries, num_q_heads, head_dim]")
num_queries, num_q_heads, head_dim = q.shape
if prefix.shape != q.shape:
raise ValueError("prefix output must have the same shape as q")
if prefix_lse.shape != (num_q_heads, num_queries):
raise ValueError("prefix_lse must have shape [num_q_heads, num_queries]")
if k_cache.ndim != 3 or v_cache.ndim != 3:
raise ValueError("flattened KV caches must have shape [slots, heads, dim]")
if k_cache.shape != v_cache.shape:
raise ValueError("K and V caches must have matching shapes")
num_kv_heads = k_cache.shape[1]
if k_cache.shape[2] != head_dim:
raise ValueError("K/V and query head dimensions must match")
if num_q_heads % num_kv_heads:
raise ValueError("query heads must be divisible by KV heads")
if suffix_page_table.ndim != 2 or suffix_page_table.shape[0] != num_queries:
raise ValueError("suffix page table must have one row per query")
if suffix_cache_seqlens.numel() != num_queries:
raise ValueError("suffix cache lengths must have one value per query")
if suffix_page_table.shape[1] == 0 or num_queries == 0:
return prefix
_fused_suffix_attention_merge_kernel[(num_queries, num_q_heads)](
q,
k_cache,
v_cache,
suffix_page_table,
suffix_cache_seqlens,
prefix,
prefix_lse,
softmax_scale,
q.stride(0),
q.stride(1),
q.stride(2),
k_cache.stride(0),
k_cache.stride(1),
k_cache.stride(2),
v_cache.stride(0),
v_cache.stride(1),
v_cache.stride(2),
suffix_page_table.stride(0),
suffix_page_table.stride(1),
prefix.stride(0),
prefix.stride(1),
prefix.stride(2),
prefix_lse.stride(0),
prefix_lse.stride(1),
NUM_Q_HEADS=num_q_heads,
NUM_KV_HEADS=num_kv_heads,
HEAD_DIM=head_dim,
BLOCK_SUFFIX=triton.next_power_of_2(suffix_page_table.shape[1]),
BLOCK_D=triton.next_power_of_2(head_dim),
num_warps=4,
num_stages=1,
)
return prefix
python/sglang/srt/arg_groups/speculative_hook.py
+165
View File
@@ -327,6 +327,171 @@ def _handle_dflash(server_args: ServerArgs) -> None:
)
def _handle_uno(server_args: ServerArgs) -> None:
cfg = resolving_view(server_args)
if not cfg.device.startswith("cuda"):
raise ValueError("UNO only supports CUDA.")
if cfg.speculative_draft_model_path is not None:
raise ValueError(
"UNO reuses the target model and does not accept "
"--speculative-draft-model-path."
)
if cfg.uno_lora_path is None:
raise ValueError("UNO requires --uno-lora-path.")
if cfg.enable_deterministic_inference:
raise ValueError(
"UNO does not support --enable-deterministic-inference because its "
"sampling path does not use per-request seeds."
)
if cfg.enable_strict_thinking:
raise ValueError(
"UNO does not support --enable-strict-thinking because it requires "
"grammar decoding."
)
verify_width = cfg.speculative_num_draft_tokens
if verify_width is None or int(verify_width) < 1:
raise ValueError(
"UNO requires --speculative-num-draft-tokens to be a positive "
"integer denoting the linear width or tree verify width Q."
)
verify_width = int(verify_width)
declare_resolution(
server_args,
"_handle_uno",
speculative_num_draft_tokens=verify_width,
)
candidate_top_k = (
1 if cfg.speculative_eagle_topk is None else int(cfg.speculative_eagle_topk)
)
if candidate_top_k < 1:
raise ValueError(
"UNO requires --speculative-eagle-topk to be at least 1, "
f"got {candidate_top_k}."
)
if candidate_top_k > 1:
if cfg.speculative_num_steps is None:
raise ValueError(
"UNO tree mode requires --speculative-num-steps so its draft "
"width F can be derived as speculative_num_steps + 1."
)
speculative_num_steps = int(cfg.speculative_num_steps)
if speculative_num_steps < 1:
raise ValueError(
"UNO tree mode requires --speculative-num-steps to be positive, "
f"got {speculative_num_steps}."
)
draft_width = speculative_num_steps + 1
if verify_width < draft_width:
raise ValueError(
f"UNO tree mode requires Q >= F; got Q={verify_width}, F={draft_width}."
)
if verify_width > 128:
raise ValueError(
"UNO tree mode currently supports at most Q=128 verify nodes, "
f"got Q={verify_width}."
)
frontier_slots = verify_width * candidate_top_k
if frontier_slots > 2048:
raise ValueError(
"UNO tree mode currently supports Q*K <= 2048; got "
f"Q*K={verify_width}*{candidate_top_k}={frontier_slots}."
)
tree_capacity = 1
nodes_at_depth = 1
for _ in range(speculative_num_steps):
nodes_at_depth *= candidate_top_k
tree_capacity += nodes_at_depth
if tree_capacity >= verify_width:
break
if verify_width > tree_capacity:
raise ValueError(
"UNO tree mode cannot build the requested Q from F and K: "
f"Q={verify_width} exceeds capacity={tree_capacity} for "
f"F={draft_width}, K={candidate_top_k}."
)
parent_width = candidate_top_k * max(speculative_num_steps - 1, 0) + 1
if verify_width - 1 > parent_width:
raise ValueError(
"UNO tree mode cannot represent the requested Q in EAGLE's "
"parent-list ABI: "
f"Q-1={verify_width - 1} exceeds "
f"K*(F-2)+1={parent_width} for "
f"F={draft_width}, K={candidate_top_k}."
)
if cfg.enable_pdmux:
raise ValueError("UNO tree mode does not yet support PDMux.")
if cfg.enable_two_batch_overlap:
raise ValueError("UNO tree mode does not yet support two-batch overlap.")
if (
cfg.speculative_accept_threshold_single != 1.0
or cfg.speculative_accept_threshold_acc != 1.0
):
raise ValueError(
"UNO tree mode reuses EAGLE target-only sampling and requires "
"both speculative accept thresholds to be 1.0."
)
declare_resolution(
server_args,
"_handle_uno",
speculative_num_steps=speculative_num_steps,
speculative_eagle_topk=candidate_top_k,
)
else:
for field in ("speculative_num_steps", "speculative_eagle_topk"):
old_value = getattr(cfg, field)
if old_value not in (None, 1):
logger.warning("UNO uses %s=1; overriding %s.", field, old_value)
declare_resolution(
server_args,
"_handle_uno",
speculative_num_steps=1,
speculative_eagle_topk=1,
)
if (cfg.tp_size, cfg.pp_size) != (1, 1):
raise ValueError("UNO requires TP=PP=1.")
if cfg.enable_dp_attention or cfg.attn_cp_size != 1:
raise ValueError("UNO does not support DP attention or context parallelism.")
if cfg.enable_lora or cfg.lora_paths:
raise ValueError("UNO does not support public Multi-LoRA serving.")
declare_resolution(
server_args,
"_handle_uno",
enable_lora_overlap_loading=False,
lora_strict_loading=True,
)
if cfg.speculative_use_rejection_sampling:
raise ValueError(
"UNO manages its own stochastic verification and does not use "
"--speculative-use-rejection-sampling."
)
if cfg.enable_mixed_chunk:
declare_resolution(
server_args,
"_handle_uno",
enable_mixed_chunk=False,
)
logger.warning(
"Mixed chunked prefill is disabled for UNO speculative decoding."
)
prefill_backend, decode_backend = attention_backends_of(resolved_view(server_args))
if (prefill_backend, decode_backend) != ("fa3", "fa3"):
raise ValueError(
"UNO requires FA3 for both prefill and decode attention; "
f"got prefill={prefill_backend!r}, decode={decode_backend!r}."
)
def _target_checkpoint_bundles_dspark_draft(server_args: ServerArgs) -> bool:
from sglang.srt.speculative.dspark_components.dspark_config import (
checkpoint_bundles_dspark_draft,
python/sglang/srt/layers/attention/flashattention_backend.py
+30 -4
View File
@@ -12,6 +12,10 @@ from sglang.kernels.ops.attention.metadata import (
prepare_swa_spec_page_table_triton,
)
from sglang.kernels.ops.attention.pa_page_table import _build_pa_page_table
from sglang.kernels.ops.attention.suffix_attention_merge import (
can_use_fused_suffix_attention_merge,
merge_suffix_attention_in_place,
)
from sglang.kernels.ops.attention.utils import assert_buffer_fits
from sglang.kernels.ops.kvcache.trtllm_mha_page_table import (
build_trtllm_mha_page_table,
@@ -1579,14 +1583,36 @@ class FlashAttentionBackend(AttentionBackend):
if use_cascade_attn:
o, softmax_lse, *rest = result
if (
use_cascade_attn
and forward_batch.spec_algorithm.is_uno()
and can_use_fused_suffix_attention_merge(
layer=layer,
q=q,
key_cache=key_cache,
value_cache=value_cache,
extra_kwargs=kwargs,
)
):
suffix_metadata = self.forward_metadata_spec_decode_expand
o = merge_suffix_attention_in_place(
q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
k_cache=key_cache.view(-1, layer.tp_k_head_num, layer.head_dim),
v_cache=value_cache.view(-1, layer.tp_v_head_num, layer.v_head_dim),
suffix_page_table=suffix_metadata.page_table,
suffix_cache_seqlens=suffix_metadata.cache_seqlens_int32,
prefix=o,
prefix_lse=softmax_lse,
softmax_scale=layer.scaling,
)
elif use_cascade_attn:
o_expand, softmax_lse_expand, *rest_expand = flash_attn_with_kvcache(
q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
# Here metadata_expand.page_table is not divided with page_size.
# This is because we loose the fine control of what token to attend,
# but has to attend to some block completely.
# The suffix table stores physical token slots, so expose
# the paged cache as page-size-one blocks.
k_cache=key_cache.view(-1, 1, layer.tp_k_head_num, layer.head_dim),
v_cache=value_cache.view(
-1, 1, layer.tp_v_head_num, layer.head_dim
-1, 1, layer.tp_v_head_num, layer.v_head_dim
),
page_table=self.forward_metadata_spec_decode_expand.page_table,
cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32,
python/sglang/srt/lora/backend/base_backend.py
+25
View File
@@ -23,6 +23,9 @@ class BaseLoRABackend(LoRABackendLmHeadMixing):
device: the device where the backend runs.
"""
supports_lora_a_overlap = False
skip_inactive_lora_batches = False
# Supporting backends implement init_prefill_cuda_graph_batch_info() and
# honor use_prefill_cuda_graph in prepare_lora_batch().
supports_prefill_cuda_graph: bool = False
@@ -52,6 +55,14 @@ class BaseLoRABackend(LoRABackendLmHeadMixing):
self.lm_head_pass_batch_infos = None
self._lm_head_pass_idx = None
def validate_lora_targets(
self,
base_model: torch.nn.Module,
target_modules: set[str],
) -> None:
"""Raise before wrapping when this backend cannot execute its targets."""
pass
def run_lora_a_embedding(
self,
input_ids: torch.Tensor,
@@ -351,6 +362,20 @@ class BaseLoRABackend(LoRABackendLmHeadMixing):
"""
pass
def prepare_lora_token_segments(
self,
*,
segment_lens: list[int],
weight_indices: list[int],
lora_ranks: list[int],
scalings: list[float],
) -> None:
"""Prepare explicit eager token-row LoRA segments."""
raise NotImplementedError(
f"LoRA backend {type(self).__name__} does not support explicit "
"token segments."
)
@triton.jit
def _compute_moe_lora_info_kernel(
python/sglang/srt/lora/backend/lora_registry.py
+7
View File
@@ -44,6 +44,13 @@ def create_torch_native_backend():
return TorchNativeLoRABackend
@register_lora_backend("uno_cublas")
def create_uno_cublas_backend():
from sglang.srt.lora.backend.uno_cublas_backend import UnoCublasLoRABackend
return UnoCublasLoRABackend
@register_lora_backend("flashinfer")
def create_flashinfer_backend():
raise ValueError(
python/sglang/srt/lora/backend/triton_backend.py
+38
View File
@@ -362,6 +362,44 @@ class TritonLoRABackend(BaseLoRABackend):
self._prepare_lm_head_batch_info(forward_batch, weight_indices, batch_info)
)
def prepare_lora_token_segments(
self,
*,
segment_lens: list[int],
weight_indices: list[int],
lora_ranks: list[int],
scalings: list[float],
) -> None:
"""Install explicit eager token-row routing metadata."""
self.reset_batch_state()
segment_lens_tensor = torch.tensor(
segment_lens, dtype=torch.int32, device=self.device
)
segment_indptr = torch.zeros(
len(segment_lens) + 1, dtype=torch.int32, device=self.device
)
segment_indptr[1:] = torch.cumsum(segment_lens_tensor, dim=0)
self.batch_info = LoRABatchInfo(
use_cuda_graph=False,
bs=len(segment_lens),
num_segments=len(segment_lens),
seg_lens=segment_lens_tensor,
seg_indptr=segment_indptr,
max_len=max(segment_lens),
weight_indices=torch.tensor(
weight_indices, dtype=torch.int32, device=self.device
),
lora_ranks=torch.tensor(lora_ranks, dtype=torch.int64, device=self.device),
scalings=torch.tensor(scalings, dtype=torch.float, device=self.device),
permutation=None,
expected_tokens=sum(segment_lens),
)
# These segments already describe physical token-row order.
self.sgemm_batch_info = None
def _prepare_lm_head_batch_info(
self,
forward_batch: ForwardBatch,
python/sglang/srt/lora/backend/uno_cublas_backend.py
+455
View File
@@ -0,0 +1,455 @@
"""Single-adapter LoRA backend for UNO draft forwards.
Parallel-linear layers overlap LoRA-A with the base GEMM on an auxiliary CUDA
stream. Other dense LoRA layers use the inherited Triton implementation.
Single-request cuBLAS batches operate only on draft rows; larger batches use
one ``mm``/``addmm_`` over all rows and zero the seed-row LoRA hidden states
between the two GEMMs.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Optional
import torch
from sglang.srt.lora.backend.triton_backend import TritonLoRABackend
from sglang.srt.lora.utils import LoRABatchInfo
@dataclass(frozen=True)
class _UnoSingleAdapterRoute:
weight_index: int
rank: int
scaling: float
batch_size: int
forward_width: int
total_rows: int
active_rows: int
@property
def lora_rows(self) -> int:
# For C=1, skipping the seed row makes both GEMMs smaller. For C>1,
# one GEMM across C*F rows is faster than C tiny (F-1)-row GEMMs.
return self.active_rows if self.batch_size == 1 else self.total_rows
@dataclass(frozen=True)
class _PendingLoRAA:
output: torch.Tensor
producer_stream: torch.cuda.Stream
class UnoCublasLoRABackend(TritonLoRABackend):
"""Fast LoRA backend for UNO's draft forwards.
Use cuBLAS and CUDA streams.
Reuse multi-LoRA batch metadata from the Triton parent.
"""
name = "uno_cublas"
supports_lora_a_overlap = True
# K2's runners prepare base-only LoRA metadata whenever any internal
# manager exists. UNO never exposes request-selectable adapters, so those
# prefill/warmup batches must stay on the plain base-model path.
skip_inactive_lora_batches = True
def __init__(
self,
max_loras_per_batch: int,
device: torch.device,
**kwargs,
):
super().__init__(max_loras_per_batch, device, **kwargs)
# Different CUDA-graph runners may capture and replay on different main
# streams. Give each main stream its own LoRA-A side stream so concurrently
# replayed graphs do not serialize or interfere through one shared stream.
self._lora_a_streams: dict[torch.cuda.Stream, torch.cuda.Stream] = {}
self._pending_lora_a: Optional[_PendingLoRAA] = None
self._use_cublas_lora_b = False
def reset_batch_state(self):
self._pending_lora_a = None
self._use_cublas_lora_b = False
super().reset_batch_state()
def validate_lora_targets(
self,
base_model: torch.nn.Module,
target_modules: set[str],
) -> None:
"""Reject target layers that cannot honor UNO's token-row routing."""
from sglang.srt.layers.linear import (
ColumnParallelLinear,
ReplicatedLinear,
RowParallelLinear,
)
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
from sglang.srt.layers.utils import get_layer_id
from sglang.srt.models.inkling_common.dense_mlp import InklingBatchDenseMLP
unsupported: list[str] = []
# Embedding and LM-head wrappers use the inherited Triton kernels and
# are handled separately by LoRAManager. Decoder-layer projections use
# either the overlapped cuBLAS path or the Triton ReplicatedLinear path.
supported = (ColumnParallelLinear, RowParallelLinear, ReplicatedLinear)
target_moe = {"gate_up_proj", "down_proj"}.issubset(target_modules)
for module_name, module in base_model.named_modules():
parts = module_name.split(".")
named_target = bool(parts) and (
parts[-1] in target_modules or ".".join(parts[-2:]) in target_modules
)
special_moe_target = target_moe and isinstance(
module, (FusedMoE, InklingBatchDenseMLP)
)
if not (named_target or special_moe_target):
continue
if get_layer_id(module_name) is None:
continue
if not isinstance(module, supported):
unsupported.append(f"{module_name} ({type(module).__name__})")
if unsupported:
raise ValueError(
"UNO's LoRA backend cannot execute these target modules: "
+ ", ".join(sorted(set(unsupported)))
)
def prepare_lora_token_segments(
self,
*,
segment_lens: list[int],
weight_indices: list[int],
lora_ranks: list[int],
scalings: list[float],
) -> None:
super().prepare_lora_token_segments(
segment_lens=segment_lens,
weight_indices=weight_indices,
lora_ranks=lora_ranks,
scalings=scalings,
)
route: Optional[_UnoSingleAdapterRoute] = None
if (
len(segment_lens) >= 2
and len(segment_lens) % 2 == 0
and len(weight_indices) == len(segment_lens)
and segment_lens[0] == 1
and segment_lens[1] > 0
):
batch_size = len(segment_lens) // 2
forward_width = segment_lens[1] + 1
base_index, adapter_index = weight_indices[:2]
adapter_rank = lora_ranks[adapter_index]
if (
base_index != adapter_index
and lora_ranks[base_index] == 0
and adapter_rank > 0
and all(
segment_lens[2 * index] == 1
and segment_lens[2 * index + 1] == forward_width - 1
and weight_indices[2 * index] == base_index
and weight_indices[2 * index + 1] == adapter_index
for index in range(batch_size)
)
):
route = _UnoSingleAdapterRoute(
weight_index=adapter_index,
rank=adapter_rank,
scaling=float(scalings[adapter_index]),
batch_size=batch_size,
forward_width=forward_width,
total_rows=sum(segment_lens),
active_rows=batch_size * (forward_width - 1),
)
# Each CUDA-graph bucket retains its own batch_info. Store the immutable
# UNO route there so switching back to a captured bucket restores the
# route corresponding to that bucket, rather than using metadata last
# written by another bucket.
self.batch_info.uno_single_adapter_route = route
def _route(self) -> _UnoSingleAdapterRoute:
route = getattr(self.batch_info, "uno_single_adapter_route", None)
if route is None:
raise RuntimeError("UNO cuBLAS execution requires an active UNO route.")
return route
@staticmethod
def _output_offsets(output_offset_cpu, output_offset) -> list[int]:
offsets = output_offset_cpu if output_offset_cpu is not None else output_offset
return [int(offset) for offset in offsets.tolist()]
@staticmethod
def _lora_a_input(
x: torch.Tensor,
route: _UnoSingleAdapterRoute,
) -> torch.Tensor:
if route.batch_size == 1:
return x[1:]
return x
def _compute_lora_a(
self,
lora_input: torch.Tensor,
active_a: torch.Tensor,
route: _UnoSingleAdapterRoute,
*,
output: Optional[torch.Tensor] = None,
) -> torch.Tensor:
hidden = torch.mm(lora_input, active_a.t(), out=output)
if route.batch_size > 1:
# LoRA must not affect each request's seed token. Zeroing C rows
# is cheaper than masking every hidden element or running C tiny
# strided-batched GEMMs.
hidden.view(route.batch_size, route.forward_width, -1)[:, 0].zero_()
return hidden
def _accumulate_lora_b(
self,
*,
hidden: torch.Tensor,
active_b: torch.Tensor,
base_output: torch.Tensor,
output_start: int,
output_end: int,
route: _UnoSingleAdapterRoute,
) -> None:
if route.batch_size == 1:
base_output[1:, output_start:output_end].addmm_(
hidden,
active_b.t(),
beta=1.0,
alpha=route.scaling,
)
return
base_output[:, output_start:output_end].addmm_(
hidden,
active_b.t(),
beta=1.0,
alpha=route.scaling,
)
def _run_lora_b(
self,
x: torch.Tensor,
weights: torch.Tensor,
base_output: torch.Tensor,
) -> torch.Tensor:
route = self._route()
active_b = weights[route.weight_index, :, : route.rank]
self._accumulate_lora_b(
hidden=x[:, : route.rank],
active_b=active_b,
base_output=base_output,
output_start=0,
output_end=active_b.shape[0],
route=route,
)
return base_output
def run_lora_a_sgemm(
self,
x: torch.Tensor,
weights: torch.Tensor,
pruned_batch_info: LoRABatchInfo = None,
stack_num: int = 1,
*args,
**kwargs,
) -> torch.Tensor:
pending = self._pending_lora_a
if pending is None:
self._use_cublas_lora_b = False
return super().run_lora_a_sgemm(
x,
weights,
pruned_batch_info,
stack_num,
*args,
**kwargs,
)
output = self._consume_lora_a_overlap(pending)
self._use_cublas_lora_b = True
return output
def run_lora_b_sgemm(
self,
x: torch.Tensor,
weights: torch.Tensor,
base_output: torch.Tensor = None,
pruned_batch_info: LoRABatchInfo = None,
*args,
**kwargs,
) -> torch.Tensor:
if not self._use_cublas_lora_b:
return super().run_lora_b_sgemm(
x,
weights,
base_output,
pruned_batch_info,
*args,
**kwargs,
)
self._use_cublas_lora_b = False
return self._run_lora_b(x, weights, base_output)
def _run_stacked_lora(
self,
*,
lora_b: torch.Tensor,
base_output: torch.Tensor,
output_offset,
output_offset_cpu,
num_slices: int,
pending: _PendingLoRAA,
) -> torch.Tensor:
route = self._route()
hidden = self._consume_lora_a_overlap(pending)
offsets = self._output_offsets(output_offset_cpu, output_offset)
for slice_index in range(num_slices):
input_start = slice_index * route.rank
input_end = input_start + route.rank
output_start = offsets[slice_index]
output_end = offsets[slice_index + 1]
active_b = lora_b[
route.weight_index,
output_start:output_end,
: route.rank,
]
self._accumulate_lora_b(
hidden=hidden[:, input_start:input_end],
active_b=active_b,
base_output=base_output,
output_start=output_start,
output_end=output_end,
route=route,
)
return base_output
def run_qkv_lora(
self,
x: torch.Tensor,
qkv_lora_a: torch.Tensor,
qkv_lora_b: torch.Tensor,
output_offset: torch.Tensor,
max_qkv_out_dim: int,
base_output: torch.Tensor = None,
n_slices: int = 3,
*args,
output_offset_cpu=None,
**kwargs,
) -> torch.Tensor:
pending = self._pending_lora_a
if pending is None:
return super().run_qkv_lora(
x,
qkv_lora_a,
qkv_lora_b,
output_offset,
max_qkv_out_dim,
base_output,
n_slices,
*args,
**kwargs,
)
return self._run_stacked_lora(
lora_b=qkv_lora_b,
base_output=base_output,
output_offset=output_offset,
output_offset_cpu=output_offset_cpu,
num_slices=n_slices,
pending=pending,
)
def run_gate_up_lora(
self,
x: torch.Tensor,
gate_up_lora_a: torch.Tensor,
gate_up_lora_b: torch.Tensor,
base_output: torch.Tensor = None,
*args,
output_offset=None,
output_offset_cpu=None,
**kwargs,
) -> torch.Tensor:
pending = self._pending_lora_a
if pending is None:
return super().run_gate_up_lora(
x,
gate_up_lora_a,
gate_up_lora_b,
base_output,
*args,
**kwargs,
)
return self._run_stacked_lora(
lora_b=gate_up_lora_b,
base_output=base_output,
output_offset=output_offset,
output_offset_cpu=output_offset_cpu,
num_slices=2,
pending=pending,
)
def start_lora_a_overlap(
self,
x: torch.Tensor,
weights: torch.Tensor,
*,
num_slices: int = 1,
) -> None:
"""Launch LoRA-A on the auxiliary stream before the base GEMM."""
if self._pending_lora_a is not None:
raise RuntimeError("Previous UNO LoRA-A overlap was not consumed.")
route = self._route()
out_dim = num_slices * route.rank
lora_input = self._lora_a_input(x, route)
active_a = weights[route.weight_index, :out_dim]
main_stream = torch.cuda.current_stream(x.device)
stream = self._lora_a_streams.get(main_stream)
if stream is None:
if torch.cuda.is_current_stream_capturing():
raise RuntimeError(
"UNO LoRA overlap side stream was not created during "
"CUDA-graph warmup."
)
stream = torch.cuda.Stream(device=x.device)
self._lora_a_streams[main_stream] = stream
# Allocate on the main/consumer stream before handing the buffer to
# the auxiliary stream. This keeps CUDA-graph allocator ownership and
# the eventual LoRA-B consumer on the same stream.
output = torch.empty(
(route.lora_rows, out_dim),
dtype=x.dtype,
device=x.device,
)
stream.wait_stream(main_stream)
with torch.cuda.stream(stream):
self._compute_lora_a(lora_input, active_a, route, output=output)
self._pending_lora_a = _PendingLoRAA(
output=output,
producer_stream=stream,
)
def _consume_lora_a_overlap(
self,
pending: _PendingLoRAA,
) -> torch.Tensor:
self._pending_lora_a = None
torch.cuda.current_stream(pending.output.device).wait_stream(
pending.producer_stream
)
return pending.output
python/sglang/srt/lora/layers.py
+36 -3
View File
@@ -66,7 +66,15 @@ class BaseLayerWithLoRA(nn.Module):
has LoRA batch metadata. batch_info is None on DP-attention idle
forwards (see LoRAManager.prepare_lora_batch), so idle forwards take
the base path."""
return self.set_lora and self.lora_backend.batch_info is not None
batch_info = self.lora_backend.batch_info
return (
self.set_lora
and batch_info is not None
and (
not self.lora_backend.skip_inactive_lora_batches
or batch_info.has_active_lora
)
)
def set_lora_info(self, *args):
pass
@@ -482,14 +490,22 @@ class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA):
)
return lora_output
def start_lora_a_overlap(self, x: torch.Tensor) -> None:
if self.lora_backend.supports_lora_a_overlap:
self.lora_backend.start_lora_a_overlap(x, self.A_buffer)
def forward(self, input_: torch.Tensor):
# duplicate the logic in ColumnParallelLinear
lora_active = self.lora_active
if lora_active:
self.start_lora_a_overlap(input_)
bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None
output_parallel = self.base_layer.quant_method.apply(
self.base_layer, input_, bias
)
if self.lora_active:
if lora_active:
output_parallel = self.apply_lora(output_parallel, input_)
if self.base_layer.gather_output:
@@ -596,6 +612,12 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
)
return lora_output
def start_lora_a_overlap(self, x: torch.Tensor) -> None:
if self.lora_backend.supports_lora_a_overlap:
self.lora_backend.start_lora_a_overlap(
x, self.A_buffer, num_slices=self._get_lora_n_slices()
)
def slice_lora_a_weights(self, A: torch.Tensor):
return A
@@ -703,6 +725,10 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
return lora_output
def start_lora_a_overlap(self, x: torch.Tensor) -> None:
if self.lora_backend.supports_lora_a_overlap:
self.lora_backend.start_lora_a_overlap(x, self.A_buffer_qkv, num_slices=3)
def slice_lora_a_weights(self, A: torch.Tensor):
return A
@@ -773,6 +799,10 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
)
return lora_output
def start_lora_a_overlap(self, x: torch.Tensor) -> None:
if self.lora_backend.supports_lora_a_overlap:
self.lora_backend.start_lora_a_overlap(x, self.A_buffer)
def forward(self, input_: torch.Tensor, skip_all_reduce=False, forward_batch=None):
if self.base_layer.input_is_parallel:
input_parallel = input_
@@ -783,6 +813,10 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
)
input_parallel = splitted_input[tp_rank].contiguous()
lora_active = self.lora_active
if lora_active:
self.start_lora_a_overlap(input_parallel)
bias_ = (
None
if (self.base_layer.tp_rank > 0 or self.base_layer.skip_bias_add)
@@ -806,7 +840,6 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
all_reduce = get_parallel().attn_tp_group.all_reduce
else:
all_reduce = tensor_model_parallel_all_reduce
lora_active = self.lora_active
if lora_active and should_reduce:
lora_a_output = self.lora_backend.run_lora_a_sgemm(
input_parallel, self.A_buffer
python/sglang/srt/lora/lora_manager.py
+48 -1
View File
@@ -17,7 +17,7 @@
import logging
import re
from typing import Dict, Iterable, List, Optional
from typing import Dict, Iterable, List, Optional, Sequence
import torch
@@ -431,6 +431,16 @@ class LoRAManager:
self.lora_backend.reset_batch_state()
def prepare_lora_batch(self, forward_batch: ForwardBatch):
# Some internal-only backends (currently UNO) use explicit token-row
# routing for their adapted forwards and want all-base batches to run
# through the plain model path. Clear any routing retained by the
# preceding adapted forward before inspecting CUDA-graph metadata.
if self.lora_backend.skip_inactive_lora_batches and not any(
uid is not None for uid in forward_batch.lora_ids
):
self.reset_lora_batch()
return
# set up batch info shared by all lora modules
bs = forward_batch.batch_size
@@ -470,6 +480,39 @@ class LoRAManager:
lora_ranks[wi] > 0 for wi in weight_indices
)
def prepare_lora_token_segments(
self,
*,
lora_ids: Sequence[Optional[str]],
segment_lens: Sequence[int],
) -> None:
"""Prepare eager LoRA routing independently of request batching."""
lora_ids = list(lora_ids)
segment_lens = list(segment_lens)
if len(lora_ids) != len(segment_lens):
raise ValueError("LoRA ids and segment lengths must have equal length.")
weight_indices = []
lora_ranks = [0] * self.max_loras_per_batch
scalings = [0.0] * self.max_loras_per_batch
for lora_id in lora_ids:
weight_index = self.memory_pool.get_buffer_id(lora_id)
weight_indices.append(weight_index)
if lora_id is not None:
lora = self.loras[lora_id]
lora_ranks[weight_index] = lora.config.r
scalings[weight_index] = lora.scaling
self.lora_backend.prepare_lora_token_segments(
segment_lens=segment_lens,
weight_indices=weight_indices,
lora_ranks=lora_ranks,
scalings=scalings,
)
self.lora_backend.batch_info.has_active_lora = any(
lora_ranks[index] > 0 for index in weight_indices
)
def update_lora_info(self):
"""
Update all LoRA modules to associate them with the latest memory buffer.
@@ -579,6 +622,10 @@ class LoRAManager:
max_lora_rank=max_lora_rank,
target_modules=target_modules,
)
self.lora_backend.validate_lora_targets(
base_model=self.base_model,
target_modules=self.target_modules,
)
if self._experts_shared_outer_override is not None:
self.experts_shared_outer_loras = self._experts_shared_outer_override
python/sglang/srt/managers/scheduler.py
+10 -1
View File
@@ -308,6 +308,7 @@ from sglang.srt.speculative.eagle_utils import (
get_draft_recurrent_hidden_state_spec_from_config,
)
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.speculative.uno_validation import validate_uno_request
from sglang.srt.utils import (
DynamicGradMode,
configure_gc_logger,
@@ -2813,6 +2814,14 @@ class Scheduler(
self._add_request_to_queue(req)
return
if self.spec_algorithm.is_uno():
error_msg = validate_uno_request(req)
if error_msg is not None:
req.set_finish_with_abort(error_msg)
self.init_req_max_new_tokens(req)
self._add_request_to_queue(req)
return
if (
req.return_sampling_mask
and self.disaggregation_mode != DisaggregationMode.NULL
@@ -4484,7 +4493,7 @@ class Scheduler(
self.decode_moment_totals,
batch_size,
step_us,
batch_size + result.num_correct_drafts,
result.get_num_generated_tokens(batch_size),
)
def maybe_send_health_check_signal(self):
python/sglang/srt/managers/scheduler_components/batch_result_processor.py
+25 -14
View File
@@ -78,6 +78,16 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
def _get_speculative_output_stride(result: GenerationBatchResult) -> int:
"""Return the padded per-request width in flattened speculative output."""
stride = result.speculative_output_stride
if stride is None:
stride = result.speculative_num_draft_tokens
if stride is None or stride < 1:
raise RuntimeError("speculative result is missing a positive output row stride")
return stride
@dataclass(kw_only=True, slots=True, frozen=True)
class SchedulerBatchResultProcessor:
is_generation: bool
@@ -711,8 +721,12 @@ class SchedulerBatchResultProcessor:
next_token_ids = result.next_token_ids.tolist()
accept_lens = result.accept_lens.tolist()
result.num_correct_drafts = sum(accept_lens) - len(batch.reqs)
result.num_correct_drafts_per_req_cpu = [x - 1 for x in accept_lens]
stride = _get_speculative_output_stride(result)
num_non_draft = result.num_non_draft_tokens_per_req
result.num_correct_drafts_per_req_cpu = [
length - num_non_draft for length in accept_lens
]
result.num_correct_drafts = sum(result.num_correct_drafts_per_req_cpu)
block_accept_lens = (
result.block_accept_lens.tolist()
@@ -740,11 +754,6 @@ class SchedulerBatchResultProcessor:
self.advance_grammar_fsm(result, batch)
predict_tokens = []
# In adaptive spec-v2, the worker state may already have switched when this
# delayed result is processed. Use the draft token count recorded on result.
stride = result.speculative_num_draft_tokens
assert stride is not None, "spec-v2 result missing speculative_num_draft_tokens"
for i, req in enumerate(batch.reqs):
accept_tokens = next_token_ids[i * stride : i * stride + accept_lens[i]]
@@ -853,8 +862,7 @@ class SchedulerBatchResultProcessor:
if result.accept_lens is None:
return
accept_lens = result.accept_lens.tolist()
stride = result.speculative_num_draft_tokens
assert stride is not None, "spec-v2 result missing speculative_num_draft_tokens"
stride = _get_speculative_output_stride(result)
retained = [None] * len(batch.reqs)
for i, req in enumerate(batch.reqs):
if req.grammar is None or req.is_retracted or req.finished():
@@ -907,11 +915,14 @@ class SchedulerBatchResultProcessor:
next_token_ids=next_token_ids,
)
self.metrics_reporter.num_generated_tokens += len(batch.reqs)
batch_size = batch.batch_size()
num_generated_tokens = result.get_num_generated_tokens(batch_size)
self.metrics_reporter.num_generated_tokens += num_generated_tokens
if not batch.spec_algorithm.is_none():
self.metrics_reporter.update_spec_metrics(
batch.batch_size(),
batch_size,
result.num_correct_drafts,
num_accept_tokens=num_generated_tokens,
num_block_accept_tokens=result.num_block_accept_tokens,
num_cap_tokens=result.num_cap_tokens,
)
@@ -982,8 +993,8 @@ class SchedulerBatchResultProcessor:
if req.return_hidden_states and logits_output.hidden_states is not None:
# hidden_states is [bs * stride, hidden_dim], one row per emitted
# token; stride = speculative_num_draft_tokens for spec, 1 for non-spec.
stride = result.speculative_num_draft_tokens or 1
# token; speculative workers record their padded row width.
stride = _get_speculative_output_stride(result) if is_spec else 1
accept_len = len(next_token_id)
start = i * stride
self._append_decode_hidden_states(
@@ -1016,7 +1027,7 @@ class SchedulerBatchResultProcessor:
self.metrics_reporter.report_decode_stats(
can_run_cuda_graph,
running_batch=batch,
num_correct_drafts=result.num_correct_drafts,
num_generated_tokens=num_generated_tokens,
)
def _normalize_decode_outputs(
python/sglang/srt/managers/scheduler_components/metrics_reporter.py
+12 -10
View File
@@ -175,9 +175,10 @@ class SchedulerMetricsReporter:
}.get(getattr(self.scheduler, "device", ""), "cuda graph")
# Cumulative spec-decoding counters (reset every decode_log_interval).
# Each update adds (num_correct_drafts + bs, bs).
# `*_accept_tokens` = drafts + bonus; `*_correct_drafts` = drafts-only.
# `*_accept_tokens` includes accepted drafts and non-draft output tokens;
# `*_correct_drafts` counts accepted draft proposals only.
self.spec_num_accept_tokens = 0 # per-log-interval
self.spec_num_correct_drafts = 0
self.spec_num_forward_ct = 0
self.spec_total_num_accept_tokens = 0 # lifetime
self.spec_total_num_forward_ct = 0
@@ -398,17 +399,16 @@ class SchedulerMetricsReporter:
self,
bs: int,
num_correct_drafts: int,
num_accept_tokens: int,
num_block_accept_tokens: int = 0,
num_cap_tokens: int = 0,
):
self.spec_num_accept_tokens += num_correct_drafts + bs
self.spec_num_accept_tokens += num_accept_tokens
self.spec_num_correct_drafts += num_correct_drafts
self.spec_num_forward_ct += bs
self.spec_num_block_accept_tokens += num_block_accept_tokens
self.spec_num_cap_tokens += num_cap_tokens
# Bonus tokens updated elsewhere
self.num_generated_tokens += num_correct_drafts
def _init_estimated_perf_constants(self) -> None:
model_config = self.scheduler.model_config
hf_text_config = model_config.hf_text_config
@@ -572,6 +572,7 @@ class SchedulerMetricsReporter:
self.forward_ct_decode = 0
self.num_generated_tokens = 0
self.spec_num_accept_tokens = 0
self.spec_num_correct_drafts = 0
self.spec_num_forward_ct = 0
self.spec_total_num_accept_tokens = 0
self.spec_total_num_forward_ct = 0
@@ -757,13 +758,13 @@ class SchedulerMetricsReporter:
self,
can_run_cuda_graph: bool,
running_batch: ScheduleBatch = None,
num_correct_drafts: int = 0,
num_generated_tokens: int = 0,
):
batch = running_batch or self.scheduler.running_batch
# Every-iteration work: realtime token counting + status logger
if self.current_scheduler_metrics_enabled:
decode_tokens = batch.batch_size() + num_correct_drafts
decode_tokens = num_generated_tokens
self.metrics_collector.increment_realtime_tokens(
# TODO unify this w/ the bumping logic in `Scheduler.num_generated_tokens` accumulator
decode_tokens=decode_tokens,
@@ -826,7 +827,7 @@ class SchedulerMetricsReporter:
spec_block_accept_length = 0
else:
spec_accept_length = self.spec_num_accept_tokens / self.spec_num_forward_ct
num_correct_drafts = self.spec_num_accept_tokens - self.spec_num_forward_ct
num_correct_drafts = self.spec_num_correct_drafts
if get_spec().speculative_num_draft_tokens:
draft_per_round = get_spec().speculative_num_draft_tokens - 1
else:
@@ -853,7 +854,8 @@ class SchedulerMetricsReporter:
)
self.spec_total_num_accept_tokens += self.spec_num_accept_tokens
self.spec_total_num_forward_ct += self.spec_num_forward_ct
self.spec_num_accept_tokens = self.spec_num_forward_ct = 0
self.spec_num_accept_tokens = self.spec_num_correct_drafts = 0
self.spec_num_forward_ct = 0
self.spec_num_block_accept_tokens = 0
self.spec_num_cap_tokens = 0
msg += f"accept len: {spec_accept_length:.2f}, accept rate: {spec_accept_rate:.2f}, "
python/sglang/srt/managers/utils.py
+11 -2
View File
@@ -22,7 +22,7 @@ from sglang.srt.state_capturer.base import TopkCaptureOutput
if TYPE_CHECKING:
from sglang.srt.managers.scheduler import GenerationBatchResult
from sglang.srt.sampling.sampling_observer import HostAuxiliaryOutput
from sglang.srt.speculative.eagle_info import EagleDraftInput
from sglang.srt.speculative.spec_info import SpecInput
logger = logging.getLogger(__name__)
@@ -71,6 +71,12 @@ class GenerationBatchResult:
delay_sample_func: Optional[callable] = None
future_indices: Optional[torch.Tensor] = None
speculative_num_draft_tokens: Optional[int] = None
# Padded row width in flattened speculative output. Existing algorithms
# default to speculative_num_draft_tokens; linear UNO emits F + 1 columns.
speculative_output_stride: Optional[int] = None
# Valid output tokens that are not accepted draft proposals. Existing
# algorithms have one bonus token; UNO also emits its clean root.
num_non_draft_tokens_per_req: int = 1
# Grammar FSM advance memoization (spec-v2 overlap). advance_grammar_fsm sets
# these once — eagerly via the scheduler's grammar barrier inside verify(), or
@@ -91,7 +97,7 @@ class GenerationBatchResult:
new_seq_lens: Optional[torch.Tensor] = None
# relay path: forward stream -> next step forward
next_draft_input: Optional[EagleDraftInput] = None
next_draft_input: Optional[SpecInput] = None
# Refs the worker wants scheduler to keep alive for the same 2-iter window
# as batch_record_buf. Used for cross-stream tensor lifetime (e.g. a spec
@@ -117,6 +123,9 @@ class GenerationBatchResult:
this rank/split (a non-last PP rank or a non-final prefill split)."""
return isinstance(self.next_token_ids, torch.Tensor)
def get_num_generated_tokens(self, batch_size: int) -> int:
return self.num_correct_drafts + batch_size * self.num_non_draft_tokens_per_req
@torch.profiler.record_function("copy_result_to_cpu")
def copy_to_cpu(self, return_logprob: bool, return_hidden_states: bool = True):
"""Copy tensors to CPU in overlap scheduling.
python/sglang/srt/mem_cache/allocation_sizing.py
+14 -5
View File
@@ -33,6 +33,11 @@ def get_alloc_len_per_decode() -> int:
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
spec_algo = SpeculativeAlgorithm.from_string(spec.speculative_algorithm)
if spec_algo.is_uno():
if spec_tokens is None:
raise RuntimeError("UNO requires speculative_num_draft_tokens")
# UNO retains an additional clean-root position beside Q/F draft slots.
return spec_tokens + 1
if page_size == 1 or spec_topk == 1 or not spec_algo.has_draft_kv():
return max(spec_steps * spec_topk, spec_tokens)
else:
@@ -88,9 +93,13 @@ def get_req_to_token_extra_context_len() -> int:
# FIXME(lsyin): temporary fix for the context length issue under spec decoding
extra = 4 + (max_speculative_num_draft_tokens() or 0)
page_size = get_alloc_page_size()
if get_spec().speculative_algorithm is not None and page_size > 1:
# kv_allocated_len is page-aligned (eagle_prepare_for_decode), so near
# the context limit the aligned reserve can overshoot by page_size - 1;
# without the headroom the row write silently lands in the neighbor row.
extra = max(extra, get_alloc_reserve_per_decode() + page_size - 1)
spec_algorithm = get_spec().speculative_algorithm
if spec_algorithm is not None:
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
spec_algo = SpeculativeAlgorithm.from_string(spec_algorithm)
if page_size > 1 or spec_algo.is_uno():
# UNO's double-buffer reserve applies at every page size. Larger
# pages may additionally round the allocation up by page_size - 1.
extra = max(extra, get_alloc_reserve_per_decode() + page_size - 1)
return extra
python/sglang/srt/model_executor/model_runner.py
+13
View File
@@ -777,6 +777,12 @@ class ModelRunner:
self.apply_torch_tp()
def maybe_init_lora_manager(self):
if self.spec_algorithm.is_uno():
from sglang.srt.speculative.uno_lora import init_uno_lora_manager
self.lora_manager, self.uno_lora_id = init_uno_lora_manager(self)
return
# Adapters apply to the target model only; the draft runs unadapted.
if get_lora().enable_lora and not self.is_draft_worker:
self.init_lora_manager()
@@ -1430,6 +1436,13 @@ class ModelRunner:
Subclasses can override this to install specialized decode graph runners.
"""
if self.spec_algorithm.is_uno():
from sglang.srt.speculative.uno_cuda_graph_runner import (
UnoDecodeCudaGraphRunner,
)
return UnoDecodeCudaGraphRunner
from sglang.srt.model_executor.runner.decode_cuda_graph_runner import (
DecodeCudaGraphRunner,
)
python/sglang/srt/server_args.py
+6 -1
View File
@@ -2067,7 +2067,12 @@ class ServerArgs:
# -------------------------------------------------------------------------
speculative_algorithm: A[
Optional[str],
"Speculative algorithm. Builtins: EAGLE, EAGLE3, NEXTN, STANDALONE, NGRAM, DFLASH, DSPARK. Or any name registered via `SpeculativeAlgorithm.register`.",
"Speculative algorithm. Builtins: EAGLE, EAGLE3, NEXTN, STANDALONE, NGRAM, DFLASH, DSPARK, UNO. Or any name registered via `SpeculativeAlgorithm.register`.",
NS("spec"),
] = None
uno_lora_path: A[
Optional[str],
"Path to the UNO draft LoRA checkpoint.",
NS("spec"),
] = None
speculative_draft_model_path: A[
python/sglang/srt/speculative/eagle_utils.py
+55
View File
@@ -663,11 +663,30 @@ def _verify_coins(
return coins, coins_for_final_sampling
def _can_use_sparse_uno_tree_target_sampling(
max_top_k: Optional[int],
sampling_info: SamplingBatchInfo,
) -> bool:
if max_top_k is None:
return False
from sglang.srt.speculative.uno_utils import _SPARSE_TOP_K_LIMIT
return bool(
_is_cuda
and max_top_k <= _SPARSE_TOP_K_LIMIT
and sampling_info.sampling_seed is None
and not sampling_info.need_min_p_sampling
and not get_spec().speculative_use_rejection_sampling
)
def eagle_sample(
verify_input: EagleVerifyInput,
batch: ScheduleBatch,
logits_output: LogitsProcessorOutput,
grammar_mask: Optional[GrammarMask] = None,
uno_target_max_top_k: Optional[int] = None,
):
"""
Verify and find accepted tokens based on logits output and batch
@@ -769,6 +788,42 @@ def eagle_sample(
tp_group.broadcast(predict, src=0)
tp_group.broadcast(accept_index, src=0)
tp_group.broadcast(num_correct_drafts, src=0)
elif _can_use_sparse_uno_tree_target_sampling(
uno_target_max_top_k,
sampling_info,
):
from sglang.srt.speculative.uno_utils import (
sample_uno_tree_target_tokens,
)
target_predict = sample_uno_tree_target_tokens(
next_token_logits=next_token_logits,
sampling_info=sampling_info,
batch_size=bs,
verify_width=verify_input.draft_token_num,
max_top_k=uno_target_max_top_k,
)
predict, accept_index, num_correct_drafts = verify_tree_greedy_func(
predicts=predict,
accept_index=accept_index,
accept_token_num=num_correct_drafts,
candidates=candidates,
retrieve_index=verify_input.retrieve_index,
retrieve_next_token=verify_input.retrieve_next_token,
retrieve_next_sibling=verify_input.retrieve_next_sibling,
target_predict=target_predict,
topk=verify_input.tree_topk,
)
tp_group = (
get_parallel().attn_tp_group
if is_dp_attention_enabled()
else get_tp_group()
)
if tp_group.world_size > 1:
tp_group.broadcast(predict, src=0)
tp_group.broadcast(accept_index, src=0)
tp_group.broadcast(num_correct_drafts, src=0)
else:
from sgl_kernel import (
top_k_renorm_prob,
python/sglang/srt/speculative/eagle_worker_common.py
+8 -1
View File
@@ -472,6 +472,7 @@ def run_eagle_verify(
metadata_ready_pre_pad: bool,
finalize_tree_path: bool,
grammar_barrier=None,
uno_target_max_top_k: Optional[int] = None,
) -> GenerationBatchResult:
"""Shared verify step: target-verify forward, sampling, acceptance bookkeeping.
@@ -584,7 +585,13 @@ def run_eagle_verify(
predict,
accept_lens,
accept_index,
) = eagle_sample(verify_input, batch, logits_output, grammar_mask)
) = eagle_sample(
verify_input,
batch,
logits_output,
grammar_mask,
uno_target_max_top_k=uno_target_max_top_k,
)
new_seq_lens = batch.seq_lens + accept_lens
clear_unaccepted_c128 = getattr(
token_to_kv_pool_allocator.get_kvcache(),
python/sglang/srt/speculative/spec_info.py
+17
View File
@@ -37,6 +37,7 @@ class SpeculativeAlgorithm(Enum):
"""
DFLASH = auto()
UNO = auto()
DSPARK = auto()
EAGLE = auto()
EAGLE3 = auto()
@@ -114,6 +115,9 @@ class SpeculativeAlgorithm(Enum):
def is_dflash(self) -> bool:
return self == SpeculativeAlgorithm.DFLASH
def is_uno(self) -> bool:
return self == SpeculativeAlgorithm.UNO
def is_dspark(self) -> bool:
return self == SpeculativeAlgorithm.DSPARK
@@ -220,6 +224,7 @@ class SpeculativeAlgorithm(Enum):
_handle_eagle_family,
_handle_frozen_kv_mtp,
_handle_ngram,
_handle_uno,
)
# Validate for every algorithm at startup: the metrics paths read the
@@ -230,6 +235,8 @@ class SpeculativeAlgorithm(Enum):
if self.is_dflash():
_handle_dflash(server_args)
elif self.is_uno():
_handle_uno(server_args)
elif self.is_dspark():
_handle_dspark(server_args)
elif self.is_frozen_kv_mtp():
@@ -304,6 +311,11 @@ class SpeculativeAlgorithm(Enum):
return DFlashWorkerV2
if self.is_uno():
from sglang.srt.speculative.uno_worker_v2 import UnoWorkerV2
return UnoWorkerV2
if self.is_dspark():
from sglang.srt.speculative.dspark_components.dspark_worker_v2 import (
DSparkWorkerV2,
@@ -356,6 +368,9 @@ class SpecInputType(IntEnum):
DFLASH_DRAFT = auto()
DFLASH_VERIFY = auto()
NGRAM_VERIFY = auto()
UNO_STATE = auto()
UNO_DRAFT = auto()
UNO_VERIFY = auto()
class SpecInput(ABC):
@@ -393,6 +408,7 @@ class SpecInput(ABC):
SpecInputType.EAGLE_DRAFT_EXTEND,
SpecInputType.FROZEN_KV_MTP_DRAFT,
SpecInputType.DFLASH_DRAFT,
SpecInputType.UNO_DRAFT,
}
def is_verify_input(self) -> bool:
@@ -401,6 +417,7 @@ class SpecInput(ABC):
SpecInputType.FROZEN_KV_MTP_VERIFY,
SpecInputType.DFLASH_VERIFY,
SpecInputType.NGRAM_VERIFY,
SpecInputType.UNO_VERIFY,
}
python/sglang/srt/speculative/spec_registry.py
+3
View File
@@ -82,6 +82,9 @@ class CustomSpecAlgo:
def is_dflash(self) -> bool:
return False
def is_uno(self) -> bool:
return False
def is_dspark(self) -> bool:
return False
python/sglang/srt/speculative/spec_utils.py
+6
View File
@@ -1051,6 +1051,12 @@ def spec_prepare_for_decode(batch: ScheduleBatch) -> None:
)
if batch.spec_algorithm.is_dflash_family():
batch.spec_info.prepare_for_decode(batch)
elif batch.spec_algorithm.is_uno():
from sglang.srt.speculative.uno_info import UnoDraftInput
if not isinstance(batch.spec_info, UnoDraftInput):
raise RuntimeError("UNO decode preparation requires UnoDraftInput")
batch.spec_info.prepare_for_decode(batch)
else:
from sglang.srt.speculative.eagle_utils import eagle_prepare_for_decode
python/sglang/srt/speculative/uno_cuda_graph_runner.py
+214
View File
@@ -0,0 +1,214 @@
import torch
from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
from sglang.srt.model_executor.runner.decode_cuda_graph_runner import (
DecodeCudaGraphRunner,
)
from sglang.srt.runtime_context import get_spec
from sglang.srt.speculative.eagle_info import EagleVerifyInput
from sglang.srt.speculative.spec_info import SpecInputType
from sglang.srt.speculative.uno_info import UnoForwardInput
from sglang.srt.speculative.uno_lora import UnoCudaGraphLoRAState
class UnoDecodeCudaGraphRunner(DecodeCudaGraphRunner):
"""Decode graph runner for linear UNO and both tree forward roles.
Linear UNO uses two-variant F-wide capture. Tree UNO uses
separate runner instances for its F-wide LoRA draft
and native Q/K EAGLE target verification.
"""
def __init__(
self,
model_runner,
*,
tree_draft_attn_backend=None,
tree_draft_width=None,
**kwargs,
):
candidate_top_k = get_spec().speculative_eagle_topk
self._tree_mode = candidate_top_k > 1
self._tree_draft_mode = tree_draft_width is not None
if self._tree_draft_mode:
self.record_nolora_graph = False
self._capture_spec_input_type = SpecInputType.UNO_DRAFT
self._lora_state = UnoCudaGraphLoRAState(
model_runner.lora_manager,
model_runner.uno_lora_id,
tree_draft_width,
)
model_runner.lora_manager.reset_lora_batch()
kwargs.update(
attn_backend=tree_draft_attn_backend,
speculative_num_steps=1,
speculative_num_draft_tokens=tree_draft_width,
)
super().__init__(model_runner, **kwargs)
model_runner.lora_manager.reset_lora_batch()
return
if self._tree_mode:
# Capture exactly one base-model target graph. The internal UNO
# adapter is active only in the rejected F-wide draft phase.
self.record_nolora_graph = False
model_runner.lora_manager.reset_lora_batch()
kwargs.update(
attn_backend=model_runner.attn_backend,
speculative_num_steps=get_spec().speculative_num_steps,
speculative_num_draft_tokens=get_spec().speculative_num_draft_tokens,
)
super().__init__(model_runner, **kwargs)
model_runner.lora_manager.reset_lora_batch()
return
forward_width = model_runner.decode_num_tokens_per_req()
self.record_nolora_graph = forward_width > 1
self._capture_spec_input_type = SpecInputType.UNO_VERIFY
self._lora_state = UnoCudaGraphLoRAState(
model_runner.lora_manager,
model_runner.uno_lora_id,
forward_width,
)
model_runner.lora_manager.reset_lora_batch()
super().__init__(model_runner, **kwargs)
def capture_prepare(self, size, stream_idx=None, num_tokens=None):
forward_batch, attn_backend, pp_proxy_tensors = super().capture_prepare(
size, stream_idx=stream_idx, num_tokens=num_tokens
)
# UNO owns token-row routing directly. K2's generic graph runner now
# keys LoRA setup off lora_manager presence, so suppress its synthetic
# request-level base-adapter routing during UNO capture.
forward_batch.lora_ids = None
return forward_batch, attn_backend, pp_proxy_tensors
def can_run_graph(self, forward_batch):
spec_info = forward_batch.spec_info
if self._tree_draft_mode:
if not isinstance(spec_info, UnoForwardInput):
return False
if spec_info.spec_input_type != SpecInputType.UNO_DRAFT:
return False
return super().can_run_graph(forward_batch)
if self._tree_mode:
if not isinstance(spec_info, EagleVerifyInput):
return False
return super().can_run_graph(forward_batch)
if not isinstance(spec_info, UnoForwardInput):
return False
# At F=1 both phases are base-only and share variant_label=None.
if spec_info.spec_input_type not in {
SpecInputType.UNO_DRAFT,
SpecInputType.UNO_VERIFY,
}:
return False
return super().can_run_graph(forward_batch)
def _resolve_lora_variant(self, forward_batch):
"""borrowed technique from multi-LoRA serving
to capture separate graph for each step."""
if self._tree_mode:
# Tree verification is always the clean target model. Keeping the
# graph keys unlabeled is sufficient because draft and verify own
# separate runners.
return None
if not self.record_nolora_graph:
return None
if forward_batch.spec_info.spec_input_type == SpecInputType.UNO_DRAFT:
return "lora"
return "nolora"
def capture_one_shape(
self,
size,
forward,
stream_idx=None,
variant_label=None,
dsa_variant=None,
):
"""capture one CUDA graph with/out UNO LoRA."""
if self._tree_draft_mode:
self._lora_state.capture_draft(size)
try:
return super().capture_one_shape(
size,
forward,
stream_idx,
None,
dsa_variant,
)
finally:
self._lora_state.reset()
if self._tree_mode:
self.model_runner.lora_manager.reset_lora_batch()
try:
return super().capture_one_shape(
size,
forward,
stream_idx,
None,
dsa_variant,
)
finally:
self.model_runner.lora_manager.reset_lora_batch()
if variant_label == "lora":
self._capture_spec_input_type = SpecInputType.UNO_DRAFT
self._lora_state.capture_draft(size)
else:
self._capture_spec_input_type = SpecInputType.UNO_VERIFY
self._lora_state.reset()
super().capture_one_shape(
size,
forward,
stream_idx,
variant_label,
dsa_variant,
)
self._lora_state.reset()
def get_spec_info(self, num_tokens: int):
if self._tree_draft_mode:
return UnoForwardInput(
spec_input_type=SpecInputType.UNO_DRAFT,
positions=self.buffers.positions[:num_tokens],
draft_token_num=self.captured_req_width,
)
if self._tree_mode:
# This deliberately mirrors DecodeCudaGraphRunner's EAGLE capture
# input. Current eagle_prepare_for_verify requests FULL hidden
# capture for every non-STANDALONE algorithm, including UNO.
spec_info = EagleVerifyInput(
draft_token=None,
custom_mask=self.buffers.custom_mask,
positions=None,
retrieve_index=None,
retrieve_next_token=None,
retrieve_next_sibling=None,
retrieve_cum_len=None,
spec_steps=self.speculative_num_steps,
topk=get_spec().speculative_eagle_topk,
draft_token_num=self.speculative_num_draft_tokens,
capture_hidden_mode=CaptureHiddenMode.FULL,
seq_lens_sum=None,
seq_lens_cpu=None,
)
spec_info.hidden_states = torch.zeros(
(num_tokens, self.model_runner.model_config.hidden_size),
dtype=self.model_runner.dtype,
device=self.model_runner.device,
)
return spec_info
return UnoForwardInput(
spec_input_type=self._capture_spec_input_type,
positions=self.buffers.positions[:num_tokens],
draft_token_num=self.captured_req_width,
)
python/sglang/srt/speculative/uno_info.py
+276
View File
@@ -0,0 +1,276 @@
from dataclasses import dataclass, field
from typing import ClassVar, List, Optional
import torch
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.mem_cache.allocation import alloc_for_spec_decode
from sglang.srt.mem_cache.allocation_sizing import (
get_alloc_reserve_per_decode,
page_aligned_decode_alloc_lens,
)
from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
@dataclass
class UnoDraftInput(SpecInput):
"""UNO state carried from one engine iteration to the next."""
# Previously emitted token at logical position C. It has no KV yet.
bonus_tokens: torch.Tensor
# Target-correct KV frontier C.
new_seq_lens: torch.Tensor
# Number of queries in each UNO forward.
forward_width: int
# FutureMap compatibility. UNO relays only bonus_tokens, but the generic
# relay payload reads these optional Eagle-shaped fields.
topk_p: ClassVar[Optional[torch.Tensor]] = None
topk_index: ClassVar[Optional[torch.Tensor]] = None
hidden_states: ClassVar[Optional[torch.Tensor]] = None
# Filled by the scheduler after an overlapped dispatch.
future_indices: Optional[torch.Tensor] = None
# Host-side upper bound for the allocator mapping prepared for this step.
reserved_seq_lens_cpu: Optional[torch.Tensor] = None
reserved_seq_lens_sum: Optional[int] = None
# Sampling metadata prepared before either internal forward.
max_top_k: int = 1
uniform_top_k_value: Optional[int] = None
def __post_init__(self):
super().__init__(SpecInputType.UNO_STATE)
if self.forward_width < 1:
raise ValueError("UNO forward_width must be positive.")
# The carried state represents one request row, not an F-row forward.
self.num_tokens_per_req = 1
self.num_tokens_for_logprob_per_req = 1
@property
def tail_width(self) -> int:
return self.forward_width + 1
@classmethod
def create_idle_input(
cls,
*,
device,
forward_width: int,
) -> "UnoDraftInput":
return cls(
bonus_tokens=torch.empty(
(0,),
dtype=torch.int64,
device=device,
),
new_seq_lens=torch.empty(
(0,),
dtype=torch.int64,
device=device,
),
forward_width=forward_width,
)
def prepare_for_decode(self, batch: ScheduleBatch) -> None:
batch.maybe_evict_swa()
batch_size = batch.batch_size()
if batch_size == 0:
return
if self.future_indices is None:
if self.bonus_tokens.numel() != batch_size:
raise RuntimeError("UNO seed count does not match decode batch size.")
if self.new_seq_lens.numel() != batch_size:
raise RuntimeError(
"UNO frontier count does not match decode batch size."
)
elif self.future_indices.numel() != batch_size:
raise RuntimeError(
"UNO future-index count does not match decode batch size."
)
committed_lengths: list[int] = []
reserve_width = int(get_alloc_reserve_per_decode())
max_top_k = 1
uniform_top_k_value = None
uniform_top_k = True
for index, req in enumerate(batch.reqs):
if req.kv is None:
raise RuntimeError("UNO decode request has no KV allocation.")
committed = int(req.kv.kv_committed_len)
allocated = int(req.kv.kv_allocated_len)
if allocated < committed:
raise RuntimeError(
"UNO encountered an invalid KV watermark: "
f"committed={committed}, allocated={allocated}."
)
committed_lengths.append(committed)
top_k = int(req.sampling_params.top_k)
max_top_k = max(max_top_k, top_k)
if index == 0:
uniform_top_k_value = top_k
elif uniform_top_k and top_k != uniform_top_k_value:
uniform_top_k = False
self.max_top_k = max_top_k
self.uniform_top_k_value = uniform_top_k_value if uniform_top_k else None
page_size = batch.token_to_kv_pool_allocator.page_size
current_lengths, next_lengths, num_needed_tokens = (
page_aligned_decode_alloc_lens(
batch.reqs,
reserve=reserve_width,
page_size=page_size,
)
)
row_width = int(batch.req_to_token_pool.req_to_token.shape[1])
if max(next_lengths) > row_width:
raise RuntimeError(
"UNO allocation exceeds the req_to_token row: "
f"needed={max(next_lengths)}, available={row_width}."
)
current_cpu = torch.tensor(
current_lengths,
dtype=torch.int32,
device="cpu",
)
next_cpu = torch.tensor(
next_lengths,
dtype=torch.int32,
device="cpu",
)
current_device = current_cpu.to(
batch.device,
non_blocking=True,
)
next_device = next_cpu.to(
batch.device,
non_blocking=True,
)
alloc_for_spec_decode(
batch.tree_cache,
batch.req_to_token_pool,
reqs=batch.reqs,
req_pool_indices=batch.req_pool_indices,
cur_kv_lens=current_device,
cur_kv_lens_cpu=current_cpu,
nxt_kv_lens=next_device,
nxt_kv_lens_cpu=next_cpu,
num_needed_tokens=num_needed_tokens,
batch=batch,
)
for req in batch.reqs:
req.decode_batch_idx += 1
batch.seq_lens_cpu = torch.tensor(
committed_lengths,
dtype=torch.int64,
device="cpu",
)
batch.seq_lens_sum = sum(committed_lengths)
self.reserved_seq_lens_cpu = next_cpu
self.reserved_seq_lens_sum = sum(next_lengths)
def filter_batch(
self,
new_indices: torch.Tensor,
new_indices_cpu: Optional[List[int]] = None,
) -> None:
if self.reserved_seq_lens_cpu is not None:
host_indices = (
new_indices_cpu if new_indices_cpu is not None else new_indices.cpu()
)
self.reserved_seq_lens_cpu = self.reserved_seq_lens_cpu[host_indices]
self.reserved_seq_lens_sum = int(self.reserved_seq_lens_cpu.sum().item())
if self.future_indices is not None:
self.future_indices = self.future_indices[new_indices]
return
self.bonus_tokens = self.bonus_tokens[new_indices]
self.new_seq_lens = self.new_seq_lens[new_indices]
def merge_batch(self, other: "UnoDraftInput") -> None:
if not isinstance(other, UnoDraftInput):
raise TypeError(f"Cannot merge UnoDraftInput with {type(other).__name__}.")
if self.forward_width != other.forward_width:
raise RuntimeError("Cannot merge UNO states with different forward widths.")
self_has_reservation = self.reserved_seq_lens_cpu is not None
other_has_reservation = other.reserved_seq_lens_cpu is not None
if self_has_reservation != other_has_reservation:
raise RuntimeError("Cannot merge prepared and unprepared UNO states.")
if self_has_reservation:
self.reserved_seq_lens_cpu = torch.cat(
(
self.reserved_seq_lens_cpu,
other.reserved_seq_lens_cpu,
)
)
self.reserved_seq_lens_sum = int(self.reserved_seq_lens_cpu.sum().item())
if self.future_indices is not None:
assert other.future_indices is not None
self.future_indices = torch.cat((self.future_indices, other.future_indices))
return
self.bonus_tokens = torch.cat((self.bonus_tokens, other.bonus_tokens))
self.new_seq_lens = torch.cat((self.new_seq_lens, other.new_seq_lens))
@dataclass
class UnoForwardInput(SpecInput):
"""Metadata for one fixed-width UNO forward."""
# constructor
spec_input_type: SpecInputType
positions: torch.Tensor
# For UNO, this is forward width F, not proposal count F - 1.
draft_token_num: int
# expected by interface
custom_mask: Optional[torch.Tensor] = None
capture_hidden_mode: CaptureHiddenMode = CaptureHiddenMode.NULL
hidden_states: Optional[torch.Tensor] = None
# derived
num_tokens_per_req: int = field(init=False)
num_tokens_for_logprob_per_req: int = field(init=False)
def __post_init__(self):
if self.spec_input_type not in {
SpecInputType.UNO_DRAFT,
SpecInputType.UNO_VERIFY,
}:
raise ValueError(f"Invalid UNO input type: {self.spec_input_type}")
if self.draft_token_num < 1:
raise ValueError("UNO forward width must be positive.")
# Dataclass-generated __init__ does not call the non-dataclass base
# initializer. This currently reassigns the same field, while also
# preserving the SpecInput initialization contract.
super().__init__(self.spec_input_type)
self.num_tokens_per_req = self.draft_token_num
self.num_tokens_for_logprob_per_req = self.draft_token_num
python/sglang/srt/speculative/uno_lora.py
+97
View File
@@ -0,0 +1,97 @@
"""Loading for UNO's draft LoRA."""
from __future__ import annotations
from typing import TYPE_CHECKING
from sglang.srt.lora.lora_manager import LoRAManager
from sglang.srt.lora.lora_registry import LoRARef
from sglang.srt.runtime_context import get_spec
if TYPE_CHECKING:
from sglang.srt.model_executor.model_runner import ModelRunner
# This name is internal to the model-execution process. It is never exposed as
# a request-selectable serving adapter.
_UNO_INTERNAL_LORA_NAME = "__uno_draft__"
# LoRA pool capacity includes the base-model slot.
_UNO_LORA_POOL_CAPACITY = 2
def init_uno_lora_manager(
model_runner: ModelRunner,
) -> tuple[LoRAManager, str]:
"""Load and pin the single UNO draft adapter."""
lora_path = get_spec().uno_lora_path
uno_ref = LoRARef(
lora_id=LoRARef.deterministic_id(
_UNO_INTERNAL_LORA_NAME,
lora_path,
),
lora_name=_UNO_INTERNAL_LORA_NAME,
lora_path=lora_path,
pinned=True,
)
manager = LoRAManager(
base_model=model_runner.model,
base_hf_config=model_runner.model_config.hf_config,
max_loras_per_batch=_UNO_LORA_POOL_CAPACITY,
load_config=model_runner.load_config,
dtype=model_runner.dtype,
server_args=model_runner.server_args,
lora_backend="uno_cublas", # fast path
tp_size=model_runner.ps.tp_size,
tp_rank=model_runner.ps.tp_rank,
# Infer these from the one trained adapter.
max_lora_rank=None,
target_modules=None,
lora_paths=[uno_ref],
)
# LoRAManager construction initially makes only the base slot resident.
# UNO always needs both fixed choices resident:
#
# None -> base model
# uno_ref.lora_id -> base model + UNO draft LoRA
manager.fetch_new_loras({None, uno_ref.lora_id})
return manager, uno_ref.lora_id
class UnoCudaGraphLoRAState:
"""Retained token-row LoRA routing for UNO draft graph buckets."""
def __init__(
self,
manager: LoRAManager,
uno_lora_id: str,
forward_width: int,
):
self.manager = manager
self.uno_lora_id = uno_lora_id
self.forward_width = forward_width
self._draft_batch_infos = {}
def capture_draft(self, batch_size: int) -> None:
if batch_size not in self._draft_batch_infos:
self.manager.prepare_lora_token_segments(
lora_ids=[None, self.uno_lora_id] * batch_size,
segment_lens=[1, self.forward_width - 1] * batch_size,
)
batch_info = self.manager.lora_backend.batch_info
batch_info.use_cuda_graph = True
self._draft_batch_infos[batch_size] = batch_info
self.activate_draft(batch_size)
def activate_draft(self, batch_size: int) -> None:
self.manager.reset_lora_batch()
self.manager.lora_backend.batch_info = self._draft_batch_infos[batch_size]
def reset(self) -> None:
self.manager.reset_lora_batch()
python/sglang/srt/speculative/uno_tree.py
+584
View File
@@ -0,0 +1,584 @@
# SPDX-License-Identifier: Apache-2.0
# Adapted from nano-vllm-uno's fixed-budget draft-tree builder for SGLang.
"""GPU-native UNO proposal-tree construction for SGLang spec-v2.
UNO owns only proposal ranking and fixed-budget best-first selection. The
result is expressed directly in EAGLE's candidate-lineage ABI so the existing
EAGLE implementation can build masks, positions, traversal links, verify the
tree, sample the accepted path, and compact KV state.
"""
from __future__ import annotations
from dataclasses import dataclass
import torch
import triton
import triton.language as tl
from flashinfer import top_k as _flashinfer_top_k
from torch import Tensor
@dataclass(frozen=True)
class UnoTreeProposal:
"""EAGLE-native representation of one fixed-width UNO tree batch.
For non-root node ``i``, ``top_scores_index[:, i - 1]`` is the implicit
candidate edge ``parent_node * candidate_top_k + candidate_rank``.
``parent_list`` maps an EAGLE candidate row back to the edge that created
that row's node. Its rows use EAGLE's native
``candidate_top_k * (max_depth - 1) + 1`` stride, including unused
padding. The tensors can therefore be passed directly to
``build_tree_kernel_efficient`` without constructing direct parent arrays.
Tensors may be backed by the caller's reusable workspace and remain valid
only until that workspace is reused.
"""
root_tokens: Tensor
draft_tokens: Tensor
parent_list: Tensor
top_scores_index: Tensor
candidate_top_k: int
max_depth: int
@property
def num_verify_tokens(self) -> int:
return int(self.draft_tokens.shape[1]) + 1
@triton.jit
def _candidate_lse_partials_kernel(
logits,
partial_max,
partial_sum,
temperature_values,
inverse_temperature,
BATCH_STRIDE: tl.constexpr,
DEPTH_STRIDE: tl.constexpr,
VOCAB_STRIDE: tl.constexpr,
NUM_DEPTHS: tl.constexpr,
VOCAB_SIZE: tl.constexpr,
NUM_BLOCKS: tl.constexpr,
BLOCK_VOCAB: tl.constexpr,
TEMPERATURE_BATCH_STRIDE: tl.constexpr,
TEMPERATURE_IS_TENSOR: tl.constexpr,
):
batch = tl.program_id(0)
depth = tl.program_id(1)
block = tl.program_id(2)
row = batch * NUM_DEPTHS + depth
offsets = block * BLOCK_VOCAB + tl.arange(0, BLOCK_VOCAB)
values = tl.load(
logits + batch * BATCH_STRIDE + depth * DEPTH_STRIDE + offsets * VOCAB_STRIDE,
mask=offsets < VOCAB_SIZE,
other=-float("inf"),
).to(tl.float32)
if TEMPERATURE_IS_TENSOR:
row_temperature = tl.load(
temperature_values + batch * TEMPERATURE_BATCH_STRIDE
).to(tl.float32)
row_inverse_temperature = tl.where(
row_temperature > 0.0,
1.0 / row_temperature,
1.0,
)
else:
row_inverse_temperature = inverse_temperature
values *= row_inverse_temperature
maximum = tl.max(values, axis=0)
total = tl.sum(tl.exp(values - maximum), axis=0)
output = row * NUM_BLOCKS + block
tl.store(partial_max + output, maximum)
tl.store(partial_sum + output, total)
@triton.jit
def _candidate_lse_finalize_kernel(
top_values,
partial_max,
partial_sum,
top_log_probs,
temperature_values,
inverse_temperature,
K: tl.constexpr,
NUM_DEPTHS: tl.constexpr,
NUM_BLOCKS: tl.constexpr,
BLOCK_K: tl.constexpr,
BLOCK_PARTIALS: tl.constexpr,
TEMPERATURE_BATCH_STRIDE: tl.constexpr,
TEMPERATURE_IS_TENSOR: tl.constexpr,
):
row = tl.program_id(0)
batch = row // NUM_DEPTHS
blocks = tl.arange(0, BLOCK_PARTIALS)
block_max = tl.load(
partial_max + row * NUM_BLOCKS + blocks,
mask=blocks < NUM_BLOCKS,
other=-float("inf"),
)
maximum = tl.max(block_max, axis=0)
block_sum = tl.load(
partial_sum + row * NUM_BLOCKS + blocks,
mask=blocks < NUM_BLOCKS,
other=0.0,
)
total = tl.sum(block_sum * tl.exp(block_max - maximum), axis=0)
normalizer = maximum + tl.log(total)
ranks = tl.arange(0, BLOCK_K)
candidates = tl.load(
top_values + row * K + ranks,
mask=ranks < K,
other=-float("inf"),
).to(tl.float32)
if TEMPERATURE_IS_TENSOR:
row_temperature = tl.load(
temperature_values + batch * TEMPERATURE_BATCH_STRIDE
).to(tl.float32)
row_inverse_temperature = tl.where(
row_temperature > 0.0,
1.0 / row_temperature,
1.0,
)
else:
row_inverse_temperature = inverse_temperature
log_probs = candidates * row_inverse_temperature - normalizer
# Proposal mass affects efficiency, not correctness. A malformed row
# must not leave the best-first frontier without a deterministic winner.
log_probs = tl.where(log_probs == log_probs, log_probs, -float("inf"))
log_probs = tl.where(log_probs > 0.0, 0.0, log_probs)
tl.store(
top_log_probs + row * K + ranks,
log_probs,
mask=ranks < K,
)
def _temperature_rows(
temperature: float | Tensor,
*,
batch_size: int,
device: torch.device,
) -> tuple[Tensor, int] | None:
if not isinstance(temperature, Tensor):
return None
if temperature.device != device:
raise ValueError("temperature and logits must share a device")
if not temperature.is_floating_point():
raise TypeError("temperature tensor must use a floating dtype")
if temperature.ndim == 0:
return temperature.reshape(1), 0
if temperature.shape not in ((batch_size,), (batch_size, 1)):
raise ValueError(
"temperature tensor must be scalar, [B], or [B, 1]; got "
f"{tuple(temperature.shape)}"
)
values = temperature.reshape(batch_size)
return values, values.stride(0)
def _build_candidate_log_probs(
logits: Tensor,
top_values: Tensor,
top_log_probs: Tensor,
partial_max: Tensor,
partial_sum: Tensor,
temperature: float | Tensor,
) -> None:
"""Normalize selected logits over the full vocabulary in FP32."""
batch_size, num_depths, vocab_size = logits.shape
num_rows = batch_size * num_depths
candidate_top_k = int(top_values.size(-1))
block_vocab = 8192
num_blocks = triton.cdiv(vocab_size, block_vocab)
temperature_rows = _temperature_rows(
temperature,
batch_size=batch_size,
device=logits.device,
)
if temperature_rows is None:
temperature_values = logits
temperature_batch_stride = 0
temperature_is_tensor = False
scalar_temperature = float(temperature)
inverse_temperature = (
1.0 / scalar_temperature if scalar_temperature > 0.0 else 1.0
)
else:
temperature_values, temperature_batch_stride = temperature_rows
temperature_is_tensor = True
inverse_temperature = 1.0
_candidate_lse_partials_kernel[(batch_size, num_depths, num_blocks)](
logits,
partial_max,
partial_sum,
temperature_values,
inverse_temperature,
BATCH_STRIDE=logits.stride(0),
DEPTH_STRIDE=logits.stride(1),
VOCAB_STRIDE=logits.stride(-1),
NUM_DEPTHS=num_depths,
VOCAB_SIZE=vocab_size,
NUM_BLOCKS=num_blocks,
BLOCK_VOCAB=block_vocab,
TEMPERATURE_BATCH_STRIDE=temperature_batch_stride,
TEMPERATURE_IS_TENSOR=temperature_is_tensor,
num_warps=4,
num_stages=1,
)
_candidate_lse_finalize_kernel[(num_rows,)](
top_values,
partial_max,
partial_sum,
top_log_probs,
temperature_values,
inverse_temperature,
K=candidate_top_k,
NUM_DEPTHS=num_depths,
NUM_BLOCKS=num_blocks,
BLOCK_K=triton.next_power_of_2(candidate_top_k),
BLOCK_PARTIALS=triton.next_power_of_2(num_blocks),
TEMPERATURE_BATCH_STRIDE=temperature_batch_stride,
TEMPERATURE_IS_TENSOR=temperature_is_tensor,
num_warps=1,
num_stages=1,
)
@triton.jit
def _build_tree_kernel(
top_token_ids,
top_log_probs,
draft_tokens,
selected_edges,
parent_list,
search_depths,
search_log_masses,
TOKEN_BATCH_STRIDE: tl.constexpr,
TOKEN_DEPTH_STRIDE: tl.constexpr,
TOKEN_RANK_STRIDE: tl.constexpr,
PROB_BATCH_STRIDE: tl.constexpr,
PROB_DEPTH_STRIDE: tl.constexpr,
PROB_RANK_STRIDE: tl.constexpr,
PARENT_BATCH_STRIDE: tl.constexpr,
NUM_DEPTHS: tl.constexpr,
K: tl.constexpr,
Q: tl.constexpr,
BLOCK_CANDIDATES: tl.constexpr,
):
batch = tl.program_id(0)
search_offset = batch * Q
edge_output_offset = batch * (Q - 1)
parent_output_offset = batch * PARENT_BATCH_STRIDE
candidate_slots = tl.arange(0, BLOCK_CANDIDATES)
candidate_parents = candidate_slots // K
candidate_ranks = candidate_slots % K
candidate_in_bounds = candidate_slots < Q * K
used = tl.zeros((BLOCK_CANDIDATES,), dtype=tl.int1)
# The root token itself is returned by reference. Only its private search
# state is stored here; EAGLE prepends it as the verify-tree root later.
tl.store(
parent_list + parent_output_offset + candidate_slots,
-1,
mask=candidate_slots < PARENT_BATCH_STRIDE,
)
tl.store(search_depths + search_offset, 0)
tl.store(search_log_masses + search_offset, 0.0)
tl.debug_barrier()
for node_index in range(1, Q):
parent_valid = candidate_in_bounds & (candidate_parents < node_index)
safe_parent = tl.where(parent_valid, candidate_parents, 0)
parent_depth = tl.load(
search_depths + search_offset + safe_parent,
mask=parent_valid,
other=NUM_DEPTHS,
)
parent_mass = tl.load(
search_log_masses + search_offset + safe_parent,
mask=parent_valid,
other=-float("inf"),
).to(tl.float32)
valid = parent_valid & (parent_depth < NUM_DEPTHS) & ~used
safe_depth = tl.where(valid, parent_depth, 0)
token = tl.load(
top_token_ids
+ batch * TOKEN_BATCH_STRIDE
+ safe_depth * TOKEN_DEPTH_STRIDE
+ candidate_ranks * TOKEN_RANK_STRIDE,
mask=valid,
other=0,
)
log_prob = tl.load(
top_log_probs
+ batch * PROB_BATCH_STRIDE
+ safe_depth * PROB_DEPTH_STRIDE
+ candidate_ranks * PROB_RANK_STRIDE,
mask=valid,
other=-float("inf"),
).to(tl.float32)
mass = parent_mass + log_prob
child_depth = parent_depth + 1
# Deterministic best-first order: mass, shallower depth, lower rank,
# lower token ID, then lower parent node.
best_mass = tl.max(tl.where(valid, mass, -float("inf")), axis=0)
winner = valid & (mass == best_mass)
best_depth = tl.min(tl.where(winner, child_depth, 1 << 30), axis=0)
winner &= child_depth == best_depth
best_rank = tl.min(tl.where(winner, candidate_ranks, 1 << 30), axis=0)
winner &= candidate_ranks == best_rank
best_token = tl.min(tl.where(winner, token, 1 << 30), axis=0)
winner &= token == best_token
best_parent = tl.min(tl.where(winner, candidate_parents, 1 << 30), axis=0)
winner &= candidate_parents == best_parent
best_slot = tl.min(tl.where(winner, candidate_slots, 1 << 30), axis=0)
selected_parent = best_slot // K
selected_rank = best_slot % K
selected_depth = tl.load(search_depths + search_offset + selected_parent)
selected_mass = tl.load(search_log_masses + search_offset + selected_parent).to(
tl.float32
) + tl.load(
top_log_probs
+ batch * PROB_BATCH_STRIDE
+ selected_depth * PROB_DEPTH_STRIDE
+ selected_rank * PROB_RANK_STRIDE
).to(tl.float32)
selected_token = tl.load(
top_token_ids
+ batch * TOKEN_BATCH_STRIDE
+ selected_depth * TOKEN_DEPTH_STRIDE
+ selected_rank * TOKEN_RANK_STRIDE
)
output_index = edge_output_offset + node_index - 1
tl.store(draft_tokens + output_index, selected_token)
# This is already EAGLE's implicit selected-edge encoding.
tl.store(selected_edges + output_index, best_slot)
if node_index < Q - 1:
# EAGLE shifts each selected edge by one candidate row. Fusing
# this write avoids separate fill/copy launches on every step.
tl.store(
parent_list + parent_output_offset + node_index,
best_slot,
)
tl.store(
search_depths + search_offset + node_index,
selected_depth + 1,
)
tl.store(
search_log_masses + search_offset + node_index,
selected_mass,
)
used |= candidate_slots == best_slot
tl.debug_barrier()
def _candidate_tree_capacity(
num_depths: int,
candidate_top_k: int,
stop_at: int,
) -> int:
capacity = 1
width = 1
for _ in range(num_depths):
width *= candidate_top_k
capacity += width
if capacity >= stop_at:
break
return capacity
@torch.inference_mode()
def build_uno_tree_proposal(
root_tokens: Tensor,
draft_logits: Tensor,
*,
max_nodes: int,
candidate_top_k: int,
temperature: float | Tensor,
workspace: dict[str, Tensor] | None = None,
) -> UnoTreeProposal:
"""Build fixed-``Q`` UNO trees directly in EAGLE's proposal ABI.
``root_tokens`` is ``[B]`` and ``draft_logits`` is ``[B, F-1, V]``.
Candidate log probabilities are normalized over the full vocabulary. No
CPU reference/fallback, direct parent array, attention mask, traversal
structure, acceptance walk, or KV operation is implemented here.
"""
if root_tokens.ndim != 1:
raise ValueError(
f"root_tokens must have shape [B], got {tuple(root_tokens.shape)}"
)
if draft_logits.ndim != 3 or draft_logits.size(0) != root_tokens.size(0):
raise ValueError(
"draft_logits must have shape [B, depth, vocab] with the same B "
f"as root_tokens; got {tuple(draft_logits.shape)}"
)
if not root_tokens.is_cuda or not draft_logits.is_cuda:
raise ValueError("UNO tree construction requires CUDA tensors")
if root_tokens.device != draft_logits.device:
raise ValueError("tree roots and draft logits must share a device")
if root_tokens.dtype not in (torch.int32, torch.int64):
raise TypeError("tree roots must use an integer dtype")
if not draft_logits.is_floating_point():
raise TypeError("draft logits must use a floating dtype")
if root_tokens.numel() == 0:
raise ValueError("UNO tree construction requires a non-empty batch")
if max_nodes < 1:
raise ValueError("max_nodes must include at least the root")
if candidate_top_k < 1:
raise ValueError("candidate_top_k must be >= 1")
batch_size = int(root_tokens.size(0))
num_depths = int(draft_logits.size(1))
vocab_size = int(draft_logits.size(2))
candidate_top_k = int(candidate_top_k)
draft_width = num_depths + 1
parent_width = candidate_top_k * max(num_depths - 1, 0) + 1
if max_nodes < draft_width:
raise ValueError(
f"max_nodes Q must be >= draft width F; got Q={max_nodes}, F={draft_width}"
)
if candidate_top_k > vocab_size:
raise ValueError(
f"candidate_top_k ({candidate_top_k}) exceeds vocabulary size "
f"({vocab_size})"
)
if max_nodes > 128 or max_nodes * candidate_top_k > 2048:
raise ValueError(
"the initial single-program UNO builder requires Q <= 128 and "
f"Q*K <= 2048; got Q={max_nodes}, K={candidate_top_k}"
)
capacity = _candidate_tree_capacity(
num_depths,
candidate_top_k,
max_nodes,
)
if capacity < max_nodes:
raise ValueError(
f"candidate set can produce only {capacity} tree nodes, but "
f"fixed tree verification requires {max_nodes}"
)
if max_nodes - 1 > parent_width:
raise ValueError(
"EAGLE's parent-list ABI cannot represent this UNO tree: "
f"Q-1={max_nodes - 1} exceeds K*(depth-1)+1={parent_width}"
)
_temperature_rows(
temperature,
batch_size=batch_size,
device=draft_logits.device,
)
def buffer(
name: str,
shape: tuple[int, ...],
dtype: torch.dtype,
) -> Tensor:
if workspace is None:
return torch.empty(
shape,
dtype=dtype,
device=draft_logits.device,
)
value = workspace.get(name)
if (
value is None
or value.shape != shape
or value.dtype != dtype
or value.device != draft_logits.device
):
value = torch.empty(
shape,
dtype=dtype,
device=draft_logits.device,
)
workspace[name] = value
return value
edge_shape = (batch_size, max_nodes - 1)
draft_tokens = buffer("draft_tokens", edge_shape, torch.long)
selected_edges = buffer("top_scores_index", edge_shape, torch.long)
parent_list = buffer(
"parent_list",
(batch_size, parent_width),
torch.long,
)
if max_nodes == 1:
parent_list.fill_(-1)
return UnoTreeProposal(
root_tokens=root_tokens,
draft_tokens=draft_tokens,
parent_list=parent_list,
top_scores_index=selected_edges,
candidate_top_k=candidate_top_k,
max_depth=num_depths,
)
candidate_shape = (batch_size, num_depths, candidate_top_k)
flat_logits = draft_logits.contiguous().view(batch_size * num_depths, vocab_size)
flat_top_values, flat_top_token_ids = _flashinfer_top_k(
flat_logits,
candidate_top_k,
sorted=True,
deterministic=False,
)
top_values = flat_top_values.view(candidate_shape)
top_token_ids = buffer("top_token_ids", candidate_shape, torch.long)
top_token_ids.copy_(flat_top_token_ids.view(candidate_shape))
top_log_probs = buffer("top_log_probs", candidate_shape, torch.float32)
num_partial_blocks = (vocab_size + 8191) // 8192
partial_shape = (batch_size * num_depths, num_partial_blocks)
_build_candidate_log_probs(
draft_logits,
top_values,
top_log_probs,
buffer("partial_lse_max", partial_shape, torch.float32),
buffer("partial_lse_sum", partial_shape, torch.float32),
temperature,
)
search_shape = (batch_size, max_nodes)
search_depths = buffer("search_depths", search_shape, torch.int32)
search_log_masses = buffer("search_log_masses", search_shape, torch.float32)
_build_tree_kernel[(batch_size,)](
top_token_ids,
top_log_probs,
draft_tokens,
selected_edges,
parent_list,
search_depths,
search_log_masses,
TOKEN_BATCH_STRIDE=top_token_ids.stride(0),
TOKEN_DEPTH_STRIDE=top_token_ids.stride(1),
TOKEN_RANK_STRIDE=top_token_ids.stride(2),
PROB_BATCH_STRIDE=top_log_probs.stride(0),
PROB_DEPTH_STRIDE=top_log_probs.stride(1),
PROB_RANK_STRIDE=top_log_probs.stride(2),
PARENT_BATCH_STRIDE=parent_list.stride(0),
NUM_DEPTHS=num_depths,
K=candidate_top_k,
Q=max_nodes,
BLOCK_CANDIDATES=triton.next_power_of_2(max_nodes * candidate_top_k),
num_warps=8,
num_stages=1,
)
return UnoTreeProposal(
root_tokens=root_tokens,
draft_tokens=draft_tokens,
parent_list=parent_list,
top_scores_index=selected_edges,
candidate_top_k=candidate_top_k,
max_depth=num_depths,
)
python/sglang/srt/speculative/uno_utils.py
+587
View File
@@ -0,0 +1,587 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Any
import torch
from flashinfer import top_k as _flashinfer_top_k
from sglang.kernels.ops.speculative.reject_sampling import (
chain_speculative_sampling_triton,
)
from sglang.srt.speculative.dflash_utils import (
_get_or_create_chain_verify_buffers,
build_dflash_verify_target_probs,
)
from sglang.srt.speculative.spec_utils import fast_sample
_SPARSE_TOP_K_LIMIT = 128
def _normalize_sparse_topk_probs(
topk_logits: torch.Tensor,
temperatures: torch.Tensor,
valid: torch.Tensor,
top_ps: torch.Tensor,
) -> torch.Tensor:
"""Normalize a compact top-k support with top-k-first top-p semantics."""
scaled = topk_logits.float() / temperatures
scaled = scaled.masked_fill(~valid, float("-inf"))
probs = torch.softmax(scaled, dim=-1)
cdf = torch.cumsum(probs, dim=-1)
probs = probs.masked_fill((cdf - probs) > top_ps, 0.0)
return probs / probs.sum(dim=-1, keepdim=True).clamp_min(1e-12)
@torch.compile(dynamic=True)
def _sparse_rejection_from_support(
candidates: torch.Tensor,
target_ids: torch.Tensor,
target_probs: torch.Tensor,
draft_ids: torch.Tensor,
draft_probs: torch.Tensor,
accept_uniforms: torch.Tensor,
final_uniforms: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Run exact p/q rejection sampling on compact linear-chain supports."""
batch_size, forward_width = candidates.shape
num_proposals = forward_width - 1
if num_proposals == 0:
final_ids = target_ids[:, 0]
final_probs = target_probs[:, 0]
accepted_counts = torch.zeros(
batch_size,
dtype=torch.int32,
device=candidates.device,
)
else:
proposal_ids = candidates[:, 1:]
p_proposal = torch.where(
target_ids[:, :num_proposals].eq(proposal_ids.unsqueeze(-1)),
target_probs[:, :num_proposals],
torch.zeros(
(),
dtype=target_probs.dtype,
device=target_probs.device,
),
).sum(dim=-1)
q_proposal = torch.where(
draft_ids.eq(proposal_ids.unsqueeze(-1)),
draft_probs,
torch.zeros(
(),
dtype=draft_probs.dtype,
device=draft_probs.device,
),
).sum(dim=-1)
ratios = torch.where(
q_proposal > 0,
p_proposal / q_proposal,
torch.zeros_like(p_proposal),
).clamp_(max=1.0)
accepted_flags = accept_uniforms[:, :num_proposals] < ratios
accepted_counts = accepted_flags.to(torch.int32).cumprod(dim=1).sum(dim=1)
batch_indices = torch.arange(batch_size, device=candidates.device)
final_rows = accepted_counts.to(torch.long)
final_ids = target_ids[batch_indices, final_rows]
final_probs = target_probs[batch_indices, final_rows]
rejected = final_rows < num_proposals
draft_rows = final_rows.clamp(max=num_proposals - 1)
final_draft_ids = draft_ids[batch_indices, draft_rows]
final_draft_probs = draft_probs[batch_indices, draft_rows]
q_on_target = torch.where(
final_ids.unsqueeze(2).eq(final_draft_ids.unsqueeze(1)),
final_draft_probs.unsqueeze(1),
torch.zeros(
(),
dtype=final_draft_probs.dtype,
device=final_draft_probs.device,
),
).sum(dim=2)
correction_probs = (final_probs - q_on_target).clamp_min_(0.0)
correction_sum = correction_probs.sum(dim=1, keepdim=True)
correction_probs = torch.where(
correction_sum > 0,
correction_probs / correction_sum.clamp_min(1e-12),
final_probs,
)
final_probs = torch.where(
rejected[:, None],
correction_probs,
final_probs,
)
cdf = torch.cumsum(final_probs, dim=-1)
thresholds = final_uniforms * final_probs.sum(dim=-1)
sampled_offsets = (cdf <= thresholds[:, None]).sum(dim=-1)
sampled_offsets.clamp_(max=final_ids.shape[-1] - 1)
bonus = (
final_ids.gather(1, sampled_offsets[:, None]).squeeze(1).to(candidates.dtype)
)
return accepted_counts, bonus
@torch.compile(dynamic=True)
def _build_sparse_target_support_tensors(
next_token_logits: torch.Tensor,
temperatures: torch.Tensor,
top_ks: torch.Tensor,
top_ps: torch.Tensor,
batch_size: int,
forward_width: int,
max_top_k: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Build compact support using the fastest available top-k primitive."""
rows = batch_size * forward_width
topk_logits, topk_ids = _flashinfer_top_k(
next_token_logits.contiguous(),
max_top_k,
sorted=True,
deterministic=False,
)
expanded_temperatures = torch.repeat_interleave(
temperatures,
forward_width,
dim=0,
).reshape(rows, -1)
expanded_top_ks = torch.repeat_interleave(
top_ks,
forward_width,
dim=0,
).reshape(rows, 1)
expanded_top_ps = torch.repeat_interleave(
top_ps,
forward_width,
dim=0,
).reshape(rows, 1)
ranks = torch.arange(
max_top_k,
dtype=expanded_top_ks.dtype,
device=next_token_logits.device,
)[None, :]
probs = _normalize_sparse_topk_probs(
topk_logits,
expanded_temperatures,
ranks < expanded_top_ks,
expanded_top_ps,
)
return (
topk_ids.view(batch_size, forward_width, max_top_k),
probs.view(batch_size, forward_width, max_top_k),
)
def _build_sparse_target_support(
*,
next_token_logits: torch.Tensor,
sampling_info: Any,
batch_size: int,
forward_width: int,
max_top_k: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Return compact target token IDs/probabilities without a dense scatter."""
if bool(getattr(sampling_info, "need_top_p_sampling", False)):
top_ps = sampling_info.top_ps
else:
top_ps = torch.ones(
(batch_size,),
dtype=torch.float32,
device=next_token_logits.device,
)
return _build_sparse_target_support_tensors(
next_token_logits,
sampling_info.temperatures,
sampling_info.top_ks,
top_ps,
batch_size,
forward_width,
max_top_k,
)
def _sample_from_support(
support_ids: torch.Tensor,
support_probs: torch.Tensor,
) -> torch.Tensor:
"""Sample one token from every compact support row."""
flat_ids = support_ids.flatten(0, 1)
flat_probs = support_probs.flatten(0, 1)
_, offsets = fast_sample(flat_probs)
return flat_ids.gather(1, offsets).view(support_ids.shape[:2])
@dataclass(frozen=True)
class UnoDraftDistribution:
"""The exact q used to sample future UNO proposal rows."""
probs: torch.Tensor
token_ids: torch.Tensor | None = None
def _run_sparse_rejection(
*,
candidates: torch.Tensor,
next_token_logits: torch.Tensor,
sampling_info: Any,
max_top_k: int,
draft_distribution: UnoDraftDistribution,
) -> tuple[torch.Tensor, torch.Tensor]:
if draft_distribution.token_ids is None:
raise RuntimeError("Sparse UNO verification requires sparse draft q.")
batch_size, forward_width = candidates.shape
support_ids, support_probs = _build_sparse_target_support(
next_token_logits=next_token_logits,
sampling_info=sampling_info,
batch_size=batch_size,
forward_width=forward_width,
max_top_k=max_top_k,
)
# Preserve the legacy path's two RNG draws and tensor shapes. The final
# uniforms select the correction/bonus; the first F - 1 acceptance coins
# are consumed by a linear chain.
accept_uniforms = torch.rand(
(batch_size, forward_width),
dtype=torch.float32,
device=next_token_logits.device,
)
final_uniforms = torch.rand(
(batch_size,),
dtype=torch.float32,
device=next_token_logits.device,
)
return _sparse_rejection_from_support(
candidates,
support_ids,
support_probs,
draft_distribution.token_ids,
draft_distribution.probs,
accept_uniforms,
final_uniforms,
)
def _build_dense_probs(
*,
next_token_logits: torch.Tensor,
sampling_info: Any,
batch_size: int,
forward_width: int,
max_top_k: int,
uniform_top_k_value: int | None,
) -> torch.Tensor:
"""Build the dense sampling distribution used by SGLang verification."""
return build_dflash_verify_target_probs(
next_token_logits=next_token_logits,
sampling_info=sampling_info,
draft_token_num=forward_width,
bs=batch_size,
max_top_k=max_top_k,
uniform_top_k_value=uniform_top_k_value,
use_sparse_topk=True,
)
def _sample_from_dense_probs(probs: torch.Tensor) -> torch.Tensor:
"""Sample one token from every dense distribution row."""
_, token_ids = fast_sample(probs.flatten(0, 1))
return token_ids.view(probs.shape[:2])
def _run_dense_rejection(
*,
candidates: torch.Tensor,
next_token_logits: torch.Tensor,
sampling_info: Any,
max_top_k: int,
uniform_top_k_value: int | None,
draft_distribution: UnoDraftDistribution,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Run SGLang's fused linear-chain p/q rejection kernel."""
if draft_distribution.token_ids is not None:
raise RuntimeError("Dense UNO verification requires dense draft q.")
batch_size, forward_width = candidates.shape
target_probs = _build_dense_probs(
next_token_logits=next_token_logits,
sampling_info=sampling_info,
batch_size=batch_size,
forward_width=forward_width,
max_top_k=max_top_k,
uniform_top_k_value=uniform_top_k_value,
)
accept_uniforms = torch.rand(
(batch_size, forward_width),
dtype=torch.float32,
device=next_token_logits.device,
)
final_uniforms = torch.rand(
(batch_size,),
dtype=torch.float32,
device=next_token_logits.device,
)
(
retrieve_index,
retrieve_next_token,
retrieve_next_sibling,
predicts,
accept_index,
accepted_counts,
) = _get_or_create_chain_verify_buffers(
bs=batch_size,
draft_token_num=forward_width,
device=next_token_logits.device,
)
chain_speculative_sampling_triton(
predicts=predicts,
accept_index=accept_index,
accept_token_num=accepted_counts,
candidates=candidates,
retrive_index=retrieve_index,
retrive_next_token=retrieve_next_token,
retrive_next_sibling=retrieve_next_sibling,
uniform_samples=accept_uniforms,
uniform_samples_for_final_sampling=final_uniforms,
target_probs=target_probs,
draft_probs=draft_distribution.probs,
threshold_single=1.0,
threshold_acc=1.0,
deterministic=True,
)
rows = torch.arange(batch_size, device=candidates.device)
bonus_positions = accept_index[
rows,
accepted_counts.to(torch.long),
].to(torch.long)
bonus = predicts[bonus_positions].to(candidates.dtype)
return accepted_counts, bonus
@dataclass
class UnoSamplingResult:
output_ids: torch.Tensor
accept_lens: torch.Tensor
new_seq_lens: torch.Tensor
next_seed_tokens: torch.Tensor
@dataclass
class UnoTreeSamplingResult:
output_ids: torch.Tensor
accept_lens: torch.Tensor
def build_uno_draft_input(
*,
seed_tokens: torch.Tensor,
forward_width: int,
vocab_size: int,
noise_tokens: torch.Tensor | None = None, # for testing
) -> torch.Tensor:
"""Build one ``[seed, uniform noise...]`` row per request.
Random noise is sampled independently from ``[0, vocab_size)``.
Supplying ``noise_tokens`` bypasses random generation. This is used
by deterministic tests and must have shape
``(batch_size, forward_width - 1)``.
``forward_width == 1`` never generates noise or consumes RNG state.
"""
seed_tokens = seed_tokens.reshape(-1).to(dtype=torch.int64)
batch_size = seed_tokens.numel()
noise_shape = (batch_size, forward_width - 1)
if forward_width == 1:
return seed_tokens[:, None]
if noise_tokens is None:
noise_tokens = torch.randint(
low=0,
high=vocab_size,
size=noise_shape,
dtype=torch.int64,
device=seed_tokens.device,
)
else:
noise_tokens = noise_tokens.to(
device=seed_tokens.device,
dtype=torch.int64,
)
draft_input_ids = seed_tokens.new_empty((batch_size, forward_width))
draft_input_ids[:, 0].copy_(seed_tokens)
draft_input_ids[:, 1:].copy_(noise_tokens)
return draft_input_ids
def sample_uno_candidates(
*,
draft_logits: torch.Tensor, # [B, F, V]
sampling_info: Any,
max_top_k: int,
uniform_top_k_value: int | None = None,
) -> tuple[torch.Tensor, UnoDraftDistribution]:
"""Sample 1 clean token from seed and F-1 draft tokens.
Depending on max_top_k value, may use sparse representations for efficiency.
candidates: [B, F] sampled tokens including clean and draft.
draft_distribution: probabilities of the draft tokens for rejection sampling.
"""
batch_size, forward_width, vocab_size = draft_logits.shape
flat_logits = draft_logits.reshape(-1, vocab_size)
if max_top_k <= _SPARSE_TOP_K_LIMIT:
support_ids, support_probs = _build_sparse_target_support(
next_token_logits=flat_logits,
sampling_info=sampling_info,
batch_size=batch_size,
forward_width=forward_width,
max_top_k=max_top_k,
)
candidates = _sample_from_support(support_ids, support_probs)
draft_distribution = UnoDraftDistribution(
token_ids=support_ids[:, 1:],
probs=support_probs[:, 1:],
)
else:
probs = _build_dense_probs(
next_token_logits=flat_logits,
sampling_info=sampling_info,
batch_size=batch_size,
forward_width=forward_width,
max_top_k=max_top_k,
uniform_top_k_value=uniform_top_k_value,
)
candidates = _sample_from_dense_probs(probs)
draft_distribution = UnoDraftDistribution(probs=probs[:, 1:])
return candidates, draft_distribution
def sample_uno_clean_root(
*,
seed_tokens: torch.Tensor,
draft_logits: torch.Tensor,
sampling_info: Any,
max_top_k: int,
uniform_top_k_value: int | None = None,
) -> torch.Tensor:
"""Sample the clean root using the current UNO sampling path."""
del seed_tokens
candidates, _ = sample_uno_candidates(
draft_logits=draft_logits[:, :1, :].contiguous(),
sampling_info=sampling_info,
max_top_k=max_top_k,
uniform_top_k_value=uniform_top_k_value,
)
return candidates[:, 0]
def sample_uno_tree_target_tokens(
*,
next_token_logits: torch.Tensor,
sampling_info: Any,
batch_size: int,
verify_width: int,
max_top_k: int,
) -> torch.Tensor:
"""Sample one target token per verify node from compact top-k support."""
support_ids, support_probs = _build_sparse_target_support(
next_token_logits=next_token_logits,
sampling_info=sampling_info,
batch_size=batch_size,
forward_width=verify_width,
max_top_k=max_top_k,
)
return _sample_from_support(support_ids, support_probs)
def pack_uno_tree_result(
*,
clean_root_tokens: torch.Tensor,
eagle_predict: torch.Tensor,
eagle_accept_lens: torch.Tensor,
draft_width: int,
) -> UnoTreeSamplingResult:
"""Convert an internal EAGLE tree result into UNO's public row."""
clean_root_tokens = clean_root_tokens.reshape(-1).to(dtype=torch.int64)
batch_size = clean_root_tokens.numel()
verify_width = eagle_predict.numel() // batch_size
predict_rows = eagle_predict.reshape(batch_size, verify_width)
output_ids = clean_root_tokens.new_zeros((batch_size, draft_width + 1))
output_ids[:, 0].copy_(clean_root_tokens)
output_ids[:, 1:].copy_(predict_rows[:, :draft_width].to(dtype=output_ids.dtype))
return UnoTreeSamplingResult(
output_ids=output_ids,
accept_lens=eagle_accept_lens + 1,
)
def pack_uno_result(
*,
candidates: torch.Tensor, # [B, F]
accepted_proposal_counts: torch.Tensor, # [B]
bonus_tokens: torch.Tensor, # [B]
committed_frontiers: torch.Tensor, # [B]
) -> UnoSamplingResult:
"""Pack acceptance into fixed-width UNO output rows."""
batch_size, forward_width = candidates.shape
output_ids = candidates.new_zeros((batch_size, forward_width + 1))
output_ids[:, :forward_width].copy_(candidates)
output_ids.scatter_(
1,
(accepted_proposal_counts.to(torch.long) + 1)[:, None],
bonus_tokens[:, None],
)
accept_lens = accepted_proposal_counts + 2
new_seq_lens = committed_frontiers + accept_lens.to(committed_frontiers.dtype)
return UnoSamplingResult(
output_ids=output_ids,
accept_lens=accept_lens,
new_seq_lens=new_seq_lens,
next_seed_tokens=bonus_tokens,
)
def run_uno_sampling(
*,
candidates: torch.Tensor, # [B, F]
next_token_logits: torch.Tensor, # [B x F, V]
sampling_info: Any,
committed_frontiers: torch.Tensor, # [B]
draft_distribution: UnoDraftDistribution,
max_top_k: int,
uniform_top_k_value: int | None = None,
) -> UnoSamplingResult:
"""Verify sampled UNO proposals against target p and pack the result."""
if draft_distribution.token_ids is not None:
accepted, bonus = _run_sparse_rejection(
candidates=candidates,
next_token_logits=next_token_logits,
sampling_info=sampling_info,
max_top_k=max_top_k,
draft_distribution=draft_distribution,
)
else:
accepted, bonus = _run_dense_rejection(
candidates=candidates,
next_token_logits=next_token_logits,
sampling_info=sampling_info,
draft_distribution=draft_distribution,
max_top_k=max_top_k,
uniform_top_k_value=uniform_top_k_value,
)
return pack_uno_result(
candidates=candidates,
accepted_proposal_counts=accepted,
bonus_tokens=bonus,
committed_frontiers=committed_frontiers,
)
python/sglang/srt/speculative/uno_validation.py
+50
View File
@@ -0,0 +1,50 @@
"""Request-admission validation for UNO speculative decoding."""
from __future__ import annotations
from typing import TYPE_CHECKING, Optional
if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req
def validate_uno_request(req: Req) -> Optional[str]:
"""Return an error for request features that UNO cannot execute."""
sampling_params = req.sampling_params
if sampling_params.min_p > 0.0:
return "UNO speculative decoding does not support min_p sampling."
has_grammar = req.grammar is not None or any(
getattr(sampling_params, field) is not None
for field in ("json_schema", "regex", "ebnf", "structural_tag")
)
if has_grammar:
return "UNO speculative decoding does not support grammar decoding."
if req.return_logprob:
return "UNO speculative decoding does not support returned logprobs."
if req.return_hidden_states_mode.need_capture():
return "UNO speculative decoding does not support return_hidden_states."
has_penalties = (
sampling_params.frequency_penalty != 0.0
or sampling_params.presence_penalty != 0.0
or sampling_params.repetition_penalty != 1.0
or sampling_params.min_new_tokens > 0
)
if has_penalties:
return "UNO speculative decoding does not support sampling penalties."
if sampling_params.logit_bias is not None:
return "UNO speculative decoding does not support logit_bias."
if req.custom_logit_processor:
return "UNO speculative decoding does not support custom logit processors."
if req.lora_id is not None:
return "UNO speculative decoding does not support request-selectable LoRA."
return None
python/sglang/srt/speculative/uno_worker_v2.py
+867
View File
@@ -0,0 +1,867 @@
from __future__ import annotations
import contextlib
import copy
import logging
import time
from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.managers.utils import GenerationBatchResult
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
ForwardBatch,
ForwardMode,
)
from sglang.srt.model_executor.forward_context import (
ForwardContext,
forward_context,
)
from sglang.srt.runtime_context import get_schedule, get_spec
from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
from sglang.srt.speculative.eagle_info import EagleDraftInput
from sglang.srt.speculative.eagle_utils import default_tree_mask_mode
from sglang.srt.speculative.eagle_worker_common import (
build_eagle_verify_input,
run_eagle_verify,
)
from sglang.srt.speculative.spec_info import SpecInputType, SpeculativeAlgorithm
from sglang.srt.speculative.spec_utils import get_plan_stream
from sglang.srt.speculative.uno_cuda_graph_runner import (
UnoDecodeCudaGraphRunner,
)
from sglang.srt.speculative.uno_info import UnoDraftInput, UnoForwardInput
from sglang.srt.speculative.uno_tree import build_uno_tree_proposal
from sglang.srt.speculative.uno_utils import (
build_uno_draft_input,
pack_uno_tree_result,
run_uno_sampling,
sample_uno_candidates,
sample_uno_clean_root,
)
from sglang.srt.utils.common import (
get_available_gpu_memory,
log_info_on_rank0,
)
if TYPE_CHECKING:
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.managers.tp_worker import TpModelWorker
from sglang.srt.server_args import ServerArgs
logger = logging.getLogger(__name__)
class UnoWorkerV2(BaseSpecWorker):
"""Single-model UNO worker with linear and native-EAGLE tree decode."""
def __init__(
self,
server_args: ServerArgs,
gpu_id: int,
ps: ParallelState,
nccl_port: int,
target_worker: TpModelWorker,
):
super().__init__()
self.server_args = server_args
self.gpu_id = gpu_id
self.ps = ps
self.nccl_port = nccl_port
self._target_worker = target_worker
self._draft_worker = None
self.model_runner = target_worker.model_runner
self.lora_manager = self.model_runner.lora_manager
self.uno_lora_id = self.model_runner.uno_lora_id
self.device = target_worker.device
self.enable_overlap = not get_schedule().disable_overlap_schedule
configured_topk = int(get_spec().speculative_eagle_topk or 1)
self.tree_mode = configured_topk > 1
# Linear UNO stores F in speculative_num_draft_tokens. Tree UNO reuses
# EAGLE's native dimensions: F=steps+1, K=eagle_topk, Q=draft_tokens.
default_forward_width = (
int(get_spec().speculative_num_steps) + 1
if self.tree_mode
else int(get_spec().speculative_num_draft_tokens)
)
self.forward_width = default_forward_width
self.verify_width = int(get_spec().speculative_num_draft_tokens)
self.candidate_top_k = (
int(get_spec().speculative_eagle_topk) if self.tree_mode else 1
)
self.tree_depth = int(get_spec().speculative_num_steps) if self.tree_mode else 1
self.num_speculative_proposals = self.forward_width - 1
self.tail_width = self.forward_width + 1
# Compatibility fields read by speculative infrastructure.
self.speculative_num_draft_tokens = self.verify_width
self.speculative_num_steps = self.tree_depth
self.topk = self.candidate_top_k
# Ordinary scheduler overlap still serializes model work on the
# forward stream, so one persistent proposal workspace is sufficient:
# the next reuse is ordered after this step's tree build and verify.
self._uno_tree_workspace = {} if self.tree_mode else None
# The scheduler constructs speculative workers before allocating the
# target KV pools. Build the private tree-draft backend later, from
# init_attention_backends(), after those pools exist.
self._uno_draft_attn_backend = None
self._uno_draft_cuda_graph_runner = None
self.plan_stream, self.plan_stream_ctx = (
get_plan_stream(self.device)
if self.tree_mode
else (None, contextlib.nullcontext())
)
self._tail_offsets = torch.arange(
self.tail_width,
dtype=torch.int64,
device=self.device,
)
def _build_uno_draft_attn_backend(self):
"""Build only the F/1 backend absent from the native Q/K target role."""
model_runner = self.model_runner
original_workspace_flag = model_runner.init_new_workspace
try:
with get_spec().override(
speculative_num_steps=1,
speculative_eagle_topk=1,
speculative_num_draft_tokens=self.forward_width,
):
return model_runner._get_attention_backend(init_new_workspace=True)
finally:
model_runner.init_new_workspace = original_workspace_flag
def init_attention_backends(self):
"""Initialize only UNO's private backend after target pool allocation."""
if self.tree_mode:
self._uno_draft_attn_backend = self._build_uno_draft_attn_backend()
def init_cuda_graphs(self):
"""Capture only the private F-wide tree-draft graph."""
self._uno_draft_cuda_graph_runner = None
if not self.tree_mode or self.model_runner.decode_cuda_graph_runner is None:
return None
if self._uno_draft_attn_backend is None:
raise RuntimeError(
"UNO tree draft graph capture requires its attention backend."
)
tic = time.perf_counter()
before_mem = get_available_gpu_memory(
self.device,
self.gpu_id,
empty_cache=False,
)
log_info_on_rank0(
logger,
"Capture UNO tree draft CUDA graph begin. "
f"num_tokens_per_req={self.forward_width}, "
f"avail mem={before_mem:.2f} GB",
)
with self._bind_uno_draft_runtime():
self._uno_draft_cuda_graph_runner = UnoDecodeCudaGraphRunner(
self.model_runner,
tree_draft_attn_backend=self._uno_draft_attn_backend,
tree_draft_width=self.forward_width,
)
after_mem = get_available_gpu_memory(
self.device,
self.gpu_id,
empty_cache=False,
)
capture_time = time.perf_counter() - tic
self._additional_graph_memory_usage["draft_decode"] = before_mem - after_mem
self._additional_graph_time_usage["draft_decode"] = capture_time
log_info_on_rank0(
logger,
"Capture UNO tree draft CUDA graph end. "
f"elapsed={capture_time:.2f} s, "
f"mem usage={(before_mem - after_mem):.2f} GB, "
f"avail mem={after_mem:.2f} GB.",
)
return None
@property
def draft_worker(self):
# Both passes use the target runner and its KV pool.
return None
@property
def last_shared_read_runner(self):
# The target verify is the final phase that reads shared scheduler
# buffers, so its runner owns the WAR-barrier completion event.
return self._target_worker.model_runner
@property
def spec_v2_attn_backends(self) -> tuple:
"""Return every attention backend touched by one UNO step.
Linear UNO uses only the target runner's native backend. Tree UNO adds
one private F-wide draft backend before finishing on the native Q-wide
target backend. The scheduler ORs these capabilities when deciding
whether FutureMap must carry a CPU sequence-length mirror.
"""
target_backend = self._target_worker.model_runner.attn_backend
if not self.tree_mode:
return (target_backend,)
return (target_backend, self._uno_draft_attn_backend)
def __getattr__(self, name):
# Scheduler-facing methods not implemented by this wrapper belong to
# the target worker. Guard initialization to avoid recursive lookup.
if name == "_target_worker":
raise AttributeError(name)
return getattr(self.target_worker, name)
def _validate_batch(self, batch: ScheduleBatch) -> None:
if batch.forward_mode.is_idle():
raise NotImplementedError("UNO does not support idle batches.")
if batch.forward_mode.is_mixed() or (
batch.forward_mode.is_decode() and batch.is_extend_in_batch
):
raise NotImplementedError(
"UNO does not support mixed extend/decode batches."
)
if not batch.spec_algorithm.is_uno():
raise RuntimeError(
"UnoWorkerV2 received a batch whose speculative algorithm is not UNO."
)
sampling_info = batch.sampling_info
if sampling_info is None:
raise RuntimeError("UNO requires sampling metadata.")
if sampling_info.need_min_p_sampling:
raise NotImplementedError("UNO does not support min-p sampling.")
if batch.has_grammar:
raise NotImplementedError("UNO does not support grammar decoding.")
if batch.return_logprob:
raise NotImplementedError("UNO does not support returned logprobs.")
if batch.return_hidden_states:
raise NotImplementedError("UNO does not support returned hidden states.")
penalizer = sampling_info.penalizer_orchestrator
penalties_active = (
(penalizer is not None and penalizer.is_required)
or sampling_info.acc_additive_penalties is not None
or sampling_info.acc_scaling_penalties is not None
)
if penalties_active:
raise NotImplementedError("UNO does not support sampling penalties.")
if sampling_info.logit_bias is not None:
raise NotImplementedError("UNO does not support logit bias.")
if sampling_info.has_custom_logit_processor:
raise NotImplementedError("UNO does not support custom logit processors.")
if any(req.lora_id is not None for req in batch.reqs):
raise NotImplementedError("UNO does not support multi-LoRA.")
def _make_forward_batch(
self,
*,
spec_input_type: SpecInputType,
input_ids: torch.Tensor,
positions: torch.Tensor,
out_cache_loc: torch.Tensor,
prefix_lens: torch.Tensor,
seq_lens_cpu: Optional[torch.Tensor],
seq_lens_sum: Optional[int],
req_pool_indices: torch.Tensor,
) -> ForwardBatch:
input_ids = input_ids.reshape(-1)
positions = positions.reshape(-1)
out_cache_loc = out_cache_loc.reshape(-1)
spec_info = UnoForwardInput(
spec_input_type=spec_input_type,
positions=positions,
draft_token_num=self.forward_width,
)
return ForwardBatch(
forward_mode=ForwardMode.TARGET_VERIFY,
batch_size=len(prefix_lens),
input_ids=input_ids,
req_pool_indices=req_pool_indices,
seq_lens=prefix_lens,
out_cache_loc=out_cache_loc,
seq_lens_sum=seq_lens_sum,
seq_lens_cpu=seq_lens_cpu,
positions=positions,
spec_algorithm=SpeculativeAlgorithm.UNO,
spec_info=spec_info,
capture_hidden_mode=CaptureHiddenMode.NULL,
return_hidden_states_before_norm=False,
)
def _run_target_block(
self,
forward_batch: ForwardBatch,
*,
need_top1: bool = True,
) -> tuple:
result = self.target_worker.forward_batch_generation(
batch=None,
forward_batch=forward_batch,
is_verify=True,
)
if result.logits_output is None:
raise RuntimeError("UNO target block returned no logits output.")
logits = result.logits_output.next_token_logits
if logits is None:
raise RuntimeError("UNO target block returned no next-token logits.")
expected_rows = forward_batch.batch_size * self.forward_width
if logits.ndim != 2 or logits.shape[0] != expected_rows:
raise RuntimeError(
"UNO target block returned an invalid logits shape: "
f"expected ({expected_rows}, vocab_size), got "
f"{tuple(logits.shape)}."
)
# DFlash consumes logits directly; only greedy acceptance needs top-1.
if not need_top1:
return result, None
predictions = torch.argmax(logits, dim=-1).view(
forward_batch.batch_size,
self.forward_width,
)
return result, predictions
@staticmethod
def _accept_and_pack(
*,
candidates: torch.Tensor,
target_top1: torch.Tensor,
committed_seq_lens: torch.Tensor,
) -> tuple:
if candidates.ndim != 2:
raise RuntimeError(
f"UNO candidates must be rank 2, got shape={tuple(candidates.shape)}."
)
if target_top1.shape != candidates.shape:
raise RuntimeError(
"UNO candidate and target shapes differ: "
f"{tuple(candidates.shape)} versus {tuple(target_top1.shape)}."
)
batch_size, forward_width = candidates.shape
device = candidates.device
# forward_width is static configuration, so this branch does not inspect
# or synchronize a device tensor.
if forward_width == 1:
accepted_specs = torch.zeros(
batch_size,
dtype=torch.int32,
device=device,
)
else:
matches = candidates[:, 1:] == target_top1[:, :-1]
accepted_specs = (
matches.to(torch.int32).cumprod(dim=1).sum(dim=1).to(torch.int32)
)
accepted_specs_long = accepted_specs.to(torch.int64)
correction = target_top1.gather(
1,
accepted_specs_long[:, None],
).squeeze(1)
output_ids = torch.zeros(
(batch_size, forward_width + 1),
dtype=torch.int64,
device=device,
)
output_ids[:, :forward_width].copy_(candidates)
output_ids.scatter_(
1,
(accepted_specs_long + 1)[:, None],
correction[:, None],
)
accept_lens = accepted_specs + 2
new_seq_lens = committed_seq_lens + accept_lens.to(committed_seq_lens.dtype)
return output_ids, accept_lens, new_seq_lens, correction
def _forward_prefill(
self,
batch: ScheduleBatch,
on_publish,
) -> GenerationBatchResult:
result = self.target_worker.forward_batch_generation(batch)
if not isinstance(result.next_token_ids, torch.Tensor):
raise RuntimeError("UNO target prefill returned no sampled seed tensor.")
seed_tokens = result.next_token_ids.reshape(-1)
if seed_tokens.shape[0] != len(batch.reqs):
raise RuntimeError(
"UNO prefill seed count does not match batch size: "
f"{seed_tokens.shape[0]} versus {len(batch.reqs)}."
)
result.new_seq_lens = batch.seq_lens
result.next_draft_input = UnoDraftInput(
bonus_tokens=seed_tokens,
new_seq_lens=batch.seq_lens,
forward_width=self.forward_width,
)
if on_publish is not None:
on_publish(result.new_seq_lens)
return result
def _forward_decode_tree(
self,
batch: ScheduleBatch,
on_publish,
) -> GenerationBatchResult:
"""Run UNO's F-wide proposal pass, then native EAGLE Q-node verify."""
draft_state = batch.spec_info
if batch.seq_lens.is_cuda:
batch.seq_lens.record_stream(
torch.get_device_module(self.device).current_stream()
)
batch_size = len(batch.seq_lens)
committed_seq_lens = batch.seq_lens.clone()
seed_tokens = draft_state.bonus_tokens.reshape(-1).to(
device=self.device,
dtype=torch.int64,
)
committed_seq_lens_cpu = None
if batch.seq_lens_cpu is not None:
committed_seq_lens_cpu = batch.seq_lens_cpu.to(
device="cpu",
dtype=torch.int64,
)
draft_seq_lens_cpu = committed_seq_lens_cpu + self.forward_width
draft_seq_lens_sum = int(draft_seq_lens_cpu.sum())
elif draft_state.reserved_seq_lens_cpu is not None:
# This host tensor is a planning upper bound only. The device
# frontier below remains the exact committed length.
draft_seq_lens_cpu = draft_state.reserved_seq_lens_cpu
draft_seq_lens_sum = draft_state.reserved_seq_lens_sum
else:
draft_seq_lens_cpu = None
draft_seq_lens_sum = None
draft_positions = (
committed_seq_lens.to(torch.int64)[:, None]
+ (self._tail_offsets[None, : self.forward_width])
)
req_pool_indices_long = batch.req_pool_indices.to(torch.int64)
req_to_token = self.model_runner.req_to_token_pool.req_to_token
draft_locs = req_to_token[
req_pool_indices_long[:, None],
draft_positions,
].to(torch.int64)
draft_input_ids = build_uno_draft_input(
seed_tokens=seed_tokens,
forward_width=self.forward_width,
vocab_size=self.model_runner.model_config.vocab_size,
)
draft_forward_batch = self._make_forward_batch(
spec_input_type=SpecInputType.UNO_DRAFT,
input_ids=draft_input_ids,
positions=draft_positions,
out_cache_loc=draft_locs,
prefix_lens=committed_seq_lens,
seq_lens_cpu=draft_seq_lens_cpu,
seq_lens_sum=draft_seq_lens_sum,
req_pool_indices=batch.req_pool_indices,
)
draft_result, _ = self._run_draft_block(
draft_forward_batch,
need_top1=False,
)
draft_logits = draft_result.logits_output.next_token_logits.reshape(
batch_size,
self.forward_width,
-1,
)
sampling_info = batch.sampling_info
if sampling_info.is_all_greedy:
clean_root_tokens = torch.argmax(
draft_logits[:, 0, :],
dim=-1,
)
else:
clean_root_tokens = sample_uno_clean_root(
seed_tokens=seed_tokens,
draft_logits=draft_logits,
sampling_info=sampling_info,
max_top_k=draft_state.max_top_k,
uniform_top_k_value=draft_state.uniform_top_k_value,
)
proposal = build_uno_tree_proposal(
clean_root_tokens,
draft_logits[:, 1:, :],
max_nodes=self.verify_width,
candidate_top_k=self.candidate_top_k,
temperature=sampling_info.temperatures,
workspace=self._uno_tree_workspace,
)
# The first pass wrote the carried seed at C. Give EAGLE a shallow
# batch whose KV-ready prefix is therefore C+1; its existing allocator
# assigns all Q tree slots at C+1 and its compactor can stay unchanged.
verify_batch = copy.copy(batch)
verify_batch.seq_lens = committed_seq_lens + 1
if committed_seq_lens_cpu is None:
verify_batch.seq_lens_cpu = None
verify_batch.seq_lens_sum = None
else:
verify_batch.seq_lens_cpu = committed_seq_lens_cpu + 1
verify_batch.seq_lens_sum = int(verify_batch.seq_lens_cpu.sum())
verify_input = build_eagle_verify_input(
verify_batch,
EagleDraftInput(bonus_tokens=proposal.root_tokens),
proposal.parent_list,
proposal.top_scores_index,
proposal.draft_tokens,
None,
target_worker=self.target_worker,
topk=self.candidate_top_k,
num_steps=self.tree_depth,
num_draft_tokens=self.verify_width,
tree_mask_mode=default_tree_mask_mode(),
device=self.device,
)
verify_batch.spec_info = verify_input
if self.plan_stream is not None:
# C+1 was produced on the forward stream immediately above. The
# generic EAGLE path receives an older, already-visible frontier;
# UNO must explicitly order its freshly derived tensor before the
# plan stream assigns Q verify cache locations from it.
self.plan_stream.wait_stream(
torch.get_device_module(self.device).current_stream()
)
eagle_result = run_eagle_verify(
verify_batch,
target_worker=self.target_worker,
req_to_token_pool=self.model_runner.req_to_token_pool,
token_to_kv_pool_allocator=(self.model_runner.token_to_kv_pool_allocator),
plan_stream=self.plan_stream,
plan_stream_ctx=self.plan_stream_ctx,
topk=self.candidate_top_k,
num_draft_tokens=self.verify_width,
device=self.device,
metadata_ready_pre_pad=False,
finalize_tree_path=True,
uno_target_max_top_k=draft_state.max_top_k,
)
packed = pack_uno_tree_result(
clean_root_tokens=clean_root_tokens,
eagle_predict=eagle_result.next_token_ids,
eagle_accept_lens=eagle_result.accept_lens,
draft_width=self.forward_width,
)
new_seq_lens = eagle_result.new_seq_lens
next_draft_input = UnoDraftInput(
bonus_tokens=eagle_result.next_draft_input.bonus_tokens,
new_seq_lens=new_seq_lens,
forward_width=self.forward_width,
)
if on_publish is not None:
on_publish(new_seq_lens)
# Preserve EAGLE's verify ForwardBatch keep-alive refs verbatim. FutureMap
# relays the wrapped bonus; on_publish above relays the new frontier.
return GenerationBatchResult(
logits_output=eagle_result.logits_output,
next_token_ids=packed.output_ids.reshape(-1),
accept_lens=packed.accept_lens,
next_draft_input=next_draft_input,
speculative_num_draft_tokens=self.forward_width,
speculative_output_stride=self.forward_width + 1,
num_non_draft_tokens_per_req=2,
new_seq_lens=new_seq_lens,
can_run_cuda_graph=eagle_result.can_run_cuda_graph,
routed_experts_output=eagle_result.routed_experts_output,
indexer_topk_output=eagle_result.indexer_topk_output,
extra_keep_alive_refs=eagle_result.extra_keep_alive_refs,
)
def _forward_decode(
self,
batch: ScheduleBatch,
on_publish,
) -> GenerationBatchResult:
if self.tree_mode:
return self._forward_decode_tree(batch, on_publish)
draft_state = batch.spec_info
if batch.seq_lens.is_cuda:
batch.seq_lens.record_stream(
torch.get_device_module(self.device).current_stream()
)
batch_size = len(batch.seq_lens)
committed_seq_lens = batch.seq_lens.clone()
seed_tokens = draft_state.bonus_tokens.reshape(-1).to(
device=self.device,
dtype=torch.int64,
)
if batch.seq_lens_cpu is not None:
committed_seq_lens_cpu = batch.seq_lens_cpu.to(
device="cpu",
dtype=torch.int32,
)
draft_seq_lens_cpu = committed_seq_lens_cpu + self.forward_width
verify_seq_lens_cpu = committed_seq_lens_cpu + self.tail_width
draft_seq_lens_sum = int(draft_seq_lens_cpu.sum())
verify_seq_lens_sum = int(verify_seq_lens_cpu.sum())
elif draft_state.reserved_seq_lens_cpu is not None:
# Triton only needs a safe host planning bound. The allocator's
# retained reservation avoids a D2H copy when FutureMap keeps the
# exact committed frontier on GPU.
draft_seq_lens_cpu = draft_state.reserved_seq_lens_cpu
verify_seq_lens_cpu = draft_state.reserved_seq_lens_cpu
draft_seq_lens_sum = draft_state.reserved_seq_lens_sum
verify_seq_lens_sum = draft_state.reserved_seq_lens_sum
else:
draft_seq_lens_cpu = None
verify_seq_lens_cpu = None
draft_seq_lens_sum = None
verify_seq_lens_sum = None
logical_positions = (
committed_seq_lens.to(torch.int64)[:, None] + (self._tail_offsets[None, :])
)
req_pool_indices_long = batch.req_pool_indices.to(torch.int64)
req_to_token = self.model_runner.req_to_token_pool.req_to_token
tail_locs = req_to_token[
req_pool_indices_long[:, None],
logical_positions,
].to(torch.int64)
draft_input_ids = build_uno_draft_input(
seed_tokens=seed_tokens,
forward_width=self.forward_width,
vocab_size=self.model_runner.model_config.vocab_size,
)
draft_forward_batch = self._make_forward_batch(
spec_input_type=SpecInputType.UNO_DRAFT,
input_ids=draft_input_ids,
positions=logical_positions[:, : self.forward_width],
out_cache_loc=tail_locs[:, : self.forward_width],
prefix_lens=committed_seq_lens,
seq_lens_cpu=draft_seq_lens_cpu,
seq_lens_sum=draft_seq_lens_sum,
req_pool_indices=batch.req_pool_indices,
)
sampling_info = batch.sampling_info
all_greedy = sampling_info.is_all_greedy
max_top_k = draft_state.max_top_k
uniform_top_k_value = draft_state.uniform_top_k_value
draft_result, candidates = self._run_draft_block(
draft_forward_batch,
need_top1=all_greedy,
)
if not all_greedy:
draft_logits = draft_result.logits_output.next_token_logits.reshape(
batch_size,
self.forward_width,
-1,
)
candidates, draft_distribution = sample_uno_candidates(
draft_logits=draft_logits,
sampling_info=sampling_info,
max_top_k=max_top_k,
uniform_top_k_value=uniform_top_k_value,
)
verify_prefix_lens = committed_seq_lens + 1
verify_forward_batch = self._make_forward_batch(
spec_input_type=SpecInputType.UNO_VERIFY,
input_ids=candidates,
positions=logical_positions[:, 1:],
out_cache_loc=tail_locs[:, 1:],
prefix_lens=verify_prefix_lens,
seq_lens_cpu=verify_seq_lens_cpu,
seq_lens_sum=verify_seq_lens_sum,
req_pool_indices=batch.req_pool_indices,
)
verify_result, target_top1 = self._run_target_block(
verify_forward_batch,
need_top1=all_greedy,
)
if all_greedy:
output_ids, accept_lens, new_seq_lens, correction = self._accept_and_pack(
candidates=candidates,
target_top1=target_top1,
committed_seq_lens=committed_seq_lens,
)
else:
sampling_result = run_uno_sampling(
candidates=candidates,
next_token_logits=verify_result.logits_output.next_token_logits,
sampling_info=sampling_info,
committed_frontiers=committed_seq_lens,
draft_distribution=draft_distribution,
max_top_k=max_top_k,
uniform_top_k_value=uniform_top_k_value,
)
output_ids = sampling_result.output_ids
accept_lens = sampling_result.accept_lens
new_seq_lens = sampling_result.new_seq_lens
correction = sampling_result.next_seed_tokens
next_draft_input = UnoDraftInput(
bonus_tokens=correction,
new_seq_lens=new_seq_lens,
forward_width=self.forward_width,
)
if on_publish is not None:
on_publish(new_seq_lens)
return GenerationBatchResult(
logits_output=verify_result.logits_output,
next_token_ids=output_ids.reshape(-1),
accept_lens=accept_lens,
next_draft_input=next_draft_input,
speculative_num_draft_tokens=self.forward_width,
speculative_output_stride=self.tail_width,
num_non_draft_tokens_per_req=2,
new_seq_lens=new_seq_lens,
can_run_cuda_graph=verify_result.can_run_cuda_graph,
routed_experts_output=verify_result.routed_experts_output,
indexer_topk_output=verify_result.indexer_topk_output,
)
def forward_batch_generation(
self,
batch: ScheduleBatch,
on_publish=None,
grammar_barrier=None,
) -> GenerationBatchResult:
del grammar_barrier
self._validate_batch(batch)
if batch.forward_mode == ForwardMode.EXTEND:
return self._forward_prefill(batch, on_publish)
if batch.forward_mode == ForwardMode.DECODE:
return self._forward_decode(batch, on_publish)
raise RuntimeError(
f"UNO expected an EXTEND or DECODE batch, got {batch.forward_mode}."
)
def update_weights_from_disk(self, recv_req):
# The scheduler updates the target worker before calling the spec worker.
return True, "UNO has no separate draft weights."
def update_weights_from_ipc(self, recv_req):
# The scheduler updates the target worker before calling the spec worker.
return True, "UNO has no separate draft weights."
def update_weights_from_tensor(self, recv_req):
# This update route selects the spec worker instead of updating both.
return self.target_worker.update_weights_from_tensor(recv_req)
@contextlib.contextmanager
def _bind_uno_draft_runtime(self):
target_attn_backend = self.model_runner.attn_backend
target_graph_runner = self.model_runner.decode_cuda_graph_runner
self.model_runner.attn_backend = self._uno_draft_attn_backend
self.model_runner.decode_cuda_graph_runner = self._uno_draft_cuda_graph_runner
try:
yield
finally:
self.model_runner.attn_backend = target_attn_backend
self.model_runner.decode_cuda_graph_runner = target_graph_runner
def _run_draft_block(
self,
forward_batch: ForwardBatch,
*,
need_top1: bool = True,
):
batch_size = forward_batch.batch_size
self.lora_manager.reset_lora_batch()
backend_context = (
self._bind_uno_draft_runtime()
if self.tree_mode
else contextlib.nullcontext()
)
attn_context = (
forward_context(ForwardContext(attn_backend=self._uno_draft_attn_backend))
if self.tree_mode
else contextlib.nullcontext()
)
# The eager runner plans through model_runner.attn_backend, while
# attention layers execute through ForwardContext. Tree UNO binds both
# to the same private F/1 backend for this draft pass.
with backend_context, attn_context:
if self.num_speculative_proposals == 0:
return self._run_target_block(
forward_batch,
need_top1=need_top1,
)
graph_runner = getattr(
self.model_runner,
"decode_cuda_graph_runner",
None,
)
# If cuda-graph is on, reuse its captured LoRA routing
# and replay it in _run_target_block.
# Else, prepare routing for eager.
if not (
forward_batch.forward_mode.is_cuda_graph()
and graph_runner is not None
and graph_runner.can_run_graph(forward_batch)
):
self.lora_manager.prepare_lora_token_segments(
lora_ids=[None, self.uno_lora_id] * batch_size,
segment_lens=[1, self.num_speculative_proposals] * batch_size,
)
result = self._run_target_block(
forward_batch,
need_top1=need_top1,
)
# Clear LoRA routing before verification
self.lora_manager.reset_lora_batch()
return result