Streamlined the code and removed unnecessary parts.

Signed-off-by: vx120 <893600387@qq.com>

Author: vx120

vllm/config/vllm.py

@@ -944,11 +944,6 @@ def _verify_sampling_replay_config(self) -> None:
             return
         if not self.use_v2_model_runner:
             raise ValueError("sampling distribution replay requires Model Runner V2")
-        if model_config.logprobs_mode != "processed_logprobs":
-            raise ValueError(
-                "sampling distribution replay requires "
-                "logprobs_mode='processed_logprobs'"
-            )
         if self.speculative_config is not None:
             raise ValueError(
                 "sampling distribution replay does not support speculative decoding"

vllm/outputs.py

@@ -186,7 +186,6 @@ def add(self, next_output: "RequestOutput", aggregate: bool) -> None:
                         if next_completion.logprobs:
                             assert completion.logprobs is not None
                             completion.logprobs.extend(next_completion.logprobs)  # type: ignore[arg-type]
-                        _merge_sampling_masks(completion, next_completion)
                         completion.cumulative_logprob = (
                             next_completion.cumulative_logprob
                         )
@@ -216,23 +215,6 @@ def __repr__(self) -> str:
         )
 
 
-def _merge_sampling_masks(
-    completion: CompletionOutput, next_completion: CompletionOutput
-) -> None:
-    current_mask = completion.sampling_mask
-    next_mask = next_completion.sampling_mask
-    if current_mask is None and next_mask is None:
-        return
-    if current_mask is None or next_mask is None:
-        raise RuntimeError("cannot merge partially missing sampling masks")
-    current_token_ids = list(current_mask.token_ids)
-    current_offsets = list(current_mask.offsets)
-    offset = len(current_token_ids)
-    current_token_ids.extend(next_mask.token_ids)
-    current_offsets.extend(offset + item for item in next_mask.offsets[1:])
-    completion.sampling_mask = SamplingMask(current_token_ids, current_offsets)
-
-
 # Sentinel to indicate request is finished, used with streaming inputs.
 STREAM_FINISHED = RequestOutput(
     request_id="",

vllm/v1/core/sched/scheduler.py

@@ -332,10 +332,6 @@ def __init__(
         self.enable_return_routed_experts = (
             vllm_config.model_config.enable_return_routed_experts
         )
-        self.enable_return_sampling_mask = (
-            vllm_config.model_config.enable_return_sampling_mask
-        )
-
         if self.enable_return_routed_experts:
             assert self.dcp_world_size == 1 and self.pcp_world_size == 1, (
                 "enable_return_routed_experts does not support context parallelism "
@@ -1825,16 +1821,10 @@ def update_from_output(
                 new_logprobs = logprobs.slice_request(req_index, len(new_token_ids))
 
             sampling_masks = model_runner_output.sampling_masks
-            if new_token_ids:
-                if sampling_masks is None:
-                    if self.enable_return_sampling_mask:
-                        raise RuntimeError(
-                            f"missing sampling mask for request {req_id}"
-                        )
-                else:
-                    new_sampling_mask = sampling_masks.slice_request(
-                        req_index, len(new_token_ids)
-                    )
+            if new_token_ids and sampling_masks is not None:
+                new_sampling_mask = sampling_masks.slice_request(
+                    req_index, len(new_token_ids)
+                )
 
             if num_nans_in_logits is not None and req_id in num_nans_in_logits:
                 request.num_nans_in_logits = num_nans_in_logits[req_id]

vllm/v1/engine/input_processor.py

@@ -326,6 +326,13 @@ def process_inputs(
             )
             if self.tokenizer is not None:
                 sampling_params.update_from_tokenizer(self.tokenizer)
+            if (
+                self.model_config.enable_return_sampling_mask
+                and sampling_params.temperature <= 0
+            ):
+                raise ValueError(
+                    "sampling distribution replay requires temperature > 0"
+                )
         else:
             pooling_params = params.clone()
 

vllm/v1/engine/output_processor.py

@@ -180,8 +180,7 @@ def __init__(
 
         # Routed experts accumulation (prompt + sample chunks)
         self.routed_experts_chunks: list[np.ndarray] = []
-        self.sampling_mask_token_ids: list[int] = []
-        self.sampling_mask_offsets: list[int] = [0]
+        self.sampling_mask_chunks: list[SamplingMaskLists] = []
 
         # Stream Interval
         self.stream_interval = stream_interval
@@ -340,33 +339,6 @@ def make_request_output(
             ec_transfer_params,
         )
 
-    def update_sampling_mask(
-        self,
-        new_token_ids: list[int],
-        sampling_mask: SamplingMaskLists | None,
-    ) -> None:
-        if sampling_mask is None:
-            if new_token_ids and len(self.sampling_mask_offsets) > 1:
-                raise RuntimeError(
-                    f"missing sampling mask for request {self.request_id}"
-                )
-            return
-        if len(sampling_mask.counts) != len(new_token_ids):
-            raise RuntimeError(
-                f"sampling mask row count does not match tokens for request "
-                f"{self.request_id}: {len(sampling_mask.counts)} rows for "
-                f"{len(new_token_ids)} tokens"
-            )
-        for row, raw_count in zip(sampling_mask.token_ids, sampling_mask.counts):
-            count = int(raw_count)
-            if count <= 0 or count > len(row):
-                raise RuntimeError(
-                    f"invalid sampling mask count {count} for request {self.request_id}"
-                )
-            kept_ids = [int(token_id) for token_id in row[:count]]
-            self.sampling_mask_token_ids.extend(kept_ids)
-            self.sampling_mask_offsets.append(len(self.sampling_mask_token_ids))
-
     def _new_request_output(
         self,
         external_req_id: str,
@@ -434,7 +406,14 @@ def _new_completion_output(
         if delta and logprobs:
             logprobs = logprobs[-len(token_ids) :]
 
-        sampling_mask = self._get_sampling_mask(token_ids, delta)
+        sampling_mask = None
+        if (
+            finished
+            and self.output_kind == RequestOutputKind.FINAL_ONLY
+            and self.sampling_mask_chunks
+        ):
+            merged = SamplingMaskLists.merge(self.sampling_mask_chunks)
+            sampling_mask = SamplingMask(merged.token_ids, merged.offsets)
 
         # Concatenate routed experts on finish
         routed_experts = None
@@ -453,27 +432,6 @@ def _new_completion_output(
             stop_reason=stop_reason if finished else None,
         )
 
-    def _get_sampling_mask(
-        self, token_ids: list[int], delta: bool
-    ) -> SamplingMask | None:
-        if len(self.sampling_mask_offsets) == 1:
-            return None
-        num_rows = len(self.sampling_mask_offsets) - 1
-        if num_rows != self.detokenizer.num_output_tokens():
-            raise RuntimeError(
-                f"sampling mask is misaligned for request {self.request_id}: "
-                f"{num_rows} rows for {self.detokenizer.num_output_tokens()} tokens"
-            )
-        start_row = num_rows - len(token_ids) if delta else 0
-        flat_start = self.sampling_mask_offsets[start_row]
-        offsets = [
-            offset - flat_start for offset in self.sampling_mask_offsets[start_row:]
-        ]
-        return SamplingMask(
-            token_ids=self.sampling_mask_token_ids[flat_start:],
-            offsets=offsets,
-        )
-
     def _new_pooling_output(self, pooling_output: torch.Tensor) -> PoolingOutput:
         return PoolingOutput(data=pooling_output)
 
@@ -704,9 +662,13 @@ def process_outputs(
             if pooling_output is None:
                 assert req_state.detokenizer is not None
                 assert req_state.logprobs_processor is not None
-                req_state.update_sampling_mask(
-                    new_token_ids, engine_core_output.new_sampling_mask
-                )
+                if (
+                    engine_core_output.new_sampling_mask is not None
+                    and req_state.output_kind == RequestOutputKind.FINAL_ONLY
+                ):
+                    req_state.sampling_mask_chunks.append(
+                        engine_core_output.new_sampling_mask
+                    )
                 # 2) Detokenize the token ids into text and perform stop checks.
                 stop_string = req_state.detokenizer.update(
                     new_token_ids, finish_reason == FinishReason.STOP

vllm/v1/outputs.py

@@ -51,33 +51,36 @@ def slice_request(self, req_idx: int, num_positions: int):
 
 
 class SamplingMaskLists(NamedTuple):
-    # [num_generated_tokens, max_kept_tokens]
+    # [num_kept_tokens]
     token_ids: np.ndarray
-    # [num_generated_tokens]
-    counts: np.ndarray
+    # [num_generated_tokens + 1]
+    offsets: np.ndarray
     # [num_reqs + 1]
     cu_num_generated_tokens: list[int] | None = None
 
     def slice_request(self, req_idx: int, num_positions: int) -> "SamplingMaskLists":
         if self.cu_num_generated_tokens is None:
             start_idx = req_idx
-            req_end_idx = req_idx + 1
         else:
             start_idx = self.cu_num_generated_tokens[req_idx]
-            req_end_idx = self.cu_num_generated_tokens[req_idx + 1]
         end_idx = start_idx + num_positions
-        if end_idx > req_end_idx:
-            raise RuntimeError(
-                "sampling mask has fewer rows than the generated tokens: "
-                f"request index {req_idx}, requested {num_positions}, "
-                f"available {req_end_idx - start_idx}"
-            )
+        flat_start = self.offsets[start_idx]
+        flat_end = self.offsets[end_idx]
         return SamplingMaskLists(
-            self.token_ids[start_idx:end_idx],
-            self.counts[start_idx:end_idx],
+            self.token_ids[flat_start:flat_end],
+            self.offsets[start_idx : end_idx + 1] - flat_start,
             None,
         )
 
+    @staticmethod
+    def merge(chunks: Sequence["SamplingMaskLists"]) -> "SamplingMaskLists":
+        token_ids = np.concatenate([chunk.token_ids for chunk in chunks])
+        counts = np.concatenate([np.diff(chunk.offsets) for chunk in chunks])
+        offsets = np.empty(len(counts) + 1, dtype=np.int64)
+        offsets[0] = 0
+        np.cumsum(counts, dtype=np.int64, out=offsets[1:])
+        return SamplingMaskLists(token_ids, offsets)
+
 
 class LogprobsTensors(NamedTuple):
     # [num_reqs x num_generated_tokens, max_num_logprobs + 1]