diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index f2b6bae08..7d9af199f 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -389,6 +389,7 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0) infer_state.is_prefill = model_input.is_prefill infer_state.is_token_healing = self.is_token_healing infer_state.return_all_prompt_logics = self.return_all_prompt_logics + infer_state.is_mtp_draft_model = 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 @@ -539,6 +540,8 @@ 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.mtp_collector = model_output.mtp_collector.unpad_decode( padded_batch_size=padded_batch_size, origin_batch_size=origin_batch_size, @@ -551,6 +554,8 @@ 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.mtp_collector = padded_model_output.mtp_collector.unpad_prefill( origin_handle_token_num=origin_handle_token_num ) @@ -742,6 +747,7 @@ 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, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), prompt_logics=infer_state.prompt_logics, ) @@ -771,6 +777,7 @@ 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, mtp_collector=infer_state.hidden_collector.finish_output(infer_state=infer_state), ) @@ -1025,11 +1032,13 @@ 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, 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, mtp_collector=infer_state1.hidden_collector.finish_output(infer_state=infer_state1), prompt_logics=infer_state1.prompt_logics, ) @@ -1074,10 +1083,12 @@ 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, 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, 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 ae645d4b7..6a06ebf84 100644 --- a/lightllm/common/basemodel/batch_objs.py +++ b/lightllm/common/basemodel/batch_objs.py @@ -199,11 +199,32 @@ class ModelOutput: # 此时 logits 依然只保存每个请求最后一个位置的 logits，prompt_logics 为可选项，仅在 # 需要返回 prompt logprobs 信息时才会非空。 prompt_logics: Optional[torch.Tensor] = None + # Sparse draft logits map each candidate column back to a token id in the + # draft model's output vocabulary. None means logits are dense and their + # column indexes are already token ids. Keep this optional field last so + # existing positional ModelOutput construction remains compatible. + logits_token_ids: Optional[torch.Tensor] = None def __post_init__(self) -> None: if self.mtp_collector is None: self.mtp_collector = ModelMtpOutputCollector() + if self.logits_token_ids is not 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 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.mtp_collector.to_no_ref_tensor() + + def index_select_logits_rows(self, index: torch.Tensor) -> "ModelOutput": + """Select vocabulary-output rows while preserving sparse token ids.""" + + 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 + ), + ) diff --git a/lightllm/common/basemodel/infer_struct.py b/lightllm/common/basemodel/infer_struct.py index 91c6e9969..6c09a8248 100755 --- a/lightllm/common/basemodel/infer_struct.py +++ b/lightllm/common/basemodel/infer_struct.py @@ -53,6 +53,14 @@ def __init__(self): self.is_token_healing: bool = False self.return_all_prompt_logics: bool = False + # Draft models only need a small candidate set from the vocabulary + # projection. The model marker enables the sparse logits path in the + # post layer without changing the target-model sampling interface. + self.is_mtp_draft_model: bool = False + # When logits are sparse, each column maps to the corresponding token + # id in the model's own output vocabulary. Dense target logits leave + # this field as None and keep the historical column-index semantics. + self.logits_token_ids: Optional[torch.Tensor] = None # 在开启 return_all_prompt_logics 模式时，保存整个 prefill 阶段每一个 # token 位置的 logits，供后续回传 prompt logprobs 信息使用。 # 仅在 prefill 阶段且需要返回 prompt logprobs 时才会被填充。 diff --git a/lightllm/models/llama/layer_infer/post_layer_infer.py b/lightllm/models/llama/layer_infer/post_layer_infer.py index bb6e4f373..f0ee2c633 100644 --- a/lightllm/models/llama/layer_infer/post_layer_infer.py +++ b/lightllm/models/llama/layer_infer/post_layer_infer.py @@ -7,15 +7,19 @@ 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.distributed.communication_op import all_gather +from lightllm.distributed.communication_op import all_gather, all_gather_into_tensor class LlamaPostLayerInfer(PostLayerInferTpl): """ """ + DEFAULT_DRAFT_LOGITS_TOPK = 128 + def __init__(self, network_config): super().__init__(network_config) self.eps_ = network_config["rms_norm_eps"] + self.draft_logits_topk_ = int(os.getenv("LIGHTLLM_DRAFT_LOGITS_TOPK", str(self.DEFAULT_DRAFT_LOGITS_TOPK))) + assert self.draft_logits_topk_ > 0 return def _norm(self, input, infer_state, layer_weight: LlamaPreAndPostLayerWeight) -> torch.Tensor: @@ -70,7 +74,13 @@ def _token_forward( input_embdings = None # 正常采样使用的 logits，始终只对应每个请求最后一个位置。 - ans_logics = self._lm_head_and_gather(last_input, token_num, layer_weight, infer_state) + ans_logics = self._lm_head_and_gather( + last_input, + token_num, + layer_weight, + infer_state, + use_sparse_logits=getattr(infer_state, "is_mtp_draft_model", False), + ) # 在 return_all_prompt_logics 模式下，prompt_logics 保存的是完整 prefill # 的 hidden state，需要在 norm/lm_head 之前取出来，避免被 input_embdings 置空。 prompt_logics_hiddens = infer_state.prompt_logics @@ -80,7 +90,11 @@ 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, + use_sparse_logits=False, ) return ans_logics @@ -91,12 +105,21 @@ def _lm_head_and_gather( token_num: int, layer_weight: LlamaPreAndPostLayerWeight, infer_state: LlamaInferStateInfo, + use_sparse_logits: bool = False, ) -> torch.Tensor: normed = self._norm(hidden, infer_state, layer_weight) normed = normed.permute(1, 0).view(-1, token_num) logic_batch = layer_weight.lm_head_weight_(input=normed, alloc_func=self.alloc_tensor) normed = None + if use_sparse_logits: + return self._gather_draft_topk_logits( + logic_batch=logic_batch, + token_num=token_num, + layer_weight=layer_weight, + infer_state=infer_state, + ) + vocab_size = layer_weight.lm_head_weight_.vocab_size if self.tp_world_size_ == 1: gather_data = logic_batch @@ -116,6 +139,66 @@ def _lm_head_and_gather( gather_data = None return ans_logics + def _gather_draft_topk_logits( + self, + logic_batch: torch.Tensor, + token_num: int, + layer_weight: LlamaPreAndPostLayerWeight, + infer_state: LlamaInferStateInfo, + ) -> torch.Tensor: + """Gather each TP rank's local top-k instead of the full vocabulary.""" + + lm_head = layer_weight.lm_head_weight_ + vocab_size = lm_head.vocab_size + local_topk = min(self.draft_logits_topk_, vocab_size) + + # Every TP vocabulary shard is expected to contain at least top-k rows, + # so every rank contributes the same fixed-shape tensors directly. + assert logic_batch.shape[0] >= local_topk, ( + f"draft logits top-k ({local_topk}) exceeds the local vocabulary shard " f"size ({logic_batch.shape[0]})" + ) + local_values, local_indexes = torch.topk(logic_batch, k=local_topk, dim=0, sorted=False) + local_values = local_values.float() + local_indexes = local_indexes.to(dtype=torch.int32) + local_token_ids = local_indexes + int(lm_head.tp_vocab_start_id) + + if self.tp_world_size_ == 1: + candidate_values = local_values.permute(1, 0) + candidate_token_ids = local_token_ids.permute(1, 0) + else: + # Pack FP32 logits and the bit representation of INT32 token ids + # into one FP32 payload. Both element types are four bytes, so the + # ids can be restored losslessly after a single all-gather. + local_payload = self.alloc_tensor( + (local_topk * 2, token_num), + dtype=torch.float32, + device=local_values.device, + ) + local_payload[:local_topk].copy_(local_values) + local_payload[local_topk:].view(torch.int32).copy_(local_token_ids) + gathered_payload = self.alloc_tensor( + (self.tp_world_size_, local_topk * 2, token_num), + dtype=torch.float32, + device=local_values.device, + ) + all_gather_into_tensor( + output_=gathered_payload, + input_=local_payload, + group=infer_state.dist_group, + async_op=False, + ) + gathered_values = gathered_payload[:, :local_topk, :] + gathered_token_ids = gathered_payload[:, local_topk:, :].view(torch.int32) + candidate_values = gathered_values.permute(2, 0, 1).reshape(token_num, -1) + candidate_token_ids = gathered_token_ids.permute(2, 0, 1).reshape(token_num, -1) + + candidate_token_ids = candidate_token_ids.to(dtype=torch.int64) + candidate_count = candidate_values.shape[1] + ans_logics = self.alloc_tensor((token_num, candidate_count), dtype=torch.float32, device=logic_batch.device) + ans_logics.copy_(candidate_values) + infer_state.logits_token_ids = candidate_token_ids.contiguous() + return ans_logics + def token_forward( self, input_embdings: torch.Tensor, infer_state: LlamaInferStateInfo, layer_weight: LlamaPreAndPostLayerWeight ): 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 5a74cd988..6583a933c 100644 --- a/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py +++ b/lightllm/models/qwen3_dspark/layer_infer/post_layer_infer.py @@ -179,9 +179,20 @@ def token_forward( # Graph unpadding still uses the leading logits dimension when token ids are returned directly. return local_logits.new_empty((token_num, 1)) - logits = self._lm_head_and_gather(last_input, token_num, layer_weight, infer_state) + logits = self._lm_head_and_gather( + last_input, + token_num, + layer_weight, + infer_state, + use_sparse_logits=getattr(infer_state, "is_mtp_draft_model", False), + ) block_logits = logits.reshape(num_reqs, self.block_size_, -1) - sampled_tokens = torch.argmax(block_logits, dim=-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) 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 e0d1b3924..bd28a0d33 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -840,14 +840,24 @@ 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) + max_probs, candidate_indexes = torch.max(probs, dim=-1) + draft_next_token_ids_gpu = self._map_logits_indexes_to_token_ids(model_output, candidate_indexes) return draft_next_token_ids_gpu, max_probs + @staticmethod + def _map_logits_indexes_to_token_ids(model_output: ModelOutput, candidate_indexes: torch.Tensor): + """Map sparse-logits columns to vocabulary ids; dense logits are identity-mapped.""" + + if model_output.logits_token_ids is None: + 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, logits: torch.Tensor, 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 6b8c23e8f..137a580ae 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 3d2c0a0e8..f2c730ccd 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 6f9477e29..be1bb0833 100644 --- a/unit_tests/common/basemodel/test_model_output.py +++ b/unit_tests/common/basemodel/test_model_output.py @@ -11,16 +11,20 @@ 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), 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) + torch.testing.assert_close(unpadded.logits_token_ids, output.logits_token_ids[: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.mtp_collector.spec_hidden.shape == (6, 3) @@ -28,6 +32,7 @@ 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), mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.arange(24).view(8, 3)), prompt_logics=torch.arange(32).view(8, 4), ) @@ -39,10 +44,27 @@ 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.mtp_collector.spec_hidden.shape == (6, 3) assert unpadded.prompt_logics.shape == (6, 4) +def test_index_select_logits_rows_preserves_sparse_token_mapping(): + output = ModelOutput( + logits=torch.tensor([[1.0, 4.0], [7.0, 2.0], [3.0, 6.0]]), + logits_token_ids=torch.tensor([[10, 40], [70, 20], [30, 60]]), + mtp_collector=ModelMtpOutputCollector(spec_hidden=torch.ones((3, 2))), + ) + + selected = output.index_select_logits_rows(torch.tensor([2, 0])) + + torch.testing.assert_close(selected.logits, torch.tensor([[3.0, 6.0], [1.0, 4.0]])) + torch.testing.assert_close(selected.logits_token_ids, torch.tensor([[30, 60], [10, 40]])) + # The helper intentionally returns only the vocabulary outputs needed by + # the proposer; hidden collection remains owned by the original output. + assert selected.mtp_collector.spec_hidden is None + + def test_decode_unpad_restores_empty_output(): model = TpPartBaseModel.__new__(TpPartBaseModel) output = ModelOutput( diff --git a/unit_tests/models/test_sparse_draft_logits.py b/unit_tests/models/test_sparse_draft_logits.py new file mode 100644 index 000000000..ea289164a --- /dev/null +++ b/unit_tests/models/test_sparse_draft_logits.py @@ -0,0 +1,157 @@ +from types import SimpleNamespace + +import torch + +from lightllm.common.basemodel.batch_objs import ModelOutput +from lightllm.models.llama.layer_infer import post_layer_infer as llama_post_layer +from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer +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_draft_logits_topk_comes_only_from_environment(monkeypatch): + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", "0") + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "1") + config = { + "rms_norm_eps": 1e-6, + "vocab_size": 1024, + "n_embed": 64, + "draft_logits_topk": 7, + } + + monkeypatch.setenv("LIGHTLLM_DRAFT_LOGITS_TOPK", "32") + assert LlamaPostLayerInfer(config).draft_logits_topk_ == 32 + + monkeypatch.delenv("LIGHTLLM_DRAFT_LOGITS_TOPK") + assert LlamaPostLayerInfer(config).draft_logits_topk_ == LlamaPostLayerInfer.DEFAULT_DRAFT_LOGITS_TOPK + + +def test_sparse_draft_logits_gather_local_candidates_and_token_ids(monkeypatch): + post = LlamaPostLayerInfer.__new__(LlamaPostLayerInfer) + post.tp_world_size_ = 2 + post.draft_logits_topk_ = 2 + post.alloc_tensor = lambda shape, dtype, device="cpu": torch.empty(shape, dtype=dtype, device=device) + + # Rank 0 owns token ids [0, 4), rank 1 owns [4, 8). Each rank contributes + # its local top-2 directly, so the returned candidate width is TP * 2. + rank0_logits = torch.tensor( + [ + [1.0, 8.0], + [9.0, 2.0], + [3.0, 4.0], + [2.0, 5.0], + ] + ) + rank1_values = torch.tensor([[10.0, 7.0], [8.0, 6.0]]) + rank1_token_ids = torch.tensor([[4, 4], [5, 5]], dtype=torch.int32) + rank1_payload = torch.empty((4, 2), dtype=torch.float32) + rank1_payload[:2].copy_(rank1_values) + rank1_payload[2:].view(torch.int32).copy_(rank1_token_ids) + gather_call = 0 + + def fake_all_gather_into_tensor(output_, input_, group=None, async_op=False): + nonlocal gather_call + assert input_.dtype == torch.float32 + assert input_.shape == (4, 2) + output_[0].copy_(input_) + output_[1].copy_(rank1_payload) + gather_call += 1 + + monkeypatch.setattr(llama_post_layer, "all_gather_into_tensor", fake_all_gather_into_tensor) + infer_state = SimpleNamespace(dist_group=None, logits_token_ids=None) + layer_weight = SimpleNamespace( + lm_head_weight_=SimpleNamespace( + vocab_size=8, + tp_vocab_start_id=0, + ) + ) + + sparse_logits = post._gather_draft_topk_logits( + logic_batch=rank0_logits, + token_num=2, + layer_weight=layer_weight, + infer_state=infer_state, + ) + + assert sparse_logits.shape == (2, 4) + assert infer_state.logits_token_ids.shape == (2, 4) + assert sparse_logits.dtype == torch.float32 + assert infer_state.logits_token_ids.dtype == torch.int64 + reconstructed = torch.full((2, 8), float("-inf")) + reconstructed.scatter_(1, infer_state.logits_token_ids.long(), sparse_logits) + expected = torch.tensor( + [ + [float("-inf"), 9.0, 3.0, float("-inf"), 10.0, 8.0, float("-inf"), float("-inf")], + [8.0, float("-inf"), float("-inf"), 5.0, 7.0, 6.0, float("-inf"), float("-inf")], + ] + ) + torch.testing.assert_close(reconstructed, expected) + assert gather_call == 1 + + +def test_draft_argmax_and_truncated_probability_restore_vocabulary_ids(): + backend = ModeBackend.__new__(ModeBackend) + output = ModelOutput( + 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) + 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) + expected_probs = torch.softmax(output.logits, dim=-1).max(dim=-1).values + torch.testing.assert_close(probs, expected_probs) + + +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_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([[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]], dtype=torch.int32) + + def gather_sparse(*args, **kwargs): + infer_state = args[3] + infer_state.logits_token_ids = sparse_token_ids + return sparse_logits + + post._lm_head_and_gather = gather_sparse + 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, + is_mtp_draft_model=True, + 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]], dtype=torch.int32))