ngram_worker.py 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169
  1. import weakref
  2. from typing import List, Optional, Set, Tuple
  3. import torch
  4. from aphrodite.common.sequence import ExecuteModelRequest, SamplerOutput
  5. from aphrodite.spec_decode.interfaces import SpeculativeProposals
  6. from aphrodite.spec_decode.proposer_worker_base import NonLLMProposerWorkerBase
  7. from aphrodite.spec_decode.top1_proposer import Top1Proposer
  8. class NGramWorker(NonLLMProposerWorkerBase):
  9. """NGramWorker provides a light drafter without need for model.
  10. Current NGramWorker only implement prompt lookup decoding,
  11. and in future we may also do RAG type drafter and other scenarios
  12. which don't rely on LLM model to give proposals.
  13. """
  14. def __init__(self, *args, **kwargs):
  15. # Get local_rank/vocab_size from kwargs attribute
  16. self.local_rank = kwargs["local_rank"]
  17. self.vocab_size = kwargs["model_config"].get_vocab_size()
  18. # Lazy initialization list.
  19. self._proposer: Top1Proposer
  20. def set_ngram_window_size(self, ngram_prompt_lookup_min: int,
  21. ngram_prompt_lookup_max: int):
  22. # Search valid candidate window between
  23. # ngram_prompt_lookup_min/ngram_prompt_lookup_max
  24. self.ngram_prompt_lookup_max = ngram_prompt_lookup_max
  25. self.ngram_prompt_lookup_min = ngram_prompt_lookup_min
  26. def init_device(self):
  27. self.device = torch.device(f"cuda:{self.local_rank}")
  28. self.load_model = lambda *args, **kwargs: None
  29. # Current only support Top1Proposer
  30. self._proposer = Top1Proposer(
  31. weakref.proxy(self),
  32. device=self.device,
  33. vocab_size=self.vocab_size,
  34. )
  35. def sampler_output(
  36. self,
  37. execute_model_req: ExecuteModelRequest,
  38. sample_len: int,
  39. # Unused parameter. NGramWorker does not use the KV Cache and
  40. # therefore does not need this parameter.
  41. seq_ids_with_bonus_token_in_last_step: Set[int],
  42. ) -> Tuple[Optional[List[SamplerOutput]], bool]:
  43. """NGram match algo to pick proposal candidate. Returns the list of
  44. sampler output, one per SequenceGroupMetadata.
  45. For ngram worker, we already done needed transposed internal, so the
  46. indicator pass to sampler_output_to_torch shall be False.
  47. """
  48. self._raise_if_unsupported(execute_model_req)
  49. has_spec_out = False
  50. token_id_list = []
  51. token_prob_list = []
  52. for idx, seq_group_metadata in enumerate(
  53. execute_model_req.seq_group_metadata_list):
  54. seq_data = next(iter(seq_group_metadata.seq_data.values()))
  55. input_ids = torch.as_tensor(seq_data.get_token_ids(),
  56. dtype=torch.long,
  57. device=self.device)
  58. input_length = seq_data.get_len()
  59. for ngram_size in range(
  60. min(self.ngram_prompt_lookup_max, input_length - 1),
  61. self.ngram_prompt_lookup_min - 1,
  62. -1,
  63. ):
  64. ngram_tensor = input_ids[-ngram_size:]
  65. if ngram_size == 1:
  66. # Do not match itself and do not use unfold and all
  67. matches = (input_ids[:-1] == ngram_tensor)
  68. else:
  69. windows = input_ids.unfold(dimension=0,
  70. size=ngram_size,
  71. step=1)
  72. # Do not match itself
  73. matches = (windows[:-1] == ngram_tensor).all(dim=-1)
  74. # first_match includes "values" (bool), indicating whether
  75. # the match is found, and "indices", indicating the index
  76. # of the first match.
  77. # Note that "first_match.values.item()" triggers GPU-CPU
  78. # sync so it is a bit inefficient, but we have not found
  79. # a better way to do this.
  80. first_match = matches.max(dim=-1)
  81. if first_match.values.item():
  82. proposal_start_idx = first_match.indices.add_(ngram_size)
  83. spec_indices = (
  84. proposal_start_idx).repeat(sample_len) + torch.arange(
  85. sample_len, device=self.device)
  86. spec_indices.clamp_(max=input_ids.shape[-1] - 1)
  87. res = input_ids.gather(dim=-1, index=spec_indices)
  88. token_id_list.append(res)
  89. token_prob_list.append(
  90. torch.nn.functional.one_hot(
  91. res,
  92. num_classes=self.vocab_size).to(torch.float32))
  93. has_spec_out = True
  94. break
  95. else:
  96. token_id_list.append(None)
  97. token_prob_list.append(None)
  98. if not has_spec_out:
  99. return None, False
  100. outputs: List[Optional[SamplerOutput]] = []
  101. for idx in range(len(execute_model_req.seq_group_metadata_list)):
  102. if token_id_list[idx] is None:
  103. outputs.append(None)
  104. else:
  105. outputs.append(
  106. SamplerOutput(
  107. outputs=None,
  108. sampled_token_probs=token_prob_list[idx],
  109. logprobs=torch.zeros((sample_len, self.vocab_size),
  110. dtype=torch.float32,
  111. device=self.device),
  112. sampled_token_ids=token_id_list[idx],
  113. ))
  114. return outputs, False
  115. def get_spec_proposals(
  116. self,
  117. execute_model_req: ExecuteModelRequest,
  118. # Unused parameter. NGramWorker does not use the KV Cache and
  119. # therefore does not need this parameter.
  120. seq_ids_with_bonus_token_in_last_step: Set[int],
  121. ) -> SpeculativeProposals:
  122. """Produce speculations given an input batch of sequences. The number of
  123. speculative tokens per sequence is determined by max_proposal_len.
  124. """
  125. return self._proposer.get_spec_proposals(
  126. execute_model_req, seq_ids_with_bonus_token_in_last_step)
  127. def _raise_if_unsupported(
  128. self,
  129. execute_model_req: ExecuteModelRequest,
  130. ) -> None:
  131. """NGramWorker does not yet implement support for cache swap
  132. operations or beam search.
  133. """
  134. if any([
  135. execute_model_req.blocks_to_swap_in,
  136. execute_model_req.blocks_to_swap_out,
  137. execute_model_req.blocks_to_copy
  138. ]):
  139. raise NotImplementedError(
  140. "NGramWorker does not support cache operations")
  141. if any(
  142. len(seq_group_metadata.seq_data.keys()) != 1
  143. for seq_group_metadata in
  144. execute_model_req.seq_group_metadata_list):
  145. raise NotImplementedError(
  146. "NGramWorker does not support beam search.")