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 28f2abf74b..3848a6bbcb 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -861,13 +861,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]]))