Support for logit bias in specdec

aphrodite/sampling_params.py

@@ -1001,10 +1001,7 @@ def _validate_spec_decode(
     ) -> None:
         if speculative_config is None:
             return
-
-        # Some sampling parameters are not yet compatible with spec decoding.
-        if self.logit_bias:
-            raise ValueError("The logit_bias sampling parameter is not yet supported with speculative decoding.")
+        return
 
     def _validate_diffusion(self, model_config: ModelConfig) -> None:
         if not model_config.is_diffusion:

aphrodite/v1/sample/logits_processor/__init__.py

@@ -183,18 +183,6 @@ def build_logitsprocs(
         logger.debug("Skipping logits processor loading because pooling models do not support logits processors.")
         return LogitsProcessors()
 
-    # Check if speculative decoding is enabled.
-    if aphrodite_config.speculative_config:
-        if custom_logitsprocs:
-            raise ValueError(STR_SPEC_DEC_REJECTS_LOGITSPROCS)
-        logger.warning("logit_bias parameter won't work with speculative decoding.")
-        return LogitsProcessors(
-            [
-                MinTokensLogitsProcessor(aphrodite_config, device, is_pin_memory),
-                MinPLogitsProcessor(aphrodite_config, device, is_pin_memory),
-            ]
-        )
-
     custom_logitsprocs_classes = _load_custom_logitsprocs(custom_logitsprocs)
     return LogitsProcessors(
         ctor(aphrodite_config, device, is_pin_memory)

aphrodite/v1/sample/logits_processor/builtin.py

@@ -155,6 +155,40 @@ def apply(self, logits: torch.Tensor) -> torch.Tensor:
             logits[self.logits_slice] += self.bias_tensor
         return logits
 
+    def apply_with_spec_decode(
+        self,
+        logits: torch.Tensor,
+        num_draft_tokens: list[int],
+    ) -> torch.Tensor:
+        if not self.biases:
+            return logits
+
+        num_draft_arr = np.asarray(num_draft_tokens, dtype=np.int64)
+        cumsum = np.concatenate([[0], np.cumsum(num_draft_arr)])
+
+        all_rows: list[np.ndarray] = []  # row indices to bias
+        all_toks: list[np.ndarray] = []  # token ids at those rows
+        all_biases: list[np.ndarray] = []  # bias values at those rows
+        for req_idx, lb in self.biases.items():
+            n_rows = int(num_draft_arr[req_idx])
+            if n_rows <= 0 or not lb:
+                continue
+            offset = cumsum[req_idx]
+            row_indices = np.arange(offset, offset + n_rows, dtype=np.int64)
+            tok_ids = np.fromiter(lb.keys(), dtype=np.int64, count=len(lb))
+            bias_vals = np.fromiter(lb.values(), dtype=np.float32, count=len(lb))
+            all_rows.append(np.repeat(row_indices, len(lb)))
+            all_toks.append(np.tile(tok_ids, n_rows))
+            all_biases.append(np.tile(bias_vals, n_rows))
+
+        if all_rows:
+            logits_slice = (
+                self._device_tensor(np.concatenate(all_rows), torch.int64),
+                self._device_tensor(np.concatenate(all_toks), torch.int64),
+            )
+            logits[logits_slice] += self._device_tensor(np.concatenate(all_biases), torch.float32)
+        return logits
+
 
 class MinTokensLogitsProcessor(LogitsProcessor):
     def __init__(self, aphrodite_config: "AphroditeConfig", device: torch.device, is_pin_memory: bool):

aphrodite/v1/sample/logits_processor/interface.py

@@ -79,6 +79,17 @@ def apply(self, logits: torch.Tensor) -> torch.Tensor:
         """
         raise NotImplementedError
 
+    def apply_with_spec_decode(
+        self,
+        logits: torch.Tensor,
+        num_draft_tokens: list[int],
+    ) -> torch.Tensor:
+        """Apply to speculative-decode target logits."""
+        raise NotImplementedError(
+            f"{type(self).__name__} does not support speculative decoding. "
+            "apply_with_spec_decode() must be implemented in order to do so."
+        )
+
     @abstractmethod
     def is_argmax_invariant(self) -> bool:
         """True if logits processor has no impact on the

aphrodite/v1/sample/rejection_sampler.py

@@ -12,10 +12,7 @@
 from aphrodite.logger import init_logger
 from aphrodite.triton_utils import tl, triton
 from aphrodite.v1.outputs import LogprobsLists, LogprobsTensors, SamplerOutput
-from aphrodite.v1.sample.logits_processor.builtin import (
-    MinPLogitsProcessor,
-    MinTokensLogitsProcessor,
-)
+from aphrodite.v1.sample.logits_processor.builtin import MinPLogitsProcessor
 from aphrodite.v1.sample.metadata import SamplingMetadata
 from aphrodite.v1.sample.ops.bad_words import apply_bad_words_with_drafts
 from aphrodite.v1.sample.ops.penalties import apply_all_penalties
@@ -308,8 +305,7 @@ def apply_logits_processors(
             apply_bad_words_with_drafts(logits, bad_words_token_ids, output_token_ids, metadata.num_draft_tokens)
 
         for processor in sampling_metadata.logitsprocs.non_argmax_invariant:
-            if isinstance(processor, MinTokensLogitsProcessor):
-                logits = processor.apply_with_spec_decode(logits, metadata.num_draft_tokens)
+            logits = processor.apply_with_spec_decode(logits, metadata.num_draft_tokens)
         holder = sampling_metadata.thinking_budget_state_holder
         if holder is not None and holder.has_tracked_requests():
             logits = holder.apply_to_logits(

tests/v1/sample/test_rejection_sampler.py

@@ -863,6 +863,58 @@ def test_allowed_token_ids(rejection_sampler):
     assert torch.equal(output.sampled_token_ids, expected)
 
 
+def test_logit_bias(rejection_sampler):
+    """Test rejection sampling with logit_bias.
+
+    The bias must reach every speculated (draft) position, not only the
+    first/bonus token. Request 0 is biased at draft position 1, request 1 at
+    position 0, and request 2 is left unbiased.
+    """
+    from aphrodite.v1.sample.logits_processor import LogitBiasLogitsProcessor
+
+    spec_tokens = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
+    output_tokens = [[1, 2, 3, 4], [4, 5, 6, 7], [7, 8, 9, 5]]
+
+    # Token 15 is the fallback the sampler falls back to when the intended
+    # token is biased below it (create_logits_tensor sets 15 to 99.0 vs 100.0).
+    logits = create_logits_tensor(output_tokens, token_idx_to_override=15)
+
+    # A -50 bias drops the target token from 100.0 to 50.0, below the 99.0
+    # fallback, flipping the greedy argmax to token 15 (i.e. a rejection).
+    logit_bias_proc = LogitBiasLogitsProcessor(None, torch.device(DEVICE_TYPE), is_pin_memory=False)
+    logit_bias_proc.biases = {
+        0: {2: -50.0},  # suppress token 2 -> reject req 0 at draft position 1
+        1: {4: -50.0},  # suppress token 4 -> reject req 1 at draft position 0
+        # request 2 is intentionally left unbiased
+    }
+
+    metadata = create_sampling_metadata(
+        all_greedy=True,
+        output_token_ids=[[], [], []],
+        spec_token_ids=spec_tokens,
+        logitsprocs=LogitsProcessors([logit_bias_proc]),
+    )
+    bonus_token_tensor = torch.tensor([output_tokens[i][-1] for i in range(len(output_tokens))], device=logits.device)
+    spec_decode_metadata = create_spec_decode_metadata(spec_tokens, logits)
+    mock_sampler_output(rejection_sampler, bonus_token_tensor)
+    output = rejection_sampler(
+        spec_decode_metadata,
+        draft_probs=None,
+        logits=logits,
+        sampling_metadata=metadata,
+    )
+
+    # Request 0: token 1 accepted, token 2 biased out -> 15, rest discarded.
+    # Request 1: token 4 biased out at the first position -> 15, rest discarded.
+    # Request 2: unbiased, all draft tokens accepted plus the bonus token.
+    expected = torch.tensor(
+        [[1, 15, -1, -1], [15, -1, -1, -1], [7, 8, 9, 5]],
+        dtype=torch.int,
+        device=logits.device,
+    )
+    assert torch.equal(output.sampled_token_ids, expected)
+
+
 @pytest.mark.parametrize("batch_size", [1, 100])
 @pytest.mark.parametrize("vocab_size", [100, 8192, 10000])
 @pytest.mark.parametrize("max_spec_len", [1, 3])