diff --git a/aphrodite/sampling_params.py b/aphrodite/sampling_params.py index c65798f41e..08be79ea41 100644 --- a/aphrodite/sampling_params.py +++ b/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: diff --git a/aphrodite/v1/sample/logits_processor/__init__.py b/aphrodite/v1/sample/logits_processor/__init__.py index 7cd36b2ed6..546a8b3032 100644 --- a/aphrodite/v1/sample/logits_processor/__init__.py +++ b/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) diff --git a/aphrodite/v1/sample/logits_processor/builtin.py b/aphrodite/v1/sample/logits_processor/builtin.py index 1fbd601d3e..6ffe191fc0 100644 --- a/aphrodite/v1/sample/logits_processor/builtin.py +++ b/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): diff --git a/aphrodite/v1/sample/logits_processor/interface.py b/aphrodite/v1/sample/logits_processor/interface.py index ba4f58b67e..c488112eb1 100644 --- a/aphrodite/v1/sample/logits_processor/interface.py +++ b/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 diff --git a/aphrodite/v1/sample/rejection_sampler.py b/aphrodite/v1/sample/rejection_sampler.py index d742779704..fcbcae53c5 100644 --- a/aphrodite/v1/sample/rejection_sampler.py +++ b/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( diff --git a/tests/v1/sample/test_rejection_sampler.py b/tests/v1/sample/test_rejection_sampler.py index 5382998416..0ee936f943 100644 --- a/tests/v1/sample/test_rejection_sampler.py +++ b/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])