From cb719c74adeef47305c595b564b269bf05f22a9f Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Wed, 24 Dec 2025 15:21:13 +0800 Subject: [PATCH] [2/N] clean duplicate code of logprob processing in spec. (#15593) --- python/sglang/srt/layers/utils/logprob.py | 105 +++++++++++++++++- python/sglang/srt/speculative/eagle_worker.py | 85 +------------- .../speculative/multi_layer_eagle_worker.py | 85 +------------- python/sglang/srt/speculative/ngram_info.py | 32 +++--- python/sglang/srt/speculative/ngram_worker.py | 87 +-------------- 5 files changed, 127 insertions(+), 267 deletions(-) diff --git a/python/sglang/srt/layers/utils/logprob.py b/python/sglang/srt/layers/utils/logprob.py index ebc3031da..6f84c15ff 100644 --- a/python/sglang/srt/layers/utils/logprob.py +++ b/python/sglang/srt/layers/utils/logprob.py @@ -2,12 +2,17 @@ from __future__ import annotations import dataclasses from enum import Enum, auto -from typing import TYPE_CHECKING, List, Optional +from typing import TYPE_CHECKING, List, Optional, Union import torch +from sglang.srt.environ import envs + if TYPE_CHECKING: - from sglang.srt.layers.logits_processor import LogitsMetadata + from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessorOutput + from sglang.srt.managers.schedule_batch import ScheduleBatch + from sglang.srt.speculative.eagle_info import EagleVerifyOutput + from sglang.srt.speculative.ngram_info import NgramVerifyInput class LogprobStage(Enum): @@ -304,3 +309,99 @@ def get_token_ids_logprobs_chunk( pt += pruned_len return next_split_pruned_len + + +def add_output_logprobs_for_spec_v1( + batch: ScheduleBatch, + res: Union[EagleVerifyOutput, NgramVerifyInput], + logits_output: Optional[LogitsProcessorOutput] = None, +): + # Extract args + if logits_output is None: + logits_output = res.logits_output + + if hasattr(res, "accept_length_per_req_cpu"): + accept_length_per_req_cpu = res.accept_length_per_req_cpu + else: + # FIXME: Get a NgramVerifyOutput class and use that instead of this hack. + accept_length_per_req_cpu = res.accept_length.tolist() + + top_logprobs_nums = batch.top_logprobs_nums + token_ids_logprobs = batch.token_ids_logprobs + accepted_indices = res.accepted_indices + assert len(accepted_indices) == len(logits_output.next_token_logits) + + temperatures = batch.sampling_info.temperatures + num_draft_tokens = batch.spec_info.draft_token_num + # acceptance indices are the indices in a "flattened" batch. + # dividing it to num_draft_tokens will yield the actual batch index. + temperatures = temperatures[accepted_indices // num_draft_tokens] + if envs.SGLANG_RETURN_ORIGINAL_LOGPROB.get(): + logprobs = torch.nn.functional.log_softmax( + logits_output.next_token_logits, dim=-1 + ) + else: + logprobs = torch.nn.functional.log_softmax( + logits_output.next_token_logits / temperatures, dim=-1 + ) + batch_next_token_ids = res.verified_id + num_tokens_per_req = [accept + 1 for accept in accept_length_per_req_cpu] + + # We should repeat top_logprobs_nums to match num_tokens_per_req. + top_logprobs_nums_repeat_interleaved = [ + num + for num, num_tokens in zip(top_logprobs_nums, num_tokens_per_req) + for _ in range(num_tokens) + ] + + token_ids_logprobs_repeat_interleaved = [ + token_ids + for token_ids, num_tokens in zip(token_ids_logprobs, num_tokens_per_req) + for _ in range(num_tokens) + ] + + # Extract logprobs + should_top_logprobs = any(x > 0 for x in top_logprobs_nums) + should_token_ids_logprobs = any(x is not None for x in token_ids_logprobs) + if should_top_logprobs: + ( + logits_output.next_token_top_logprobs_val, + logits_output.next_token_top_logprobs_idx, + ) = get_top_logprobs( + logprobs, + top_logprobs_nums_repeat_interleaved, + ) + + if should_token_ids_logprobs: + ( + logits_output.next_token_token_ids_logprobs_val, + logits_output.next_token_token_ids_logprobs_idx, + ) = get_token_ids_logprobs( + logprobs, + token_ids_logprobs_repeat_interleaved, + ) + + logits_output.next_token_logprobs = logprobs[ + torch.arange(len(batch_next_token_ids), device=batch.sampling_info.device), + batch_next_token_ids, + ] + + # Add output logprobs to the request + pt = 0 + next_token_logprobs = logits_output.next_token_logprobs.tolist() + verified_ids = batch_next_token_ids.tolist() + token_top_logprobs_val = logits_output.next_token_top_logprobs_val + token_top_logprobs_idx = logits_output.next_token_top_logprobs_idx + for req, num_tokens in zip(batch.reqs, num_tokens_per_req, strict=True): + for _ in range(num_tokens): + # TODO: add token_ids_logprobs to each request + if req.return_logprob: + req.output_token_logprobs_val.append(next_token_logprobs[pt]) + req.output_token_logprobs_idx.append(verified_ids[pt]) + if req.top_logprobs_num > 0: + assert ( + should_top_logprobs + ), "Inconsistent state: should_top_logprobs is False" + req.output_top_logprobs_val.append(token_top_logprobs_val[pt]) + req.output_top_logprobs_idx.append(token_top_logprobs_idx[pt]) + pt += 1 diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index ed9e752c5..7514d7338 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/python/sglang/srt/speculative/eagle_worker.py @@ -14,7 +14,7 @@ from sglang.srt.layers.moe.utils import ( speculative_moe_a2a_backend_context, speculative_moe_backend_context, ) -from sglang.srt.layers.utils.logprob import get_token_ids_logprobs, get_top_logprobs +from sglang.srt.layers.utils.logprob import add_output_logprobs_for_spec_v1 from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult @@ -62,7 +62,6 @@ from sglang.srt.utils import ( MultiprocessingSerializer, empty_context, get_available_gpu_memory, - get_bool_env_var, is_cuda, is_npu, next_power_of_2, @@ -75,7 +74,6 @@ if is_cuda(): from sgl_kernel import segment_packbits # noqa: F401 logger = logging.getLogger(__name__) -SGLANG_RETURN_ORIGINAL_LOGPROB = get_bool_env_var("SGLANG_RETURN_ORIGINAL_LOGPROB") class EAGLEWorker(TpModelWorker): @@ -756,7 +754,7 @@ class EAGLEWorker(TpModelWorker): ) if batch.return_logprob: - self.add_logprob_values(batch, res, logits_output) + add_output_logprobs_for_spec_v1(batch, res, logits_output) # Prepare the batch for the next draft forwards. batch.forward_mode = ( @@ -845,85 +843,6 @@ class EAGLEWorker(TpModelWorker): model=self.target_worker.model_runner.model, ) - def add_logprob_values( - self, - batch: ScheduleBatch, - res: EagleVerifyOutput, - logits_output: LogitsProcessorOutput, - ): - # Extract args - logits_output = res.logits_output - top_logprobs_nums = batch.top_logprobs_nums - token_ids_logprobs = batch.token_ids_logprobs - accepted_indices = res.accepted_indices - assert len(accepted_indices) == len(logits_output.next_token_logits) - - temperatures = batch.sampling_info.temperatures - num_draft_tokens = batch.spec_info.draft_token_num - # acceptance indices are the indices in a "flattened" batch. - # dividing it to num_draft_tokens will yield the actual batch index. - temperatures = temperatures[accepted_indices // num_draft_tokens] - if SGLANG_RETURN_ORIGINAL_LOGPROB: - logprobs = torch.nn.functional.log_softmax( - logits_output.next_token_logits, dim=-1 - ) - else: - logprobs = torch.nn.functional.log_softmax( - logits_output.next_token_logits / temperatures, dim=-1 - ) - batch_next_token_ids = res.verified_id - num_tokens_per_req = [accept + 1 for accept in res.accept_length_per_req_cpu] - - # We should repeat top_logprobs_nums to match num_tokens_per_req. - top_logprobs_nums_repeat_interleaved = [] - token_ids_logprobs_repeat_interleaved = [] - for num, num_tokens in zip(top_logprobs_nums, num_tokens_per_req): - top_logprobs_nums_repeat_interleaved.extend([num] * num_tokens) - for token_ids, num_tokens in zip(token_ids_logprobs, num_tokens_per_req): - token_ids_logprobs_repeat_interleaved.extend([token_ids] * num_tokens) - - # Extract logprobs - if any(x > 0 for x in top_logprobs_nums): - ( - logits_output.next_token_top_logprobs_val, - logits_output.next_token_top_logprobs_idx, - ) = get_top_logprobs( - logprobs, - top_logprobs_nums_repeat_interleaved, - ) - - if any(x is not None for x in token_ids_logprobs): - ( - logits_output.next_token_token_ids_logprobs_val, - logits_output.next_token_token_ids_logprobs_idx, - ) = get_token_ids_logprobs( - logprobs, - token_ids_logprobs_repeat_interleaved, - ) - - logits_output.next_token_logprobs = logprobs[ - torch.arange(len(batch_next_token_ids), device=batch.sampling_info.device), - batch_next_token_ids, - ] - - # Add output logprobs to the request - pt = 0 - next_token_logprobs = logits_output.next_token_logprobs.tolist() - verified_ids = batch_next_token_ids.tolist() - for req, num_tokens in zip(batch.reqs, num_tokens_per_req, strict=True): - for _ in range(num_tokens): - if req.return_logprob: - req.output_token_logprobs_val.append(next_token_logprobs[pt]) - req.output_token_logprobs_idx.append(verified_ids[pt]) - if req.top_logprobs_num > 0: - req.output_top_logprobs_val.append( - res.logits_output.next_token_top_logprobs_val[pt] - ) - req.output_top_logprobs_idx.append( - res.logits_output.next_token_top_logprobs_idx[pt] - ) - pt += 1 - def forward_draft_extend( self, batch: ScheduleBatch, diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker.py b/python/sglang/srt/speculative/multi_layer_eagle_worker.py index 455a8e289..9dc3c5fd7 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker.py @@ -22,7 +22,7 @@ from sglang.srt.distributed import get_tp_group from sglang.srt.layers.dp_attention import get_attention_tp_group from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.moe.utils import speculative_moe_backend_context -from sglang.srt.layers.utils.logprob import get_token_ids_logprobs, get_top_logprobs +from sglang.srt.layers.utils.logprob import add_output_logprobs_for_spec_v1 from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker @@ -65,7 +65,6 @@ from sglang.srt.speculative.spec_utils import ( from sglang.srt.utils import ( empty_context, get_available_gpu_memory, - get_bool_env_var, is_cuda, is_npu, next_power_of_2, @@ -77,7 +76,6 @@ if is_cuda(): from sgl_kernel import segment_packbits # noqa: F401 logger = logging.getLogger(__name__) -SGLANG_RETURN_ORIGINAL_LOGPROB = get_bool_env_var("SGLANG_RETURN_ORIGINAL_LOGPROB") class MultiLayerEagleWorker(TpModelWorker): @@ -727,7 +725,7 @@ class MultiLayerEagleWorker(TpModelWorker): ) if batch.return_logprob: - self.add_logprob_values(batch, res, logits_output) + add_output_logprobs_for_spec_v1(batch, res, logits_output) # Prepare the batch for the next draft forwards. batch.forward_mode = ( @@ -737,85 +735,6 @@ class MultiLayerEagleWorker(TpModelWorker): return logits_output, res, model_worker_batch, can_run_cuda_graph - def add_logprob_values( - self, - batch: ScheduleBatch, - res: EagleVerifyOutput, - logits_output: LogitsProcessorOutput, - ): - # Extract args - logits_output = res.logits_output - top_logprobs_nums = batch.top_logprobs_nums - token_ids_logprobs = batch.token_ids_logprobs - accepted_indices = res.accepted_indices - assert len(accepted_indices) == len(logits_output.next_token_logits) - - temperatures = batch.sampling_info.temperatures - num_draft_tokens = batch.spec_info.draft_token_num - # acceptance indices are the indices in a "flattened" batch. - # dividing it to num_draft_tokens will yield the actual batch index. - temperatures = temperatures[accepted_indices // num_draft_tokens] - if SGLANG_RETURN_ORIGINAL_LOGPROB: - logprobs = torch.nn.functional.log_softmax( - logits_output.next_token_logits, dim=-1 - ) - else: - logprobs = torch.nn.functional.log_softmax( - logits_output.next_token_logits / temperatures, dim=-1 - ) - batch_next_token_ids = res.verified_id - num_tokens_per_req = [accept + 1 for accept in res.accept_length_per_req_cpu] - - # We should repeat top_logprobs_nums to match num_tokens_per_req. - top_logprobs_nums_repeat_interleaved = [] - token_ids_logprobs_repeat_interleaved = [] - for num, num_tokens in zip(top_logprobs_nums, num_tokens_per_req): - top_logprobs_nums_repeat_interleaved.extend([num] * num_tokens) - for token_ids, num_tokens in zip(token_ids_logprobs, num_tokens_per_req): - token_ids_logprobs_repeat_interleaved.extend([token_ids] * num_tokens) - - # Extract logprobs - if any(x > 0 for x in top_logprobs_nums): - ( - logits_output.next_token_top_logprobs_val, - logits_output.next_token_top_logprobs_idx, - ) = get_top_logprobs( - logprobs, - top_logprobs_nums_repeat_interleaved, - ) - - if any(x is not None for x in token_ids_logprobs): - ( - logits_output.next_token_token_ids_logprobs_val, - logits_output.next_token_token_ids_logprobs_idx, - ) = get_token_ids_logprobs( - logprobs, - token_ids_logprobs_repeat_interleaved, - ) - - logits_output.next_token_logprobs = logprobs[ - torch.arange(len(batch_next_token_ids), device=batch.sampling_info.device), - batch_next_token_ids, - ] - - # Add output logprobs to the request - pt = 0 - next_token_logprobs = logits_output.next_token_logprobs.tolist() - verified_ids = batch_next_token_ids.tolist() - for req, num_tokens in zip(batch.reqs, num_tokens_per_req, strict=True): - for _ in range(num_tokens): - if req.return_logprob: - req.output_token_logprobs_val.append(next_token_logprobs[pt]) - req.output_token_logprobs_idx.append(verified_ids[pt]) - if req.top_logprobs_num > 0: - req.output_top_logprobs_val.append( - res.logits_output.next_token_top_logprobs_val[pt] - ) - req.output_top_logprobs_idx.append( - res.logits_output.next_token_top_logprobs_idx[pt] - ) - pt += 1 - def forward_draft_extend( self, batch: ScheduleBatch, diff --git a/python/sglang/srt/speculative/ngram_info.py b/python/sglang/srt/speculative/ngram_info.py index 637f51c97..67194aef6 100644 --- a/python/sglang/srt/speculative/ngram_info.py +++ b/python/sglang/srt/speculative/ngram_info.py @@ -155,7 +155,7 @@ class NgramVerifyInput(SpecInput): batch: ScheduleBatch, logits_output: torch.Tensor, ): - accept_index_cpu = self.accept_index.tolist() + accept_index_cpu = self.accepted_indices.tolist() predict_cpu = self.predict.tolist() has_finished = False @@ -171,7 +171,7 @@ class NgramVerifyInput(SpecInput): if req.finished(): has_finished = True # set all tokens after finished token to -1 and break - self.accept_index[i, j + 1 :] = -1 + self.accepted_indices[i, j + 1 :] = -1 break else: if req.grammar is not None: @@ -180,7 +180,7 @@ class NgramVerifyInput(SpecInput): except ValueError as e: logger.info( f"{i=}, {req=}\n" - f"{self.accept_index=}\n" + f"{self.accepted_indices=}\n" f"{self.predict=}\n" ) raise e @@ -190,15 +190,17 @@ class NgramVerifyInput(SpecInput): ) if has_finished: - self.accept_length = (self.accept_index != -1).sum(dim=1) - 1 - self.accept_index = self.accept_index[self.accept_index != -1] + self.accept_length = (self.accepted_indices != -1).sum(dim=1) - 1 + self.accepted_indices = self.accepted_indices[self.accepted_indices != -1] logits_output.next_token_logits = logits_output.next_token_logits[ - self.accept_index + self.accepted_indices ] if logits_output.hidden_states: - logits_output.hidden_states = logits_output.hidden_states[self.accept_index] - self.verified_id = self.predict[self.accept_index] + logits_output.hidden_states = logits_output.hidden_states[ + self.accepted_indices + ] + self.verified_id = self.predict[self.accepted_indices] def _free_cache( self, batch: ScheduleBatch, page_size: int, accept_length_cpu: torch.Tensor @@ -208,16 +210,16 @@ class NgramVerifyInput(SpecInput): if page_size == 1: # TODO: boolean array index leads to a device sync. Remove it. evict_mask = torch.full_like(self.draft_token, True, dtype=torch.bool) - evict_mask[self.accept_index] = False + evict_mask[self.accepted_indices] = False batch.token_to_kv_pool_allocator.free(batch.out_cache_loc[evict_mask]) - batch.out_cache_loc = batch.out_cache_loc[self.accept_index] + batch.out_cache_loc = batch.out_cache_loc[self.accepted_indices] else: # Shift the accepted tokens to the beginning. # Only evict the last part src_cache_loc, tgt_cache_loc, to_free_num_slots = get_src_tgt_cache_loc( batch.seq_lens, batch.out_cache_loc, - self.accept_index, + self.accepted_indices, self.accept_length, self.draft_token_num, page_size, @@ -285,14 +287,14 @@ class NgramVerifyInput(SpecInput): predict_shape = list(logits_output.next_token_logits.shape)[:-1] predict_shape[-1] += 1 self.predict = torch.empty(predict_shape, dtype=torch.int32, device=self.device) - self.accept_index = torch.full( + self.accepted_indices = torch.full( (bs, self.draft_token_num), -1, dtype=torch.int32, device=self.device ) self.accept_length = torch.empty((bs,), dtype=torch.int32, device=self.device) verify_tree_greedy( predicts=self.predict, # mutable - accept_index=self.accept_index, # mutable + accept_index=self.accepted_indices, # mutable accept_token_num=self.accept_length, # mutable candidates=candidates, retrive_index=self.retrive_index, @@ -312,7 +314,7 @@ class NgramVerifyInput(SpecInput): predict_shape = list(logits_output.next_token_logits.shape)[:-1] predict_shape[-1] += 1 self.predict = torch.empty(predict_shape, dtype=torch.int32, device=self.device) - self.accept_index = torch.full( + self.accepted_indices = torch.full( (bs, self.draft_token_num), -1, dtype=torch.int32, device=self.device ) self.accept_length = torch.empty((bs,), dtype=torch.int32, device=self.device) @@ -354,7 +356,7 @@ class NgramVerifyInput(SpecInput): ) tree_speculative_sampling_target_only( predicts=self.predict, # mutable - accept_index=self.accept_index, # mutable + accept_index=self.accepted_indices, # mutable accept_token_num=self.accept_length, # mutable candidates=candidates.to(torch.int64), retrive_index=self.retrive_index.to(torch.int64), diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py index 4296c6924..fdeec130d 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/python/sglang/srt/speculative/ngram_worker.py @@ -5,9 +5,7 @@ import numpy as np import torch from sgl_kernel.speculative import reconstruct_indices_from_tree_mask -from sglang.srt.environ import envs -from sglang.srt.layers.logits_processor import LogitsProcessorOutput -from sglang.srt.layers.utils.logprob import get_token_ids_logprobs, get_top_logprobs +from sglang.srt.layers.utils.logprob import add_output_logprobs_for_spec_v1 from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker @@ -198,85 +196,6 @@ class NGRAMWorker: ) batch.spec_info.prepare_for_verify(batch, self.page_size) - def add_logprob_values( - self, - batch: ScheduleBatch, - res: NgramVerifyInput, - logits_output: LogitsProcessorOutput, - ): - # Extract args - top_logprobs_nums = batch.top_logprobs_nums - token_ids_logprobs = batch.token_ids_logprobs - accepted_indices = res.accept_index - assert len(accepted_indices) == len(logits_output.next_token_logits) - - temperatures = batch.sampling_info.temperatures - num_draft_tokens = batch.spec_info.draft_token_num - # acceptance indices are the indices in a "flattened" batch. - # dividing it to num_draft_tokens will yield the actual batch index. - temperatures = temperatures[accepted_indices // num_draft_tokens] - if envs.SGLANG_RETURN_ORIGINAL_LOGPROB.get(): - logprobs = torch.nn.functional.log_softmax( - logits_output.next_token_logits, dim=-1 - ) - else: - logprobs = torch.nn.functional.log_softmax( - logits_output.next_token_logits / temperatures, dim=-1 - ) - batch_next_token_ids = res.verified_id - accept_length_per_req_cpu = res.accept_length.tolist() - num_tokens_per_req = [accept + 1 for accept in accept_length_per_req_cpu] - - # We should repeat top_logprobs_nums to match num_tokens_per_req. - top_logprobs_nums_repeat_interleaved = [] - token_ids_logprobs_repeat_interleaved = [] - for num, num_tokens in zip(top_logprobs_nums, num_tokens_per_req): - top_logprobs_nums_repeat_interleaved.extend([num] * num_tokens) - for token_ids, num_tokens in zip(token_ids_logprobs, num_tokens_per_req): - token_ids_logprobs_repeat_interleaved.extend([token_ids] * num_tokens) - - # Extract logprobs - if any(x > 0 for x in top_logprobs_nums): - ( - logits_output.next_token_top_logprobs_val, - logits_output.next_token_top_logprobs_idx, - ) = get_top_logprobs( - logprobs, - top_logprobs_nums_repeat_interleaved, - ) - - if any(x is not None for x in token_ids_logprobs): - ( - logits_output.next_token_token_ids_logprobs_val, - logits_output.next_token_token_ids_logprobs_idx, - ) = get_token_ids_logprobs( - logprobs, - token_ids_logprobs_repeat_interleaved, - ) - - logits_output.next_token_logprobs = logprobs[ - torch.arange(len(batch_next_token_ids), device=batch.sampling_info.device), - batch_next_token_ids, - ] - - # Add output logprobs to the request - pt = 0 - next_token_logprobs = logits_output.next_token_logprobs.tolist() - verified_ids = batch_next_token_ids.tolist() - for req, num_tokens in zip(batch.reqs, num_tokens_per_req, strict=True): - for _ in range(num_tokens): - if req.return_logprob: - req.output_token_logprobs_val.append(next_token_logprobs[pt]) - req.output_token_logprobs_idx.append(verified_ids[pt]) - if req.top_logprobs_num > 0: - req.output_top_logprobs_val.append( - logits_output.next_token_top_logprobs_val[pt] - ) - req.output_top_logprobs_idx.append( - logits_output.next_token_top_logprobs_idx[pt] - ) - pt += 1 - def _update_ngram_cache(self, batch: ScheduleBatch): batch_tokens = [] for req in batch.reqs: @@ -305,14 +224,14 @@ class NGRAMWorker: batch_result.logits_output, batch_result.can_run_cuda_graph, ) - verify_input = model_worker_batch.spec_info + verify_input: NgramVerifyInput = model_worker_batch.spec_info logits_output, next_token_ids, num_accepted_tokens = verify_input.verify( batch, logits_output, self.page_size ) # Store accept_lens for per-request metrics accept_lens = verify_input.accept_length if batch.return_logprob: - self.add_logprob_values(batch, verify_input, logits_output) + add_output_logprobs_for_spec_v1(batch, verify_input, logits_output) self._update_ngram_cache(batch) batch.forward_mode = ForwardMode.DECODE