from __future__ import annotations import logging import os import time from contextlib import contextmanager from typing import TYPE_CHECKING, List, Optional import torch import triton import triton.language as tl from huggingface_hub import snapshot_download from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject from sglang.srt.distributed.parallel_state import ( GroupCoordinator, patch_tensor_parallel_group, ) from sglang.srt.environ import envs from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.common import get_last_loc from sglang.srt.server_args import ServerArgs, get_global_server_args from sglang.srt.utils import is_cuda, is_hip, is_npu, next_power_of_2 _is_cuda = is_cuda() _is_hip = is_hip() _is_npu = is_npu() if TYPE_CHECKING: from sglang.srt.speculative.eagle_info import EagleVerifyInput if _is_cuda: from sgl_kernel import fast_topk elif _is_hip: from sgl_kernel import fast_topk else: from sglang.srt.utils.common import fast_topk logger = logging.getLogger(__name__) # Simulate acceptance length for benchmarking purposes SIMULATE_ACC_LEN = envs.SGLANG_SIMULATE_ACC_LEN.get() # turn off if < 0 SIMULATE_ACC_METHOD = envs.SGLANG_SIMULATE_ACC_METHOD.get() TREE_TRAVERSE_TIME_THRESHOLD = 1 # TODO: set this properly TREE_SPEC_KERNEL_AVAILABLE = _is_cuda # This kernel is only available for CUDA now def spec_need_hidden_states(server_args: Optional[ServerArgs] = None) -> bool: if server_args is None: server_args = get_global_server_args() # TODO(lsyin): also skip when 1) step = 1 or 2) standalone draft model return not server_args.enable_multi_layer_eagle @triton.jit def create_extend_after_decode_spec_info( verified_id, seq_lens, accept_lens, positions, new_verified_id, bs_upper: tl.constexpr, ): pid = tl.program_id(axis=0) offsets = tl.arange(0, bs_upper) seq_length = tl.load(seq_lens + pid) accept_length = tl.load(accept_lens + pid) accept_len_cumsum = tl.sum( tl.load(accept_lens + offsets, mask=offsets < pid, other=0) ) positions_ptr = positions + accept_len_cumsum mask = offsets < accept_length tl.store(positions_ptr + offsets, seq_length - accept_length + offsets, mask) accept_len_cumsum += accept_length - 1 verified_id_data = tl.load(verified_id + accept_len_cumsum) tl.store(new_verified_id + pid, verified_id_data) @triton.jit def assign_req_to_token_pool( req_pool_indices, req_to_token, start_offset, end_offset, out_cache_loc, pool_len: tl.constexpr, bs_upper: tl.constexpr, ): BLOCK_SIZE: tl.constexpr = 32 pid = tl.program_id(axis=0) kv_start = tl.load(start_offset + pid) kv_end = tl.load(end_offset + pid) token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len length_offset = tl.arange(0, bs_upper) start = tl.load(start_offset + length_offset, mask=length_offset < pid, other=0) end = tl.load(end_offset + length_offset, mask=length_offset < pid, other=0) out_offset = tl.sum(end - start, axis=0) out_cache_ptr = out_cache_loc + out_offset save_offset = tl.arange(0, BLOCK_SIZE) + kv_start load_offset = tl.arange(0, BLOCK_SIZE) num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE) for _ in range(num_loop): mask = save_offset < kv_end data = tl.load(out_cache_ptr + load_offset, mask=mask) tl.store(token_pool + save_offset, data, mask=mask) save_offset += BLOCK_SIZE load_offset += BLOCK_SIZE def assign_req_to_token_pool_func( req_pool_indices: torch.Tensor, req_to_token: torch.Tensor, start_offset: torch.Tensor, end_offset: torch.Tensor, out_cache_loc: torch.Tensor, batch_size: int, ): assign_req_to_token_pool[(batch_size,)]( req_pool_indices, req_to_token, start_offset, end_offset, out_cache_loc, req_to_token.shape[1], next_power_of_2(batch_size), ) @triton.jit def assign_draft_cache_locs( req_pool_indices, req_to_token, seq_lens, extend_lens, num_new_pages_per_topk, out_cache_loc, source_cache_loc, target_cache_loc, last_page_lens_cumsum, duplicate_cache_len: tl.constexpr, pool_len: tl.constexpr, topk: tl.constexpr, speculative_num_steps: tl.constexpr, page_size: tl.constexpr, bs_upper: tl.constexpr, iter_upper: tl.constexpr, ): BLOCK_SIZE: tl.constexpr = 128 pid = tl.program_id(axis=0) if page_size == 1 or topk == 1: copy_len = topk * speculative_num_steps out_cache_ptr = out_cache_loc + pid * topk * speculative_num_steps else: bs_offset = tl.arange(0, bs_upper) copy_len = tl.load(extend_lens + pid) cum_copy_len = tl.sum(tl.load(extend_lens + bs_offset, mask=bs_offset < pid)) out_cache_ptr = out_cache_loc + cum_copy_len # Part 1: Copy from out_cache_loc to req_to_token kv_start = tl.load(seq_lens + pid) token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len num_loop = tl.cdiv(copy_len, BLOCK_SIZE) for i in range(num_loop): copy_offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE mask = copy_offset < copy_len data = tl.load(out_cache_ptr + copy_offset, mask=mask) tl.store(token_pool + kv_start + copy_offset, data, mask=mask) if page_size != 1 and topk != 1 and duplicate_cache_len > 0: # Part 2: Copy indices into source_cache_loc and target_cache_loc # Expected output: src:[8,9,10,8,9,10...] tgt:[16,17,18,24,25,26...] prefix_len = tl.load(seq_lens + pid) last_page_len = prefix_len % page_size offsets = tl.arange(0, page_size) mask = offsets < last_page_len num_new_pages_per_topk_ = tl.load(num_new_pages_per_topk + pid) prefix_base = token_pool + prefix_len - last_page_len src_indices = tl.load(prefix_base + offsets, mask=mask) last_page_lens_cumsum_ = tl.load(last_page_lens_cumsum + pid) # Skip the first one since no copy is needed for topk_id in range(1, topk): tl.store( source_cache_loc + (topk - 1) * (last_page_lens_cumsum_ - last_page_len) + (topk_id - 1) * last_page_len + offsets, src_indices, mask=mask, ) tgt_indices = tl.load( prefix_base + topk_id * num_new_pages_per_topk_ * page_size + offsets, mask=mask, ) tl.store( target_cache_loc + (topk - 1) * (last_page_lens_cumsum_ - last_page_len) + (topk_id - 1) * last_page_len + offsets, tgt_indices, mask=mask, ) # Part 3: Copy and remove the used indices for duplication # speculative_num_steps=5, page_size=4, num_new_pages_per_topk_=2, last_page_len=1 # - xxxxx .. | - xxxxx .. | # topk=0 topk=1 # "-" means prefix tokens # "x" means speculative draft tokens # "." means padded tokens # we only want to copy the "x" part. iter_offset = tl.arange(0, iter_upper) for topk_id in range(topk): mask_upper = iter_offset < (speculative_num_steps + last_page_len) mask_lower = iter_offset >= last_page_len combined_mask = mask_upper & mask_lower indices = tl.load( prefix_base + topk_id * num_new_pages_per_topk_ * page_size + iter_offset, mask=combined_mask, other=0, ) # Shift from previous batches ptr_offset = pid * speculative_num_steps * topk # Subtract last_page_len to fill the gap of duplicated last page tokens. # For example, token pool is (1, 2, 3, 4 ,5) and last page is 1, # we write 2, 3, 4 to the front of out_cache_loc. tl.store( out_cache_loc + ptr_offset + topk_id * speculative_num_steps - last_page_len + iter_offset, indices, mask=combined_mask, ) @triton.jit def generate_draft_decode_kv_indices( req_pool_indices, req_to_token, paged_kernel_lens, kv_indices, kv_indptr, positions, pool_len: tl.constexpr, kv_indices_stride: tl.constexpr, kv_indptr_stride: tl.constexpr, bs_upper: tl.constexpr, iter_upper: tl.constexpr, num_tokens_upper: tl.constexpr, page_size: tl.constexpr, ): BLOCK_SIZE: tl.constexpr = 128 iters = tl.program_id(axis=0) bid = tl.program_id(axis=1) topk_id = tl.program_id(axis=2) num_steps = tl.num_programs(axis=0) num_seqs = tl.num_programs(axis=1) topk = tl.num_programs(axis=2) kv_indices += kv_indices_stride * iters kv_indptr += kv_indptr_stride * iters iters += 1 load_offset = tl.arange(0, bs_upper) seq_lens = tl.load(paged_kernel_lens + load_offset, mask=load_offset < bid, other=0) seq_len = tl.load(paged_kernel_lens + bid) cum_seq_len = tl.sum(seq_lens) # Update kv_indices kv_offset = cum_seq_len * topk + bid * iters * topk + topk_id * (seq_len + iters) kv_ptr = kv_indices + kv_offset token_pool_ptr = req_to_token + tl.load(req_pool_indices + bid) * pool_len kv_offset = tl.arange(0, BLOCK_SIZE) num_loop = tl.cdiv(seq_len, BLOCK_SIZE) for _ in range(num_loop): mask = kv_offset < seq_len data = tl.load(token_pool_ptr + kv_offset, mask=mask) tl.store(kv_ptr + kv_offset, data, mask=mask) kv_offset += BLOCK_SIZE extend_offset = tl.arange(0, iter_upper) if page_size == 1 or topk == 1: extend_data = tl.load( token_pool_ptr + seq_len + topk_id * num_steps + tl.arange(0, iter_upper), mask=extend_offset < iters, ) else: prefix_len = seq_len last_page_len = prefix_len % page_size num_new_pages_per_topk = ( last_page_len + num_steps + page_size - 1 ) // page_size prefix_base = seq_len // page_size * page_size start = ( prefix_base + topk_id * num_new_pages_per_topk * page_size + last_page_len ) extend_data = tl.load( token_pool_ptr + start + extend_offset, mask=extend_offset < iters, ) tl.store(kv_ptr + seq_len + extend_offset, extend_data, mask=extend_offset < iters) # Update kv_indptr bs_offset = tl.arange(0, num_tokens_upper) zid = bid * topk + topk_id if zid == 0: zid = num_seqs * topk positions = tl.load(positions + bs_offset, mask=bs_offset < zid, other=0) base = tl.sum(positions) tl.store(kv_indptr + zid, base + zid * iters) @triton.jit def align_evict_mask_to_page_size( seq_lens, evict_mask, page_size: tl.constexpr, num_draft_tokens: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): t_range = tl.arange(0, BLOCK_SIZE) bid = tl.program_id(axis=0) seq_len = tl.load(seq_lens + bid) io_mask = t_range < num_draft_tokens mask_row = tl.load( evict_mask + bid * num_draft_tokens + t_range, mask=io_mask, other=0 ) num_trues = tl.sum(mask_row) num_false = num_draft_tokens - num_trues start = (seq_len + num_false - 1) // page_size * page_size - seq_len for i in range(max(start, 0), min(start + page_size, num_draft_tokens)): tl.store(evict_mask + bid * num_draft_tokens + i, False) @triton.jit def get_target_cache_loc( tgt_cache_loc, to_free_slots, accept_length, to_free_num_slots, out_cache_loc, num_verify_tokens: tl.constexpr, num_verify_tokens_upper: tl.constexpr, bs_upper: tl.constexpr, ): bid = tl.program_id(axis=0) offset = tl.arange(0, num_verify_tokens_upper) bs_offset = tl.arange(0, bs_upper) # write the first part to tgt_cache_loc accept_len_all = tl.load(accept_length + bs_offset, mask=bs_offset < bid) tgt_cache_loc_start = tl.sum(accept_len_all) + bid copy_len = tl.load(accept_length + bid) + 1 out_cache_loc_row = tl.load( out_cache_loc + bid * num_verify_tokens + offset, mask=offset < copy_len ) tl.store( tgt_cache_loc + tgt_cache_loc_start + offset, out_cache_loc_row, mask=offset < copy_len, ) # write the second part to to_free_num_pages to_free_num_slots_all = tl.load(to_free_num_slots + bs_offset, mask=bs_offset < bid) to_free_num_slots_cur = tl.load(to_free_num_slots + bid) out_cache_loc_start = num_verify_tokens - to_free_num_slots_cur to_free_slots_start = tl.sum(to_free_num_slots_all) copy_len = to_free_num_slots_cur out_cache_loc_row = tl.load( out_cache_loc + bid * num_verify_tokens + out_cache_loc_start + offset, mask=offset < copy_len, ) tl.store( to_free_slots + to_free_slots_start + offset, out_cache_loc_row, mask=offset < copy_len, ) @torch.compile(dynamic=True, disable=_is_npu) def get_src_tgt_cache_loc( seq_lens: torch.Tensor, out_cache_loc: torch.Tensor, accept_index: torch.Tensor, accept_length: torch.Tensor, draft_token_num: int, page_size: int, ): src_cache_loc = out_cache_loc[accept_index] tgt_cache_loc = torch.empty_like(src_cache_loc) extended_len = seq_lens + draft_token_num keep_len = torch.minimum( (seq_lens + accept_length + 1 + page_size - 1) // page_size * page_size, extended_len, ) to_free_num_slots = extended_len - keep_len return src_cache_loc, tgt_cache_loc, to_free_num_slots @triton.jit def filter_finished_cache_loc_kernel( out_cache_loc, tgt_cache_loc, accept_length, accept_length_filter, bs_upper: tl.constexpr, num_verify_tokens_upper: tl.constexpr, ): bid = tl.program_id(0) bs_offset = tl.arange(0, bs_upper) accept_length_all = tl.load(accept_length + bs_offset, mask=bs_offset < bid) old_start = tl.sum(accept_length_all) + bid accept_length_filter_all = tl.load( accept_length_filter + bs_offset, mask=bs_offset < bid ) new_start = tl.sum(accept_length_filter_all) copy_len = tl.load(accept_length_filter + bid) copy_offset = tl.arange(0, num_verify_tokens_upper) value = tl.load( tgt_cache_loc + old_start + copy_offset, mask=copy_offset < copy_len ) tl.store( out_cache_loc + new_start + copy_offset, value, mask=copy_offset < copy_len ) @torch.compile(dynamic=True, disable=_is_npu) def create_accept_length_filter( accept_length: torch.Tensor, unfinished_index_device: torch.Tensor, seq_lens: torch.Tensor, ): accept_length_filter = torch.zeros_like(accept_length) accept_length_filter[unfinished_index_device] = ( accept_length[unfinished_index_device] + 1 ) seq_lens.add_(accept_length + 1) return accept_length_filter @torch.compile(dynamic=True, disable=_is_npu) def select_top_k_tokens( i: int, topk_p: torch.Tensor, topk_index: torch.Tensor, hidden_states: torch.Tensor, scores: torch.Tensor, topk: int, ): if i == 0: # The first step after extend input_ids = topk_index.flatten() if hidden_states is not None: hidden_states = hidden_states.repeat_interleave(topk, dim=0) scores = topk_p # shape: (b, topk) tree_info = ( topk_p.unsqueeze(1), # shape: (b, 1, topk) topk_index, # shape: (b, topk) torch.arange(-1, topk, dtype=torch.long, device=input_ids.device) .unsqueeze(0) .repeat(topk_p.shape[0], 1), # shape: (b, topk + 1) ) else: # The later decode steps expand_scores = torch.mul( scores.unsqueeze(2), topk_p.reshape(-1, topk, topk) ) # (b, topk, 1) x (b, topk ,topk) -> (b, topk, topk) topk_cs_p, topk_cs_index = fast_topk( expand_scores.flatten(start_dim=1), topk, dim=-1 ) # (b, topk) scores = topk_cs_p # shape: (b, topk) topk_index = topk_index.reshape(-1, topk**2) input_ids = torch.gather(topk_index, index=topk_cs_index, dim=1).flatten() if hidden_states.shape[0] > 0: selected_input_index = topk_cs_index.flatten() // topk + torch.arange( 0, hidden_states.shape[0], step=topk, device=topk_index.device ).repeat_interleave(topk) hidden_states = hidden_states[selected_input_index, :] tree_info = ( expand_scores, # shape: (b, topk, topk) topk_index, # shape: (b, topk * topk) topk_cs_index + (topk**2 * (i - 1) + topk), # shape: (b, topk) ) return input_ids, hidden_states, scores, tree_info def generate_simulated_accept_index( accept_index, predict, accept_length, bs, spec_steps, simulate_acc_len: float = SIMULATE_ACC_LEN, simulate_acc_method: str = SIMULATE_ACC_METHOD, ): assert simulate_acc_len > 0.0 if simulate_acc_method == "multinomial": simulated_values = torch.normal( mean=simulate_acc_len, std=1.0, size=(1,), device="cpu", ) # clamp simulated values to be between 1 and self.spec_steps simulated_values = torch.clamp(simulated_values, min=1.0, max=spec_steps + 1) simulate_acc_len = int(simulated_values.round().item()) elif simulate_acc_method == "match-expected": # multinomial sampling does not match the expected length # we keep it for the sake of compatibility of existing tests # but it's better to use "match-expected" for the cases that need to # match the expected length, One caveat is that this will only sample # either round down or round up of the expected length simulate_acc_len = max(1.0, min(spec_steps + 1, simulate_acc_len)) lower = int(simulate_acc_len // 1) upper = lower + 1 if lower < spec_steps + 1 else lower if lower == upper: simulate_acc_len = lower else: weight_upper = simulate_acc_len - lower weight_lower = 1.0 - weight_upper probs = torch.tensor([weight_lower, weight_upper], device="cpu") sampled_index = torch.multinomial(probs, num_samples=1) simulate_acc_len = lower if sampled_index == 0 else upper else: raise ValueError(f"Invalid simulate_acc_method: {SIMULATE_ACC_METHOD}") accept_indx_first_col = accept_index[:, 0].view(-1, 1) sim_accept_index = torch.full( (bs, spec_steps + 1), -1, dtype=torch.int32, device="cuda" ) sim_accept_index[:, :simulate_acc_len] = accept_indx_first_col + torch.arange( simulate_acc_len, device=accept_index.device ) accept_length.fill_(simulate_acc_len - 1) predict.fill_(100) # some legit token id return sim_accept_index def traverse_tree( retrieve_next_token: torch.Tensor, retrieve_next_sibling: torch.Tensor, draft_tokens: torch.Tensor, grammar: BaseGrammarObject, allocate_token_bitmask: torch.Tensor, vocab_size: Optional[int] = None, ): """ Traverse the tree constructed by the draft model to generate the logits mask. """ assert ( retrieve_next_token.shape == retrieve_next_sibling.shape == draft_tokens.shape ) def dfs( curr: int, retrieve_next_token: torch.Tensor, retrieve_next_sibling: torch.Tensor, parent_pos: int, ): if curr == 0: # the first token generated by the target model, and thus it is always # accepted from the previous iteration accepted = True else: parent_bitmask = allocate_token_bitmask[parent_pos] curr_token_id = draft_tokens[curr] if vocab_size and curr_token_id >= vocab_size: accepted = False else: # 32 boolean bitmask values are packed into 32-bit integers accepted = ( parent_bitmask[curr_token_id // 32] & (1 << (curr_token_id % 32)) ) != 0 if accepted: if curr != 0: # Accept the current token grammar.accept_token(draft_tokens[curr]) if not grammar.is_terminated(): # Generate the bitmask for the current token grammar.fill_vocab_mask(allocate_token_bitmask, curr) if retrieve_next_token[curr] != -1: # Visit the child node dfs( retrieve_next_token[curr], retrieve_next_token, retrieve_next_sibling, curr, ) if curr != 0: # Rollback the current token grammar.rollback(1) if retrieve_next_sibling[curr] != -1: # Visit the sibling node dfs( retrieve_next_sibling[curr], retrieve_next_token, retrieve_next_sibling, parent_pos, ) dfs(0, retrieve_next_token, retrieve_next_sibling, -1) def generate_token_bitmask( reqs: List[Req], verify_input: EagleVerifyInput, retrieve_next_token_cpu: torch.Tensor, retrieve_next_sibling_cpu: torch.Tensor, draft_tokens_cpu: torch.Tensor, vocab_size: int, ): """ Generate the logit mask for structured output. Draft model's token can be either valid or invalid with respect to the grammar. We need to perform DFS to 1. figure out which tokens are accepted by the grammar. 2. if so, what is the corresponding logit mask. """ num_draft_tokens = draft_tokens_cpu.shape[-1] allocate_token_bitmask = None assert len(reqs) == retrieve_next_token_cpu.shape[0] grammar = None for i, req in enumerate(reqs): if req.grammar is not None: if allocate_token_bitmask is None: allocate_token_bitmask = req.grammar.allocate_vocab_mask( vocab_size=vocab_size, batch_size=draft_tokens_cpu.numel(), device="cpu", ) grammar = req.grammar s = time.perf_counter() traverse_tree( retrieve_next_token_cpu[i], retrieve_next_sibling_cpu[i], draft_tokens_cpu[i], req.grammar, allocate_token_bitmask[ i * num_draft_tokens : (i + 1) * num_draft_tokens ], vocab_size=vocab_size, ) tree_traverse_time = time.perf_counter() - s if tree_traverse_time > TREE_TRAVERSE_TIME_THRESHOLD: logger.warning( f"Bit mask generation took {tree_traverse_time} seconds with " f"grammar: {req.grammar}" ) verify_input.grammar = grammar return allocate_token_bitmask def load_token_map(token_map_path: str) -> List[int]: if not os.path.exists(token_map_path): cache_dir = snapshot_download( os.path.dirname(token_map_path), ignore_patterns=["*.bin", "*.safetensors"], ) token_map_path = os.path.join(cache_dir, os.path.basename(token_map_path)) hot_token_id = torch.load(token_map_path, weights_only=True) return torch.tensor(hot_token_id, dtype=torch.int64) @contextmanager def draft_tp_context(tp_group: GroupCoordinator): # Draft model doesn't use dp and has its own tp group. # We disable mscclpp now because it doesn't support 2 comm groups. with patch_tensor_parallel_group(tp_group): yield def detect_nan(logits_output: LogitsProcessorOutput): logits = logits_output.next_token_logits if torch.any(torch.isnan(logits)): logger.error("Detected errors during sampling! NaN in the logits.") raise ValueError("Detected errors during sampling! NaN in the logits.") # Disable torch.compile for this function because it will be # even slower. # @torch.compile(dynamic=True) def get_last_loc_large_page_size_large_top_k( req_to_token: torch.Tensor, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor, speculative_num_steps: int, topk: int, page_size: int, ): prefix_lens = seq_lens last_page_lens = prefix_lens % page_size num_new_pages_per_topk = ( last_page_lens + speculative_num_steps + page_size - 1 ) // page_size seq_lens = prefix_lens // page_size * page_size + num_new_pages_per_topk * ( page_size * topk ) extend_lens = seq_lens - prefix_lens last_loc = get_last_loc( req_to_token, req_pool_indices, prefix_lens, ) return ( prefix_lens, seq_lens, last_loc, num_new_pages_per_topk, extend_lens, last_page_lens, )