Add tests for PLessLogitsProcessor

Signed-off-by: Runyan <blackbird7@hotmail.com>

Author: ryttry

tests/v1/logits_processors/test_correctness.py

@@ -30,6 +30,7 @@
     MinPLogitsProcessor,
     MinTokensLogitsProcessor,
     MoveDirectionality,
+    PLessLogitsProcessor,
     build_logitsprocs,
 )
 from vllm.v1.sample.metadata import SamplingMetadata
@@ -190,7 +191,7 @@ def _generate_test_fakes(batch_size: int, device: str) -> LogitsprocsTestFakes:
 def _sampling_params_from_logitproc(logitproc_type: LogitprocType) -> SamplingParams:
     """Customize request SamplingParams for a specified logitproc"""
     # SamplingParams for req with no logitproc
-    kwargs = {"min_p": 0.0, "logit_bias": None, "min_tokens": 0}
+    kwargs = {"p_less": False, "min_p": 0.0, "logit_bias": None, "min_tokens": 0}
     if fxn := logitsprocs_test_mapping[logitproc_type].gen_request_fxn:
         fxn(kwargs)
     return SamplingParams(**kwargs)
@@ -296,6 +297,71 @@ def _logit_bias_validate(
                 )
 
 
+def _p_less_params(kwargs: dict) -> None:
+    """
+    The p-less logits processor configuration.
+    """
+    kwargs["p_less"] = True
+
+
+def _p_less_validate(
+    test_fakes: LogitsprocsTestFakes,
+    persistent_batch: list[LogitsProcsRequestParams],
+    logits_new: torch.Tensor,
+    batch_index: int,
+    request_params: LogitsProcsRequestParams,
+    step_idx: int,
+) -> None:
+    """
+    Validate that the p-less logits processor is applied correctly.
+    """
+    logits_old_for_batch_id = test_fakes.logits[
+        persistent_batch[batch_index].workload_index
+    ].cpu()
+    logits_new_for_batch_id = logits_new[batch_index].cpu()
+    if request_params.params.p_less:
+        # p-less decoding is used when p_less is set to True
+        # Convert logits to probabilities
+        probabilities = logits_old_for_batch_id.softmax(dim=-1)
+        # Calculate threshold probabilities for the batch for filtering out invalid
+        # tokens
+        threshold_probabilities = probabilities.square().sum(dim=-1, keepdim=True)
+        # Create boolean mask for invalid tokens
+        invalid_token_mask = probabilities < threshold_probabilities
+        # Apply mask to convert logits of invalid tokens to negative infinity
+        logits_new_exp = logits_old_for_batch_id.masked_fill(
+            invalid_token_mask, -float("inf")
+        )
+        for token_id in range(VOCAB_SIZE):
+            logit_new = logits_new_for_batch_id[token_id]
+            logit_new_exp = logits_new_exp[token_id]
+            if logit_new_exp != pytest.approx(logit_new):
+                _raise_error_invalid(
+                    msg_suffix=(
+                        f"Token {token_id} logit value {logit_new} "
+                        f"does not match expected value {logit_new_exp}"
+                    ),
+                    batch_index=batch_index,
+                    request_params=request_params,
+                    step_idx=step_idx,
+                )
+    else:
+        # p-less decoding is not used when p_less is set to False
+        for token_id in range(VOCAB_SIZE):
+            logit_old = logits_old_for_batch_id[token_id]
+            logit_new = logits_new_for_batch_id[token_id]
+            if logit_new != pytest.approx(logit_old):
+                _raise_error_invalid(
+                    msg_suffix=(
+                        f"Token {token_id} logit value {logit_old} "
+                        f"does not match expected value {logit_new}"
+                    ),
+                    batch_index=batch_index,
+                    request_params=request_params,
+                    step_idx=step_idx,
+                )
+
+
 def _min_p_params(kwargs: dict) -> None:
     """Min-p logitproc config"""
     kwargs["min_p"] = 0.1
@@ -582,6 +648,9 @@ class LogitsprocTestHelpers(NamedTuple):
     LogitBiasLogitsProcessor: LogitsprocTestHelpers(
         gen_request_fxn=_logit_bias_params, eval_fxn=_logit_bias_validate
     ),
+    PLessLogitsProcessor: LogitsprocTestHelpers(
+        gen_request_fxn=_p_less_params, eval_fxn=_p_less_validate
+    ),
     MinPLogitsProcessor: LogitsprocTestHelpers(
         gen_request_fxn=_min_p_params, eval_fxn=_min_p_validate
     ),

tests/v1/logits_processors/test_custom_offline.py

@@ -269,7 +269,7 @@ def test_rejects_custom_logitsprocs(
         llm = LLM(**llm_kwargs)
         # Require that no custom logitsprocs have been loaded
         # (built-in processors may exist: MinTokensLogitsProcessor,
-        # LogitBiasLogitsProcessor, MinPLogitsProcessor)
+        # LogitBiasLogitsProcessor, MinPLogitsProcessor, PLessLogitsProcessor)
         worker = llm.llm_engine.model_executor.driver_worker.worker
         for proc in worker.model_runner.input_batch.logitsprocs.all:
             assert not isinstance(proc, DummyLogitsProcessor)