From 9a6150792b6b50e9c9853c1aea9a45d6f47cb08e Mon Sep 17 00:00:00 2001 From: sufubao Date: Thu, 27 Aug 2026 12:17:17 +0800 Subject: [PATCH 1/8] perf: avoid full-vocab all-gather for draft greedy sampling --- lightllm/common/basemodel/basemodel.py | 19 +++ lightllm/common/basemodel/batch_objs.py | 36 ++++++ lightllm/common/basemodel/infer_struct.py | 3 + .../post_process/greedy_sample.py | 109 ++++++++++++++++ .../post_process/vocab_parallel_greedy.py | 118 ++++++++++++++++++ .../triton_kernel/transpose_convert.py | 65 ++++++++++ .../llama/layer_infer/post_layer_infer.py | 18 ++- .../layer_infer/post_layer_infer.py | 6 +- .../model_infer/mode_backend/base_backend.py | 18 ++- .../dp_overlap_proposers/eagle_with_att.py | 2 +- .../proposers/eagle_with_att.py | 2 +- .../common/basemodel/test_model_output.py | 24 ++++ .../test_vocab_parallel_greedy.py | 65 ++++++++++ .../test_vocab_parallel_greedy_output.py | 74 +++++++++++ 14 files changed, 551 insertions(+), 8 deletions(-) create mode 100644 lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py create mode 100644 lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py create mode 100644 lightllm/common/basemodel/triton_kernel/transpose_convert.py create mode 100644 unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py create mode 100644 unit_tests/models/test_vocab_parallel_greedy_output.py diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index e80f2b552f..b561595b98 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -384,6 +384,7 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) infer_state.input_ids = model_input.input_ids infer_state.is_prefill = model_input.is_prefill infer_state.return_all_prompt_logics = self.return_all_prompt_logics + infer_state.use_vocab_parallel_greedy = self.is_mtp_draft_model infer_state.batch_size = model_input.batch_size infer_state.total_token_num = model_input.total_token_num infer_state.max_q_seq_len = model_input.max_q_seq_len @@ -534,6 +535,9 @@ def _create_unpad_decode_model_output(self, model_output: ModelOutput, origin_ba return model_output new_model_output = copy.copy(model_output) new_model_output.logits = new_model_output.logits[0:origin_batch_size] + if new_model_output.logits_token_ids is not None: + new_model_output.logits_token_ids = new_model_output.logits_token_ids[0:origin_batch_size] + new_model_output.logits_logsumexp = new_model_output.logits_logsumexp[0:origin_batch_size] new_model_output.mtp_collector = model_output.mtp_collector.unpad_decode( padded_batch_size=padded_batch_size, origin_batch_size=origin_batch_size, @@ -546,6 +550,9 @@ def _create_unpad_prefill_model_output( new_model_output = copy.copy(padded_model_output) # logits 始终只对应每个请求最后一个位置，移除 padding 的 req 对应的行。 new_model_output.logits = new_model_output.logits[0:origin_batch_size] + if new_model_output.logits_token_ids is not None: + new_model_output.logits_token_ids = new_model_output.logits_token_ids[0:origin_batch_size] + new_model_output.logits_logsumexp = new_model_output.logits_logsumexp[0:origin_batch_size] new_model_output.mtp_collector = padded_model_output.mtp_collector.unpad_prefill( origin_handle_token_num=origin_handle_token_num ) @@ -737,6 +744,8 @@ def prefill_func(input_tensors, _infer_state): hidden_collector.add_final_hidden(last_input_embs) model_output = ModelOutput( logits=predict_logits.contiguous(), + logits_token_ids=infer_state.logits_token_ids, + logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, ) @@ -766,6 +775,8 @@ def _token_forward(self, infer_state: InferStateInfo): hidden_collector.add_final_hidden(last_input_embs) model_output = ModelOutput( logits=predict_logits.contiguous(), + logits_token_ids=infer_state.logits_token_ids, + logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), ) @@ -1020,11 +1031,15 @@ def _overlap_tpsp_context_forward(self, infer_state: InferStateInfo, infer_state hidden_collector1.add_final_hidden(last_input_embs1) model_output = ModelOutput( logits=predict_logits.contiguous(), + logits_token_ids=infer_state.logits_token_ids, + logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, ) model_output1 = ModelOutput( logits=predict_logits1.contiguous(), + logits_token_ids=infer_state1.logits_token_ids, + logits_logsumexp=infer_state1.logits_logsumexp, mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1), prompt_logics=infer_state1.prompt_logics, ) @@ -1069,10 +1084,14 @@ def _overlap_tpsp_token_forward(self, infer_state: InferStateInfo, infer_state1: hidden_collector1.add_final_hidden(last_input_embs1) model_output = ModelOutput( logits=predict_logits.contiguous(), + logits_token_ids=infer_state.logits_token_ids, + logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), ) model_output1 = ModelOutput( logits=predict_logits1.contiguous(), + logits_token_ids=infer_state1.logits_token_ids, + logits_logsumexp=infer_state1.logits_logsumexp, mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1), ) diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index ae645d4b7b..edddb749e5 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -200,10 +200,46 @@ class ModelOutput: # 需要返回 prompt logprobs 信息时才会非空。 prompt_logics: Optional[torch.Tensor] = None + # Vocab-parallel outputs keep logits as logits while mapping each sparse + # column back to its global token id. logits_logsumexp is computed over the + # complete vocabulary, so sparse argmax probabilities remain exact. + # Both fields are None for historical dense logits. + logits_token_ids: Optional[torch.Tensor] = None + logits_logsumexp: Optional[torch.Tensor] = None + def __post_init__(self) -> None: if self.mtp_collector is None: self.mtp_collector = ModelMtpOutputCollector() + assert (self.logits_token_ids is None) == (self.logits_logsumexp is None) + if self.logits_token_ids is not None: + assert self.logits.ndim == 2 + assert self.logits_token_ids.shape == self.logits.shape + assert self.logits_token_ids.dtype in (torch.int32, torch.int64) + assert self.logits_token_ids.device == self.logits.device + assert self.logits_logsumexp.shape == (self.logits.shape[0],) + assert self.logits_logsumexp.dtype == torch.float32 + assert self.logits_logsumexp.device == self.logits.device def to_no_ref_tensor(self): self.logits = tensor_to_no_ref_tensor(self.logits) + if self.logits_token_ids is not None: + self.logits_token_ids = tensor_to_no_ref_tensor(self.logits_token_ids) + self.logits_logsumexp = tensor_to_no_ref_tensor(self.logits_logsumexp) self.mtp_collector.to_no_ref_tensor() + + @property + def has_vocab_parallel_logits(self) -> bool: + return self.logits_token_ids is not None + + def index_select_logits_rows(self, index: torch.Tensor) -> "ModelOutput": + """Select logit rows without dropping their vocabulary metadata.""" + + return ModelOutput( + logits=self.logits.index_select(0, index), + logits_token_ids=( + self.logits_token_ids.index_select(0, index) if self.logits_token_ids is not None else None + ), + logits_logsumexp=( + self.logits_logsumexp.index_select(0, index) if self.logits_logsumexp is not None else None + ), + ) diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 29648aa78e..5c1e361729 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -52,6 +52,9 @@ def __init__(self): self.mem_index: torch.Tensor = None self.return_all_prompt_logics: bool = False + self.use_vocab_parallel_greedy: bool = False + self.logits_token_ids: Optional[torch.Tensor] = None + self.logits_logsumexp: Optional[torch.Tensor] = None # 在开启 return_all_prompt_logics 模式时，保存整个 prefill 阶段每一个 # token 位置的 logits，供后续回传 prompt logprobs 信息使用。 # 仅在 prefill 阶段且需要返回 prompt logprobs 时才会被填充。 diff --git a/lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py b/lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py new file mode 100644 index 0000000000..8e7e99e1ab --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py @@ -0,0 +1,109 @@ +"""Local greedy statistics for distributed vocabulary shards.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _greedy_sample_stage1_kernel( + logits, + partial_max, + partial_sum, + partial_argmax, + stride_row, + stride_col, + vocab_size: tl.constexpr, + num_blocks: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + block = tl.program_id(1) + offsets = block * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + values = tl.load( + logits + row * stride_row + offsets * stride_col, + mask=offsets < vocab_size, + other=-float("inf"), + ) + values = values.to(tl.float32) + + block_max = tl.max(values, axis=0) + block_sum = tl.sum(tl.exp(values - block_max), axis=0) + block_argmax = tl.argmax(values, axis=0) + block * BLOCK_SIZE + output_offset = row * num_blocks + block + tl.store(partial_max + output_offset, block_max) + tl.store(partial_sum + output_offset, block_sum) + tl.store(partial_argmax + output_offset, block_argmax) + + +@triton.jit +def _greedy_sample_stage2_stats_kernel( + partial_max, + partial_sum, + partial_argmax, + output_stats, + output_argmax, + num_blocks: tl.constexpr, + batch_size: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < num_blocks + input_offset = row * num_blocks + offsets + block_max = tl.load(partial_max + input_offset, mask=mask, other=-float("inf")) + block_sum = tl.load(partial_sum + input_offset, mask=mask, other=0.0) + block_argmax = tl.load(partial_argmax + input_offset, mask=mask, other=0x7FFFFFFF) + + global_max = tl.max(block_max, axis=0) + global_sum = tl.sum(block_sum * tl.exp(block_max - global_max), axis=0) + candidate_ids = tl.where(block_max == global_max, block_argmax, 0x7FFFFFFF) + global_argmax = tl.min(candidate_ids, axis=0) + tl.store(output_stats + row, global_max) + tl.store(output_stats + batch_size + row, global_max + tl.log(global_sum)) + tl.store(output_argmax + row, global_argmax) + + +def _launch_stage1(logits: torch.Tensor, scratch: torch.Tensor, block_size: int, num_blocks: int) -> None: + batch_size, vocab_size = logits.shape + _greedy_sample_stage1_kernel[(batch_size, num_blocks)]( + logits, + scratch[0], + scratch[1], + scratch[2], + logits.stride(0), + logits.stride(1), + vocab_size=vocab_size, + num_blocks=num_blocks, + BLOCK_SIZE=block_size, + num_warps=8, + ) + + +@torch.no_grad() +def greedy_sample_local_stats(logits: torch.Tensor, alloc_func=torch.empty) -> torch.Tensor: + """Return local max, logsumexp and argmax rows for distributed greedy sampling.""" + + assert logits.ndim == 2 and logits.is_cuda and logits.is_contiguous() + batch_size, vocab_size = logits.shape + block_size = 4096 + num_blocks = triton.cdiv(vocab_size, block_size) + scratch = alloc_func((3, batch_size, num_blocks), dtype=torch.float32, device=logits.device) + # The third FP32 row carries INT32 argmax bits. Keeping one fixed-size + # payload gives the distributed reducer a single collective without losing + # token-id precision through a numeric int-to-float conversion. + output_stats = alloc_func((3, batch_size), dtype=torch.float32, device=logits.device) + + _launch_stage1(logits, scratch, block_size, num_blocks) + _greedy_sample_stage2_stats_kernel[(batch_size,)]( + scratch[0], + scratch[1], + scratch[2], + output_stats, + output_stats[2].view(torch.int32), + num_blocks=num_blocks, + batch_size=batch_size, + BLOCK_SIZE=triton.next_power_of_2(num_blocks), + num_warps=4, + ) + return output_stats diff --git a/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py b/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py new file mode 100644 index 0000000000..fbfa4191e6 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py @@ -0,0 +1,118 @@ +"""Greedy sampling directly from tensor-parallel vocabulary shards.""" + +import torch +import triton +import triton.language as tl + +from lightllm.common.basemodel.triton_kernel.post_process.greedy_sample import ( + greedy_sample_local_stats, +) +from lightllm.common.basemodel.triton_kernel.transpose_convert import ( + transpose_convert_2d, +) +from lightllm.distributed.communication_op import all_gather_into_tensor + + +@triton.jit +def _combine_vocab_parallel_stats_kernel( + gathered_stats, + gathered_argmax, + output_logits, + output_token_ids, + output_logsumexp, + token_num, + vocab_size: tl.constexpr, + tp_world_size: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + token_offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + token_mask = token_offsets < token_num + rank_stride = 3 * token_num + + global_max = tl.full((BLOCK_SIZE,), -float("inf"), tl.float32) + global_id = tl.full((BLOCK_SIZE,), 0x7FFFFFFF, tl.int32) + for rank in tl.static_range(tp_world_size): + rank_base = rank * rank_stride + local_max = tl.load( + gathered_stats + rank_base + token_offsets, + mask=token_mask, + other=-float("inf"), + ) + local_id = tl.load( + gathered_argmax + rank_base + 2 * token_num + token_offsets, + mask=token_mask, + other=0x7FFFFFFF, + ) + local_id += (rank * vocab_size) // tp_world_size + wins = (local_max > global_max) | ((local_max == global_max) & (local_id < global_id)) + global_max = tl.where(wins, local_max, global_max) + global_id = tl.where(wins, local_id, global_id) + + global_sum = tl.zeros((BLOCK_SIZE,), tl.float32) + for rank in tl.static_range(tp_world_size): + rank_base = rank * rank_stride + local_lse = tl.load( + gathered_stats + rank_base + token_num + token_offsets, + mask=token_mask, + other=-float("inf"), + ) + global_sum += tl.exp(local_lse - global_max) + + tl.store(output_logits + token_offsets, global_max, mask=token_mask) + tl.store(output_token_ids + token_offsets, global_id, mask=token_mask) + tl.store(output_logsumexp + token_offsets, global_max + tl.log(global_sum), mask=token_mask) + + +@torch.no_grad() +def vocab_parallel_greedy( + local_logits: torch.Tensor, + *, + vocab_size: int, + tp_world_size: int, + group, + alloc_func, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return exact sparse logits, global token ids and full-vocab logsumexp.""" + + assert local_logits.ndim == 2 and local_logits.is_cuda and local_logits.is_contiguous() + local_vocab_size, token_num = local_logits.shape + assert local_vocab_size in { + vocab_size // tp_world_size, + (vocab_size + tp_world_size - 1) // tp_world_size, + } + + transposed_logits = alloc_func( + (token_num, local_vocab_size), + dtype=local_logits.dtype, + device=local_logits.device, + ) + transpose_convert_2d(local_logits, transposed_logits) + local_stats = greedy_sample_local_stats(transposed_logits, alloc_func=alloc_func) + + if tp_world_size == 1: + gathered_stats = local_stats.view(1, 3, token_num) + else: + gathered_stats = alloc_func((tp_world_size, 3, token_num), dtype=torch.float32, device=local_logits.device) + all_gather_into_tensor( + output_=gathered_stats, + input_=local_stats, + group=group, + async_op=False, + ) + + output_logits = alloc_func((token_num, 1), dtype=torch.float32, device=local_logits.device) + output_token_ids = alloc_func((token_num, 1), dtype=torch.int64, device=local_logits.device) + output_logsumexp = alloc_func((token_num,), dtype=torch.float32, device=local_logits.device) + _combine_vocab_parallel_stats_kernel[(triton.cdiv(token_num, 256),)]( + gathered_stats, + gathered_stats.view(torch.int32), + output_logits, + output_token_ids, + output_logsumexp, + token_num, + vocab_size=vocab_size, + tp_world_size=tp_world_size, + BLOCK_SIZE=256, + num_warps=4, + ) + return output_logits, output_token_ids, output_logsumexp diff --git a/lightllm/common/basemodel/triton_kernel/transpose_convert.py b/lightllm/common/basemodel/triton_kernel/transpose_convert.py new file mode 100644 index 0000000000..b618d69526 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/transpose_convert.py @@ -0,0 +1,65 @@ +"""Tiled transpose kernels used by the post-layer logits path.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _transpose_convert_2d_kernel( + input_ptr, + output_ptr, + rows, + cols, + input_stride_0, + input_stride_1, + output_stride_0, + output_stride_1, + BLOCK_ROWS: tl.constexpr, + BLOCK_COLS: tl.constexpr, +): + row_offsets = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + col_offsets = tl.program_id(1) * BLOCK_COLS + tl.arange(0, BLOCK_COLS) + input_offsets = row_offsets[:, None] * input_stride_0 + col_offsets[None, :] * input_stride_1 + mask = (row_offsets[:, None] < rows) & (col_offsets[None, :] < cols) + values = tl.load(input_ptr + input_offsets, mask=mask) + + output_offsets = col_offsets[:, None] * output_stride_0 + row_offsets[None, :] * output_stride_1 + tl.store(output_ptr + output_offsets, tl.trans(values), mask=tl.trans(mask)) + + +@torch.no_grad() +def transpose_convert_2d( + input_tensor: torch.Tensor, + output_tensor: torch.Tensor, + *, + block_rows: int = 64, + block_cols: int = 64, + num_warps: int = 8, + num_stages: int = 1, +) -> torch.Tensor: + """Transpose a contiguous 2-D CUDA tensor while converting its dtype.""" + + assert input_tensor.is_cuda and output_tensor.is_cuda + assert input_tensor.device == output_tensor.device + assert input_tensor.ndim == 2 and output_tensor.ndim == 2 + assert output_tensor.shape == (input_tensor.shape[1], input_tensor.shape[0]) + assert input_tensor.is_contiguous() and output_tensor.is_contiguous() + + rows, cols = input_tensor.shape + grid = (triton.cdiv(rows, block_rows), triton.cdiv(cols, block_cols)) + _transpose_convert_2d_kernel[grid]( + input_tensor, + output_tensor, + rows, + cols, + input_tensor.stride(0), + input_tensor.stride(1), + output_tensor.stride(0), + output_tensor.stride(1), + BLOCK_ROWS=block_rows, + BLOCK_COLS=block_cols, + num_warps=num_warps, + num_stages=num_stages, + ) + return output_tensor diff --git a/lightllm/models/llama/layer_infer/post_layer_infer.py b/lightllm/models/llama/layer_infer/post_layer_infer.py index 6e4b15a55d..11d916b9e7 100644 --- a/lightllm/models/llama/layer_infer/post_layer_infer.py +++ b/lightllm/models/llama/layer_infer/post_layer_infer.py @@ -7,6 +7,9 @@ from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight from lightllm.models.llama.infer_struct import LlamaInferStateInfo from lightllm.common.basemodel import PostLayerInferTpl +from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy import ( + vocab_parallel_greedy, +) from lightllm.distributed.communication_op import all_gather @@ -64,7 +67,7 @@ def _token_forward( if prompt_logics_hiddens is not None: prompt_token_num = prompt_logics_hiddens.shape[0] infer_state.prompt_logics = self._lm_head_and_gather( - prompt_logics_hiddens, prompt_token_num, layer_weight, infer_state + prompt_logics_hiddens, prompt_token_num, layer_weight, infer_state, force_full_logits=True ) return ans_logics @@ -75,6 +78,7 @@ def _lm_head_and_gather( token_num: int, layer_weight: LlamaPreAndPostLayerWeight, infer_state: LlamaInferStateInfo, + force_full_logits: bool = False, ) -> torch.Tensor: normed = self._norm(hidden, infer_state, layer_weight) normed = normed.permute(1, 0).view(-1, token_num) @@ -82,6 +86,18 @@ def _lm_head_and_gather( normed = None vocab_size = layer_weight.lm_head_weight_.vocab_size + if infer_state.use_vocab_parallel_greedy and not force_full_logits: + logits, token_ids, logsumexp = vocab_parallel_greedy( + logic_batch, + vocab_size=vocab_size, + tp_world_size=self.tp_world_size_, + group=infer_state.dist_group, + alloc_func=self.alloc_tensor, + ) + infer_state.logits_token_ids = token_ids + infer_state.logits_logsumexp = logsumexp + return logits + if self.tp_world_size_ == 1: gather_data = logic_batch else: diff --git a/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py b/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py index 5a74cd988e..eb4481fd6b 100644 --- a/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py +++ b/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py @@ -181,7 +181,11 @@ def token_forward( logits = self._lm_head_and_gather(last_input, token_num, layer_weight, infer_state) block_logits = logits.reshape(num_reqs, self.block_size_, -1) - sampled_tokens = torch.argmax(block_logits, dim=-1) + if infer_state.logits_token_ids is None: + sampled_tokens = torch.argmax(block_logits, dim=-1) + else: + assert block_logits.shape[-1] == 1 + sampled_tokens = infer_state.logits_token_ids.reshape(num_reqs, self.block_size_) confidence_logits = self.predict_confidence_logits( block_hidden, anchor_token_ids=anchor_token_ids, diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 07b88471f0..1913f82547 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -876,13 +876,23 @@ def _trans_req_ids_to_req_objs(self, req_ids: List[int]) -> List[InferReq]: def _gen_argmax_token_ids(self, model_output: ModelOutput): logits = model_output.logits - return torch.argmax(logits, dim=-1) + candidate_indexes = torch.argmax(logits, dim=-1) + return self._map_logits_indexes_to_token_ids(model_output, candidate_indexes) def _gen_argmax_token_ids_and_prob(self, model_output: ModelOutput): logits = model_output.logits - probs = torch.softmax(logits, dim=-1) - max_probs, draft_next_token_ids_gpu = torch.max(probs, dim=-1) - return draft_next_token_ids_gpu, max_probs + if model_output.has_vocab_parallel_logits: + max_logits, candidate_indexes = torch.max(logits, dim=-1) + token_ids = self._map_logits_indexes_to_token_ids(model_output, candidate_indexes) + return token_ids, torch.exp(max_logits - model_output.logits_logsumexp) + max_probs, token_ids = torch.max(torch.softmax(logits, dim=-1), dim=-1) + return token_ids, max_probs + + @staticmethod + def _map_logits_indexes_to_token_ids(model_output: ModelOutput, candidate_indexes: torch.Tensor): + if not model_output.has_vocab_parallel_logits: + return candidate_indexes + return model_output.logits_token_ids.gather(1, candidate_indexes.long().view(-1, 1)).view(-1).long() def _sample_and_scatter_token( self, diff --git a/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/eagle_with_att.py b/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/eagle_with_att.py index 6b8c23e8fd..137a580ae6 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/eagle_with_att.py +++ b/lightllm/server/router/model_infer/mtp_speculative/dp_overlap_proposers/eagle_with_att.py @@ -141,7 +141,7 @@ def propose_next_overlap( req_num_by_batch, ) ): - accepted_tail_output = ModelOutput(logits=extend_output.logits.index_select(0, accepted_tail_rows)) + accepted_tail_output = extend_output.index_select_logits_rows(accepted_tail_rows) if self.enable_dynmaic_mtp: draft_token_ids, draft_token_probs = self._gen_argmax_token_ids_and_prob(accepted_tail_output) draft_token_probs = draft_token_probs.float() diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py index 3d2c0a0e86..f2c730ccd5 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/eagle_with_att.py @@ -81,7 +81,7 @@ def propose_next( # 只在 req_num 行 logits 上进行 argmax，避免为未接受的 verify 行执行 # vocabulary reduction。第一列 proposal 来自每个请求的 accepted tail。 - accepted_tail_output = ModelOutput(logits=extend_output.logits.index_select(0, accepted_tail_rows)) + accepted_tail_output = extend_output.index_select_logits_rows(accepted_tail_rows) if self.enable_dynmaic_mtp: draft_token_ids, draft_token_probs = self._gen_argmax_token_ids_and_prob(accepted_tail_output) schedule_scores_by_step.append(draft_token_probs.float().unsqueeze(1)) diff --git a/unit_tests/common/basemodel/test_model_output.py b/unit_tests/common/basemodel/test_model_output.py index 6f9477e294..8ec94ae5dc 100644 --- a/unit_tests/common/basemodel/test_model_output.py +++ b/unit_tests/common/basemodel/test_model_output.py @@ -7,20 +7,40 @@ from lightllm.common.basemodel.batch_objs import ModelInput, ModelMtpOutputCollector, ModelOutput +def test_vocab_parallel_metadata_follows_row_selection(): + output = ModelOutput( + logits=torch.tensor([[8.0], [7.0], [6.0]]), + logits_token_ids=torch.tensor([[4], [9], [2]]), + logits_logsumexp=torch.tensor([8.5, 7.25, 6.75]), + ) + + selected = output.index_select_logits_rows(torch.tensor([2, 0])) + + torch.testing.assert_close(selected.logits.view(-1), torch.tensor([6.0, 8.0])) + torch.testing.assert_close(selected.logits_token_ids.view(-1), torch.tensor([2, 4])) + torch.testing.assert_close(selected.logits_logsumexp, torch.tensor([6.75, 8.5])) + + def test_decode_unpad_slices_spec_output_with_logits(): model = TpPartBaseModel.__new__(TpPartBaseModel) output = ModelOutput( logits=torch.arange(24).view(6, 4), + logits_token_ids=torch.arange(100, 124).view(6, 4), + logits_logsumexp=torch.arange(6, dtype=torch.float32), mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(18).view(6, 3)), ) unpadded = model._create_unpad_decode_model_output(output, origin_batch_size=4) assert unpadded.logits.shape == (4, 4) + assert unpadded.logits_token_ids.shape == (4, 4) + assert unpadded.logits_logsumexp.shape == (4,) assert unpadded.mtp_collector.spec_hidden.shape == (4, 3) # Unpadding returns a shallow output copy and leaves the graph-owned # tensors on the original ModelOutput intact. assert output.logits.shape == (6, 4) + assert output.logits_token_ids.shape == (6, 4) + assert output.logits_logsumexp.shape == (6,) assert output.mtp_collector.spec_hidden.shape == (6, 3) @@ -28,6 +48,8 @@ def test_prefill_unpad_uses_token_rows_for_spec_hidden(): model = TpPartBaseModel.__new__(TpPartBaseModel) output = ModelOutput( logits=torch.arange(20).view(5, 4), + logits_token_ids=torch.arange(100, 120).view(5, 4), + logits_logsumexp=torch.arange(5, dtype=torch.float32), mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(24).view(8, 3)), prompt_logics=torch.arange(32).view(8, 4), ) @@ -39,6 +61,8 @@ def test_prefill_unpad_uses_token_rows_for_spec_hidden(): ) assert unpadded.logits.shape == (3, 4) + assert unpadded.logits_token_ids.shape == (3, 4) + assert unpadded.logits_logsumexp.shape == (3,) assert unpadded.mtp_collector.spec_hidden.shape == (6, 3) assert unpadded.prompt_logics.shape == (6, 4) diff --git a/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py b/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py new file mode 100644 index 0000000000..279268a85c --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py @@ -0,0 +1,65 @@ +import importlib + +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.post_process.greedy_sample import ( + greedy_sample_local_stats, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for Triton kernels") + + +@pytest.mark.parametrize("token_num", [1, 7, 64]) +def test_vocab_parallel_greedy_matches_full_logits(monkeypatch, token_num): + module = importlib.import_module("lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy") + tp_world_size = 4 + local_vocab_size = 8192 + vocab_size = tp_world_size * local_vocab_size + generator = torch.Generator(device="cuda").manual_seed(20260826 + token_num) + local_logits_by_rank = [ + torch.randn( + (local_vocab_size, token_num), + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ) + for _ in range(tp_world_size) + ] + + # Exercise deterministic tie-breaking both between local reduction blocks + # and across tensor-parallel ranks. The smallest global token id must win. + local_logits_by_rank[0][4097, 0] = 20.0 + local_logits_by_rank[0][3, 0] = 20.0 + local_logits_by_rank[3][2, 0] = 20.0 + + local_stats_by_rank = [ + greedy_sample_local_stats(local_logits.transpose(0, 1).contiguous()) for local_logits in local_logits_by_rank + ] + + def fake_all_gather_into_tensor(output_, input_, **_kwargs): + for output, local_stats in zip(output_, local_stats_by_rank): + output.copy_(local_stats) + + monkeypatch.setattr(module, "all_gather_into_tensor", fake_all_gather_into_tensor) + actual_logits, actual_ids, actual_logsumexp = module.vocab_parallel_greedy( + local_logits_by_rank[0], + vocab_size=vocab_size, + tp_world_size=tp_world_size, + group=None, + alloc_func=torch.empty, + ) + actual_ids = actual_ids.view(-1) + actual_logprobs = actual_logits.view(-1) - actual_logsumexp + + full_logits = torch.cat(local_logits_by_rank, dim=0).transpose(0, 1).float() + expected_ids = full_logits.argmax(dim=1) + expected_logits = full_logits.gather(1, expected_ids[:, None]).view(-1) + expected_logsumexp = torch.logsumexp(full_logits, dim=1) + expected_logprobs = torch.log_softmax(full_logits, dim=1).gather(1, expected_ids[:, None]).squeeze(1) + + torch.testing.assert_close(actual_ids, expected_ids, rtol=0, atol=0) + torch.testing.assert_close(actual_logits.view(-1), expected_logits, rtol=0, atol=0) + torch.testing.assert_close(actual_logsumexp, expected_logsumexp, rtol=2e-4, atol=2e-4) + torch.testing.assert_close(actual_logprobs, expected_logprobs, rtol=2e-4, atol=2e-4) diff --git a/unit_tests/models/test_vocab_parallel_greedy_output.py b/unit_tests/models/test_vocab_parallel_greedy_output.py new file mode 100644 index 0000000000..88e5a7c282 --- /dev/null +++ b/unit_tests/models/test_vocab_parallel_greedy_output.py @@ -0,0 +1,74 @@ +from types import SimpleNamespace + +import torch + +from lightllm.common.basemodel.batch_objs import ModelOutput +from lightllm.models.qwen3_dspark.layer_infer.post_layer_infer import Qwen3DSparkPostLayerInfer +from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend + + +def test_argmax_restores_global_token_ids_and_exact_probabilities(): + backend = ModeBackend.__new__(ModeBackend) + output = ModelOutput( + logits=torch.tensor([[3.0, 1.0], [0.0, 5.0]]), + logits_token_ids=torch.tensor([[30, 10], [100, 500]]), + logits_logsumexp=torch.tensor([4.0, 5.25]), + ) + + token_ids = backend._gen_argmax_token_ids(output) + token_ids_with_prob, probs = backend._gen_argmax_token_ids_and_prob(output) + + torch.testing.assert_close(token_ids, torch.tensor([30, 500])) + torch.testing.assert_close(token_ids_with_prob, token_ids) + torch.testing.assert_close(probs, torch.exp(torch.tensor([-1.0, -0.25]))) + + +def test_dense_argmax_keeps_column_index_semantics(): + backend = ModeBackend.__new__(ModeBackend) + output = ModelOutput(logits=torch.tensor([[1.0, 4.0, 2.0]])) + + torch.testing.assert_close(backend._gen_argmax_token_ids(output), torch.tensor([1])) + + +def test_dspark_confidence_path_receives_global_token_ids(): + post = Qwen3DSparkPostLayerInfer.__new__(Qwen3DSparkPostLayerInfer) + post.block_size_ = 2 + post.markov_rank_ = 0 + post._slice_get_last_input = lambda input_embeddings, infer_state: (input_embeddings, 4) + sparse_logits = torch.tensor([[4.0], [5.0], [7.0], [9.0]]) + sparse_token_ids = torch.tensor([[40], [50], [70], [90]]) + + def gather_vocab_parallel(*args, **kwargs): + infer_state = args[3] + infer_state.logits_token_ids = sparse_token_ids + return sparse_logits + + post._lm_head_and_gather = gather_vocab_parallel + observed = {} + + def predict_confidence(block_hidden, anchor_token_ids, sampled_tokens, layer_weight): + observed["sampled_tokens"] = sampled_tokens + return None + + post.predict_confidence_logits = predict_confidence + + class Collector: + def add_mtp_outputs(self, **kwargs): + self.outputs = kwargs + + collector = Collector() + infer_state = SimpleNamespace( + is_prefill=False, + input_ids=torch.tensor([1, 0, 2, 0]), + logits_token_ids=None, + hidden_collector=collector, + ) + + returned_logits = post.token_forward( + input_embdings=torch.ones((4, 3)), + infer_state=infer_state, + layer_weight=object(), + ) + + torch.testing.assert_close(returned_logits, sparse_logits) + torch.testing.assert_close(observed["sampled_tokens"], torch.tensor([[40, 50], [70, 90]])) From ac8a45629dc6064568a2db73de438b11f47e68c0 Mon Sep 17 00:00:00 2001 From: sufubao Date: Wed, 2 Sep 2026 16:30:07 +0800 Subject: [PATCH 2/8] perf: share vocab-parallel top-k logits --- lightllm/common/basemodel/basemodel.py | 39 +++--- lightllm/common/basemodel/batch_objs.py | 31 +++-- lightllm/common/basemodel/cuda_graph.py | 5 + lightllm/common/basemodel/infer_struct.py | 3 +- .../common/basemodel/prefill_cuda_graph.py | 5 + .../post_process/greedy_sample.py | 109 ---------------- .../post_process/vocab_parallel_greedy.py | 118 ------------------ .../post_process/vocab_parallel_topk.py | 94 ++++++++++++++ .../triton_kernel/transpose_convert.py | 65 ---------- .../gemma4/layer_infer/post_layer_infer.py | 16 ++- .../llama/layer_infer/post_layer_infer.py | 27 ++-- .../layer_infer/post_layer_infer.py | 7 +- .../model_infer/mode_backend/base_backend.py | 30 +++-- .../mode_backend/chunked_prefill/impl.py | 30 +++-- .../chunked_prefill/impl_for_reward_model.py | 2 +- .../mode_backend/diverse_backend/impl.py | 9 +- .../mode_backend/dp_backend/impl.py | 52 +++----- .../mode_backend/generic_post_process.py | 43 ++++++- .../mode_backend/generic_pre_process.py | 11 ++ .../common/basemodel/test_model_output.py | 61 +++++++-- .../test_vocab_parallel_greedy.py | 65 ---------- .../triton_kernel/test_vocab_parallel_topk.py | 81 ++++++++++++ .../models/test_gemma4_vocab_parallel_topk.py | 67 ++++++++++ ....py => test_vocab_parallel_topk_output.py} | 18 ++- .../test_vocab_parallel_topk_sampling.py | 91 ++++++++++++++ 25 files changed, 585 insertions(+), 494 deletions(-) delete mode 100644 lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py delete mode 100644 lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py create mode 100644 lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_topk.py delete mode 100644 lightllm/common/basemodel/triton_kernel/transpose_convert.py delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py create mode 100644 unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_topk.py create mode 100644 unit_tests/models/test_gemma4_vocab_parallel_topk.py rename unit_tests/models/{test_vocab_parallel_greedy_output.py => test_vocab_parallel_topk_output.py} (78%) create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_topk_sampling.py diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index b561595b98..bebf2671cc 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -23,6 +23,9 @@ from lightllm.common.basemodel.prefill_cuda_graph import PrefillCudaGraph from lightllm.common.quantization import Quantcfg from lightllm.common.basemodel.triton_kernel.gather_token_id import gather_token, gather_token_prefill_decode_mixed +from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk import ( + is_vocab_parallel_topk_enabled, +) from lightllm.utils.log_utils import init_logger from lightllm.utils.dist_utils import get_dp_world_size from lightllm.utils.profile_max_tokens import profile_mtp_weight_memory @@ -378,13 +381,22 @@ def forward(self, model_input: ModelInput): else: return self._decode(model_input) + def _is_cuda_graph_output_compatible(self, *model_inputs: ModelInput) -> bool: + """Whether inputs match the dense/sparse contract captured at startup.""" + + return ( + self.is_mtp_draft_model + or not is_vocab_parallel_topk_enabled() + or all(model_input.use_vocab_parallel_topk for model_input in model_inputs) + ) + def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0): infer_state = self.infer_state_class() infer_state.hidden_collector = self.hidden_collector_prototype.new_instance() infer_state.input_ids = model_input.input_ids infer_state.is_prefill = model_input.is_prefill infer_state.return_all_prompt_logics = self.return_all_prompt_logics - infer_state.use_vocab_parallel_greedy = self.is_mtp_draft_model + infer_state.use_vocab_parallel_topk = self.is_mtp_draft_model or model_input.use_vocab_parallel_topk infer_state.batch_size = model_input.batch_size infer_state.total_token_num = model_input.total_token_num infer_state.max_q_seq_len = model_input.max_q_seq_len @@ -537,7 +549,6 @@ def _create_unpad_decode_model_output(self, model_output: ModelOutput, origin_ba new_model_output.logits = new_model_output.logits[0:origin_batch_size] if new_model_output.logits_token_ids is not None: new_model_output.logits_token_ids = new_model_output.logits_token_ids[0:origin_batch_size] - new_model_output.logits_logsumexp = new_model_output.logits_logsumexp[0:origin_batch_size] new_model_output.mtp_collector = model_output.mtp_collector.unpad_decode( padded_batch_size=padded_batch_size, origin_batch_size=origin_batch_size, @@ -552,7 +563,6 @@ def _create_unpad_prefill_model_output( new_model_output.logits = new_model_output.logits[0:origin_batch_size] if new_model_output.logits_token_ids is not None: new_model_output.logits_token_ids = new_model_output.logits_token_ids[0:origin_batch_size] - new_model_output.logits_logsumexp = new_model_output.logits_logsumexp[0:origin_batch_size] new_model_output.mtp_collector = padded_model_output.mtp_collector.unpad_prefill( origin_handle_token_num=origin_handle_token_num ) @@ -650,9 +660,13 @@ def _decode( # CUDA Graph 可能继续向上对齐 batch size，并因此加入 seq_len=2 的 # dummy request。先用最终可能出现的 KV 长度判断 graph，再统一 padding 一次。 infer_max_kv_seq_len = max(2, model_input.max_kv_seq_len) - use_cuda_graph = self.graph is not None and self.graph.can_run( - batch_size=infer_batch_size, - max_len_in_batch=infer_max_kv_seq_len, + use_cuda_graph = ( + self._is_cuda_graph_output_compatible(model_input) + and self.graph is not None + and self.graph.can_run( + batch_size=infer_batch_size, + max_len_in_batch=infer_max_kv_seq_len, + ) ) need_capture = False if use_cuda_graph: @@ -685,7 +699,6 @@ def _decode( @final def _context_forward(self, infer_state: InferStateInfo): - input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight) if self.args.enable_dp_prefill_balance: assert not self.args.enable_prefill_cudagraph, "not support now" @@ -745,7 +758,6 @@ def prefill_func(input_tensors, _infer_state): model_output = ModelOutput( logits=predict_logits.contiguous(), logits_token_ids=infer_state.logits_token_ids, - logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, ) @@ -776,7 +788,6 @@ def _token_forward(self, infer_state: InferStateInfo): model_output = ModelOutput( logits=predict_logits.contiguous(), logits_token_ids=infer_state.logits_token_ids, - logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), ) @@ -908,7 +919,11 @@ def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1 infer_batch_size = max(1, origin_batch_size0, origin_batch_size1) infer_batch_size = triton.cdiv(infer_batch_size, self.tp_world_size_) * self.tp_world_size_ - if self.graph is not None and self.graph.can_run(infer_batch_size, max_len_in_batch): + if ( + self._is_cuda_graph_output_compatible(model_input0, model_input1) + and self.graph is not None + and self.graph.can_run(infer_batch_size, max_len_in_batch) + ): infer_batch_size = self.graph.find_closest_graph_batch_size(infer_batch_size) need_capture = self.graph.need_capture(infer_batch_size) padded_model_input0 = self._create_padded_decode_model_input(model_input0, infer_batch_size) @@ -1032,14 +1047,12 @@ def _overlap_tpsp_context_forward(self, infer_state: InferStateInfo, infer_state model_output = ModelOutput( logits=predict_logits.contiguous(), logits_token_ids=infer_state.logits_token_ids, - logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, ) model_output1 = ModelOutput( logits=predict_logits1.contiguous(), logits_token_ids=infer_state1.logits_token_ids, - logits_logsumexp=infer_state1.logits_logsumexp, mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1), prompt_logics=infer_state1.prompt_logics, ) @@ -1085,13 +1098,11 @@ def _overlap_tpsp_token_forward(self, infer_state: InferStateInfo, infer_state1: model_output = ModelOutput( logits=predict_logits.contiguous(), logits_token_ids=infer_state.logits_token_ids, - logits_logsumexp=infer_state.logits_logsumexp, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), ) model_output1 = ModelOutput( logits=predict_logits1.contiguous(), logits_token_ids=infer_state1.logits_token_ids, - logits_logsumexp=infer_state1.logits_logsumexp, mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1), ) diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index edddb749e5..77ed1d400d 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -55,6 +55,10 @@ class ModelInput: # 的 draft 模型的输入 mtp_draft_input_hiddens: Optional[torch.Tensor] = None + # The router sets this only when a target-model batch can sample directly + # from sparse candidates. Draft models always enable the same output form. + use_vocab_parallel_topk: bool = False + def to_cuda(self): self.check_input() @@ -200,31 +204,23 @@ class ModelOutput: # 需要返回 prompt logprobs 信息时才会非空。 prompt_logics: Optional[torch.Tensor] = None - # Vocab-parallel outputs keep logits as logits while mapping each sparse - # column back to its global token id. logits_logsumexp is computed over the - # complete vocabulary, so sparse argmax probabilities remain exact. - # Both fields are None for historical dense logits. + # Sparse vocab-parallel outputs map every candidate column back to its + # global token id. None means logits are dense and column indexes are ids. logits_token_ids: Optional[torch.Tensor] = None - logits_logsumexp: Optional[torch.Tensor] = None def __post_init__(self) -> None: if self.mtp_collector is None: self.mtp_collector = ModelMtpOutputCollector() - assert (self.logits_token_ids is None) == (self.logits_logsumexp is None) if self.logits_token_ids is not None: assert self.logits.ndim == 2 assert self.logits_token_ids.shape == self.logits.shape assert self.logits_token_ids.dtype in (torch.int32, torch.int64) assert self.logits_token_ids.device == self.logits.device - assert self.logits_logsumexp.shape == (self.logits.shape[0],) - assert self.logits_logsumexp.dtype == torch.float32 - assert self.logits_logsumexp.device == self.logits.device def to_no_ref_tensor(self): self.logits = tensor_to_no_ref_tensor(self.logits) if self.logits_token_ids is not None: self.logits_token_ids = tensor_to_no_ref_tensor(self.logits_token_ids) - self.logits_logsumexp = tensor_to_no_ref_tensor(self.logits_logsumexp) self.mtp_collector.to_no_ref_tensor() @property @@ -239,7 +235,18 @@ def index_select_logits_rows(self, index: torch.Tensor) -> "ModelOutput": logits_token_ids=( self.logits_token_ids.index_select(0, index) if self.logits_token_ids is not None else None ), - logits_logsumexp=( - self.logits_logsumexp.index_select(0, index) if self.logits_logsumexp is not None else None + ) + + @classmethod + def concat_logits_rows(cls, outputs: List["ModelOutput"]) -> "ModelOutput": + """Concatenate outputs that share the same dense or sparse layout.""" + + assert outputs + has_vocab_parallel_logits = outputs[0].has_vocab_parallel_logits + assert all(output.has_vocab_parallel_logits == has_vocab_parallel_logits for output in outputs) + return cls( + logits=torch.cat([output.logits for output in outputs], dim=0), + logits_token_ids=( + torch.cat([output.logits_token_ids for output in outputs], dim=0) if has_vocab_parallel_logits else None ), ) diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index 5849cccf54..c080c6f49e 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -9,6 +9,9 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.distributed import dist_group_manager from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput +from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk import ( + is_vocab_parallel_topk_enabled, +) from lightllm.utils.torch_memory_saver_utils import ( TorchMemorySaverWrapper, MemoryTag, @@ -279,6 +282,7 @@ def warmup(self, model): b_position_delta=torch.zeros(batch_size, dtype=torch.int32, device="cuda"), is_prefill=False, multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], + use_vocab_parallel_topk=is_vocab_parallel_topk_enabled(), **model._gen_special_model_input(batch_size), ) model_output: ModelOutput = model.forward(model_input) @@ -340,6 +344,7 @@ def warmup_overlap(self, model): b_shared_radix_node_id=b_shared_radix_node_id, b_position_delta=torch.zeros(batch_size, dtype=torch.int32, device="cuda"), multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], + use_vocab_parallel_topk=is_vocab_parallel_topk_enabled(), **model._gen_special_model_input(batch_size), ) decode_batches.append(micro_batch) diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 5c1e361729..3091c878a9 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -52,9 +52,8 @@ def __init__(self): self.mem_index: torch.Tensor = None self.return_all_prompt_logics: bool = False - self.use_vocab_parallel_greedy: bool = False + self.use_vocab_parallel_topk: bool = False self.logits_token_ids: Optional[torch.Tensor] = None - self.logits_logsumexp: Optional[torch.Tensor] = None # 在开启 return_all_prompt_logics 模式时，保存整个 prefill 阶段每一个 # token 位置的 logits，供后续回传 prompt logprobs 信息使用。 # 仅在 prefill 阶段且需要返回 prompt logprobs 时才会被填充。 diff --git a/lightllm/common/basemodel/prefill_cuda_graph.py b/lightllm/common/basemodel/prefill_cuda_graph.py index bf6039a48f..594dd1ef94 100644 --- a/lightllm/common/basemodel/prefill_cuda_graph.py +++ b/lightllm/common/basemodel/prefill_cuda_graph.py @@ -10,6 +10,9 @@ from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor from lightllm.distributed import dist_group_manager from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput +from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk import ( + is_vocab_parallel_topk_enabled, +) from .infer_struct import InferStateInfo from .cuda_graph import CudaGraph @@ -220,6 +223,7 @@ def warmup(self, model): is_prefill=True, b_prefill_has_output_cpu=[False], multimodal_params=[{"images": [], "audios": []}], + use_vocab_parallel_topk=is_vocab_parallel_topk_enabled(), **model._gen_special_model_input(token_num=total_token_num), ) model_output: ModelOutput = model.forward(model_input) @@ -281,6 +285,7 @@ def warmup_overlap(self, model): is_prefill=True, b_prefill_has_output_cpu=[False], multimodal_params=[{"images": [], "audios": []}], + use_vocab_parallel_topk=is_vocab_parallel_topk_enabled(), **model._gen_special_model_input(token_num=total_token_num), ) diff --git a/lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py b/lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py deleted file mode 100644 index 8e7e99e1ab..0000000000 --- a/lightllm/common/basemodel/triton_kernel/post_process/greedy_sample.py +++ /dev/null @@ -1,109 +0,0 @@ -"""Local greedy statistics for distributed vocabulary shards.""" - -import torch -import triton -import triton.language as tl - - -@triton.jit -def _greedy_sample_stage1_kernel( - logits, - partial_max, - partial_sum, - partial_argmax, - stride_row, - stride_col, - vocab_size: tl.constexpr, - num_blocks: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - row = tl.program_id(0) - block = tl.program_id(1) - offsets = block * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - values = tl.load( - logits + row * stride_row + offsets * stride_col, - mask=offsets < vocab_size, - other=-float("inf"), - ) - values = values.to(tl.float32) - - block_max = tl.max(values, axis=0) - block_sum = tl.sum(tl.exp(values - block_max), axis=0) - block_argmax = tl.argmax(values, axis=0) + block * BLOCK_SIZE - output_offset = row * num_blocks + block - tl.store(partial_max + output_offset, block_max) - tl.store(partial_sum + output_offset, block_sum) - tl.store(partial_argmax + output_offset, block_argmax) - - -@triton.jit -def _greedy_sample_stage2_stats_kernel( - partial_max, - partial_sum, - partial_argmax, - output_stats, - output_argmax, - num_blocks: tl.constexpr, - batch_size: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - row = tl.program_id(0) - offsets = tl.arange(0, BLOCK_SIZE) - mask = offsets < num_blocks - input_offset = row * num_blocks + offsets - block_max = tl.load(partial_max + input_offset, mask=mask, other=-float("inf")) - block_sum = tl.load(partial_sum + input_offset, mask=mask, other=0.0) - block_argmax = tl.load(partial_argmax + input_offset, mask=mask, other=0x7FFFFFFF) - - global_max = tl.max(block_max, axis=0) - global_sum = tl.sum(block_sum * tl.exp(block_max - global_max), axis=0) - candidate_ids = tl.where(block_max == global_max, block_argmax, 0x7FFFFFFF) - global_argmax = tl.min(candidate_ids, axis=0) - tl.store(output_stats + row, global_max) - tl.store(output_stats + batch_size + row, global_max + tl.log(global_sum)) - tl.store(output_argmax + row, global_argmax) - - -def _launch_stage1(logits: torch.Tensor, scratch: torch.Tensor, block_size: int, num_blocks: int) -> None: - batch_size, vocab_size = logits.shape - _greedy_sample_stage1_kernel[(batch_size, num_blocks)]( - logits, - scratch[0], - scratch[1], - scratch[2], - logits.stride(0), - logits.stride(1), - vocab_size=vocab_size, - num_blocks=num_blocks, - BLOCK_SIZE=block_size, - num_warps=8, - ) - - -@torch.no_grad() -def greedy_sample_local_stats(logits: torch.Tensor, alloc_func=torch.empty) -> torch.Tensor: - """Return local max, logsumexp and argmax rows for distributed greedy sampling.""" - - assert logits.ndim == 2 and logits.is_cuda and logits.is_contiguous() - batch_size, vocab_size = logits.shape - block_size = 4096 - num_blocks = triton.cdiv(vocab_size, block_size) - scratch = alloc_func((3, batch_size, num_blocks), dtype=torch.float32, device=logits.device) - # The third FP32 row carries INT32 argmax bits. Keeping one fixed-size - # payload gives the distributed reducer a single collective without losing - # token-id precision through a numeric int-to-float conversion. - output_stats = alloc_func((3, batch_size), dtype=torch.float32, device=logits.device) - - _launch_stage1(logits, scratch, block_size, num_blocks) - _greedy_sample_stage2_stats_kernel[(batch_size,)]( - scratch[0], - scratch[1], - scratch[2], - output_stats, - output_stats[2].view(torch.int32), - num_blocks=num_blocks, - batch_size=batch_size, - BLOCK_SIZE=triton.next_power_of_2(num_blocks), - num_warps=4, - ) - return output_stats diff --git a/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py b/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py deleted file mode 100644 index fbfa4191e6..0000000000 --- a/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_greedy.py +++ /dev/null @@ -1,118 +0,0 @@ -"""Greedy sampling directly from tensor-parallel vocabulary shards.""" - -import torch -import triton -import triton.language as tl - -from lightllm.common.basemodel.triton_kernel.post_process.greedy_sample import ( - greedy_sample_local_stats, -) -from lightllm.common.basemodel.triton_kernel.transpose_convert import ( - transpose_convert_2d, -) -from lightllm.distributed.communication_op import all_gather_into_tensor - - -@triton.jit -def _combine_vocab_parallel_stats_kernel( - gathered_stats, - gathered_argmax, - output_logits, - output_token_ids, - output_logsumexp, - token_num, - vocab_size: tl.constexpr, - tp_world_size: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - token_offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - token_mask = token_offsets < token_num - rank_stride = 3 * token_num - - global_max = tl.full((BLOCK_SIZE,), -float("inf"), tl.float32) - global_id = tl.full((BLOCK_SIZE,), 0x7FFFFFFF, tl.int32) - for rank in tl.static_range(tp_world_size): - rank_base = rank * rank_stride - local_max = tl.load( - gathered_stats + rank_base + token_offsets, - mask=token_mask, - other=-float("inf"), - ) - local_id = tl.load( - gathered_argmax + rank_base + 2 * token_num + token_offsets, - mask=token_mask, - other=0x7FFFFFFF, - ) - local_id += (rank * vocab_size) // tp_world_size - wins = (local_max > global_max) | ((local_max == global_max) & (local_id < global_id)) - global_max = tl.where(wins, local_max, global_max) - global_id = tl.where(wins, local_id, global_id) - - global_sum = tl.zeros((BLOCK_SIZE,), tl.float32) - for rank in tl.static_range(tp_world_size): - rank_base = rank * rank_stride - local_lse = tl.load( - gathered_stats + rank_base + token_num + token_offsets, - mask=token_mask, - other=-float("inf"), - ) - global_sum += tl.exp(local_lse - global_max) - - tl.store(output_logits + token_offsets, global_max, mask=token_mask) - tl.store(output_token_ids + token_offsets, global_id, mask=token_mask) - tl.store(output_logsumexp + token_offsets, global_max + tl.log(global_sum), mask=token_mask) - - -@torch.no_grad() -def vocab_parallel_greedy( - local_logits: torch.Tensor, - *, - vocab_size: int, - tp_world_size: int, - group, - alloc_func, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Return exact sparse logits, global token ids and full-vocab logsumexp.""" - - assert local_logits.ndim == 2 and local_logits.is_cuda and local_logits.is_contiguous() - local_vocab_size, token_num = local_logits.shape - assert local_vocab_size in { - vocab_size // tp_world_size, - (vocab_size + tp_world_size - 1) // tp_world_size, - } - - transposed_logits = alloc_func( - (token_num, local_vocab_size), - dtype=local_logits.dtype, - device=local_logits.device, - ) - transpose_convert_2d(local_logits, transposed_logits) - local_stats = greedy_sample_local_stats(transposed_logits, alloc_func=alloc_func) - - if tp_world_size == 1: - gathered_stats = local_stats.view(1, 3, token_num) - else: - gathered_stats = alloc_func((tp_world_size, 3, token_num), dtype=torch.float32, device=local_logits.device) - all_gather_into_tensor( - output_=gathered_stats, - input_=local_stats, - group=group, - async_op=False, - ) - - output_logits = alloc_func((token_num, 1), dtype=torch.float32, device=local_logits.device) - output_token_ids = alloc_func((token_num, 1), dtype=torch.int64, device=local_logits.device) - output_logsumexp = alloc_func((token_num,), dtype=torch.float32, device=local_logits.device) - _combine_vocab_parallel_stats_kernel[(triton.cdiv(token_num, 256),)]( - gathered_stats, - gathered_stats.view(torch.int32), - output_logits, - output_token_ids, - output_logsumexp, - token_num, - vocab_size=vocab_size, - tp_world_size=tp_world_size, - BLOCK_SIZE=256, - num_warps=4, - ) - return output_logits, output_token_ids, output_logsumexp diff --git a/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_topk.py b/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_topk.py new file mode 100644 index 0000000000..bcb48c0090 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_topk.py @@ -0,0 +1,94 @@ +"""Collect sparse candidates directly from tensor-parallel vocabulary shards.""" + +import os + +import torch + +from lightllm.distributed.communication_op import all_gather_into_tensor +from lightllm.utils.envs_utils import enable_env_vars + + +VOCAB_PARALLEL_TOPK_ENV = "LIGHTLLM_VOCAB_PARALLEL_TOPK" +VOCAB_PARALLEL_TOPK_SIZE_ENV = "LIGHTLLM_VOCAB_PARALLEL_TOPK_SIZE" +DEFAULT_VOCAB_PARALLEL_TOPK = 128 + + +def is_vocab_parallel_topk_enabled() -> bool: + """Whether target-model greedy batches may use sparse vocabulary output.""" + + return enable_env_vars(VOCAB_PARALLEL_TOPK_ENV) + + +def get_vocab_parallel_topk_size() -> int: + topk = int(os.getenv(VOCAB_PARALLEL_TOPK_SIZE_ENV, str(DEFAULT_VOCAB_PARALLEL_TOPK))) + assert topk > 0, f"{VOCAB_PARALLEL_TOPK_SIZE_ENV} must be positive, got {topk}" + return topk + + +@torch.no_grad() +def vocab_parallel_topk( + local_logits: torch.Tensor, + *, + vocab_size: int, + vocab_start_id: int, + topk: int, + tp_world_size: int, + group, + alloc_func, +) -> tuple[torch.Tensor, torch.Tensor]: + """Gather each TP rank's local top-k logits and their global token ids. + + The returned width is ``tp_world_size * topk``. It intentionally keeps the + union of local candidates: greedy selection remains exact, while probability + calculations over the sparse result are an inexpensive approximation. + """ + + assert local_logits.ndim == 2 and local_logits.is_cuda and local_logits.is_contiguous() + local_vocab_size, token_num = local_logits.shape + # Collectives require every rank to contribute the same shape. Vocabulary + # shards can differ by one row, so cap against the smallest possible shard. + local_topk = min(topk, vocab_size // tp_world_size) + assert local_topk > 0 + assert local_vocab_size >= local_topk + + local_values, local_indexes = torch.topk(local_logits, k=local_topk, dim=0, sorted=False) + local_values = local_values.float() + local_token_ids = local_indexes.to(torch.int32).add_(int(vocab_start_id)) + + if tp_world_size == 1: + candidate_values = local_values.permute(1, 0).contiguous() + candidate_token_ids = local_token_ids.permute(1, 0).contiguous() + else: + # Values and ids are both four bytes. Bit-packing the ids into the FP32 + # payload keeps the operation to one fixed-shape collective. + local_payload = alloc_func( + (local_topk * 2, token_num), + dtype=torch.float32, + device=local_logits.device, + ) + local_payload[:local_topk].copy_(local_values) + local_payload[local_topk:].view(torch.int32).copy_(local_token_ids) + + gathered_payload = alloc_func( + (tp_world_size, local_topk * 2, token_num), + dtype=torch.float32, + device=local_logits.device, + ) + all_gather_into_tensor( + output_=gathered_payload, + input_=local_payload, + group=group, + async_op=False, + ) + candidate_values = gathered_payload[:, :local_topk, :].permute(2, 0, 1).reshape(token_num, -1) + candidate_token_ids = ( + gathered_payload[:, local_topk:, :].view(torch.int32).permute(2, 0, 1).reshape(token_num, -1) + ) + + output_logits = alloc_func( + candidate_values.shape, + dtype=torch.float32, + device=local_logits.device, + ) + output_logits.copy_(candidate_values) + return output_logits, candidate_token_ids.to(torch.int64).contiguous() diff --git a/lightllm/common/basemodel/triton_kernel/transpose_convert.py b/lightllm/common/basemodel/triton_kernel/transpose_convert.py deleted file mode 100644 index b618d69526..0000000000 --- a/lightllm/common/basemodel/triton_kernel/transpose_convert.py +++ /dev/null @@ -1,65 +0,0 @@ -"""Tiled transpose kernels used by the post-layer logits path.""" - -import torch -import triton -import triton.language as tl - - -@triton.jit -def _transpose_convert_2d_kernel( - input_ptr, - output_ptr, - rows, - cols, - input_stride_0, - input_stride_1, - output_stride_0, - output_stride_1, - BLOCK_ROWS: tl.constexpr, - BLOCK_COLS: tl.constexpr, -): - row_offsets = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) - col_offsets = tl.program_id(1) * BLOCK_COLS + tl.arange(0, BLOCK_COLS) - input_offsets = row_offsets[:, None] * input_stride_0 + col_offsets[None, :] * input_stride_1 - mask = (row_offsets[:, None] < rows) & (col_offsets[None, :] < cols) - values = tl.load(input_ptr + input_offsets, mask=mask) - - output_offsets = col_offsets[:, None] * output_stride_0 + row_offsets[None, :] * output_stride_1 - tl.store(output_ptr + output_offsets, tl.trans(values), mask=tl.trans(mask)) - - -@torch.no_grad() -def transpose_convert_2d( - input_tensor: torch.Tensor, - output_tensor: torch.Tensor, - *, - block_rows: int = 64, - block_cols: int = 64, - num_warps: int = 8, - num_stages: int = 1, -) -> torch.Tensor: - """Transpose a contiguous 2-D CUDA tensor while converting its dtype.""" - - assert input_tensor.is_cuda and output_tensor.is_cuda - assert input_tensor.device == output_tensor.device - assert input_tensor.ndim == 2 and output_tensor.ndim == 2 - assert output_tensor.shape == (input_tensor.shape[1], input_tensor.shape[0]) - assert input_tensor.is_contiguous() and output_tensor.is_contiguous() - - rows, cols = input_tensor.shape - grid = (triton.cdiv(rows, block_rows), triton.cdiv(cols, block_cols)) - _transpose_convert_2d_kernel[grid]( - input_tensor, - output_tensor, - rows, - cols, - input_tensor.stride(0), - input_tensor.stride(1), - output_tensor.stride(0), - output_tensor.stride(1), - BLOCK_ROWS=block_rows, - BLOCK_COLS=block_cols, - num_warps=num_warps, - num_stages=num_stages, - ) - return output_tensor diff --git a/lightllm/models/gemma4/layer_infer/post_layer_infer.py b/lightllm/models/gemma4/layer_infer/post_layer_infer.py index b736a2d6c1..354d91705d 100644 --- a/lightllm/models/gemma4/layer_infer/post_layer_infer.py +++ b/lightllm/models/gemma4/layer_infer/post_layer_infer.py @@ -5,18 +5,16 @@ class Gemma4PostLayerInfer(LlamaPostLayerInfer): """ Same final RMSNorm + tied lm_head path as Llama, with an extra tanh-based - logit softcap at the end: logits = softcap * tanh(logits / softcap). + transform before sampling: logits = softcap * tanh(logits / softcap). """ def __init__(self, network_config): super().__init__(network_config) self.final_logit_softcapping = float(network_config.get("final_logit_softcapping")) - def token_forward(self, input_embdings, infer_state, layer_weight): - logits = super().token_forward(input_embdings, infer_state, layer_weight) - if self.final_logit_softcapping is not None and self.final_logit_softcapping > 0: - cap = self.final_logit_softcapping - logits = torch.tanh(logits / cap) * cap - if infer_state.prompt_logics is not None: - infer_state.prompt_logics = torch.tanh(infer_state.prompt_logics / cap) * cap - return logits + def _apply_logit_postprocessing(self, logits: torch.Tensor) -> torch.Tensor: + if self.final_logit_softcapping is None or self.final_logit_softcapping <= 0: + return logits + cap = self.final_logit_softcapping + # The historical path materializes FP32 logits before applying softcap. + return torch.tanh(logits.float() / cap) * cap diff --git a/lightllm/models/llama/layer_infer/post_layer_infer.py b/lightllm/models/llama/layer_infer/post_layer_infer.py index 11d916b9e7..496de58ef1 100644 --- a/lightllm/models/llama/layer_infer/post_layer_infer.py +++ b/lightllm/models/llama/layer_infer/post_layer_infer.py @@ -1,4 +1,3 @@ -import os import torch import torch.functional as F import torch.distributed as dist @@ -7,8 +6,9 @@ from lightllm.models.llama.layer_weights.pre_and_post_layer_weight import LlamaPreAndPostLayerWeight from lightllm.models.llama.infer_struct import LlamaInferStateInfo from lightllm.common.basemodel import PostLayerInferTpl -from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy import ( - vocab_parallel_greedy, +from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk import ( + get_vocab_parallel_topk_size, + vocab_parallel_topk, ) from lightllm.distributed.communication_op import all_gather @@ -19,11 +19,17 @@ class LlamaPostLayerInfer(PostLayerInferTpl): def __init__(self, network_config): super().__init__(network_config) self.eps_ = network_config["rms_norm_eps"] + self.vocab_parallel_topk_ = get_vocab_parallel_topk_size() return def _norm(self, input, infer_state, layer_weight: LlamaPreAndPostLayerWeight) -> torch.Tensor: return layer_weight.final_norm_weight_(input=input, eps=self.eps_, alloc_func=self.alloc_tensor) + def _apply_logit_postprocessing(self, logits: torch.Tensor) -> torch.Tensor: + """Apply model-specific transforms while the tensor still contains logits.""" + + return logits + def _slice_get_last_input(self, input_embdings: torch.Tensor, infer_state: LlamaInferStateInfo): embed_dim_ = input_embdings.shape[1] if infer_state.is_prefill: @@ -85,17 +91,20 @@ def _lm_head_and_gather( logic_batch = layer_weight.lm_head_weight_(input=normed, alloc_func=self.alloc_tensor) normed = None - vocab_size = layer_weight.lm_head_weight_.vocab_size - if infer_state.use_vocab_parallel_greedy and not force_full_logits: - logits, token_ids, logsumexp = vocab_parallel_greedy( + lm_head = layer_weight.lm_head_weight_ + vocab_size = lm_head.vocab_size + if infer_state.use_vocab_parallel_topk and not force_full_logits: + logic_batch = self._apply_logit_postprocessing(logic_batch) + logits, token_ids = vocab_parallel_topk( logic_batch, vocab_size=vocab_size, + vocab_start_id=lm_head.tp_vocab_start_id, + topk=self.vocab_parallel_topk_, tp_world_size=self.tp_world_size_, group=infer_state.dist_group, alloc_func=self.alloc_tensor, ) infer_state.logits_token_ids = token_ids - infer_state.logits_logsumexp = logsumexp return logits if self.tp_world_size_ == 1: @@ -114,12 +123,11 @@ def _lm_head_and_gather( ans_logics = self.alloc_tensor((token_num, vocab_size), dtype=torch.float32) ans_logics[:, :] = gather_data.permute(1, 0) gather_data = None - return ans_logics + return self._apply_logit_postprocessing(ans_logics) def token_forward( self, input_embdings: torch.Tensor, infer_state: LlamaInferStateInfo, layer_weight: LlamaPreAndPostLayerWeight ): - return self._token_forward(input_embdings=input_embdings, infer_state=infer_state, layer_weight=layer_weight) def overlap_tpsp_token_forward( @@ -130,7 +138,6 @@ def overlap_tpsp_token_forward( infer_state1: LlamaInferStateInfo, layer_weight: BaseLayerWeight, ): - logics = self.token_forward(input_embdings, infer_state, layer_weight=layer_weight) logics1 = self.token_forward(input_embdings1, infer_state1, layer_weight=layer_weight) diff --git a/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py b/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py index eb4481fd6b..2c89610f13 100644 --- a/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py +++ b/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py @@ -181,11 +181,12 @@ def token_forward( logits = self._lm_head_and_gather(last_input, token_num, layer_weight, infer_state) block_logits = logits.reshape(num_reqs, self.block_size_, -1) + candidate_indexes = torch.argmax(block_logits, dim=-1) if infer_state.logits_token_ids is None: - sampled_tokens = torch.argmax(block_logits, dim=-1) + sampled_tokens = candidate_indexes else: - assert block_logits.shape[-1] == 1 - sampled_tokens = infer_state.logits_token_ids.reshape(num_reqs, self.block_size_) + block_token_ids = infer_state.logits_token_ids.reshape(num_reqs, self.block_size_, -1) + sampled_tokens = block_token_ids.gather(-1, candidate_indexes.unsqueeze(-1)).squeeze(-1) confidence_logits = self.predict_confidence_logits( block_hidden, anchor_token_ids=anchor_token_ids, diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 1913f82547..0357a25770 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -385,12 +385,19 @@ def _async_copy_next_token_infos_to_pin_mem( ) return next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu - def _get_next_token_ranks(self, logits: torch.Tensor, next_token_ids: torch.Tensor) -> torch.Tensor: + def _get_next_token_ranks(self, model_output: ModelOutput, next_token_ids: torch.Tensor) -> torch.Tensor: """计算（或占位）每个 next token 在 vocab 上的 1-based rank（GPU tensor）。 仅 ``--enable_rl`` 时做真实 rank；否则返回 GPU 常量 ``-1``，避免 O(batch * vocab) 比较。 下游 async_copy 在同样条件下会忽略该返回值。 """ + if model_output.has_vocab_parallel_logits: + return g_pin_mem_manager.get_const_gpu_tensor( + key="next_token_ranks", + shape=next_token_ids.shape, + fill_value=1 if self.args.enable_rl else -1, + dtype=torch.int32, + ) if not self.args.enable_rl: return g_pin_mem_manager.get_const_gpu_tensor( key="next_token_ranks", @@ -398,8 +405,8 @@ def _get_next_token_ranks(self, logits: torch.Tensor, next_token_ids: torch.Tens fill_value=-1, dtype=torch.int32, ) - selected_logits = logits.gather(1, next_token_ids.long().view(-1, 1)) - return (logits > selected_logits).sum(dim=-1, dtype=torch.int32) + 1 + selected_logits = model_output.logits.gather(1, next_token_ids.long().view(-1, 1)) + return (model_output.logits > selected_logits).sum(dim=-1, dtype=torch.int32) + 1 def _capture_prompt_logprobs_if_needed( self, @@ -696,7 +703,6 @@ def _get_classed_reqs( can_alloc_token_num = g_infer_context.get_can_alloc_token_num() for req_obj in ready_reqs: - if req_obj.filter_mark: finished_reqs.append(req_obj) continue @@ -881,11 +887,8 @@ def _gen_argmax_token_ids(self, model_output: ModelOutput): def _gen_argmax_token_ids_and_prob(self, model_output: ModelOutput): logits = model_output.logits - if model_output.has_vocab_parallel_logits: - max_logits, candidate_indexes = torch.max(logits, dim=-1) - token_ids = self._map_logits_indexes_to_token_ids(model_output, candidate_indexes) - return token_ids, torch.exp(max_logits - model_output.logits_logsumexp) - max_probs, token_ids = torch.max(torch.softmax(logits, dim=-1), dim=-1) + max_probs, candidate_indexes = torch.max(torch.softmax(logits, dim=-1), dim=-1) + token_ids = self._map_logits_indexes_to_token_ids(model_output, candidate_indexes) return token_ids, max_probs @staticmethod @@ -896,7 +899,7 @@ def _map_logits_indexes_to_token_ids(model_output: ModelOutput, candidate_indexe def _sample_and_scatter_token( self, - logits: torch.Tensor, + model_output: ModelOutput, b_req_idx: torch.Tensor, b_mtp_index: torch.Tensor, run_reqs: List[InferReq], @@ -904,13 +907,14 @@ def _sample_and_scatter_token( b_prefill_has_output_cpu: torch.Tensor = None, mask_func: Optional[Callable] = None, ): - + logits = model_output.logits if mask_func is not None: + assert not model_output.has_vocab_parallel_logits, "constrained sampling requires dense logits" assert len(run_reqs) == logits.shape[0] mask_func(run_reqs, logits) - next_token_ids, next_token_logprobs = sample(logits, run_reqs, self.eos_id) - next_token_ranks = self._get_next_token_ranks(logits, next_token_ids) + next_token_ids, next_token_logprobs = sample(model_output, run_reqs, self.eos_id) + next_token_ranks = self._get_next_token_ranks(model_output, next_token_ids) b_has_out = None if is_prefill: b_has_out = g_pin_mem_manager.gen_from_list( diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 4d09476849..ae7ecde684 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -107,11 +107,18 @@ def prefill_normal( ): # 第一阶段: 模型推理 model_input, run_reqs = prepare_prefill_inputs(prefill_reqs, is_chuncked_mode=not self.disable_chunked_prefill) + if self.prefill_mask_func is not None: + model_input.use_vocab_parallel_topk = False with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) self._capture_prompt_logprobs_if_needed(model_input, run_reqs, model_output.prompt_logics) - (_, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu,) = self._sample_and_scatter_token( - logits=model_output.logits, + ( + _, + next_token_ids_cpu, + next_token_logprobs_cpu, + next_token_ranks_cpu, + ) = self._sample_and_scatter_token( + model_output=model_output, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, run_reqs=run_reqs, @@ -152,10 +159,17 @@ def decode_normal( decode_reqs: List[InferReq], ): model_input, run_reqs = prepare_decode_inputs(decode_reqs) + if self.decode_mask_func is not None: + model_input.use_vocab_parallel_topk = False with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) - (_, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu,) = self._sample_and_scatter_token( - logits=model_output.logits, + ( + _, + next_token_ids_cpu, + next_token_logprobs_cpu, + next_token_ranks_cpu, + ) = self._sample_and_scatter_token( + model_output=model_output, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, run_reqs=run_reqs, @@ -191,6 +205,8 @@ def prefill_mtp( prefill_reqs: List[InferReq], ): model_input, run_reqs = prepare_prefill_inputs(prefill_reqs, is_chuncked_mode=not self.disable_chunked_prefill) + if self.prefill_mask_func is not None: + model_input.use_vocab_parallel_topk = False with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) self._capture_prompt_logprobs_if_needed(model_input, run_reqs, model_output.prompt_logics) @@ -200,7 +216,7 @@ def prefill_mtp( next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=model_output.logits, + model_output=model_output, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, run_reqs=run_reqs, @@ -272,11 +288,11 @@ def decode_mtp( selected_rows = async_selected_row_mask_cpu.tensor.tolist() run_reqs = [req for req, selected in zip(run_reqs, selected_rows) if selected] next_token_ids, next_token_logprobs = sample( - model_output.logits, + model_output, run_reqs, self.eos_id, ) - next_token_ranks = self._get_next_token_ranks(model_output.logits, next_token_ids) + next_token_ranks = self._get_next_token_ranks(model_output, next_token_ids) b_req_mtp_start_loc = gen_b_req_mtp_start_loc(model_input.b_mtp_index, num_reqs=req_num) mtp_accept_len, accepted_index = mtp_utils.verify_mtp_tokens( diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_reward_model.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_reward_model.py index dfb1020820..cabf4f7de0 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_reward_model.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_reward_model.py @@ -14,9 +14,9 @@ def __init__(self) -> None: return def reward_prefill(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq]): - assert self.disable_chunked_prefill is True model_input, run_reqs = prepare_prefill_inputs(prefill_reqs, is_chuncked_mode=not self.disable_chunked_prefill) + model_input.use_vocab_parallel_topk = False model_output = self.model.forward(model_input) scores: torch.Tensor = model_output.logits diff --git a/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py index 21979cbef0..204e1312e3 100644 --- a/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/diverse_backend/impl.py @@ -40,10 +40,7 @@ def beam_prefill(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq ) with torch.cuda.stream(g_infer_context.get_overlap_stream()): - model_output = self.model.forward(model_input) - logits = model_output.logits - batch_idx, run_reqs = self._diverse_copy( master_reqs=group_reqs, b_prefill_has_out=model_input.b_prefill_has_output_cpu ) @@ -60,11 +57,11 @@ def beam_prefill(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq non_blocking=True ) - logits = logits[batch_idx] + sampled_output = model_output.index_select_logits_rows(batch_idx) b_mtp_index = model_input.b_mtp_index[batch_idx] - next_token_ids, next_token_logprobs = sample(logits, run_reqs, self.eos_id) - next_token_ranks = self._get_next_token_ranks(logits, next_token_ids) + next_token_ids, next_token_logprobs = sample(sampled_output, run_reqs, self.eos_id) + next_token_ranks = self._get_next_token_ranks(sampled_output, next_token_ids) scatter_token( next_token_ids=next_token_ids, diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index 9a81927bc1..fb5aea673a 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -187,7 +187,7 @@ def prefill_normal( next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=model_output.logits, + model_output=model_output, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, run_reqs=run_reqs, @@ -240,7 +240,7 @@ def decode_normal(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=model_output.logits, + model_output=model_output, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, run_reqs=run_reqs, @@ -287,13 +287,8 @@ def prefill_overlap(self, event_pack: OverlapEventPack, prefill_reqs: List[Infer model_output0, model_output1 = self.model.microbatch_overlap_prefill(model_input0, model_input1) self._capture_prompt_logprobs_if_needed(model_input0, run_reqs0, model_output0.prompt_logics) self._capture_prompt_logprobs_if_needed(model_input1, run_reqs1, model_output1.prompt_logics) - logits0 = model_output0.logits - logits1 = model_output1.logits - req_num0, req_num1 = len(run_reqs0), len(run_reqs1) - logits = torch.empty((req_num0 + req_num1, logits0.shape[1]), dtype=logits0.dtype, device=logits0.device) - logits[0:req_num0, :].copy_(logits0, non_blocking=True) - logits[req_num0 : req_num0 + req_num1, :].copy_(logits1, non_blocking=True) + sampled_output = ModelOutput.concat_logits_rows([model_output0, model_output1]) run_reqs = run_reqs0 + run_reqs1 b_has_out_cpu = model_input0.b_prefill_has_output_cpu + model_input1.b_prefill_has_output_cpu @@ -307,7 +302,7 @@ def prefill_overlap(self, event_pack: OverlapEventPack, prefill_reqs: List[Infer next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=logits, + model_output=sampled_output, b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, run_reqs=run_reqs, @@ -356,7 +351,7 @@ def decode_overlap(self, event_pack: OverlapEventPack, decode_reqs: List[InferRe with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output0, model_output1 = self.model.microbatch_overlap_decode(model_input0, model_input1) if req_num0 + req_num1 > 0: - logits = torch.cat((model_output0.logits, model_output1.logits), dim=0) + sampled_output = ModelOutput.concat_logits_rows([model_output0, model_output1]) b_req_idx = torch.cat((model_input0.b_req_idx, model_input1.b_req_idx), dim=0) b_mtp_index = torch.cat((model_input0.b_mtp_index, model_input1.b_mtp_index), dim=0) ( @@ -365,7 +360,7 @@ def decode_overlap(self, event_pack: OverlapEventPack, decode_reqs: List[InferRe next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=logits, + model_output=sampled_output, b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, run_reqs=run_reqs, @@ -421,7 +416,7 @@ def prefill_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq] next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=model_output.logits, + model_output=model_output, b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, run_reqs=run_reqs, @@ -446,7 +441,6 @@ def prefill_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq] sync_event.record() if req_num > 0: - # 第二阶段 event_pack.notify_post_handle_and_wait_pre_post_handle() update_packs = self._pre_post_handle(run_reqs, is_chuncked_mode=not self.disable_chunked_prefill) @@ -499,11 +493,11 @@ def decode_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[InferReq]): if req_num > 0: next_token_ids, next_token_logprobs = sample( - model_output.logits, + model_output, run_reqs, self.eos_id, ) - next_token_ranks = self._get_next_token_ranks(model_output.logits, next_token_ids) + next_token_ranks = self._get_next_token_ranks(model_output, next_token_ids) b_req_mtp_start_loc = gen_b_req_mtp_start_loc( b_mtp_index=model_input.b_mtp_index, @@ -649,17 +643,9 @@ def prefill_overlap_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[I model_output0, model_output1 = self.model.microbatch_overlap_prefill(model_input0, model_input1) self._capture_prompt_logprobs_if_needed(model_input0, run_reqs0, model_output0.prompt_logics) self._capture_prompt_logprobs_if_needed(model_input1, run_reqs1, model_output1.prompt_logics) - logits0 = model_output0.logits - logits1 = model_output1.logits req_num0, req_num1 = len(run_reqs0), len(run_reqs1) req_num = req_num0 + req_num1 - logits = torch.empty( - (req_num0 + req_num1, logits0.shape[1]), - dtype=logits0.dtype, - device=logits0.device, - ) - logits[0:req_num0, :].copy_(logits0, non_blocking=True) - logits[req_num0 : (req_num0 + req_num1), :].copy_(logits1, non_blocking=True) + sampled_output = ModelOutput.concat_logits_rows([model_output0, model_output1]) run_reqs = run_reqs0 + run_reqs1 b_has_out_cpu = model_input0.b_prefill_has_output_cpu + model_input1.b_prefill_has_output_cpu @@ -673,7 +659,7 @@ def prefill_overlap_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[I next_token_logprobs_cpu, next_token_ranks_cpu, ) = self._sample_and_scatter_token( - logits=logits, + model_output=sampled_output, run_reqs=run_reqs, b_req_idx=b_req_idx, b_mtp_index=b_mtp_index, @@ -681,7 +667,7 @@ def prefill_overlap_mtp(self, event_pack: OverlapEventPack, prefill_reqs: List[I b_prefill_has_output_cpu=b_has_out_cpu, ) else: - next_token_ids = torch.empty((0,), dtype=torch.int64, device=logits.device) + next_token_ids = torch.empty((0,), dtype=torch.int64, device=sampled_output.logits.device) target_next_token_ids_gpu0 = next_token_ids[:req_num0] target_next_token_ids_gpu1 = next_token_ids[req_num0:] @@ -769,20 +755,12 @@ def decode_overlap_mtp(self, event_pack: OverlapEventPack, decode_reqs: List[Inf verify_row_num0 = model_input0.batch_size verify_row_num1 = model_input1.batch_size verify_row_num = verify_row_num0 + verify_row_num1 - logits0 = model_output0.logits - logits1 = model_output1.logits run_reqs = run_reqs0 + run_reqs1 if req_num > 0: assert len(run_reqs) == verify_row_num - logits = torch.empty( - (verify_row_num, logits0.shape[1]), - dtype=logits0.dtype, - device=logits0.device, - ) - logits[:verify_row_num0, :].copy_(logits0, non_blocking=True) - logits[verify_row_num0:, :].copy_(logits1, non_blocking=True) - next_token_ids, next_token_logprobs = sample(logits, run_reqs, self.eos_id) - next_token_ranks = self._get_next_token_ranks(logits, next_token_ids) + sampled_output = ModelOutput.concat_logits_rows([model_output0, model_output1]) + next_token_ids, next_token_logprobs = sample(sampled_output, run_reqs, self.eos_id) + next_token_ranks = self._get_next_token_ranks(sampled_output, next_token_ids) ( next_token_ids_cpu, next_token_logprobs_cpu, diff --git a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py index 5b29ea0510..bd20f50d1f 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py @@ -1,5 +1,9 @@ import torch from typing import List, Tuple +from lightllm.common.basemodel.batch_objs import ModelOutput +from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk import ( + is_vocab_parallel_topk_enabled, +) from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty import apply_penalty from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty_gpu_cache import apply_penalty_gpu_cache from lightllm.common.basemodel.triton_kernel.post_process.apply_invalid_token import apply_invalid_token_ids @@ -8,7 +12,44 @@ from lightllm.utils.envs_utils import get_env_start_args -def sample(logits: torch.Tensor, reqs: List[InferReq], eos_id: List[int] = [2]): +def _can_use_unmodified_greedy_logits(reqs: List[InferReq]) -> bool: + """Whether sampling is exactly argmax over the incoming logits.""" + + for req_obj in reqs: + sample_param = req_obj.sampling_param + shm_param = sample_param.shm_param + if shm_param.top_k != 1 or shm_param.temperature != 1.0: + return False + if ( + shm_param.presence_penalty != 0.0 + or shm_param.frequency_penalty != 0.0 + or shm_param.repetition_penalty != 1.0 + ): + return False + if shm_param.exponential_decay_length_penalty.to_tuple()[1] != 1.0: + return False + out_token_len = req_obj.get_cur_total_len() - req_obj.shm_req.input_len + if out_token_len < shm_param.min_new_tokens - 1: + return False + if sample_param.invalid_token_ids: + return False + return True + + +def can_use_vocab_parallel_topk(reqs: List[InferReq]) -> bool: + return is_vocab_parallel_topk_enabled() and _can_use_unmodified_greedy_logits(reqs) + + +def sample(model_output: ModelOutput, reqs: List[InferReq], eos_id: List[int] = [2]): + logits = model_output.logits + if model_output.has_vocab_parallel_logits: + if not _can_use_unmodified_greedy_logits(reqs): + raise RuntimeError("vocab-parallel top-k logits require unmodified greedy requests") + probs = torch.softmax(logits, dim=-1) + max_probs, candidate_indexes = torch.max(probs, dim=-1) + token_ids = model_output.logits_token_ids.gather(1, candidate_indexes.view(-1, 1)).view(-1).long() + return token_ids, torch.log(max_probs) + ( b_req_idx, b_temperatures, diff --git a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py index 22731439c4..ee95ff24c5 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py @@ -3,6 +3,9 @@ from typing import List, Tuple from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context from lightllm.common.basemodel.batch_objs import ModelInput +from lightllm.server.router.model_infer.mode_backend.generic_post_process import ( + can_use_vocab_parallel_topk, +) INT64_MAX = torch.iinfo(torch.int64).max @@ -87,6 +90,7 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> is_prefill=True, b_prefill_has_output_cpu=b_prefill_has_output, multimodal_params=batch_multimodal_params, + use_vocab_parallel_topk=can_use_vocab_parallel_topk(run_reqs), ) return model_input, run_reqs @@ -160,6 +164,7 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In b_shared_radix_node_id=b_shared_radix_node_id, is_prefill=False, multimodal_params=multimodal_params, + use_vocab_parallel_topk=can_use_vocab_parallel_topk(run_reqs), ) return model_input, run_reqs @@ -176,6 +181,9 @@ def overlap_prepare_decode_inputs(req_objs: List[InferReq]): model_input1, run_reqs1 = prepare_decode_inputs( req_objs=decode_reqs1, ) + use_vocab_parallel_topk = can_use_vocab_parallel_topk(run_reqs0 + run_reqs1) + model_input0.use_vocab_parallel_topk = use_vocab_parallel_topk + model_input1.use_vocab_parallel_topk = use_vocab_parallel_topk return model_input0, run_reqs0, decode_reqs0, model_input1, run_reqs1, decode_reqs1 @@ -211,6 +219,9 @@ def overlap_prepare_prefill_inputs(req_objs: List[InferReq]): req_objs=right_reqs, is_chuncked_mode=True, ) + use_vocab_parallel_topk = can_use_vocab_parallel_topk(run_reqs0 + run_reqs1) + model_input0.use_vocab_parallel_topk = use_vocab_parallel_topk + model_input1.use_vocab_parallel_topk = use_vocab_parallel_topk return model_input0, run_reqs0, model_input1, run_reqs1 diff --git a/unit_tests/common/basemodel/test_model_output.py b/unit_tests/common/basemodel/test_model_output.py index 8ec94ae5dc..68a5286774 100644 --- a/unit_tests/common/basemodel/test_model_output.py +++ b/unit_tests/common/basemodel/test_model_output.py @@ -7,18 +7,17 @@ from lightllm.common.basemodel.batch_objs import ModelInput, ModelMtpOutputCollector, ModelOutput -def test_vocab_parallel_metadata_follows_row_selection(): +def test_vocab_parallel_metadata_follows_row_operations(): output = ModelOutput( - logits=torch.tensor([[8.0], [7.0], [6.0]]), - logits_token_ids=torch.tensor([[4], [9], [2]]), - logits_logsumexp=torch.tensor([8.5, 7.25, 6.75]), + logits=torch.tensor([[8.0, 1.0], [7.0, 2.0], [6.0, 3.0]]), + logits_token_ids=torch.tensor([[4, 14], [9, 19], [2, 12]]), ) selected = output.index_select_logits_rows(torch.tensor([2, 0])) + combined = ModelOutput.concat_logits_rows([selected, output.index_select_logits_rows(torch.tensor([1]))]) - torch.testing.assert_close(selected.logits.view(-1), torch.tensor([6.0, 8.0])) - torch.testing.assert_close(selected.logits_token_ids.view(-1), torch.tensor([2, 4])) - torch.testing.assert_close(selected.logits_logsumexp, torch.tensor([6.75, 8.5])) + torch.testing.assert_close(combined.logits, torch.tensor([[6.0, 3.0], [8.0, 1.0], [7.0, 2.0]])) + torch.testing.assert_close(combined.logits_token_ids, torch.tensor([[2, 12], [4, 14], [9, 19]])) def test_decode_unpad_slices_spec_output_with_logits(): @@ -26,7 +25,6 @@ def test_decode_unpad_slices_spec_output_with_logits(): output = ModelOutput( logits=torch.arange(24).view(6, 4), logits_token_ids=torch.arange(100, 124).view(6, 4), - logits_logsumexp=torch.arange(6, dtype=torch.float32), mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(18).view(6, 3)), ) @@ -34,13 +32,11 @@ def test_decode_unpad_slices_spec_output_with_logits(): assert unpadded.logits.shape == (4, 4) assert unpadded.logits_token_ids.shape == (4, 4) - assert unpadded.logits_logsumexp.shape == (4,) assert unpadded.mtp_collector.spec_hidden.shape == (4, 3) # Unpadding returns a shallow output copy and leaves the graph-owned # tensors on the original ModelOutput intact. assert output.logits.shape == (6, 4) assert output.logits_token_ids.shape == (6, 4) - assert output.logits_logsumexp.shape == (6,) assert output.mtp_collector.spec_hidden.shape == (6, 3) @@ -49,7 +45,6 @@ def test_prefill_unpad_uses_token_rows_for_spec_hidden(): output = ModelOutput( logits=torch.arange(20).view(5, 4), logits_token_ids=torch.arange(100, 120).view(5, 4), - logits_logsumexp=torch.arange(5, dtype=torch.float32), mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(24).view(8, 3)), prompt_logics=torch.arange(32).view(8, 4), ) @@ -62,7 +57,6 @@ def test_prefill_unpad_uses_token_rows_for_spec_hidden(): assert unpadded.logits.shape == (3, 4) assert unpadded.logits_token_ids.shape == (3, 4) - assert unpadded.logits_logsumexp.shape == (3,) assert unpadded.mtp_collector.spec_hidden.shape == (6, 3) assert unpadded.prompt_logics.shape == (6, 4) @@ -118,6 +112,48 @@ def _create_empty_decode_input(): ) +def test_infer_state_enables_vocab_parallel_topk_for_draft_or_requested_target(monkeypatch): + model = TpPartBaseModel.__new__(TpPartBaseModel) + model.infer_state_class = basemodel.InferStateInfo + model.hidden_collector_prototype = SimpleNamespace(new_instance=lambda: object()) + model.is_token_healing = False + model.return_all_prompt_logics = False + model.is_mtp_draft_model = False + model.mem_manager = object() + model.req_manager = object() + model.decode_att_backend = SimpleNamespace(create_att_decode_state=lambda infer_state: None) + model.decode_att_backend1 = None + monkeypatch.setattr(basemodel.dist_group_manager, "get_group", lambda _: None) + + model_input = _create_empty_decode_input() + infer_state = model._create_inferstate(model_input) + assert not infer_state.use_vocab_parallel_topk + + model_input.use_vocab_parallel_topk = True + infer_state = model._create_inferstate(model_input) + assert infer_state.use_vocab_parallel_topk + + model.is_mtp_draft_model = True + model_input.use_vocab_parallel_topk = False + infer_state = model._create_inferstate(model_input) + assert infer_state.use_vocab_parallel_topk + + +def test_cuda_graph_contract_falls_back_for_dense_target_batch(monkeypatch): + model = TpPartBaseModel.__new__(TpPartBaseModel) + model.is_mtp_draft_model = False + sparse_input = SimpleNamespace(use_vocab_parallel_topk=True) + dense_input = SimpleNamespace(use_vocab_parallel_topk=False) + + monkeypatch.setattr(basemodel, "is_vocab_parallel_topk_enabled", lambda: True) + assert model._is_cuda_graph_output_compatible(sparse_input) + assert not model._is_cuda_graph_output_compatible(dense_input) + assert not model._is_cuda_graph_output_compatible(sparse_input, dense_input) + + model.is_mtp_draft_model = True + assert model._is_cuda_graph_output_compatible(dense_input) + + @torch.no_grad() def test_decode_pads_only_once_after_selecting_execution_path(monkeypatch): monkeypatch.setattr(basemodel, "copy_kv_index_to_req", lambda *args: None) @@ -141,6 +177,7 @@ def test_decode_pads_only_once_after_selecting_execution_path(monkeypatch): ) in execution_configs: model = TpPartBaseModel.__new__(TpPartBaseModel) model.args = SimpleNamespace(enable_tpsp_mix_mode=enable_tpsp_mix_mode) + model.is_mtp_draft_model = False model.tp_world_size_ = tp_world_size model.mem_manager = SimpleNamespace(HOLD_TOKEN_MEMINDEX=99) model.req_manager = SimpleNamespace(HOLD_REQUEST_ID=88, req_to_token_indexs=object()) diff --git a/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py b/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py deleted file mode 100644 index 279268a85c..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_greedy.py +++ /dev/null @@ -1,65 +0,0 @@ -import importlib - -import pytest -import torch - -from lightllm.common.basemodel.triton_kernel.post_process.greedy_sample import ( - greedy_sample_local_stats, -) - - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for Triton kernels") - - -@pytest.mark.parametrize("token_num", [1, 7, 64]) -def test_vocab_parallel_greedy_matches_full_logits(monkeypatch, token_num): - module = importlib.import_module("lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_greedy") - tp_world_size = 4 - local_vocab_size = 8192 - vocab_size = tp_world_size * local_vocab_size - generator = torch.Generator(device="cuda").manual_seed(20260826 + token_num) - local_logits_by_rank = [ - torch.randn( - (local_vocab_size, token_num), - dtype=torch.bfloat16, - device="cuda", - generator=generator, - ) - for _ in range(tp_world_size) - ] - - # Exercise deterministic tie-breaking both between local reduction blocks - # and across tensor-parallel ranks. The smallest global token id must win. - local_logits_by_rank[0][4097, 0] = 20.0 - local_logits_by_rank[0][3, 0] = 20.0 - local_logits_by_rank[3][2, 0] = 20.0 - - local_stats_by_rank = [ - greedy_sample_local_stats(local_logits.transpose(0, 1).contiguous()) for local_logits in local_logits_by_rank - ] - - def fake_all_gather_into_tensor(output_, input_, **_kwargs): - for output, local_stats in zip(output_, local_stats_by_rank): - output.copy_(local_stats) - - monkeypatch.setattr(module, "all_gather_into_tensor", fake_all_gather_into_tensor) - actual_logits, actual_ids, actual_logsumexp = module.vocab_parallel_greedy( - local_logits_by_rank[0], - vocab_size=vocab_size, - tp_world_size=tp_world_size, - group=None, - alloc_func=torch.empty, - ) - actual_ids = actual_ids.view(-1) - actual_logprobs = actual_logits.view(-1) - actual_logsumexp - - full_logits = torch.cat(local_logits_by_rank, dim=0).transpose(0, 1).float() - expected_ids = full_logits.argmax(dim=1) - expected_logits = full_logits.gather(1, expected_ids[:, None]).view(-1) - expected_logsumexp = torch.logsumexp(full_logits, dim=1) - expected_logprobs = torch.log_softmax(full_logits, dim=1).gather(1, expected_ids[:, None]).squeeze(1) - - torch.testing.assert_close(actual_ids, expected_ids, rtol=0, atol=0) - torch.testing.assert_close(actual_logits.view(-1), expected_logits, rtol=0, atol=0) - torch.testing.assert_close(actual_logsumexp, expected_logsumexp, rtol=2e-4, atol=2e-4) - torch.testing.assert_close(actual_logprobs, expected_logprobs, rtol=2e-4, atol=2e-4) diff --git a/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_topk.py b/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_topk.py new file mode 100644 index 0000000000..1b7c2b4dc9 --- /dev/null +++ b/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_topk.py @@ -0,0 +1,81 @@ +import importlib + +import pytest +import torch + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("token_num", [1, 7, 64]) +def test_vocab_parallel_topk_collects_candidates_and_preserves_global_argmax(monkeypatch, token_num): + module = importlib.import_module("lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk") + tp_world_size = 4 + local_vocab_size = 1024 + vocab_size = tp_world_size * local_vocab_size + local_topk = 16 + generator = torch.Generator(device="cuda").manual_seed(20260902 + token_num) + local_logits_by_rank = [ + torch.randn( + (local_vocab_size, token_num), + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ) + for _ in range(tp_world_size) + ] + local_logits_by_rank[3][2, 0] = 20.0 + + def pack_rank(rank): + values, indexes = torch.topk(local_logits_by_rank[rank], k=local_topk, dim=0, sorted=False) + payload = torch.empty((local_topk * 2, token_num), dtype=torch.float32, device="cuda") + payload[:local_topk].copy_(values.float()) + payload[local_topk:].view(torch.int32).copy_(indexes.to(torch.int32) + rank * local_vocab_size) + return payload + + payloads = [pack_rank(rank) for rank in range(tp_world_size)] + + def fake_all_gather_into_tensor(output_, input_, **_kwargs): + assert output_.shape == (tp_world_size, local_topk * 2, token_num) + assert input_.shape == (local_topk * 2, token_num) + for output, payload in zip(output_, payloads): + output.copy_(payload) + + monkeypatch.setattr(module, "all_gather_into_tensor", fake_all_gather_into_tensor) + actual_logits, actual_ids = module.vocab_parallel_topk( + local_logits_by_rank[0], + vocab_size=vocab_size, + vocab_start_id=0, + topk=local_topk, + tp_world_size=tp_world_size, + group=None, + alloc_func=torch.empty, + ) + + assert actual_logits.shape == (token_num, tp_world_size * local_topk) + assert actual_ids.shape == actual_logits.shape + assert actual_logits.dtype == torch.float32 + assert actual_ids.dtype == torch.int64 + + full_logits = torch.cat(local_logits_by_rank, dim=0).transpose(0, 1).float() + candidate_indexes = actual_logits.argmax(dim=1, keepdim=True) + actual_argmax_ids = actual_ids.gather(1, candidate_indexes).view(-1) + torch.testing.assert_close(actual_argmax_ids, full_logits.argmax(dim=1), rtol=0, atol=0) + + +def test_vocab_parallel_topk_single_rank_caps_width_to_vocabulary(): + module = importlib.import_module("lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk") + local_logits = torch.tensor([[1.0], [3.0], [2.0]], device="cuda") + + logits, token_ids = module.vocab_parallel_topk( + local_logits, + vocab_size=3, + vocab_start_id=0, + topk=128, + tp_world_size=1, + group=None, + alloc_func=torch.empty, + ) + + assert logits.shape == (1, 3) + torch.testing.assert_close(token_ids.sort(dim=1).values, torch.tensor([[0, 1, 2]], device="cuda")) diff --git a/unit_tests/models/test_gemma4_vocab_parallel_topk.py b/unit_tests/models/test_gemma4_vocab_parallel_topk.py new file mode 100644 index 0000000000..13ff5c3c24 --- /dev/null +++ b/unit_tests/models/test_gemma4_vocab_parallel_topk.py @@ -0,0 +1,67 @@ +from types import SimpleNamespace + +import torch + +import lightllm.models.llama.layer_infer.post_layer_infer as llama_post_layer +from lightllm.models.gemma4.layer_infer.post_layer_infer import Gemma4PostLayerInfer + + +def test_vocab_parallel_topk_softcaps_local_logits_before_candidate_selection(monkeypatch): + post = Gemma4PostLayerInfer.__new__(Gemma4PostLayerInfer) + post.final_logit_softcapping = 2.0 + post.vocab_parallel_topk_ = 2 + post.tp_world_size_ = 1 + post.alloc_tensor = torch.empty + post._norm = lambda hidden, infer_state, layer_weight: hidden + + local_logits = torch.tensor( + [[4.0, -4.0], [2.0, -2.0], [1.0, -1.0]], + dtype=torch.bfloat16, + ) + + class LMHead: + vocab_size = 3 + tp_vocab_start_id = 0 + + def __call__(self, input, alloc_func): + return local_logits + + sparse_logits = torch.tensor([[1.5, 1.0], [-1.0, -1.5]]) + token_ids = torch.tensor([[0, 1], [2, 1]]) + captured = {} + + def fake_vocab_parallel_topk(logits, **kwargs): + captured["logits"] = logits + return sparse_logits, token_ids + + monkeypatch.setattr(llama_post_layer, "vocab_parallel_topk", fake_vocab_parallel_topk) + infer_state = SimpleNamespace( + dist_group=None, + use_vocab_parallel_topk=True, + logits_token_ids=None, + ) + + result = post._lm_head_and_gather( + hidden=torch.empty((2, 3)), + token_num=2, + layer_weight=SimpleNamespace(lm_head_weight_=LMHead()), + infer_state=infer_state, + ) + + expected = torch.tanh(local_logits.float() / 2.0) * 2.0 + assert captured["logits"].dtype == torch.float32 + torch.testing.assert_close(captured["logits"], expected) + assert result is sparse_logits + assert infer_state.logits_token_ids is token_ids + + +def test_full_logits_softcap_after_float32_conversion(): + post = Gemma4PostLayerInfer.__new__(Gemma4PostLayerInfer) + post.final_logit_softcapping = 2.0 + logits = torch.tensor([[1.234375, -3.140625]], dtype=torch.bfloat16) + + actual = post._apply_logit_postprocessing(logits) + expected = torch.tanh(logits.float() / 2.0) * 2.0 + + assert actual.dtype == torch.float32 + torch.testing.assert_close(actual, expected) diff --git a/unit_tests/models/test_vocab_parallel_greedy_output.py b/unit_tests/models/test_vocab_parallel_topk_output.py similarity index 78% rename from unit_tests/models/test_vocab_parallel_greedy_output.py rename to unit_tests/models/test_vocab_parallel_topk_output.py index 88e5a7c282..e8dc0adb2d 100644 --- a/unit_tests/models/test_vocab_parallel_greedy_output.py +++ b/unit_tests/models/test_vocab_parallel_topk_output.py @@ -7,12 +7,11 @@ from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend -def test_argmax_restores_global_token_ids_and_exact_probabilities(): +def test_argmax_restores_global_token_ids_and_uses_candidate_probability(): backend = ModeBackend.__new__(ModeBackend) output = ModelOutput( - logits=torch.tensor([[3.0, 1.0], [0.0, 5.0]]), - logits_token_ids=torch.tensor([[30, 10], [100, 500]]), - logits_logsumexp=torch.tensor([4.0, 5.25]), + logits=torch.tensor([[3.0, 1.0, 2.0], [0.0, 5.0, 4.0]]), + logits_token_ids=torch.tensor([[30, 10, 20], [100, 500, 400]]), ) token_ids = backend._gen_argmax_token_ids(output) @@ -20,7 +19,7 @@ def test_argmax_restores_global_token_ids_and_exact_probabilities(): torch.testing.assert_close(token_ids, torch.tensor([30, 500])) torch.testing.assert_close(token_ids_with_prob, token_ids) - torch.testing.assert_close(probs, torch.exp(torch.tensor([-1.0, -0.25]))) + torch.testing.assert_close(probs, torch.softmax(output.logits, dim=-1).max(dim=-1).values) def test_dense_argmax_keeps_column_index_semantics(): @@ -30,13 +29,13 @@ def test_dense_argmax_keeps_column_index_semantics(): torch.testing.assert_close(backend._gen_argmax_token_ids(output), torch.tensor([1])) -def test_dspark_confidence_path_receives_global_token_ids(): +def test_dspark_confidence_path_receives_mapped_sparse_token_ids(): post = Qwen3DSparkPostLayerInfer.__new__(Qwen3DSparkPostLayerInfer) post.block_size_ = 2 post.markov_rank_ = 0 post._slice_get_last_input = lambda input_embeddings, infer_state: (input_embeddings, 4) - sparse_logits = torch.tensor([[4.0], [5.0], [7.0], [9.0]]) - sparse_token_ids = torch.tensor([[40], [50], [70], [90]]) + sparse_logits = torch.tensor([[1.0, 4.0], [5.0, 2.0], [3.0, 7.0], [9.0, 8.0]]) + sparse_token_ids = torch.tensor([[10, 40], [50, 20], [30, 70], [90, 80]]) def gather_vocab_parallel(*args, **kwargs): infer_state = args[3] @@ -56,12 +55,11 @@ class Collector: def add_mtp_outputs(self, **kwargs): self.outputs = kwargs - collector = Collector() infer_state = SimpleNamespace( is_prefill=False, input_ids=torch.tensor([1, 0, 2, 0]), logits_token_ids=None, - hidden_collector=collector, + hidden_collector=Collector(), ) returned_logits = post.token_forward( diff --git a/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_topk_sampling.py b/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_topk_sampling.py new file mode 100644 index 0000000000..09f74e3b07 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_topk_sampling.py @@ -0,0 +1,91 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.batch_objs import ModelOutput +from lightllm.server.router.model_infer.mode_backend.generic_post_process import ( + _can_use_unmodified_greedy_logits, + can_use_vocab_parallel_topk, + sample, +) +from lightllm.utils.envs_utils import enable_env_vars + + +def make_req(**overrides): + values = { + "top_k": 1, + "temperature": 1.0, + "presence_penalty": 0.0, + "frequency_penalty": 0.0, + "repetition_penalty": 1.0, + "min_new_tokens": 1, + "decay_factor": 1.0, + "invalid_token_ids": [], + "output_len": 0, + } + values.update(overrides) + shm_param = SimpleNamespace( + top_k=values["top_k"], + temperature=values["temperature"], + presence_penalty=values["presence_penalty"], + frequency_penalty=values["frequency_penalty"], + repetition_penalty=values["repetition_penalty"], + min_new_tokens=values["min_new_tokens"], + exponential_decay_length_penalty=SimpleNamespace(to_tuple=lambda: (1, values["decay_factor"])), + ) + input_len = 10 + return SimpleNamespace( + sampling_param=SimpleNamespace( + shm_param=shm_param, + invalid_token_ids=values["invalid_token_ids"], + ), + shm_req=SimpleNamespace(input_len=input_len), + get_cur_total_len=lambda: input_len + values["output_len"], + ) + + +def test_accepts_unmodified_greedy_requests(): + assert _can_use_unmodified_greedy_logits([make_req(), make_req(output_len=5)]) + + +def test_feature_gate_requires_environment_and_eligible_batch(monkeypatch): + monkeypatch.delenv("LIGHTLLM_VOCAB_PARALLEL_TOPK", raising=False) + enable_env_vars.cache_clear() + assert not can_use_vocab_parallel_topk([make_req()]) + + monkeypatch.setenv("LIGHTLLM_VOCAB_PARALLEL_TOPK", "1") + enable_env_vars.cache_clear() + assert can_use_vocab_parallel_topk([make_req()]) + assert not can_use_vocab_parallel_topk([make_req(top_k=2)]) + enable_env_vars.cache_clear() + + +def test_samples_sparse_candidates_and_maps_global_token_ids(): + model_output = ModelOutput( + logits=torch.tensor([[9.0, 7.0, 1.0], [4.0, 6.0, 5.0]], dtype=torch.float32), + logits_token_ids=torch.tensor([[17, 3, 8], [3, 20, 11]], dtype=torch.int64), + ) + + token_ids, token_logprobs = sample(model_output, [make_req(), make_req()]) + + expected_probs = torch.softmax(model_output.logits, dim=-1).max(dim=-1).values + torch.testing.assert_close(token_ids, torch.tensor([17, 20])) + torch.testing.assert_close(token_logprobs, torch.log(expected_probs)) + + +@pytest.mark.parametrize( + "override", + [ + {"top_k": 2}, + {"temperature": 0.5}, + {"presence_penalty": 0.1}, + {"frequency_penalty": 0.1}, + {"repetition_penalty": 1.1}, + {"decay_factor": 1.1}, + {"invalid_token_ids": [7]}, + {"min_new_tokens": 2}, + ], +) +def test_rejects_logits_modifiers(override): + assert not _can_use_unmodified_greedy_logits([make_req(**override)]) From f99737d113166e76f20ff15954c8b2675f63213e Mon Sep 17 00:00:00 2001 From: sufubao Date: Wed, 2 Sep 2026 16:33:19 +0800 Subject: [PATCH 3/8] style: apply pre-commit formatting --- .../mode_backend/chunked_prefill/impl.py | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index ae7ecde684..6db9246385 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -112,12 +112,7 @@ def prefill_normal( with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) self._capture_prompt_logprobs_if_needed(model_input, run_reqs, model_output.prompt_logics) - ( - _, - next_token_ids_cpu, - next_token_logprobs_cpu, - next_token_ranks_cpu, - ) = self._sample_and_scatter_token( + (_, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu,) = self._sample_and_scatter_token( model_output=model_output, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, @@ -163,12 +158,7 @@ def decode_normal( model_input.use_vocab_parallel_topk = False with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) - ( - _, - next_token_ids_cpu, - next_token_logprobs_cpu, - next_token_ranks_cpu, - ) = self._sample_and_scatter_token( + (_, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu,) = self._sample_and_scatter_token( model_output=model_output, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, From f9c3443be1ba99567909a47720f705a7ab78a0a7 Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 7 Sep 2026 10:31:59 +0800 Subject: [PATCH 4/8] fix: preserve deterministic vocab-parallel greedy selection --- docs/CN/source/tutorial/api_server_args.rst | 16 ++++ docs/EN/source/tutorial/api_server_args.rst | 18 ++++ .../post_process/vocab_parallel_topk.py | 40 +++++---- .../model_infer/mode_backend/base_backend.py | 3 +- .../mode_backend/generic_post_process.py | 5 +- .../triton_kernel/test_vocab_parallel_topk.py | 89 +++++++++++++++++++ .../models/test_vocab_parallel_topk_output.py | 13 +++ .../test_vocab_parallel_topk_sampling.py | 11 +++ 8 files changed, 177 insertions(+), 18 deletions(-) diff --git a/docs/CN/source/tutorial/api_server_args.rst b/docs/CN/source/tutorial/api_server_args.rst index fbb63d09f0..f04645c8c6 100644 --- a/docs/CN/source/tutorial/api_server_args.rst +++ b/docs/CN/source/tutorial/api_server_args.rst @@ -641,6 +641,22 @@ MTP 多预测参数 增加此值允许更多预测，但确保模型与指定的步数兼容。 目前 deepseekv3/r1 模型仅支持 1 步 +词表并行 Top-k 环境变量 +----------------------- + +``LIGHTLLM_VOCAB_PARALLEL_TOPK=1`` 为满足 ``top_k=1``、``temperature=1`` +且无需修改 logits 的主模型批次启用稀疏输出。惩罚项、EOS 屏蔽、无效 token +屏蔽及约束采样会回退到完整词表输出。主模型默认关闭此功能；使用共享 Llama +后处理层的草稿模型始终使用稀疏候选。 + +``LIGHTLLM_VOCAB_PARALLEL_TOPK_SIZE`` 设置每个 TP rank 的本地候选数， +必须为正整数，默认 ``128``，实际取值不超过最小词表分片的大小。 +服务启动前，各 rank 应设置相同的值。 + +贪心采样与完整词表的 argmax 一致，分数相同时选择最小 token ID。 +返回的概率及动态 MTP 使用的草稿概率仅在候选集内归一化，属于近似值。 +回退到完整词表的批次会绕过为稀疏主模型输出捕获的 decode CUDA Graph。 + DeepSeek 冗余专家参数 --------------------- diff --git a/docs/EN/source/tutorial/api_server_args.rst b/docs/EN/source/tutorial/api_server_args.rst index 69edf50a86..fc17d1a7f0 100644 --- a/docs/EN/source/tutorial/api_server_args.rst +++ b/docs/EN/source/tutorial/api_server_args.rst @@ -644,6 +644,24 @@ MTP Multi-Prediction Parameters Increasing this value allows more predictions, but ensure the model is compatible with the specified number of steps. Currently deepseekv3/r1 models only support 1 step +Vocabulary-Parallel Top-k Environment Variables +----------------------------------------------- + +``LIGHTLLM_VOCAB_PARALLEL_TOPK=1`` enables sparse output for target-model +batches with ``top_k=1``, ``temperature=1``, and no logit modifiers. Penalties, +EOS masking, invalid-token masking, and constrained sampling use dense output. +The target-model feature is disabled by default; draft models using the shared +Llama post layer always use sparse candidates. + +``LIGHTLLM_VOCAB_PARALLEL_TOPK_SIZE`` sets the positive number of local +candidates per TP rank (default: ``128``), capped at the smallest vocabulary +shard size. Set the same value on all ranks before starting the service. + +Greedy token selection preserves dense argmax, including choosing the smallest +token ID on ties. Returned probabilities and dynamic-MTP draft probabilities +are normalized over the candidate set and are approximate. Dense fallback +batches bypass decode CUDA Graphs captured for sparse target output. + DeepSeek Redundant Expert Parameters ------------------------------------ diff --git a/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_topk.py b/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_topk.py index bcb48c0090..9d8d258fcd 100644 --- a/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_topk.py +++ b/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_topk.py @@ -38,9 +38,10 @@ def vocab_parallel_topk( ) -> tuple[torch.Tensor, torch.Tensor]: """Gather each TP rank's local top-k logits and their global token ids. - The returned width is ``tp_world_size * topk``. It intentionally keeps the - union of local candidates: greedy selection remains exact, while probability - calculations over the sparse result are an inexpensive approximation. + The returned width is ``tp_world_size * min(topk, vocab_size // tp_world_size)``. + Each shard's first candidate is its lowest-id maximum. Together with rank + order, this preserves dense argmax tie-breaking. Probabilities over the + candidate union remain an approximation. """ assert local_logits.ndim == 2 and local_logits.is_cuda and local_logits.is_contiguous() @@ -51,13 +52,19 @@ def vocab_parallel_topk( assert local_topk > 0 assert local_vocab_size >= local_topk - local_values, local_indexes = torch.topk(local_logits, k=local_topk, dim=0, sorted=False) - local_values = local_values.float() - local_token_ids = local_indexes.to(torch.int32).add_(int(vocab_start_id)) - + if local_topk == 1: + local_values, local_indexes = torch.max(local_logits, dim=0, keepdim=True) + else: + local_values, local_indexes = torch.topk(local_logits, k=local_topk, dim=0, sorted=True) + # topk does not promise stable ties, and can omit the first maximum. + # Swap it into the first slot if present, otherwise replace a tied max. + # Both operations preserve candidate values and avoid duplicate ids. + first_max = torch.argmax(local_logits, dim=0, keepdim=True) + local_indexes = torch.where(local_indexes == first_max, local_indexes[:1], local_indexes) + local_indexes[:1].copy_(first_max) if tp_world_size == 1: - candidate_values = local_values.permute(1, 0).contiguous() - candidate_token_ids = local_token_ids.permute(1, 0).contiguous() + candidate_values = local_values.permute(1, 0) + candidate_token_ids = local_indexes.add_(int(vocab_start_id)).permute(1, 0) else: # Values and ids are both four bytes. Bit-packing the ids into the FP32 # payload keeps the operation to one fixed-shape collective. @@ -67,7 +74,7 @@ def vocab_parallel_topk( device=local_logits.device, ) local_payload[:local_topk].copy_(local_values) - local_payload[local_topk:].view(torch.int32).copy_(local_token_ids) + local_payload[local_topk:].view(torch.int32).copy_(local_indexes).add_(int(vocab_start_id)) gathered_payload = alloc_func( (tp_world_size, local_topk * 2, token_num), @@ -80,15 +87,18 @@ def vocab_parallel_topk( group=group, async_op=False, ) - candidate_values = gathered_payload[:, :local_topk, :].permute(2, 0, 1).reshape(token_num, -1) - candidate_token_ids = ( - gathered_payload[:, local_topk:, :].view(torch.int32).permute(2, 0, 1).reshape(token_num, -1) - ) + candidate_values = gathered_payload[:, :local_topk, :].permute(2, 0, 1) + candidate_token_ids = gathered_payload[:, local_topk:, :].view(torch.int32).permute(2, 0, 1) + # copy_ converts dtype and materializes the final layout in one pass. output_logits = alloc_func( candidate_values.shape, dtype=torch.float32, device=local_logits.device, ) output_logits.copy_(candidate_values) - return output_logits, candidate_token_ids.to(torch.int64).contiguous() + output_token_ids = alloc_func(candidate_token_ids.shape, dtype=torch.int64, device=local_logits.device) + output_token_ids.copy_(candidate_token_ids) + return output_logits.view(token_num, tp_world_size * local_topk), output_token_ids.view( + token_num, tp_world_size * local_topk + ) diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 0357a25770..fbb5b9c1c3 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -887,7 +887,8 @@ def _gen_argmax_token_ids(self, model_output: ModelOutput): def _gen_argmax_token_ids_and_prob(self, model_output: ModelOutput): logits = model_output.logits - max_probs, candidate_indexes = torch.max(torch.softmax(logits, dim=-1), dim=-1) + candidate_indexes = torch.argmax(logits, dim=-1) + max_probs = torch.softmax(logits, dim=-1).gather(1, candidate_indexes.view(-1, 1)).view(-1) token_ids = self._map_logits_indexes_to_token_ids(model_output, candidate_indexes) return token_ids, max_probs diff --git a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py index bd20f50d1f..84bd268172 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py @@ -45,9 +45,10 @@ def sample(model_output: ModelOutput, reqs: List[InferReq], eos_id: List[int] = if model_output.has_vocab_parallel_logits: if not _can_use_unmodified_greedy_logits(reqs): raise RuntimeError("vocab-parallel top-k logits require unmodified greedy requests") + candidate_indexes = torch.argmax(logits, dim=-1, keepdim=True) probs = torch.softmax(logits, dim=-1) - max_probs, candidate_indexes = torch.max(probs, dim=-1) - token_ids = model_output.logits_token_ids.gather(1, candidate_indexes.view(-1, 1)).view(-1).long() + max_probs = probs.gather(1, candidate_indexes).view(-1) + token_ids = model_output.logits_token_ids.gather(1, candidate_indexes).view(-1).long() return token_ids, torch.log(max_probs) ( diff --git a/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_topk.py b/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_topk.py index 1b7c2b4dc9..2e069ab26e 100644 --- a/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_topk.py +++ b/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_topk.py @@ -1,7 +1,10 @@ import importlib +from datetime import timedelta import pytest import torch +import torch.distributed as dist +import torch.multiprocessing as mp pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") @@ -79,3 +82,89 @@ def test_vocab_parallel_topk_single_rank_caps_width_to_vocabulary(): assert logits.shape == (1, 3) torch.testing.assert_close(token_ids.sort(dim=1).values, torch.tensor([[0, 1, 2]], device="cuda")) + + +@pytest.mark.parametrize("topk", [1, 4, 16]) +def test_vocab_parallel_topk_preserves_first_argmax_without_duplicate_candidates(monkeypatch, topk): + module = importlib.import_module("lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk") + original_topk = torch.topk + + def reversed_topk(input, **kwargs): + # A valid top-k implementation may choose the largest tied token ids. + values, indexes = original_topk(input.flip(0), **kwargs) + return values, input.shape[0] - 1 - indexes + + monkeypatch.setattr(torch, "topk", reversed_topk) + local_logits = torch.ones((16, 3), device="cuda") + local_logits[:2, 1:] = 0 + local_logits[7, 2] = 2 + + logits, token_ids = module.vocab_parallel_topk( + local_logits, + vocab_size=16, + vocab_start_id=0, + topk=topk, + tp_world_size=1, + group=None, + alloc_func=torch.empty, + ) + + selected = token_ids.gather(1, logits.argmax(dim=1, keepdim=True)).view(-1) + torch.testing.assert_close(selected, local_logits.argmax(dim=0)) + torch.testing.assert_close(logits, local_logits.T.gather(1, token_ids)) + assert (token_ids.sort(dim=1).values.diff(dim=1) > 0).all() + torch.testing.assert_close(logits.sort(dim=1).values, local_logits.topk(topk, dim=0).values.T.sort(dim=1).values) + + +def _check_nccl_topk(rank, init_method): + module = importlib.import_module("lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk") + torch.cuda.set_device(rank) + dist.init_process_group("nccl", rank=rank, world_size=2, init_method=init_method, timeout=timedelta(seconds=60)) + graph = None + try: + # Uneven shards, ties within/across ranks, and top-k capped by the smaller shard. + full_logits = torch.ones((17, 3), device="cuda") + full_logits[:2, 1] = 0 + full_logits[12, 2] = 3 + start, end = (0, 8) if rank == 0 else (8, 17) + local_logits = full_logits[start:end].clone() + for topk in (1, 4, 128): + + def forward(): + return module.vocab_parallel_topk( + local_logits, + vocab_size=17, + vocab_start_id=start, + topk=topk, + tp_world_size=2, + group=dist.group.WORLD, + alloc_func=torch.empty, + ) + + forward() + torch.cuda.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + logits, token_ids = forward() + + for winner in (12, 3): + full_logits[:, 2] = 0 + full_logits[winner, 2] = 3 + local_logits.copy_(full_logits[start:end]) + graph.replay() + torch.cuda.synchronize() + assert logits.shape == (3, 2 * min(topk, 8)) + selected = token_ids.gather(1, logits.argmax(dim=1, keepdim=True)).view(-1) + torch.testing.assert_close(selected, full_logits.argmax(dim=0)) + torch.testing.assert_close(logits, full_logits.T.gather(1, token_ids)) + assert (token_ids.sort(dim=1).values.diff(dim=1) > 0).all() + finally: + # NCCL shutdown waits for communicators retained by captured graphs. + del graph + dist.destroy_process_group() + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Two CUDA devices are required for NCCL") +def test_vocab_parallel_topk_nccl_and_cuda_graph_replay(tmp_path): + mp.spawn(_check_nccl_topk, args=(f"file://{tmp_path / 'nccl_init'}",), nprocs=2, join=True) diff --git a/unit_tests/models/test_vocab_parallel_topk_output.py b/unit_tests/models/test_vocab_parallel_topk_output.py index e8dc0adb2d..069347741a 100644 --- a/unit_tests/models/test_vocab_parallel_topk_output.py +++ b/unit_tests/models/test_vocab_parallel_topk_output.py @@ -29,6 +29,19 @@ def test_dense_argmax_keeps_column_index_semantics(): torch.testing.assert_close(backend._gen_argmax_token_ids(output), torch.tensor([1])) +def test_argmax_probability_keeps_logit_winner_when_softmax_rounds_to_a_tie(): + backend = ModeBackend.__new__(ModeBackend) + output = ModelOutput( + logits=torch.tensor([[0.0, 1e-8]]), + logits_token_ids=torch.tensor([[1, 101]]), + ) + + token_ids, probs = backend._gen_argmax_token_ids_and_prob(output) + + torch.testing.assert_close(token_ids, backend._gen_argmax_token_ids(output)) + torch.testing.assert_close(probs, torch.tensor([0.5])) + + def test_dspark_confidence_path_receives_mapped_sparse_token_ids(): post = Qwen3DSparkPostLayerInfer.__new__(Qwen3DSparkPostLayerInfer) post.block_size_ = 2 diff --git a/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_topk_sampling.py b/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_topk_sampling.py index 09f74e3b07..c482464622 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_topk_sampling.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_topk_sampling.py @@ -74,6 +74,17 @@ def test_samples_sparse_candidates_and_maps_global_token_ids(): torch.testing.assert_close(token_logprobs, torch.log(expected_probs)) +def test_sparse_sampling_selects_logits_before_softmax_rounding(): + output = ModelOutput( + logits=torch.tensor([[0.0, 1e-8]]), + logits_token_ids=torch.tensor([[1, 101]]), + ) + + token_ids, _ = sample(output, [make_req()]) + + torch.testing.assert_close(token_ids, torch.tensor([101])) + + @pytest.mark.parametrize( "override", [ From 570ee1dae1ee9b53b3a07ed828125e9657a63662 Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 7 Sep 2026 13:17:01 +0800 Subject: [PATCH 5/8] refactor: simplify vocab-parallel sampling and batch eligibility --- lightllm/common/basemodel/batch_objs.py | 16 +++++++------ .../post_process/vocab_parallel_topk.py | 15 +----------- .../llama/layer_infer/post_layer_infer.py | 6 +++-- .../layer_infer/post_layer_infer.py | 7 ++---- .../model_infer/mode_backend/base_backend.py | 23 ++++-------------- .../chunked_prefill/impl_for_reward_model.py | 1 - .../mode_backend/generic_post_process.py | 24 +++++++------------ .../mode_backend/generic_pre_process.py | 12 +++++----- .../models/test_vocab_parallel_topk_output.py | 23 ++++++++++++++++++ .../mode_backend/test_generic_pre_process.py | 24 +++++++++++++++++++ .../test_vocab_parallel_topk_sampling.py | 9 ++++--- 11 files changed, 89 insertions(+), 71 deletions(-) diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index 77ed1d400d..d6379412e2 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -7,6 +7,12 @@ from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor +def logits_indexes_to_token_ids(indexes: torch.Tensor, token_ids: Optional[torch.Tensor]) -> torch.Tensor: + if token_ids is None: + return indexes + return token_ids.gather(1, indexes.reshape(-1, 1)).reshape_as(indexes).long() + + @dataclass class ModelInput: # 通用变量 @@ -223,10 +229,6 @@ def to_no_ref_tensor(self): self.logits_token_ids = tensor_to_no_ref_tensor(self.logits_token_ids) self.mtp_collector.to_no_ref_tensor() - @property - def has_vocab_parallel_logits(self) -> bool: - return self.logits_token_ids is not None - def index_select_logits_rows(self, index: torch.Tensor) -> "ModelOutput": """Select logit rows without dropping their vocabulary metadata.""" @@ -242,11 +244,11 @@ def concat_logits_rows(cls, outputs: List["ModelOutput"]) -> "ModelOutput": """Concatenate outputs that share the same dense or sparse layout.""" assert outputs - has_vocab_parallel_logits = outputs[0].has_vocab_parallel_logits - assert all(output.has_vocab_parallel_logits == has_vocab_parallel_logits for output in outputs) + has_token_ids = outputs[0].logits_token_ids is not None + assert all((output.logits_token_ids is not None) == has_token_ids for output in outputs) return cls( logits=torch.cat([output.logits for output in outputs], dim=0), logits_token_ids=( - torch.cat([output.logits_token_ids for output in outputs], dim=0) if has_vocab_parallel_logits else None + torch.cat([output.logits_token_ids for output in outputs], dim=0) if has_token_ids else None ), ) diff --git a/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_topk.py b/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_topk.py index 9d8d258fcd..bc7ef4ce5d 100644 --- a/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_topk.py +++ b/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_topk.py @@ -1,28 +1,15 @@ """Collect sparse candidates directly from tensor-parallel vocabulary shards.""" -import os - import torch from lightllm.distributed.communication_op import all_gather_into_tensor from lightllm.utils.envs_utils import enable_env_vars -VOCAB_PARALLEL_TOPK_ENV = "LIGHTLLM_VOCAB_PARALLEL_TOPK" -VOCAB_PARALLEL_TOPK_SIZE_ENV = "LIGHTLLM_VOCAB_PARALLEL_TOPK_SIZE" -DEFAULT_VOCAB_PARALLEL_TOPK = 128 - - def is_vocab_parallel_topk_enabled() -> bool: """Whether target-model greedy batches may use sparse vocabulary output.""" - return enable_env_vars(VOCAB_PARALLEL_TOPK_ENV) - - -def get_vocab_parallel_topk_size() -> int: - topk = int(os.getenv(VOCAB_PARALLEL_TOPK_SIZE_ENV, str(DEFAULT_VOCAB_PARALLEL_TOPK))) - assert topk > 0, f"{VOCAB_PARALLEL_TOPK_SIZE_ENV} must be positive, got {topk}" - return topk + return enable_env_vars("LIGHTLLM_VOCAB_PARALLEL_TOPK") @torch.no_grad() diff --git a/lightllm/models/llama/layer_infer/post_layer_infer.py b/lightllm/models/llama/layer_infer/post_layer_infer.py index 496de58ef1..376600b252 100644 --- a/lightllm/models/llama/layer_infer/post_layer_infer.py +++ b/lightllm/models/llama/layer_infer/post_layer_infer.py @@ -1,3 +1,5 @@ +import os + import torch import torch.functional as F import torch.distributed as dist @@ -7,7 +9,6 @@ from lightllm.models.llama.infer_struct import LlamaInferStateInfo from lightllm.common.basemodel import PostLayerInferTpl from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk import ( - get_vocab_parallel_topk_size, vocab_parallel_topk, ) from lightllm.distributed.communication_op import all_gather @@ -19,7 +20,8 @@ class LlamaPostLayerInfer(PostLayerInferTpl): def __init__(self, network_config): super().__init__(network_config) self.eps_ = network_config["rms_norm_eps"] - self.vocab_parallel_topk_ = get_vocab_parallel_topk_size() + self.vocab_parallel_topk_ = int(os.getenv("LIGHTLLM_VOCAB_PARALLEL_TOPK_SIZE", 128)) + assert self.vocab_parallel_topk_ > 0, "LIGHTLLM_VOCAB_PARALLEL_TOPK_SIZE must be positive" return def _norm(self, input, infer_state, layer_weight: LlamaPreAndPostLayerWeight) -> torch.Tensor: diff --git a/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py b/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py index 2c89610f13..5a445032b3 100644 --- a/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py +++ b/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py @@ -1,5 +1,6 @@ import torch +from lightllm.common.basemodel.batch_objs import logits_indexes_to_token_ids from lightllm.distributed.communication_op import all_gather_into_tensor from lightllm.models.qwen3_dflash.infer_struct import Qwen3DFlashInferStateInfo from lightllm.models.qwen3_dflash.layer_infer.post_layer_infer import Qwen3DFlashPostLayerInfer @@ -182,11 +183,7 @@ def token_forward( logits = self._lm_head_and_gather(last_input, token_num, layer_weight, infer_state) block_logits = logits.reshape(num_reqs, self.block_size_, -1) candidate_indexes = torch.argmax(block_logits, dim=-1) - if infer_state.logits_token_ids is None: - sampled_tokens = candidate_indexes - else: - block_token_ids = infer_state.logits_token_ids.reshape(num_reqs, self.block_size_, -1) - sampled_tokens = block_token_ids.gather(-1, candidate_indexes.unsqueeze(-1)).squeeze(-1) + sampled_tokens = logits_indexes_to_token_ids(candidate_indexes, infer_state.logits_token_ids) confidence_logits = self.predict_confidence_logits( block_hidden, anchor_token_ids=anchor_token_ids, diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index fbb5b9c1c3..6bba6a515f 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -19,7 +19,7 @@ from lightllm.common.linear_att_cache_manager import LinearAttCacheManager from lightllm.server.router.dynamic_prompt.linear_att_radix_cache import LinearAttPagedRadixCache from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache -from lightllm.common.basemodel.batch_objs import ModelOutput, ModelInput +from lightllm.common.basemodel.batch_objs import ModelOutput, ModelInput, logits_indexes_to_token_ids from lightllm.utils.dist_utils import init_distributed_env from lightllm.utils.envs_utils import get_unique_server_name from lightllm.server.core.objs import ShmReqManager, StartArgs @@ -391,20 +391,13 @@ def _get_next_token_ranks(self, model_output: ModelOutput, next_token_ids: torch 仅 ``--enable_rl`` 时做真实 rank；否则返回 GPU 常量 ``-1``，避免 O(batch * vocab) 比较。 下游 async_copy 在同样条件下会忽略该返回值。 """ - if model_output.has_vocab_parallel_logits: + if not self.args.enable_rl or model_output.logits_token_ids is not None: return g_pin_mem_manager.get_const_gpu_tensor( key="next_token_ranks", shape=next_token_ids.shape, fill_value=1 if self.args.enable_rl else -1, dtype=torch.int32, ) - if not self.args.enable_rl: - return g_pin_mem_manager.get_const_gpu_tensor( - key="next_token_ranks", - shape=next_token_ids.shape, - fill_value=-1, - dtype=torch.int32, - ) selected_logits = model_output.logits.gather(1, next_token_ids.long().view(-1, 1)) return (model_output.logits > selected_logits).sum(dim=-1, dtype=torch.int32) + 1 @@ -883,21 +876,15 @@ def _trans_req_ids_to_req_objs(self, req_ids: List[int]) -> List[InferReq]: def _gen_argmax_token_ids(self, model_output: ModelOutput): logits = model_output.logits candidate_indexes = torch.argmax(logits, dim=-1) - return self._map_logits_indexes_to_token_ids(model_output, candidate_indexes) + return logits_indexes_to_token_ids(candidate_indexes, model_output.logits_token_ids) def _gen_argmax_token_ids_and_prob(self, model_output: ModelOutput): logits = model_output.logits candidate_indexes = torch.argmax(logits, dim=-1) max_probs = torch.softmax(logits, dim=-1).gather(1, candidate_indexes.view(-1, 1)).view(-1) - token_ids = self._map_logits_indexes_to_token_ids(model_output, candidate_indexes) + token_ids = logits_indexes_to_token_ids(candidate_indexes, model_output.logits_token_ids) return token_ids, max_probs - @staticmethod - def _map_logits_indexes_to_token_ids(model_output: ModelOutput, candidate_indexes: torch.Tensor): - if not model_output.has_vocab_parallel_logits: - return candidate_indexes - return model_output.logits_token_ids.gather(1, candidate_indexes.long().view(-1, 1)).view(-1).long() - def _sample_and_scatter_token( self, model_output: ModelOutput, @@ -910,7 +897,7 @@ def _sample_and_scatter_token( ): logits = model_output.logits if mask_func is not None: - assert not model_output.has_vocab_parallel_logits, "constrained sampling requires dense logits" + assert model_output.logits_token_ids is None, "constrained sampling requires dense logits" assert len(run_reqs) == logits.shape[0] mask_func(run_reqs, logits) diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_reward_model.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_reward_model.py index cabf4f7de0..95aeaa3e9a 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_reward_model.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl_for_reward_model.py @@ -16,7 +16,6 @@ def __init__(self) -> None: def reward_prefill(self, event_pack: OverlapEventPack, prefill_reqs: List[InferReq]): assert self.disable_chunked_prefill is True model_input, run_reqs = prepare_prefill_inputs(prefill_reqs, is_chuncked_mode=not self.disable_chunked_prefill) - model_input.use_vocab_parallel_topk = False model_output = self.model.forward(model_input) scores: torch.Tensor = model_output.logits diff --git a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py index 84bd268172..4fd1e294dd 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py @@ -1,6 +1,6 @@ import torch from typing import List, Tuple -from lightllm.common.basemodel.batch_objs import ModelOutput +from lightllm.common.basemodel.batch_objs import ModelOutput, logits_indexes_to_token_ids from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk import ( is_vocab_parallel_topk_enabled, ) @@ -18,21 +18,15 @@ def _can_use_unmodified_greedy_logits(reqs: List[InferReq]) -> bool: for req_obj in reqs: sample_param = req_obj.sampling_param shm_param = sample_param.shm_param - if shm_param.top_k != 1 or shm_param.temperature != 1.0: - return False if ( - shm_param.presence_penalty != 0.0 - or shm_param.frequency_penalty != 0.0 - or shm_param.repetition_penalty != 1.0 + shm_param.top_k != 1 + or shm_param.temperature != 1.0 + or req_obj.need_out_token_id_statistics + or shm_param.exponential_decay_length_penalty.to_tuple()[1] != 1.0 + or req_obj.cur_output_len < shm_param.min_new_tokens - 1 + or sample_param.invalid_token_ids ): return False - if shm_param.exponential_decay_length_penalty.to_tuple()[1] != 1.0: - return False - out_token_len = req_obj.get_cur_total_len() - req_obj.shm_req.input_len - if out_token_len < shm_param.min_new_tokens - 1: - return False - if sample_param.invalid_token_ids: - return False return True @@ -42,13 +36,13 @@ def can_use_vocab_parallel_topk(reqs: List[InferReq]) -> bool: def sample(model_output: ModelOutput, reqs: List[InferReq], eos_id: List[int] = [2]): logits = model_output.logits - if model_output.has_vocab_parallel_logits: + if model_output.logits_token_ids is not None: if not _can_use_unmodified_greedy_logits(reqs): raise RuntimeError("vocab-parallel top-k logits require unmodified greedy requests") candidate_indexes = torch.argmax(logits, dim=-1, keepdim=True) probs = torch.softmax(logits, dim=-1) max_probs = probs.gather(1, candidate_indexes).view(-1) - token_ids = model_output.logits_token_ids.gather(1, candidate_indexes).view(-1).long() + token_ids = logits_indexes_to_token_ids(candidate_indexes.view(-1), model_output.logits_token_ids) return token_ids, torch.log(max_probs) ( diff --git a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py index ee95ff24c5..eadf50ae78 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py @@ -181,9 +181,9 @@ def overlap_prepare_decode_inputs(req_objs: List[InferReq]): model_input1, run_reqs1 = prepare_decode_inputs( req_objs=decode_reqs1, ) - use_vocab_parallel_topk = can_use_vocab_parallel_topk(run_reqs0 + run_reqs1) - model_input0.use_vocab_parallel_topk = use_vocab_parallel_topk - model_input1.use_vocab_parallel_topk = use_vocab_parallel_topk + model_input0.use_vocab_parallel_topk = model_input1.use_vocab_parallel_topk = ( + model_input0.use_vocab_parallel_topk and model_input1.use_vocab_parallel_topk + ) return model_input0, run_reqs0, decode_reqs0, model_input1, run_reqs1, decode_reqs1 @@ -219,9 +219,9 @@ def overlap_prepare_prefill_inputs(req_objs: List[InferReq]): req_objs=right_reqs, is_chuncked_mode=True, ) - use_vocab_parallel_topk = can_use_vocab_parallel_topk(run_reqs0 + run_reqs1) - model_input0.use_vocab_parallel_topk = use_vocab_parallel_topk - model_input1.use_vocab_parallel_topk = use_vocab_parallel_topk + model_input0.use_vocab_parallel_topk = model_input1.use_vocab_parallel_topk = ( + model_input0.use_vocab_parallel_topk and model_input1.use_vocab_parallel_topk + ) return model_input0, run_reqs0, model_input1, run_reqs1 diff --git a/unit_tests/models/test_vocab_parallel_topk_output.py b/unit_tests/models/test_vocab_parallel_topk_output.py index 069347741a..664d67e2b5 100644 --- a/unit_tests/models/test_vocab_parallel_topk_output.py +++ b/unit_tests/models/test_vocab_parallel_topk_output.py @@ -1,10 +1,33 @@ from types import SimpleNamespace import torch +import pytest from lightllm.common.basemodel.batch_objs import ModelOutput from lightllm.models.qwen3_dspark.layer_infer.post_layer_infer import Qwen3DSparkPostLayerInfer from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend +from lightllm.server.router.model_infer.mode_backend import base_backend + + +@pytest.mark.parametrize("enable_rl", [False, True]) +@pytest.mark.parametrize("sparse", [False, True]) +def test_token_ranks_preserve_dense_and_sparse_semantics(monkeypatch, enable_rl, sparse): + monkeypatch.setattr( + base_backend.g_pin_mem_manager, + "get_const_gpu_tensor", + lambda key, shape, fill_value, dtype: torch.full(shape, fill_value, dtype=dtype), + ) + backend = ModeBackend.__new__(ModeBackend) + backend.args = SimpleNamespace(enable_rl=enable_rl) + output = ModelOutput( + logits=torch.tensor([[3.0, 3.0, 1.0]]), + logits_token_ids=torch.tensor([[30, 50, 10]]) if sparse else None, + ) + selected = torch.tensor([30 if sparse else 2]) + expected = (1 if sparse else 3) if enable_rl else -1 + torch.testing.assert_close( + backend._get_next_token_ranks(output, selected), torch.tensor([expected], dtype=torch.int32) + ) def test_argmax_restores_global_token_ids_and_uses_candidate_probability(): diff --git a/unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py b/unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py index 62296634a9..a5264ff3cf 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py @@ -1,6 +1,7 @@ from types import SimpleNamespace import torch +import pytest from lightllm.server.router.model_infer.mode_backend import generic_pre_process @@ -170,3 +171,26 @@ def test_overlap_decode_preserves_empty_microbatch(monkeypatch): assert model_input1.batch_size == 0 assert model_input1.b_req_idx.shape == (0,) assert model_input1.mem_indexes_cpu.shape == (0,) + + +@pytest.mark.parametrize("prefill", [False, True]) +@pytest.mark.parametrize("eligible", [[], [True], [False], [True, True], [True, False], [False, True]]) +def test_overlap_reuses_batch_eligibility(monkeypatch, prefill, eligible): + _patch_empty_input_context(monkeypatch) + calls = [] + + def can_use(reqs): + calls.append(reqs) + return all(req.eligible for req in reqs) + + monkeypatch.setattr(generic_pre_process, "can_use_vocab_parallel_topk", can_use) + reqs = [_make_prefill_req(i, 1) if prefill else _make_decode_req(i) for i in range(len(eligible))] + for req, flag in zip(reqs, eligible): + req.eligible = flag + if prefill: + first, _, second, _ = generic_pre_process.overlap_prepare_prefill_inputs(reqs) + else: + first, _, _, second, _, _ = generic_pre_process.overlap_prepare_decode_inputs(reqs) + + assert first.use_vocab_parallel_topk == second.use_vocab_parallel_topk == all(eligible) + assert len(calls) == 2 diff --git a/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_topk_sampling.py b/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_topk_sampling.py index c482464622..e1acf78dc9 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_topk_sampling.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_topk_sampling.py @@ -34,14 +34,17 @@ def make_req(**overrides): min_new_tokens=values["min_new_tokens"], exponential_decay_length_penalty=SimpleNamespace(to_tuple=lambda: (1, values["decay_factor"])), ) - input_len = 10 return SimpleNamespace( sampling_param=SimpleNamespace( shm_param=shm_param, invalid_token_ids=values["invalid_token_ids"], ), - shm_req=SimpleNamespace(input_len=input_len), - get_cur_total_len=lambda: input_len + values["output_len"], + cur_output_len=values["output_len"], + need_out_token_id_statistics=( + values["presence_penalty"] != 0.0 + or values["frequency_penalty"] != 0.0 + or values["repetition_penalty"] != 1.0 + ), ) From acf3238eed79263f398b3d993010d30742cdd20d Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 7 Sep 2026 15:02:30 +0800 Subject: [PATCH 6/8] Refactor vocab-parallel sampling into exact top1 and opt-in approximate topk --- docs/CN/source/tutorial/api_server_args.rst | 44 +++++++--- docs/EN/source/tutorial/api_server_args.rst | 59 +++++++++---- lightllm/common/basemodel/basemodel.py | 64 ++++++++------ lightllm/common/basemodel/batch_objs.py | 13 ++- lightllm/common/basemodel/cuda_graph.py | 40 ++++++--- lightllm/common/basemodel/infer_struct.py | 4 +- .../common/basemodel/prefill_cuda_graph.py | 5 -- .../post_process/vocab_parallel_topk.py | 85 +++++++++++++++++-- .../llama/layer_infer/post_layer_infer.py | 16 ++-- lightllm/server/api_cli.py | 16 ++++ lightllm/server/api_start.py | 2 + lightllm/server/core/objs/start_args_type.py | 3 + .../model_infer/mode_backend/base_backend.py | 5 +- .../mode_backend/chunked_prefill/impl.py | 20 +++-- .../mode_backend/generic_post_process.py | 54 +++++++++--- .../mode_backend/generic_pre_process.py | 26 ++++-- .../basemodel/test_cuda_graph_layout.py | 50 +++++++++++ .../common/basemodel/test_model_output.py | 57 +++++++++---- .../triton_kernel/test_vocab_parallel_topk.py | 80 ++++++++++++++++- .../models/test_gemma4_vocab_parallel_topk.py | 36 +++++++- .../models/test_vocab_parallel_topk_output.py | 5 +- .../mode_backend/test_generic_pre_process.py | 44 +++++++++- .../test_vocab_parallel_topk_sampling.py | 84 ++++++++++++++---- unit_tests/server/test_mtp_start_args.py | 21 +++++ 24 files changed, 673 insertions(+), 160 deletions(-) diff --git a/docs/CN/source/tutorial/api_server_args.rst b/docs/CN/source/tutorial/api_server_args.rst index f04645c8c6..0ba15801f9 100644 --- a/docs/CN/source/tutorial/api_server_args.rst +++ b/docs/CN/source/tutorial/api_server_args.rst @@ -641,21 +641,41 @@ MTP 多预测参数 增加此值允许更多预测，但确保模型与指定的步数兼容。 目前 deepseekv3/r1 模型仅支持 1 步 -词表并行 Top-k 环境变量 ------------------------ +词表并行采样 +------------ + +.. option:: --disable_vocab_parallel_top1 + + 关闭主模型和草稿模型的自动贪心归约。默认对满足 ``top_k=1``、 + ``temperature=1`` 且无 logits 修改需求的主模型批次，汇总每卡最大值及 + 一个归一化统计量。token 与完整词表 argmax 一致，包括并列值规则； + 主模型 logprob 使用完整词表分母，允许浮点归约的舍入误差。 + 固定贪心 draft 只取 top-1，不计算概率。 + +.. option:: --enable_vocab_parallel_topk + + 显式接受主模型的近似采样，默认关闭。先截取每卡候选，再在汇总候选上 + 应用请求的 temperature、top-k 和 top-p。返回 logprob 使用请求 top-k/top-p + 过滤前的候选 softmax，既不是完整词表 logprob，也不是最终过滤分布的 + logprob。符合条件的 greedy 批次仍优先使用准确 top-1，除非已关闭它。 + +.. option:: --vocab_parallel_topk_size + + 每个 TP rank 的候选数，正整数，默认 ``128``，不超过最小词表分片大小。 + 与请求的 ``top_k`` 不同；改变 TP 数量可能改变近似分布。 + 允许设置为 ``1``，但在 TP=1 时无法提供有区分度的概率置信度。 -``LIGHTLLM_VOCAB_PARALLEL_TOPK=1`` 为满足 ``top_k=1``、``temperature=1`` -且无需修改 logits 的主模型批次启用稀疏输出。惩罚项、EOS 屏蔽、无效 token -屏蔽及约束采样会回退到完整词表输出。主模型默认关闭此功能；使用共享 Llama -后处理层的草稿模型始终使用稀疏候选。 +动态 Vanilla/EAGLE/DFlash draft 默认使用上述候选数计算近似调度置信度； +DSpark 使用独立 confidence head，可以保留 top-1。 +``--disable_vocab_parallel_top1`` 同时关闭 draft 的自动归约；配合不开启 +``--enable_vocab_parallel_topk``，即可让主模型和 draft 全部恢复完整输出。 -``LIGHTLLM_VOCAB_PARALLEL_TOPK_SIZE`` 设置每个 TP rank 的本地候选数， -必须为正整数，默认 ``128``，实际取值不超过最小词表分片的大小。 -服务启动前，各 rank 应设置相同的值。 +惩罚项、EOS/无效 token 屏蔽、约束采样、非 greedy 的显式 seed 及 RL rank +需求回退完整 logits。不支持的模型 head、硬件平台，以及候选通信量不小于 +完整 logits 的情况也回退。Prompt logits 始终完整。Overlap 两侧使用兼容布局； +decode 在启动时分别捕获启用布局的 graph，因此启动时间和 graph 显存可能增加。 +近似模式不保证与完整模式同 seed 的输出一致。 -贪心采样与完整词表的 argmax 一致，分数相同时选择最小 token ID。 -返回的概率及动态 MTP 使用的草稿概率仅在候选集内归一化，属于近似值。 -回退到完整词表的批次会绕过为稀疏主模型输出捕获的 decode CUDA Graph。 DeepSeek 冗余专家参数 --------------------- diff --git a/docs/EN/source/tutorial/api_server_args.rst b/docs/EN/source/tutorial/api_server_args.rst index fc17d1a7f0..b64a02866a 100644 --- a/docs/EN/source/tutorial/api_server_args.rst +++ b/docs/EN/source/tutorial/api_server_args.rst @@ -644,23 +644,48 @@ MTP Multi-Prediction Parameters Increasing this value allows more predictions, but ensure the model is compatible with the specified number of steps. Currently deepseekv3/r1 models only support 1 step -Vocabulary-Parallel Top-k Environment Variables ------------------------------------------------ - -``LIGHTLLM_VOCAB_PARALLEL_TOPK=1`` enables sparse output for target-model -batches with ``top_k=1``, ``temperature=1``, and no logit modifiers. Penalties, -EOS masking, invalid-token masking, and constrained sampling use dense output. -The target-model feature is disabled by default; draft models using the shared -Llama post layer always use sparse candidates. - -``LIGHTLLM_VOCAB_PARALLEL_TOPK_SIZE`` sets the positive number of local -candidates per TP rank (default: ``128``), capped at the smallest vocabulary -shard size. Set the same value on all ranks before starting the service. - -Greedy token selection preserves dense argmax, including choosing the smallest -token ID on ties. Returned probabilities and dynamic-MTP draft probabilities -are normalized over the candidate set and are approximate. Dense fallback -batches bypass decode CUDA Graphs captured for sparse target output. +Vocabulary-Parallel Sampling +---------------------------- + +.. option:: --disable_vocab_parallel_top1 + + Disable automatic greedy reduction for target and draft models. By default, + eligible target batches with ``top_k=1`` and ``temperature=1`` gather one + maximum per rank and one normalization statistic. Selected tokens preserve + dense argmax tie-breaking; target logprobs use the full vocabulary denominator + (subject to floating-point reduction rounding). Fixed greedy drafts collect + top-1 without computing probabilities. + +.. option:: --enable_vocab_parallel_topk + + Opt into approximate target sampling, disabled by default. Each rank retains + local candidates before applying request temperature, top-k and top-p to the + gathered candidate set. Returned logprobs use the candidate softmax before + request top-k/top-p filtering; they are not full-vocabulary logprobs or the + final filtered sampling distribution's logprobs. Eligible greedy batches + still prefer exact top-1 unless it is disabled. + +.. option:: --vocab_parallel_topk_size + + Positive candidate count per TP rank, default ``128``, capped at the smallest + vocabulary shard. This is distinct from request ``top_k``. Changing TP size + can change the approximate distribution. Size ``1`` is allowed, but provides + no useful confidence signal at TP=1. + +Dynamic Vanilla/EAGLE/DFlash drafts use this candidate count for approximate +scheduling confidence by default. DSpark uses its independent confidence head +and can retain top-1. ``--disable_vocab_parallel_top1`` also disables automatic +draft reduction; combined with leaving ``--enable_vocab_parallel_topk`` unset, +it restores dense output for both target and draft models. + +Penalties, EOS/invalid-token masking, constrained sampling, non-greedy explicit +seeds and non-greedy RL ranks retain dense output. Unsupported heads/platforms +and candidate payloads no smaller than dense logits also retain the dense path. +Prompt logits remain dense. Overlap microbatches share a compatible layout; +decode graphs for enabled layouts are captured separately at startup, increasing +capture time and potentially graph memory usage. Approximate mode does not +promise seed-equivalent output to dense mode. + DeepSeek Redundant Expert Parameters ------------------------------------ diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index bebf2671cc..b6716b9892 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -23,9 +23,6 @@ from lightllm.common.basemodel.prefill_cuda_graph import PrefillCudaGraph from lightllm.common.quantization import Quantcfg from lightllm.common.basemodel.triton_kernel.gather_token_id import gather_token, gather_token_prefill_decode_mixed -from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk import ( - is_vocab_parallel_topk_enabled, -) from lightllm.utils.log_utils import init_logger from lightllm.utils.dist_utils import get_dp_world_size from lightllm.utils.profile_max_tokens import profile_mtp_weight_memory @@ -381,22 +378,42 @@ def forward(self, model_input: ModelInput): else: return self._decode(model_input) - def _is_cuda_graph_output_compatible(self, *model_inputs: ModelInput) -> bool: - """Whether inputs match the dense/sparse contract captured at startup.""" - + def _vocab_parallel_output_mode(self, model_input: ModelInput) -> tuple[int, bool]: + if self.args.hardware_platform != "cuda": + return 0, False + if self.is_mtp_draft_model: + if self.args.disable_vocab_parallel_top1: + return 0, False + needs_probs = self.args.mtp_dynamic_verify and self.args.mtp_mode != "dspark" + return (self.args.vocab_parallel_topk_size if needs_probs else 1), False return ( - self.is_mtp_draft_model - or not is_vocab_parallel_topk_enabled() - or all(model_input.use_vocab_parallel_topk for model_input in model_inputs) + model_input.vocab_parallel_topk, + model_input.vocab_parallel_topk == 1 + and model_input.vocab_parallel_greedy + and not self.args.disable_vocab_parallel_top1, ) + def vocab_parallel_graph_modes(self): + if self.args.hardware_platform != "cuda": + return [(0, False)] + if self.is_mtp_draft_model: + return [self._vocab_parallel_output_mode(None)] + modes = [(0, False)] + if not self.args.disable_vocab_parallel_top1: + modes.append((1, True)) + if self.args.enable_vocab_parallel_topk: + modes.append((self.args.vocab_parallel_topk_size, False)) + return modes + def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0): infer_state = self.infer_state_class() infer_state.hidden_collector = self.hidden_collector_prototype.new_instance() infer_state.input_ids = model_input.input_ids infer_state.is_prefill = model_input.is_prefill infer_state.return_all_prompt_logics = self.return_all_prompt_logics - infer_state.use_vocab_parallel_topk = self.is_mtp_draft_model or model_input.use_vocab_parallel_topk + infer_state.vocab_parallel_topk, infer_state.vocab_parallel_greedy = self._vocab_parallel_output_mode( + model_input + ) infer_state.batch_size = model_input.batch_size infer_state.total_token_num = model_input.total_token_num infer_state.max_q_seq_len = model_input.max_q_seq_len @@ -660,18 +677,14 @@ def _decode( # CUDA Graph 可能继续向上对齐 batch size，并因此加入 seq_len=2 的 # dummy request。先用最终可能出现的 KV 长度判断 graph，再统一 padding 一次。 infer_max_kv_seq_len = max(2, model_input.max_kv_seq_len) - use_cuda_graph = ( - self._is_cuda_graph_output_compatible(model_input) - and self.graph is not None - and self.graph.can_run( - batch_size=infer_batch_size, - max_len_in_batch=infer_max_kv_seq_len, - ) + use_cuda_graph = self.graph is not None and self.graph.can_run( + batch_size=infer_batch_size, + max_len_in_batch=infer_max_kv_seq_len, ) need_capture = False if use_cuda_graph: infer_batch_size = self.graph.find_closest_graph_batch_size(batch_size=infer_batch_size) - need_capture = self.graph.need_capture(infer_batch_size) + need_capture = self.graph.need_capture(infer_batch_size, self._vocab_parallel_output_mode(model_input)) model_input = self._create_padded_decode_model_input(model_input=model_input, new_batch_size=infer_batch_size) infer_state = self._create_inferstate(model_input) @@ -758,6 +771,7 @@ def prefill_func(input_tensors, _infer_state): model_output = ModelOutput( logits=predict_logits.contiguous(), logits_token_ids=infer_state.logits_token_ids, + logits_are_logprobs=infer_state.logits_are_logprobs, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, ) @@ -788,6 +802,7 @@ def _token_forward(self, infer_state: InferStateInfo): model_output = ModelOutput( logits=predict_logits.contiguous(), logits_token_ids=infer_state.logits_token_ids, + logits_are_logprobs=infer_state.logits_are_logprobs, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), ) @@ -919,13 +934,10 @@ def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1 infer_batch_size = max(1, origin_batch_size0, origin_batch_size1) infer_batch_size = triton.cdiv(infer_batch_size, self.tp_world_size_) * self.tp_world_size_ - if ( - self._is_cuda_graph_output_compatible(model_input0, model_input1) - and self.graph is not None - and self.graph.can_run(infer_batch_size, max_len_in_batch) - ): + if self.graph is not None and self.graph.can_run(infer_batch_size, max_len_in_batch): infer_batch_size = self.graph.find_closest_graph_batch_size(infer_batch_size) - need_capture = self.graph.need_capture(infer_batch_size) + assert self._vocab_parallel_output_mode(model_input0) == self._vocab_parallel_output_mode(model_input1) + need_capture = self.graph.need_capture(infer_batch_size, self._vocab_parallel_output_mode(model_input0)) padded_model_input0 = self._create_padded_decode_model_input(model_input0, infer_batch_size) padded_model_input1 = self._create_padded_decode_model_input(model_input1, infer_batch_size) infer_state0 = self._create_inferstate(padded_model_input0, 0) @@ -1047,12 +1059,14 @@ def _overlap_tpsp_context_forward(self, infer_state: InferStateInfo, infer_state model_output = ModelOutput( logits=predict_logits.contiguous(), logits_token_ids=infer_state.logits_token_ids, + logits_are_logprobs=infer_state.logits_are_logprobs, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, ) model_output1 = ModelOutput( logits=predict_logits1.contiguous(), logits_token_ids=infer_state1.logits_token_ids, + logits_are_logprobs=infer_state1.logits_are_logprobs, mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1), prompt_logics=infer_state1.prompt_logics, ) @@ -1098,11 +1112,13 @@ def _overlap_tpsp_token_forward(self, infer_state: InferStateInfo, infer_state1: model_output = ModelOutput( logits=predict_logits.contiguous(), logits_token_ids=infer_state.logits_token_ids, + logits_are_logprobs=infer_state.logits_are_logprobs, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), ) model_output1 = ModelOutput( logits=predict_logits1.contiguous(), logits_token_ids=infer_state1.logits_token_ids, + logits_are_logprobs=infer_state1.logits_are_logprobs, mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1), ) diff --git a/lightllm/common/basemodel/batch_objs.py b/lightllm/common/basemodel/batch_objs.py index d6379412e2..8c67518be6 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -61,9 +61,9 @@ class ModelInput: # 的 draft 模型的输入 mtp_draft_input_hiddens: Optional[torch.Tensor] = None - # The router sets this only when a target-model batch can sample directly - # from sparse candidates. Draft models always enable the same output form. - use_vocab_parallel_topk: bool = False + # 0: dense logits; positive: per-rank candidate count. + vocab_parallel_topk: int = 0 + vocab_parallel_greedy: bool = False def to_cuda(self): self.check_input() @@ -213,6 +213,8 @@ class ModelOutput: # Sparse vocab-parallel outputs map every candidate column back to its # global token id. None means logits are dense and column indexes are ids. logits_token_ids: Optional[torch.Tensor] = None + # Exact target top-1 output already contains the selected full-vocab logprob. + logits_are_logprobs: bool = False def __post_init__(self) -> None: if self.mtp_collector is None: @@ -222,6 +224,8 @@ def __post_init__(self) -> None: assert self.logits_token_ids.shape == self.logits.shape assert self.logits_token_ids.dtype in (torch.int32, torch.int64) assert self.logits_token_ids.device == self.logits.device + if self.logits_are_logprobs: + assert self.logits_token_ids is not None and self.logits.shape[1] == 1 def to_no_ref_tensor(self): self.logits = tensor_to_no_ref_tensor(self.logits) @@ -233,6 +237,7 @@ def index_select_logits_rows(self, index: torch.Tensor) -> "ModelOutput": """Select logit rows without dropping their vocabulary metadata.""" return ModelOutput( + logits_are_logprobs=self.logits_are_logprobs, logits=self.logits.index_select(0, index), logits_token_ids=( self.logits_token_ids.index_select(0, index) if self.logits_token_ids is not None else None @@ -246,7 +251,9 @@ def concat_logits_rows(cls, outputs: List["ModelOutput"]) -> "ModelOutput": assert outputs has_token_ids = outputs[0].logits_token_ids is not None assert all((output.logits_token_ids is not None) == has_token_ids for output in outputs) + assert all(output.logits_are_logprobs == outputs[0].logits_are_logprobs for output in outputs) return cls( + logits_are_logprobs=outputs[0].logits_are_logprobs, logits=torch.cat([output.logits for output in outputs], dim=0), logits_token_ids=( torch.cat([output.logits_token_ids for output in outputs], dim=0) if has_token_ids else None diff --git a/lightllm/common/basemodel/cuda_graph.py b/lightllm/common/basemodel/cuda_graph.py index c080c6f49e..1cdb006267 100644 --- a/lightllm/common/basemodel/cuda_graph.py +++ b/lightllm/common/basemodel/cuda_graph.py @@ -9,9 +9,6 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.distributed import dist_group_manager from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput -from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk import ( - is_vocab_parallel_topk_enabled, -) from lightllm.utils.torch_memory_saver_utils import ( TorchMemorySaverWrapper, MemoryTag, @@ -86,10 +83,10 @@ def __init__( def can_run(self, batch_size, max_len_in_batch): return batch_size <= self.max_batch_size and max_len_in_batch <= self.graph_max_len_in_batch - def need_capture(self, batch_size): + def need_capture(self, batch_size, output_mode=(0, False)): find_batch_size = self.find_closest_graph_batch_size(batch_size) if find_batch_size is not None: - return find_batch_size not in self.graph + return (find_batch_size, *output_mode) not in self.graph else: assert False, "dead code" @@ -128,7 +125,11 @@ def _capture_decode(self, decode_func, infer_state: InferStateInfo): with self.torch_memory_saver.cuda_graph(graph_obj, pool=self.mempool): model_output = decode_func(infer_state) - self.graph[batch_size] = (graph_obj, infer_state, model_output) + self.graph[(batch_size, infer_state.vocab_parallel_topk, infer_state.vocab_parallel_greedy)] = ( + graph_obj, + infer_state, + model_output, + ) graph_obj.replay() self._measure_replay_cost(graph_obj=graph_obj, batch_size=batch_size) return model_output @@ -163,7 +164,7 @@ def _capture_decode_overlap( with self.torch_memory_saver.cuda_graph(graph_obj, pool=self.mempool): model_output, model_output1 = decode_func(infer_state, infer_state1) - self.graph[batch_size] = ( + self.graph[(batch_size, infer_state.vocab_parallel_topk, infer_state.vocab_parallel_greedy)] = ( graph_obj, infer_state, infer_state1, @@ -215,7 +216,9 @@ def capture_decode( def _replay(self, infer_state: InferStateInfo): batch_size = infer_state.input_ids.shape[0] - graph_obj, graph_infer_state, graph_output = self.graph[batch_size] + graph_obj, graph_infer_state, graph_output = self.graph[ + (batch_size, infer_state.vocab_parallel_topk, infer_state.vocab_parallel_greedy) + ] graph_infer_state.copy_for_cuda_graph(infer_state) graph_obj.replay() return graph_output @@ -232,7 +235,11 @@ def _replay_overlap( graph_infer_state1, graph_model_output, graph_model_output1, - ) = self.graph[batch_size] + ) = self.graph[(batch_size, infer_state.vocab_parallel_topk, infer_state.vocab_parallel_greedy)] + assert (infer_state.vocab_parallel_topk, infer_state.vocab_parallel_greedy) == ( + infer_state1.vocab_parallel_topk, + infer_state1.vocab_parallel_greedy, + ) graph_infer_state.copy_for_cuda_graph(infer_state) graph_infer_state1.copy_for_cuda_graph(infer_state1) graph_obj.replay() @@ -282,11 +289,13 @@ def warmup(self, model): b_position_delta=torch.zeros(batch_size, dtype=torch.int32, device="cuda"), is_prefill=False, multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], - use_vocab_parallel_topk=is_vocab_parallel_topk_enabled(), **model._gen_special_model_input(batch_size), ) - model_output: ModelOutput = model.forward(model_input) - del model_output + for topk, greedy in model.vocab_parallel_graph_modes(): + model_input.vocab_parallel_topk = topk + model_input.vocab_parallel_greedy = greedy + model_output: ModelOutput = model.forward(model_input) + del model_output del input_ids del mem_indexes del b_req_idx @@ -344,7 +353,6 @@ def warmup_overlap(self, model): b_shared_radix_node_id=b_shared_radix_node_id, b_position_delta=torch.zeros(batch_size, dtype=torch.int32, device="cuda"), multimodal_params=[{"images": [], "audios": []} for _ in range(batch_size)], - use_vocab_parallel_topk=is_vocab_parallel_topk_enabled(), **model._gen_special_model_input(batch_size), ) decode_batches.append(micro_batch) @@ -355,7 +363,11 @@ def warmup_overlap(self, model): del locals()[var_name] torch.cuda.empty_cache() - _, _ = model.microbatch_overlap_decode(decode_batches[0], decode_batches[1]) + for topk, greedy in model.vocab_parallel_graph_modes(): + for batch in decode_batches: + batch.vocab_parallel_topk = topk + batch.vocab_parallel_greedy = greedy + _, _ = model.microbatch_overlap_decode(decode_batches[0], decode_batches[1]) model.mem_manager.free_all() model.req_manager.free_all() diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 3091c878a9..00f68314d5 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -52,7 +52,9 @@ def __init__(self): self.mem_index: torch.Tensor = None self.return_all_prompt_logics: bool = False - self.use_vocab_parallel_topk: bool = False + self.vocab_parallel_topk: int = 0 + self.vocab_parallel_greedy: bool = False + self.logits_are_logprobs: bool = False self.logits_token_ids: Optional[torch.Tensor] = None # 在开启 return_all_prompt_logics 模式时，保存整个 prefill 阶段每一个 # token 位置的 logits，供后续回传 prompt logprobs 信息使用。 diff --git a/lightllm/common/basemodel/prefill_cuda_graph.py b/lightllm/common/basemodel/prefill_cuda_graph.py index 594dd1ef94..bf6039a48f 100644 --- a/lightllm/common/basemodel/prefill_cuda_graph.py +++ b/lightllm/common/basemodel/prefill_cuda_graph.py @@ -10,9 +10,6 @@ from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor from lightllm.distributed import dist_group_manager from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput -from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk import ( - is_vocab_parallel_topk_enabled, -) from .infer_struct import InferStateInfo from .cuda_graph import CudaGraph @@ -223,7 +220,6 @@ def warmup(self, model): is_prefill=True, b_prefill_has_output_cpu=[False], multimodal_params=[{"images": [], "audios": []}], - use_vocab_parallel_topk=is_vocab_parallel_topk_enabled(), **model._gen_special_model_input(token_num=total_token_num), ) model_output: ModelOutput = model.forward(model_input) @@ -285,7 +281,6 @@ def warmup_overlap(self, model): is_prefill=True, b_prefill_has_output_cpu=[False], multimodal_params=[{"images": [], "audios": []}], - use_vocab_parallel_topk=is_vocab_parallel_topk_enabled(), **model._gen_special_model_input(token_num=total_token_num), ) diff --git a/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_topk.py b/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_topk.py index bc7ef4ce5d..5086028292 100644 --- a/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_topk.py +++ b/lightllm/common/basemodel/triton_kernel/post_process/vocab_parallel_topk.py @@ -1,15 +1,75 @@ """Collect sparse candidates directly from tensor-parallel vocabulary shards.""" import torch +import triton +import triton.language as tl from lightllm.distributed.communication_op import all_gather_into_tensor -from lightllm.utils.envs_utils import enable_env_vars -def is_vocab_parallel_topk_enabled() -> bool: - """Whether target-model greedy batches may use sparse vocabulary output.""" +@triton.jit +def _top1_stats( + Logits, Stats, V: tl.constexpr, B: tl.constexpr, START: tl.constexpr, BV: tl.constexpr, BB: tl.constexpr +): + part = tl.program_id(0) + v = part * BV + tl.arange(0, BV) + b = tl.program_id(1) * BB + tl.arange(0, BB) + x = tl.load(Logits + v[:, None] * B + b[None, :], (v[:, None] < V) & (b[None, :] < B), other=-float("inf")).to( + tl.float32 + ) + maximum = tl.max(x, axis=0) + ids = tl.min(tl.where((x == maximum[None, :]) & (v[:, None] < V), v[:, None] + START, 2147483647), axis=0) + mass = tl.sum(tl.where(maximum[None, :] == -float("inf"), 0.0, tl.exp(x - maximum[None, :])), axis=0) + tl.store(Stats + (part * 3) * B + b, maximum, b < B) + tl.store(Stats + (part * 3 + 1) * B + b, ids.to(tl.float32, bitcast=True), b < B) + tl.store(Stats + (part * 3 + 2) * B + b, mass, b < B) + + +@triton.jit +def _merge_top1_stats( + Stats, Output, Ids, N: tl.constexpr, B: tl.constexpr, FINAL: tl.constexpr, BN: tl.constexpr, BB: tl.constexpr +): + n = tl.arange(0, BN) + b = tl.program_id(0) * BB + tl.arange(0, BB) + mask = (n[:, None] < N) & (b[None, :] < B) + maxima = tl.load(Stats + (n[:, None] * 3) * B + b[None, :], mask, other=-float("inf")) + ids = tl.load(Stats + (n[:, None] * 3 + 1) * B + b[None, :], mask, other=0).to(tl.int32, bitcast=True) + mass = tl.load(Stats + (n[:, None] * 3 + 2) * B + b[None, :], mask, other=0.0) + maximum = tl.max(maxima, axis=0) + winner = tl.min(tl.where((maxima == maximum[None, :]) & (n[:, None] < N), ids, 2147483647), axis=0) + # Keep the sum relative to its maximum; storing max + log(sum) would + # lose the normalizer when logits have a large common offset. + denominator = tl.sum(tl.where(maxima == -float("inf"), 0.0, mass * tl.exp(maxima - maximum[None, :])), axis=0) + if FINAL: + tl.store(Output + b, -tl.log(denominator), b < B) + tl.store(Ids + b, winner, b < B) + else: + tl.store(Output + b, maximum, b < B) + tl.store(Output + B + b, winner.to(tl.float32, bitcast=True), b < B) + tl.store(Output + 2 * B + b, denominator, b < B) + - return enable_env_vars("LIGHTLLM_VOCAB_PARALLEL_TOPK") +def _exact_top1(local_logits, vocab_start_id, tp_world_size, group, alloc_func): + vocab, batch = local_logits.shape + parts = triton.cdiv(vocab, 1024) + kwargs = dict(dtype=torch.float32, device=local_logits.device) + partial = alloc_func((parts, 3, batch), **kwargs) + payload = alloc_func((3, batch), **kwargs) + _top1_stats[(parts, triton.cdiv(batch, 8))](local_logits, partial, vocab, batch, vocab_start_id, 1024, 8) + _merge_top1_stats[(triton.cdiv(batch, 8),)]( + partial, payload, payload, parts, batch, False, triton.next_power_of_2(parts), 8 + ) + if tp_world_size > 1: + gathered = alloc_func((tp_world_size, 3, batch), **kwargs) + all_gather_into_tensor(output_=gathered, input_=payload, group=group, async_op=False) + else: + gathered = payload + logprobs = alloc_func((batch, 1), **kwargs) + ids = alloc_func((batch, 1), dtype=torch.int64, device=local_logits.device) + _merge_top1_stats[(triton.cdiv(batch, 8),)]( + gathered, logprobs, ids, tp_world_size, batch, True, triton.next_power_of_2(tp_world_size), 8 + ) + return logprobs, ids @torch.no_grad() @@ -22,6 +82,7 @@ def vocab_parallel_topk( tp_world_size: int, group, alloc_func, + normalize_top1: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """Gather each TP rank's local top-k logits and their global token ids. @@ -29,6 +90,10 @@ def vocab_parallel_topk( Each shard's first candidate is its lowest-id maximum. Together with rank order, this preserves dense argmax tie-breaking. Probabilities over the candidate union remain an approximation. + + With ``normalize_top1``, gather one extra shifted exponential sum per rank in the same + collective, select the global winner, and return its full-vocab logprob and + token id as one column. Winner selection precedes normalization rounding. """ assert local_logits.ndim == 2 and local_logits.is_cuda and local_logits.is_contiguous() @@ -38,6 +103,9 @@ def vocab_parallel_topk( local_topk = min(topk, vocab_size // tp_world_size) assert local_topk > 0 assert local_vocab_size >= local_topk + assert not normalize_top1 or local_topk == 1 + if normalize_top1: + return _exact_top1(local_logits, vocab_start_id, tp_world_size, group, alloc_func) if local_topk == 1: local_values, local_indexes = torch.max(local_logits, dim=0, keepdim=True) @@ -61,7 +129,7 @@ def vocab_parallel_topk( device=local_logits.device, ) local_payload[:local_topk].copy_(local_values) - local_payload[local_topk:].view(torch.int32).copy_(local_indexes).add_(int(vocab_start_id)) + local_payload[local_topk : local_topk * 2].view(torch.int32).copy_(local_indexes).add_(int(vocab_start_id)) gathered_payload = alloc_func( (tp_world_size, local_topk * 2, token_num), @@ -75,7 +143,7 @@ def vocab_parallel_topk( async_op=False, ) candidate_values = gathered_payload[:, :local_topk, :].permute(2, 0, 1) - candidate_token_ids = gathered_payload[:, local_topk:, :].view(torch.int32).permute(2, 0, 1) + candidate_token_ids = gathered_payload[:, local_topk : local_topk * 2, :].view(torch.int32).permute(2, 0, 1) # copy_ converts dtype and materializes the final layout in one pass. output_logits = alloc_func( @@ -86,6 +154,5 @@ def vocab_parallel_topk( output_logits.copy_(candidate_values) output_token_ids = alloc_func(candidate_token_ids.shape, dtype=torch.int64, device=local_logits.device) output_token_ids.copy_(candidate_token_ids) - return output_logits.view(token_num, tp_world_size * local_topk), output_token_ids.view( - token_num, tp_world_size * local_topk - ) + width = tp_world_size * local_topk + return output_logits.view(token_num, width), output_token_ids.view(token_num, width) diff --git a/lightllm/models/llama/layer_infer/post_layer_infer.py b/lightllm/models/llama/layer_infer/post_layer_infer.py index 376600b252..4edaf2ba2c 100644 --- a/lightllm/models/llama/layer_infer/post_layer_infer.py +++ b/lightllm/models/llama/layer_infer/post_layer_infer.py @@ -1,5 +1,3 @@ -import os - import torch import torch.functional as F import torch.distributed as dist @@ -20,8 +18,6 @@ class LlamaPostLayerInfer(PostLayerInferTpl): def __init__(self, network_config): super().__init__(network_config) self.eps_ = network_config["rms_norm_eps"] - self.vocab_parallel_topk_ = int(os.getenv("LIGHTLLM_VOCAB_PARALLEL_TOPK_SIZE", 128)) - assert self.vocab_parallel_topk_ > 0, "LIGHTLLM_VOCAB_PARALLEL_TOPK_SIZE must be positive" return def _norm(self, input, infer_state, layer_weight: LlamaPreAndPostLayerWeight) -> torch.Tensor: @@ -95,18 +91,26 @@ def _lm_head_and_gather( lm_head = layer_weight.lm_head_weight_ vocab_size = lm_head.vocab_size - if infer_state.use_vocab_parallel_topk and not force_full_logits: + topk = min(infer_state.vocab_parallel_topk, vocab_size // self.tp_world_size_) + normalize_top1 = infer_state.vocab_parallel_greedy and topk == 1 + # FP32 value + INT32 id (+ a normalizer for exact target top-1). + # Use the same minimum shard size on every rank when deciding to gather. + payload_bytes = (2 * topk + int(normalize_top1)) * 4 + dense_bytes = (vocab_size // self.tp_world_size_) * logic_batch.element_size() + if topk and payload_bytes < dense_bytes and not force_full_logits: logic_batch = self._apply_logit_postprocessing(logic_batch) logits, token_ids = vocab_parallel_topk( logic_batch, vocab_size=vocab_size, vocab_start_id=lm_head.tp_vocab_start_id, - topk=self.vocab_parallel_topk_, + topk=topk, tp_world_size=self.tp_world_size_, group=infer_state.dist_group, alloc_func=self.alloc_tensor, + normalize_top1=normalize_top1, ) infer_state.logits_token_ids = token_ids + infer_state.logits_are_logprobs = normalize_top1 return logits if self.tp_world_size_ == 1: diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 60b5fad4e8..e4d40e03d2 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -700,6 +700,22 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: help="""sampling used impl. 'triton' is use torch and triton kernel, flashinfer use flashinfer sampling impl""", ) + parser.add_argument( + "--disable_vocab_parallel_top1", + action="store_true", + help="Disable automatic vocab-parallel greedy selection for target and draft models.", + ) + parser.add_argument( + "--enable_vocab_parallel_topk", + action="store_true", + help="Accept approximate target sampling and logprobs over per-rank top-k candidates.", + ) + parser.add_argument( + "--vocab_parallel_topk_size", + type=int, + default=128, + help="Positive candidate count per TP rank for approximate sampling and dynamic MTP confidence.", + ) parser.add_argument( "--penalty_counter_mode", type=str, diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 37fe837ad1..fb883f4db0 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -34,6 +34,8 @@ def _set_envs_and_config(args: StartArgs): def _launch_subprocesses(args: StartArgs): + if args.vocab_parallel_topk_size < 1: + raise ValueError("--vocab_parallel_topk_size must be positive") _set_envs_and_config(args) if args.mtp_mode is not None: diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index a9aef608bd..50492c20ba 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -174,6 +174,9 @@ class StartArgs: ) llm_kv_quant_group_size: int = field(default=8) sampling_backend: str = field(default="triton", metadata={"choices": ["triton", "flashinfer"]}) + disable_vocab_parallel_top1: bool = field(default=False) + enable_vocab_parallel_topk: bool = field(default=False) + vocab_parallel_topk_size: int = field(default=128) penalty_counter_mode: str = field( default="gpu_counter", metadata={"choices": ["cpu_counter", "pin_mem_counter", "gpu_counter"]} ) diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 6bba6a515f..43a0fc1c28 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -391,13 +391,14 @@ def _get_next_token_ranks(self, model_output: ModelOutput, next_token_ids: torch 仅 ``--enable_rl`` 时做真实 rank；否则返回 GPU 常量 ``-1``，避免 O(batch * vocab) 比较。 下游 async_copy 在同样条件下会忽略该返回值。 """ - if not self.args.enable_rl or model_output.logits_token_ids is not None: + if not self.args.enable_rl or model_output.logits_are_logprobs: return g_pin_mem_manager.get_const_gpu_tensor( key="next_token_ranks", shape=next_token_ids.shape, fill_value=1 if self.args.enable_rl else -1, dtype=torch.int32, ) + assert model_output.logits_token_ids is None, "RL ranks require dense logits or exact greedy output" selected_logits = model_output.logits.gather(1, next_token_ids.long().view(-1, 1)) return (model_output.logits > selected_logits).sum(dim=-1, dtype=torch.int32) + 1 @@ -880,6 +881,8 @@ def _gen_argmax_token_ids(self, model_output: ModelOutput): def _gen_argmax_token_ids_and_prob(self, model_output: ModelOutput): logits = model_output.logits + if model_output.logits_are_logprobs: + return model_output.logits_token_ids[:, 0].long(), logits[:, 0].exp() candidate_indexes = torch.argmax(logits, dim=-1) max_probs = torch.softmax(logits, dim=-1).gather(1, candidate_indexes.view(-1, 1)).view(-1) token_ids = logits_indexes_to_token_ids(candidate_indexes, model_output.logits_token_ids) diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index 6db9246385..ad86b9d7c1 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -108,11 +108,16 @@ def prefill_normal( # 第一阶段: 模型推理 model_input, run_reqs = prepare_prefill_inputs(prefill_reqs, is_chuncked_mode=not self.disable_chunked_prefill) if self.prefill_mask_func is not None: - model_input.use_vocab_parallel_topk = False + model_input.vocab_parallel_topk = 0 with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) self._capture_prompt_logprobs_if_needed(model_input, run_reqs, model_output.prompt_logics) - (_, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu,) = self._sample_and_scatter_token( + ( + _, + next_token_ids_cpu, + next_token_logprobs_cpu, + next_token_ranks_cpu, + ) = self._sample_and_scatter_token( model_output=model_output, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, @@ -155,10 +160,15 @@ def decode_normal( ): model_input, run_reqs = prepare_decode_inputs(decode_reqs) if self.decode_mask_func is not None: - model_input.use_vocab_parallel_topk = False + model_input.vocab_parallel_topk = 0 with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) - (_, next_token_ids_cpu, next_token_logprobs_cpu, next_token_ranks_cpu,) = self._sample_and_scatter_token( + ( + _, + next_token_ids_cpu, + next_token_logprobs_cpu, + next_token_ranks_cpu, + ) = self._sample_and_scatter_token( model_output=model_output, b_req_idx=model_input.b_req_idx, b_mtp_index=model_input.b_mtp_index, @@ -196,7 +206,7 @@ def prefill_mtp( ): model_input, run_reqs = prepare_prefill_inputs(prefill_reqs, is_chuncked_mode=not self.disable_chunked_prefill) if self.prefill_mask_func is not None: - model_input.use_vocab_parallel_topk = False + model_input.vocab_parallel_topk = 0 with torch.cuda.stream(g_infer_context.get_overlap_stream()): model_output = self.model.forward(model_input) self._capture_prompt_logprobs_if_needed(model_input, run_reqs, model_output.prompt_logics) diff --git a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py index 4fd1e294dd..9e13eede6b 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py @@ -1,9 +1,6 @@ import torch from typing import List, Tuple from lightllm.common.basemodel.batch_objs import ModelOutput, logits_indexes_to_token_ids -from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk import ( - is_vocab_parallel_topk_enabled, -) from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty import apply_penalty from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty_gpu_cache import apply_penalty_gpu_cache from lightllm.common.basemodel.triton_kernel.post_process.apply_invalid_token import apply_invalid_token_ids @@ -12,16 +9,15 @@ from lightllm.utils.envs_utils import get_env_start_args -def _can_use_unmodified_greedy_logits(reqs: List[InferReq]) -> bool: - """Whether sampling is exactly argmax over the incoming logits.""" +def _can_use_unmodified_logits(reqs: List[InferReq]) -> bool: + """Modifiers that require full-vocabulary processing retain the dense path.""" for req_obj in reqs: sample_param = req_obj.sampling_param shm_param = sample_param.shm_param if ( - shm_param.top_k != 1 - or shm_param.temperature != 1.0 - or req_obj.need_out_token_id_statistics + req_obj.need_out_token_id_statistics + or (shm_param.top_k == 1 and shm_param.temperature != 1.0) or shm_param.exponential_decay_length_penalty.to_tuple()[1] != 1.0 or req_obj.cur_output_len < shm_param.min_new_tokens - 1 or sample_param.invalid_token_ids @@ -30,17 +26,51 @@ def _can_use_unmodified_greedy_logits(reqs: List[InferReq]) -> bool: return True -def can_use_vocab_parallel_topk(reqs: List[InferReq]) -> bool: - return is_vocab_parallel_topk_enabled() and _can_use_unmodified_greedy_logits(reqs) +def get_vocab_parallel_topk(reqs: List[InferReq]) -> int: + args = get_env_start_args() + if args.hardware_platform != "cuda": + return 0 + if not _can_use_unmodified_logits(reqs): + return 0 + if all(req.sampling_param.shm_param.top_k == 1 and req.sampling_param.shm_param.temperature == 1.0 for req in reqs): + if not args.disable_vocab_parallel_top1: + return 1 + # RL ranks require the complete vocabulary for non-greedy output. Explicit + # seeds retain the existing full-vocabulary RNG mapping. + if args.enable_vocab_parallel_topk and not args.enable_rl and all(req.generator is None for req in reqs): + return args.vocab_parallel_topk_size + return 0 def sample(model_output: ModelOutput, reqs: List[InferReq], eos_id: List[int] = [2]): logits = model_output.logits if model_output.logits_token_ids is not None: - if not _can_use_unmodified_greedy_logits(reqs): - raise RuntimeError("vocab-parallel top-k logits require unmodified greedy requests") + if not _can_use_unmodified_logits(reqs): + raise RuntimeError("vocab-parallel logits do not support logits modifiers") + if model_output.logits_are_logprobs: + assert all(req.sampling_param.shm_param.top_k == 1 for req in reqs) + return model_output.logits_token_ids[:, 0].long(), logits[:, 0] + # Allocate only the small candidate-shaped sampling tensors. Existing + # samplers still see column indexes; restore global ids exactly once. + params = [req.sampling_param.shm_param for req in reqs] + all_greedy = all(p.top_k == 1 for p in params) + if not all_greedy: + temperatures = g_pin_mem_manager.gen_from_list( + key="temperatures", data=[p.temperature for p in params], dtype=torch.float32 + ).cuda(non_blocking=True) + logits = logits / temperatures[:, None] candidate_indexes = torch.argmax(logits, dim=-1, keepdim=True) probs = torch.softmax(logits, dim=-1) + if not all_greedy: + top_ks = g_pin_mem_manager.gen_from_list( + key="top_ks", data=[min(p.top_k, logits.shape[1]) for p in params], dtype=torch.int32 + ).cuda(non_blocking=True) + top_ps = g_pin_mem_manager.gen_from_list( + key="top_ps", data=[p.top_p for p in params], dtype=torch.float32 + ).cuda(non_blocking=True) + sampled_indexes, _ = _top_p_top_k_sample(reqs, probs, top_ps, top_ks, False) + # Greedy rows in mixed batches must keep raw-logit tie-breaking. + candidate_indexes = torch.where(top_ks[:, None] == 1, candidate_indexes, sampled_indexes[:, None]) max_probs = probs.gather(1, candidate_indexes).view(-1) token_ids = logits_indexes_to_token_ids(candidate_indexes.view(-1), model_output.logits_token_ids) return token_ids, torch.log(max_probs) diff --git a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py index eadf50ae78..18a8a22ddb 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py @@ -4,7 +4,7 @@ from lightllm.server.router.model_infer.infer_batch import InferReq, g_infer_context from lightllm.common.basemodel.batch_objs import ModelInput from lightllm.server.router.model_infer.mode_backend.generic_post_process import ( - can_use_vocab_parallel_topk, + get_vocab_parallel_topk, ) INT64_MAX = torch.iinfo(torch.int64).max @@ -90,7 +90,8 @@ def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> is_prefill=True, b_prefill_has_output_cpu=b_prefill_has_output, multimodal_params=batch_multimodal_params, - use_vocab_parallel_topk=can_use_vocab_parallel_topk(run_reqs), + vocab_parallel_topk=get_vocab_parallel_topk(run_reqs), + vocab_parallel_greedy=all(req.sampling_param.shm_param.top_k == 1 for req in run_reqs), ) return model_input, run_reqs @@ -164,7 +165,8 @@ def prepare_decode_inputs(req_objs: List[InferReq]) -> Tuple[ModelInput, List[In b_shared_radix_node_id=b_shared_radix_node_id, is_prefill=False, multimodal_params=multimodal_params, - use_vocab_parallel_topk=can_use_vocab_parallel_topk(run_reqs), + vocab_parallel_topk=get_vocab_parallel_topk(run_reqs), + vocab_parallel_greedy=all(req.sampling_param.shm_param.top_k == 1 for req in run_reqs), ) return model_input, run_reqs @@ -181,8 +183,13 @@ def overlap_prepare_decode_inputs(req_objs: List[InferReq]): model_input1, run_reqs1 = prepare_decode_inputs( req_objs=decode_reqs1, ) - model_input0.use_vocab_parallel_topk = model_input1.use_vocab_parallel_topk = ( - model_input0.use_vocab_parallel_topk and model_input1.use_vocab_parallel_topk + model_input0.vocab_parallel_topk = model_input1.vocab_parallel_topk = ( + max(model_input0.vocab_parallel_topk, model_input1.vocab_parallel_topk) + if model_input0.vocab_parallel_topk and model_input1.vocab_parallel_topk + else 0 + ) + model_input0.vocab_parallel_greedy = model_input1.vocab_parallel_greedy = ( + model_input0.vocab_parallel_greedy and model_input1.vocab_parallel_greedy ) return model_input0, run_reqs0, decode_reqs0, model_input1, run_reqs1, decode_reqs1 @@ -219,8 +226,13 @@ def overlap_prepare_prefill_inputs(req_objs: List[InferReq]): req_objs=right_reqs, is_chuncked_mode=True, ) - model_input0.use_vocab_parallel_topk = model_input1.use_vocab_parallel_topk = ( - model_input0.use_vocab_parallel_topk and model_input1.use_vocab_parallel_topk + model_input0.vocab_parallel_topk = model_input1.vocab_parallel_topk = ( + max(model_input0.vocab_parallel_topk, model_input1.vocab_parallel_topk) + if model_input0.vocab_parallel_topk and model_input1.vocab_parallel_topk + else 0 + ) + model_input0.vocab_parallel_greedy = model_input1.vocab_parallel_greedy = ( + model_input0.vocab_parallel_greedy and model_input1.vocab_parallel_greedy ) return model_input0, run_reqs0, model_input1, run_reqs1 diff --git a/unit_tests/common/basemodel/test_cuda_graph_layout.py b/unit_tests/common/basemodel/test_cuda_graph_layout.py index 17ddcd98da..1cf1850ef0 100644 --- a/unit_tests/common/basemodel/test_cuda_graph_layout.py +++ b/unit_tests/common/basemodel/test_cuda_graph_layout.py @@ -1,6 +1,7 @@ from types import SimpleNamespace import pytest +import torch import lightllm.common.basemodel.cuda_graph as cuda_graph_module from lightllm.common.basemodel.cuda_graph import CudaGraph @@ -32,6 +33,55 @@ def test_dynamic_schedule_uses_compacted_physical_rows(_graph_args): assert _batch_sizes(max_batch_size=128) == [1, 2, 3, 4, *range(6, 129, 2)] +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_decode_replays_separate_dense_top1_topk_graphs(_graph_args): + from lightllm.common.basemodel.batch_objs import ModelOutput + from lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk import vocab_parallel_topk + + graph = CudaGraph(1, 1, 1, max_batch_size=1) + + class State: + def __init__(self, topk, greedy): + self.input_ids = torch.zeros(1, dtype=torch.int64, device="cuda") + self.local_logits = torch.zeros((32, 1), device="cuda") + self.vocab_parallel_topk = topk + self.vocab_parallel_greedy = greedy + + def copy_for_cuda_graph(self, new): + self.local_logits.copy_(new.local_logits) + + def forward(state): + if state.vocab_parallel_topk == 0: + return ModelOutput(logits=state.local_logits.T.clone()) + logits, ids = vocab_parallel_topk( + state.local_logits, + vocab_size=32, + vocab_start_id=0, + topk=state.vocab_parallel_topk, + tp_world_size=1, + group=None, + alloc_func=torch.empty, + normalize_top1=state.vocab_parallel_greedy, + ) + return ModelOutput(logits=logits, logits_token_ids=ids, logits_are_logprobs=state.vocab_parallel_greedy) + + modes = [(0, False), (1, True), (4, False), (1, False)] + for mode in modes: + assert graph.need_capture(1, mode) + graph.capture_decode(forward, State(*mode)) + assert not graph.need_capture(1, mode) + for winner, mode in enumerate(modes[::-1] + modes): + state = State(*mode) + state.local_logits[winner, 0] = 5 + expected = forward(state) + actual = graph.replay(state) + torch.cuda.synchronize() + torch.testing.assert_close(actual.logits, expected.logits) + if actual.logits_token_ids is not None: + torch.testing.assert_close(actual.logits_token_ids, expected.logits_token_ids) + assert actual.logits_are_logprobs == expected.logits_are_logprobs + + def test_public_static_schedule_preserves_original_static_mtp_default(_graph_args): assert CudaGraph.gen_cuda_graph_batch_sizes( batch_step_size_before_split=8, diff --git a/unit_tests/common/basemodel/test_model_output.py b/unit_tests/common/basemodel/test_model_output.py index 68a5286774..347e34c4a6 100644 --- a/unit_tests/common/basemodel/test_model_output.py +++ b/unit_tests/common/basemodel/test_model_output.py @@ -1,6 +1,7 @@ from types import SimpleNamespace import torch +from lightllm.server.core.objs import StartArgs from lightllm.common.basemodel import basemodel from lightllm.common.basemodel.basemodel import TpPartBaseModel @@ -20,6 +21,16 @@ def test_vocab_parallel_metadata_follows_row_operations(): torch.testing.assert_close(combined.logits_token_ids, torch.tensor([[2, 12], [4, 14], [9, 19]])) +def test_exact_logprob_metadata_follows_row_operations(): + output = ModelOutput( + logits=torch.tensor([[-2.0], [-3.0]]), logits_token_ids=torch.tensor([[20], [30]]), logits_are_logprobs=True + ) + selected = output.index_select_logits_rows(torch.tensor([1])) + combined = ModelOutput.concat_logits_rows([output, selected]) + assert combined.logits_are_logprobs + torch.testing.assert_close(combined.logits[:, 0], torch.tensor([-2.0, -3.0, -3.0])) + + def test_decode_unpad_slices_spec_output_with_logits(): model = TpPartBaseModel.__new__(TpPartBaseModel) output = ModelOutput( @@ -114,6 +125,7 @@ def _create_empty_decode_input(): def test_infer_state_enables_vocab_parallel_topk_for_draft_or_requested_target(monkeypatch): model = TpPartBaseModel.__new__(TpPartBaseModel) + model.args = StartArgs() model.infer_state_class = basemodel.InferStateInfo model.hidden_collector_prototype = SimpleNamespace(new_instance=lambda: object()) model.is_token_healing = False @@ -127,31 +139,42 @@ def test_infer_state_enables_vocab_parallel_topk_for_draft_or_requested_target(m model_input = _create_empty_decode_input() infer_state = model._create_inferstate(model_input) - assert not infer_state.use_vocab_parallel_topk + assert not infer_state.vocab_parallel_topk - model_input.use_vocab_parallel_topk = True + model_input.vocab_parallel_topk = True infer_state = model._create_inferstate(model_input) - assert infer_state.use_vocab_parallel_topk + assert infer_state.vocab_parallel_topk model.is_mtp_draft_model = True - model_input.use_vocab_parallel_topk = False + model_input.vocab_parallel_topk = False infer_state = model._create_inferstate(model_input) - assert infer_state.use_vocab_parallel_topk + assert infer_state.vocab_parallel_topk -def test_cuda_graph_contract_falls_back_for_dense_target_batch(monkeypatch): +def test_output_modes_distinguish_target_draft_and_confidence(): model = TpPartBaseModel.__new__(TpPartBaseModel) + model.args = StartArgs() model.is_mtp_draft_model = False - sparse_input = SimpleNamespace(use_vocab_parallel_topk=True) - dense_input = SimpleNamespace(use_vocab_parallel_topk=False) - - monkeypatch.setattr(basemodel, "is_vocab_parallel_topk_enabled", lambda: True) - assert model._is_cuda_graph_output_compatible(sparse_input) - assert not model._is_cuda_graph_output_compatible(dense_input) - assert not model._is_cuda_graph_output_compatible(sparse_input, dense_input) - + assert model.vocab_parallel_graph_modes() == [(0, False), (1, True)] + model.args.enable_vocab_parallel_topk = True + assert model.vocab_parallel_graph_modes() == [(0, False), (1, True), (128, False)] + model.args.disable_vocab_parallel_top1 = True + model.args.vocab_parallel_topk_size = 1 + assert model.vocab_parallel_graph_modes() == [(0, False), (1, False)] + assert model._vocab_parallel_output_mode(SimpleNamespace(vocab_parallel_topk=1, vocab_parallel_greedy=True)) == ( + 1, + False, + ) + model.args = StartArgs() model.is_mtp_draft_model = True - assert model._is_cuda_graph_output_compatible(dense_input) + assert model._vocab_parallel_output_mode(None) == (1, False) + model.args.mtp_dynamic_verify = True + model.args.mtp_mode = "eagle3" + assert model._vocab_parallel_output_mode(None) == (128, False) + model.args.mtp_mode = "dspark" + assert model._vocab_parallel_output_mode(None) == (1, False) + model.args.disable_vocab_parallel_top1 = True + assert model._vocab_parallel_output_mode(None) == (0, False) @torch.no_grad() @@ -176,7 +199,7 @@ def test_decode_pads_only_once_after_selecting_execution_path(monkeypatch): expected_batch_size, ) in execution_configs: model = TpPartBaseModel.__new__(TpPartBaseModel) - model.args = SimpleNamespace(enable_tpsp_mix_mode=enable_tpsp_mix_mode) + model.args = StartArgs(enable_tpsp_mix_mode=enable_tpsp_mix_mode) model.is_mtp_draft_model = False model.tp_world_size_ = tp_world_size model.mem_manager = SimpleNamespace(HOLD_TOKEN_MEMINDEX=99) @@ -211,7 +234,7 @@ def create_infer_state(model_input): graph = SimpleNamespace() graph.can_run = lambda **kwargs: True graph.find_closest_graph_batch_size = lambda batch_size: graph_batch_size - graph.need_capture = lambda batch_size: need_capture + graph.need_capture = lambda batch_size, output_mode: need_capture graph.capture_decode = lambda decode_func, infer_state: ModelOutput( logits=torch.ones((infer_state.b_req_idx.shape[0], 4)) ) diff --git a/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_topk.py b/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_topk.py index 2e069ab26e..3a82cee483 100644 --- a/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_topk.py +++ b/unit_tests/common/basemodel/triton_kernel/test_vocab_parallel_topk.py @@ -128,7 +128,7 @@ def _check_nccl_topk(rank, init_method): full_logits[12, 2] = 3 start, end = (0, 8) if rank == 0 else (8, 17) local_logits = full_logits[start:end].clone() - for topk in (1, 4, 128): + for topk, normalize_top1 in ((1, False), (4, False), (128, False), (1, True)): def forward(): return module.vocab_parallel_topk( @@ -139,6 +139,7 @@ def forward(): tp_world_size=2, group=dist.group.WORLD, alloc_func=torch.empty, + normalize_top1=normalize_top1, ) forward() @@ -154,11 +155,44 @@ def forward(): local_logits.copy_(full_logits[start:end]) graph.replay() torch.cuda.synchronize() - assert logits.shape == (3, 2 * min(topk, 8)) + assert logits.shape == (3, 1 if normalize_top1 else 2 * min(topk, 8)) selected = token_ids.gather(1, logits.argmax(dim=1, keepdim=True)).view(-1) torch.testing.assert_close(selected, full_logits.argmax(dim=0)) - torch.testing.assert_close(logits, full_logits.T.gather(1, token_ids)) + reference = full_logits.T.log_softmax(-1) if normalize_top1 else full_logits.T + torch.testing.assert_close(logits, reference.gather(1, token_ids)) assert (token_ids.sort(dim=1).values.diff(dim=1) > 0).all() + if not normalize_top1: + # The target sampler must choose the same global token on each + # TP rank, including approximate random rows beside greedy rows. + from types import SimpleNamespace + from lightllm.server.core.objs import StartArgs + from lightllm.common.basemodel.batch_objs import ModelOutput + from lightllm.server.router.model_infer.mode_backend import generic_post_process as sampling + + sampling.get_env_start_args = lambda: StartArgs(enable_vocab_parallel_topk=True) + reqs = [ + SimpleNamespace( + sampling_param=SimpleNamespace( + shm_param=SimpleNamespace( + top_k=k, + top_p=0.9, + temperature=1.0, + min_new_tokens=1, + exponential_decay_length_penalty=SimpleNamespace(to_tuple=lambda: (1, 1.0)), + ), + invalid_token_ids=[], + ), + cur_output_len=0, + need_out_token_id_statistics=False, + generator=None, + ) + for k in (3, 1, 3) + ] + torch.manual_seed(1234) + sampled, _ = sampling.sample(ModelOutput(logits=logits, logits_token_ids=token_ids), reqs) + rank_ids = [torch.empty_like(sampled) for _ in range(2)] + dist.all_gather(rank_ids, sampled) + torch.testing.assert_close(rank_ids[0], rank_ids[1]) finally: # NCCL shutdown waits for communicators retained by captured graphs. del graph @@ -168,3 +202,43 @@ def forward(): @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Two CUDA devices are required for NCCL") def test_vocab_parallel_topk_nccl_and_cuda_graph_replay(tmp_path): mp.spawn(_check_nccl_topk, args=(f"file://{tmp_path / 'nccl_init'}",), nprocs=2, join=True) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +def test_exact_top1_keeps_full_vocab_logprob_and_logit_winner(dtype): + module = importlib.import_module("lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk") + local_logits = torch.zeros((1024, 3), device="cuda", dtype=dtype) + local_logits[17, 0] = 10 + local_logits[1, 1] = 1e-7 + actual, ids = module.vocab_parallel_topk( + local_logits, + vocab_size=1024, + vocab_start_id=0, + topk=1, + tp_world_size=1, + group=None, + alloc_func=torch.empty, + normalize_top1=True, + ) + torch.testing.assert_close(ids[:, 0], local_logits.argmax(dim=0)) + torch.testing.assert_close(actual, local_logits.T.float().log_softmax(-1).gather(1, ids)) + + +@pytest.mark.parametrize("offset", [0.0, -10000.0, 10000.0, 1e20, -float("inf")]) +def test_exact_top1_normalizer_handles_partial_blocks_and_large_offsets(offset): + module = importlib.import_module("lightllm.common.basemodel.triton_kernel.post_process.vocab_parallel_topk") + local_logits = torch.full((1537, 9), offset, device="cuda") + if offset == -float("inf"): + local_logits[0] = 0 + actual, ids = module.vocab_parallel_topk( + local_logits, + vocab_size=1537, + vocab_start_id=0, + topk=1, + tp_world_size=1, + group=None, + alloc_func=torch.empty, + normalize_top1=True, + ) + torch.testing.assert_close(ids, torch.zeros((9, 1), dtype=torch.int64, device="cuda")) + torch.testing.assert_close(actual, local_logits.T.log_softmax(-1)[:, :1]) diff --git a/unit_tests/models/test_gemma4_vocab_parallel_topk.py b/unit_tests/models/test_gemma4_vocab_parallel_topk.py index 13ff5c3c24..90e1047a0d 100644 --- a/unit_tests/models/test_gemma4_vocab_parallel_topk.py +++ b/unit_tests/models/test_gemma4_vocab_parallel_topk.py @@ -1,6 +1,7 @@ from types import SimpleNamespace import torch +import pytest import lightllm.models.llama.layer_infer.post_layer_infer as llama_post_layer from lightllm.models.gemma4.layer_infer.post_layer_infer import Gemma4PostLayerInfer @@ -9,7 +10,6 @@ def test_vocab_parallel_topk_softcaps_local_logits_before_candidate_selection(monkeypatch): post = Gemma4PostLayerInfer.__new__(Gemma4PostLayerInfer) post.final_logit_softcapping = 2.0 - post.vocab_parallel_topk_ = 2 post.tp_world_size_ = 1 post.alloc_tensor = torch.empty post._norm = lambda hidden, infer_state, layer_weight: hidden @@ -17,10 +17,10 @@ def test_vocab_parallel_topk_softcaps_local_logits_before_candidate_selection(mo local_logits = torch.tensor( [[4.0, -4.0], [2.0, -2.0], [1.0, -1.0]], dtype=torch.bfloat16, - ) + ).repeat(8, 1) class LMHead: - vocab_size = 3 + vocab_size = 24 tp_vocab_start_id = 0 def __call__(self, input, alloc_func): @@ -37,7 +37,8 @@ def fake_vocab_parallel_topk(logits, **kwargs): monkeypatch.setattr(llama_post_layer, "vocab_parallel_topk", fake_vocab_parallel_topk) infer_state = SimpleNamespace( dist_group=None, - use_vocab_parallel_topk=True, + vocab_parallel_topk=2, + vocab_parallel_greedy=False, logits_token_ids=None, ) @@ -65,3 +66,30 @@ def test_full_logits_softcap_after_float32_conversion(): assert actual.dtype == torch.float32 torch.testing.assert_close(actual, expected) + + +@pytest.mark.parametrize("topk,force_full", [(128, False), (1, True)]) +def test_large_payload_and_prompt_logits_retain_dense_output(monkeypatch, topk, force_full): + post = llama_post_layer.LlamaPostLayerInfer.__new__(llama_post_layer.LlamaPostLayerInfer) + post.tp_world_size_ = 1 + post.alloc_tensor = torch.empty + post._norm = lambda hidden, *_: hidden + local = torch.arange(16, dtype=torch.float32).reshape(8, 2) + + class Head: + vocab_size = 8 + tp_vocab_start_id = 0 + + def __call__(self, **kwargs): + return local + + def unexpected_sparse(*args, **kwargs): + raise AssertionError("dense fallback must not collect candidates") + + monkeypatch.setattr(llama_post_layer, "vocab_parallel_topk", unexpected_sparse) + state = SimpleNamespace(vocab_parallel_topk=topk, vocab_parallel_greedy=False, logits_token_ids=None) + actual = post._lm_head_and_gather( + torch.empty(2, 3), 2, SimpleNamespace(lm_head_weight_=Head()), state, force_full_logits=force_full + ) + torch.testing.assert_close(actual, local.T) + assert state.logits_token_ids is None diff --git a/unit_tests/models/test_vocab_parallel_topk_output.py b/unit_tests/models/test_vocab_parallel_topk_output.py index 664d67e2b5..665d950485 100644 --- a/unit_tests/models/test_vocab_parallel_topk_output.py +++ b/unit_tests/models/test_vocab_parallel_topk_output.py @@ -20,8 +20,9 @@ def test_token_ranks_preserve_dense_and_sparse_semantics(monkeypatch, enable_rl, backend = ModeBackend.__new__(ModeBackend) backend.args = SimpleNamespace(enable_rl=enable_rl) output = ModelOutput( - logits=torch.tensor([[3.0, 3.0, 1.0]]), - logits_token_ids=torch.tensor([[30, 50, 10]]) if sparse else None, + logits=torch.tensor([[-0.5]]) if sparse else torch.tensor([[3.0, 3.0, 1.0]]), + logits_token_ids=torch.tensor([[30]]) if sparse else None, + logits_are_logprobs=sparse, ) selected = torch.tensor([30 if sparse else 2]) expected = (1 if sparse else 3) if enable_rl else -1 diff --git a/unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py b/unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py index a5264ff3cf..2eadc2cae5 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_generic_pre_process.py @@ -2,10 +2,13 @@ import torch import pytest +from lightllm.server.core.objs import StartArgs +from lightllm.server.router.model_infer.mode_backend import generic_post_process from lightllm.server.router.model_infer.mode_backend import generic_pre_process def _patch_empty_input_context(monkeypatch): + monkeypatch.setattr(generic_post_process, "get_env_start_args", lambda: StartArgs()) mem_manager = SimpleNamespace( HOLD_TOKEN_MEMINDEX=-1, alloc=lambda size: torch.empty((size,), dtype=torch.int32), @@ -25,6 +28,7 @@ def _patch_overlap_input_context(monkeypatch): def _make_prefill_req(req_idx: int, token_num: int): input_token_ids = [req_idx] * token_num return SimpleNamespace( + **_sampling_fields(), req_idx=req_idx, cur_kv_len=0, multimodal_params={"images": [], "audios": []}, @@ -36,6 +40,7 @@ def _make_prefill_req(req_idx: int, token_num: int): def _make_decode_req(req_idx: int): return SimpleNamespace( + **_sampling_fields(), req_idx=req_idx, cur_kv_len=3, mtp_step=0, @@ -46,6 +51,23 @@ def _make_decode_req(req_idx: int): ) +def _sampling_fields(): + return dict( + sampling_param=SimpleNamespace( + shm_param=SimpleNamespace( + top_k=1, + temperature=1.0, + min_new_tokens=1, + exponential_decay_length_penalty=SimpleNamespace(to_tuple=lambda: (1, 1.0)), + ), + invalid_token_ids=[], + ), + need_out_token_id_statistics=False, + cur_output_len=0, + generator=None, + ) + + def test_prepare_prefill_inputs_allows_empty_batch(monkeypatch): _patch_empty_input_context(monkeypatch) @@ -183,7 +205,7 @@ def can_use(reqs): calls.append(reqs) return all(req.eligible for req in reqs) - monkeypatch.setattr(generic_pre_process, "can_use_vocab_parallel_topk", can_use) + monkeypatch.setattr(generic_pre_process, "get_vocab_parallel_topk", can_use) reqs = [_make_prefill_req(i, 1) if prefill else _make_decode_req(i) for i in range(len(eligible))] for req, flag in zip(reqs, eligible): req.eligible = flag @@ -192,5 +214,23 @@ def can_use(reqs): else: first, _, _, second, _, _ = generic_pre_process.overlap_prepare_decode_inputs(reqs) - assert first.use_vocab_parallel_topk == second.use_vocab_parallel_topk == all(eligible) + assert first.vocab_parallel_topk == second.vocab_parallel_topk == all(eligible) assert len(calls) == 2 + + +@pytest.mark.parametrize("prefill", [False, True]) +@pytest.mark.parametrize("masked", [False, True]) +def test_overlap_unifies_greedy_random_and_fallback_layouts(monkeypatch, prefill, masked): + _patch_empty_input_context(monkeypatch) + args = StartArgs(enable_vocab_parallel_topk=True) + monkeypatch.setattr(generic_post_process, "get_env_start_args", lambda: args) + reqs = [_make_prefill_req(i, 1) if prefill else _make_decode_req(i) for i in range(2)] + reqs[1].sampling_param.shm_param.top_k = 20 + if masked: + reqs[1].sampling_param.invalid_token_ids = [4] + if prefill: + first, _, second, _ = generic_pre_process.overlap_prepare_prefill_inputs(reqs) + else: + first, _, _, second, _, _ = generic_pre_process.overlap_prepare_decode_inputs(reqs) + assert first.vocab_parallel_topk == second.vocab_parallel_topk == (0 if masked else 128) + assert not first.vocab_parallel_greedy and not second.vocab_parallel_greedy diff --git a/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_topk_sampling.py b/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_topk_sampling.py index e1acf78dc9..d48d1e7a0c 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_topk_sampling.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_vocab_parallel_topk_sampling.py @@ -1,20 +1,30 @@ from types import SimpleNamespace +import importlib.util import pytest import torch from lightllm.common.basemodel.batch_objs import ModelOutput from lightllm.server.router.model_infer.mode_backend.generic_post_process import ( - _can_use_unmodified_greedy_logits, - can_use_vocab_parallel_topk, + _can_use_unmodified_logits, + get_vocab_parallel_topk, sample, ) -from lightllm.utils.envs_utils import enable_env_vars +from lightllm.server.core.objs import StartArgs +from lightllm.server.router.model_infer.mode_backend import generic_post_process + + +@pytest.fixture(autouse=True) +def args(monkeypatch): + args = StartArgs() + monkeypatch.setattr(generic_post_process, "get_env_start_args", lambda: args) + return args def make_req(**overrides): values = { "top_k": 1, + "top_p": 1.0, "temperature": 1.0, "presence_penalty": 0.0, "frequency_penalty": 0.0, @@ -27,6 +37,7 @@ def make_req(**overrides): values.update(overrides) shm_param = SimpleNamespace( top_k=values["top_k"], + top_p=values["top_p"], temperature=values["temperature"], presence_penalty=values["presence_penalty"], frequency_penalty=values["frequency_penalty"], @@ -40,6 +51,7 @@ def make_req(**overrides): invalid_token_ids=values["invalid_token_ids"], ), cur_output_len=values["output_len"], + generator=None, need_out_token_id_statistics=( values["presence_penalty"] != 0.0 or values["frequency_penalty"] != 0.0 @@ -49,19 +61,18 @@ def make_req(**overrides): def test_accepts_unmodified_greedy_requests(): - assert _can_use_unmodified_greedy_logits([make_req(), make_req(output_len=5)]) - + assert _can_use_unmodified_logits([make_req(), make_req(output_len=5)]) -def test_feature_gate_requires_environment_and_eligible_batch(monkeypatch): - monkeypatch.delenv("LIGHTLLM_VOCAB_PARALLEL_TOPK", raising=False) - enable_env_vars.cache_clear() - assert not can_use_vocab_parallel_topk([make_req()]) - monkeypatch.setenv("LIGHTLLM_VOCAB_PARALLEL_TOPK", "1") - enable_env_vars.cache_clear() - assert can_use_vocab_parallel_topk([make_req()]) - assert not can_use_vocab_parallel_topk([make_req(top_k=2)]) - enable_env_vars.cache_clear() +def test_feature_gate_requires_explicit_approximation(args): + assert get_vocab_parallel_topk([make_req()]) == 1 + assert get_vocab_parallel_topk([make_req(top_k=2)]) == 0 + args.enable_vocab_parallel_topk = True + assert get_vocab_parallel_topk([make_req(top_k=2)]) == 128 + args.disable_vocab_parallel_top1 = True + assert get_vocab_parallel_topk([make_req()]) == 128 + args.enable_vocab_parallel_topk = False + assert get_vocab_parallel_topk([make_req()]) == 0 def test_samples_sparse_candidates_and_maps_global_token_ids(): @@ -91,7 +102,6 @@ def test_sparse_sampling_selects_logits_before_softmax_rounding(): @pytest.mark.parametrize( "override", [ - {"top_k": 2}, {"temperature": 0.5}, {"presence_penalty": 0.1}, {"frequency_penalty": 0.1}, @@ -102,4 +112,46 @@ def test_sparse_sampling_selects_logits_before_softmax_rounding(): ], ) def test_rejects_logits_modifiers(override): - assert not _can_use_unmodified_greedy_logits([make_req(**override)]) + assert not _can_use_unmodified_logits([make_req(**override)]) + + +def test_seed_and_random_rl_fall_back(args): + args.enable_vocab_parallel_topk = True + req = make_req(top_k=10) + req.generator = object() + assert get_vocab_parallel_topk([req]) == 0 + req.generator = None + args.enable_rl = True + assert get_vocab_parallel_topk([req]) == 0 + assert get_vocab_parallel_topk([make_req()]) == 1 + + +def test_exact_greedy_logprob_is_not_renormalized(): + output = ModelOutput( + logits=torch.tensor([[-2.0]]), logits_token_ids=torch.tensor([[123]]), logits_are_logprobs=True + ) + ids, logprobs = sample(output, [make_req()]) + assert ids.item() == 123 + assert logprobs.item() == -2.0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("backend", ["triton", "flashinfer"]) +def test_approximate_sampling_matches_candidate_reference(args, backend): + if backend == "flashinfer" and importlib.util.find_spec("flashinfer") is None: + pytest.skip("FlashInfer not installed") + args.sampling_backend = backend + args.enable_vocab_parallel_topk = True + reqs = [make_req(top_k=3, top_p=0.85, temperature=0.7), make_req()] + logits = torch.tensor([[1.0, 4.0, 3.0, 2.0], [1.0, 5.0, 5.0, 0.0]], device="cuda") + ids = torch.tensor([[100, 40, 30, 20], [10, 50, 60, 0]], device="cuda") + probs = (logits / torch.tensor([[0.7], [1.0]], device="cuda")).softmax(-1) + torch.manual_seed(42) + indexes, _ = generic_post_process._top_p_top_k_sample( + reqs, probs, torch.tensor([0.85, 1.0], device="cuda"), torch.tensor([3, 1], device="cuda"), False + ) + indexes[1] = logits[1].argmax() + torch.manual_seed(42) + actual_ids, actual_logprobs = sample(ModelOutput(logits=logits, logits_token_ids=ids), reqs) + torch.testing.assert_close(actual_ids, ids.gather(1, indexes[:, None]).flatten()) + torch.testing.assert_close(actual_logprobs, probs.gather(1, indexes[:, None]).flatten().log()) diff --git a/unit_tests/server/test_mtp_start_args.py b/unit_tests/server/test_mtp_start_args.py index 87e12479d4..0d2e7e4199 100644 --- a/unit_tests/server/test_mtp_start_args.py +++ b/unit_tests/server/test_mtp_start_args.py @@ -1,9 +1,30 @@ import pytest +import argparse +from lightllm.server.api_cli import add_cli_args from lightllm.server.api_start import _launch_subprocesses from lightllm.server.core.objs.start_args_type import StartArgs +def test_vocab_parallel_cli_defaults_and_overrides(): + parser = add_cli_args(argparse.ArgumentParser()) + defaults = parser.parse_args([]) + assert not defaults.disable_vocab_parallel_top1 + assert not defaults.enable_vocab_parallel_topk + assert defaults.vocab_parallel_topk_size == StartArgs().vocab_parallel_topk_size == 128 + configured = parser.parse_args( + ["--disable_vocab_parallel_top1", "--enable_vocab_parallel_topk", "--vocab_parallel_topk_size", "32"] + ) + assert configured.disable_vocab_parallel_top1 and configured.enable_vocab_parallel_topk + assert configured.vocab_parallel_topk_size == 32 + + +@pytest.mark.parametrize("size", [0, -1]) +def test_invalid_candidate_count_rejected_before_launch(size): + with pytest.raises(ValueError, match="must be positive"): + _launch_subprocesses(StartArgs(vocab_parallel_topk_size=size)) + + def test_mtp_requires_cuda_graph(monkeypatch): monkeypatch.setattr("lightllm.server.api_start._set_envs_and_config", lambda args: None) args = StartArgs(mtp_mode="vanilla_no_att", disable_cudagraph=True) From ae55a4676cc0fbe907105f572763723e9b4cf0fd Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 7 Sep 2026 15:28:15 +0800 Subject: [PATCH 7/8] refactor: deduplicate vocab-parallel layout alignment --- .../mode_backend/generic_post_process.py | 9 ++---- .../mode_backend/generic_pre_process.py | 30 +++++++++---------- 2 files changed, 17 insertions(+), 22 deletions(-) diff --git a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py index 9e13eede6b..5d6c8cb09f 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_post_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_post_process.py @@ -28,13 +28,10 @@ def _can_use_unmodified_logits(reqs: List[InferReq]) -> bool: def get_vocab_parallel_topk(reqs: List[InferReq]) -> int: args = get_env_start_args() - if args.hardware_platform != "cuda": + if args.hardware_platform != "cuda" or not _can_use_unmodified_logits(reqs): return 0 - if not _can_use_unmodified_logits(reqs): - return 0 - if all(req.sampling_param.shm_param.top_k == 1 and req.sampling_param.shm_param.temperature == 1.0 for req in reqs): - if not args.disable_vocab_parallel_top1: - return 1 + if not args.disable_vocab_parallel_top1 and all(req.sampling_param.shm_param.top_k == 1 for req in reqs): + return 1 # RL ranks require the complete vocabulary for non-greedy output. Explicit # seeds retain the existing full-vocabulary RNG mapping. if args.enable_vocab_parallel_topk and not args.enable_rl and all(req.generator is None for req in reqs): diff --git a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py index 18a8a22ddb..3b0e9d331c 100644 --- a/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py +++ b/lightllm/server/router/model_infer/mode_backend/generic_pre_process.py @@ -10,6 +10,18 @@ INT64_MAX = torch.iinfo(torch.int64).max +def _align_vocab_parallel_outputs(left: ModelInput, right: ModelInput): + # Both microbatches must share one layout for graph replay and output concatenation. + left.vocab_parallel_topk = right.vocab_parallel_topk = ( + max(left.vocab_parallel_topk, right.vocab_parallel_topk) + if left.vocab_parallel_topk and right.vocab_parallel_topk + else 0 + ) + left.vocab_parallel_greedy = right.vocab_parallel_greedy = ( + left.vocab_parallel_greedy and right.vocab_parallel_greedy + ) + + def prepare_prefill_inputs(req_objs: List[InferReq], is_chuncked_mode: bool) -> Tuple[ModelInput, List[InferReq]]: run_reqs = [] total_token_num = 0 @@ -183,14 +195,7 @@ def overlap_prepare_decode_inputs(req_objs: List[InferReq]): model_input1, run_reqs1 = prepare_decode_inputs( req_objs=decode_reqs1, ) - model_input0.vocab_parallel_topk = model_input1.vocab_parallel_topk = ( - max(model_input0.vocab_parallel_topk, model_input1.vocab_parallel_topk) - if model_input0.vocab_parallel_topk and model_input1.vocab_parallel_topk - else 0 - ) - model_input0.vocab_parallel_greedy = model_input1.vocab_parallel_greedy = ( - model_input0.vocab_parallel_greedy and model_input1.vocab_parallel_greedy - ) + _align_vocab_parallel_outputs(model_input0, model_input1) return model_input0, run_reqs0, decode_reqs0, model_input1, run_reqs1, decode_reqs1 @@ -226,14 +231,7 @@ def overlap_prepare_prefill_inputs(req_objs: List[InferReq]): req_objs=right_reqs, is_chuncked_mode=True, ) - model_input0.vocab_parallel_topk = model_input1.vocab_parallel_topk = ( - max(model_input0.vocab_parallel_topk, model_input1.vocab_parallel_topk) - if model_input0.vocab_parallel_topk and model_input1.vocab_parallel_topk - else 0 - ) - model_input0.vocab_parallel_greedy = model_input1.vocab_parallel_greedy = ( - model_input0.vocab_parallel_greedy and model_input1.vocab_parallel_greedy - ) + _align_vocab_parallel_outputs(model_input0, model_input1) return model_input0, run_reqs0, model_input1, run_reqs1 From 297bbf4f15573ec49191a7198bf31641ed57e454 Mon Sep 17 00:00:00 2001 From: sufubao Date: Mon, 7 Sep 2026 15:54:45 +0800 Subject: [PATCH 8/8] style: align vocab topk validation with startup checks --- lightllm/server/api_start.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index fb883f4db0..1f9979c97a 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -34,9 +34,10 @@ def _set_envs_and_config(args: StartArgs): def _launch_subprocesses(args: StartArgs): + _set_envs_and_config(args) + if args.vocab_parallel_topk_size < 1: raise ValueError("--vocab_parallel_topk_size must be positive") - _set_envs_and_config(args) if args.mtp_mode is not None: assert (