Take positions as input for .forward

EricLBuehler · Dec 17, 2023 · 9c08e04 · 9c08e04
1 parent a9fdd0e
commit 9c08e04
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 4 deletions.
diff --git a/src/openai/pipelines/llama.rs b/src/openai/pipelines/llama.rs
@@ -180,6 +180,7 @@ impl<'s> ModulePipeline<'s> for LlamaPipeline {
  fn forward(
  &mut self,
  _input_tokens: Tensor,
+ _input_positions: Tensor,
  _kv_cache: Option<Arc<Vec<(Tensor, Tensor)>>>,
  _input_metadata: InputMetadata,
  ) -> Result<Vec<TokenOrFinishReason>, APIError> {

diff --git a/src/openai/pipelines/llm_engine.rs b/src/openai/pipelines/llm_engine.rs
@@ -128,7 +128,7 @@ impl<'a> LLMEngine<'a> {
 
  let PreparedInputs {
  tokens,
- positions: _,
+ positions,
  metadata,
  } = if scheduled
  .front()
@@ -147,9 +147,12 @@ impl<'a> LLMEngine<'a> {
  self.prepare_decode(scheduled)
  }?;
 
- let result =
- self.pipeline
- .forward(tokens, Some(self.cache_engine.get_kv_cache()), metadata)?;
+ let result = self.pipeline.forward(
+ tokens,
+ positions,
+ Some(self.cache_engine.get_kv_cache()),
+ metadata,
+ )?;
 
  for (result, (_, seq)) in zip(result, seqs) {
  match result {

diff --git a/src/openai/pipelines/mod.rs b/src/openai/pipelines/mod.rs
@@ -21,6 +21,7 @@ pub trait ModulePipeline<'s>: Send + Sync {
  fn forward(
  &mut self,
  input_tokens: Tensor,
+ input_positions: Tensor,
  kv_cache: Option<Arc<Vec<(Tensor, Tensor)>>>,
  input_metadata: InputMetadata,
  ) -> Result<Vec<TokenOrFinishReason>, APIError>;