Optimize generation loop performance

src/hugging_llama/model_manager.py

@@ -677,16 +677,13 @@ async def evict_expired(self) -> None:
 
 
 def run_generation(
-    manager: ModelManager,
     request_options: GenerateOptions,
     input_ids: torch.LongTensor,
     attention_mask: torch.LongTensor,
-    tokenizer,
     model,
-    prompt_text: str,
     streamer: TextIteratorStreamer,
 ) -> dict[str, Any]:
-    prompt_tokens = len(tokenizer(prompt_text, return_tensors="pt")["input_ids"][0])
+    prompt_tokens = int(attention_mask.sum().item())
     generation_kwargs: dict[str, Any] = {
         "input_ids": input_ids.to(model.device),
         "attention_mask": attention_mask.to(model.device),
@@ -720,5 +717,6 @@ def run_generation(
         torch.manual_seed(request_options.seed)
         if torch.cuda.is_available():
             torch.cuda.manual_seed_all(request_options.seed)
-    model.generate(**generation_kwargs)
+    with torch.inference_mode():
+        model.generate(**generation_kwargs)
     return {"prompt_tokens": prompt_tokens}

src/hugging_llama/server.py

@@ -187,13 +187,10 @@ async def _stream_generate(
             None,
             partial(
                 run_generation,
-                manager,
                 request_options,
                 input_ids,
                 attention_mask,
-                tokenizer,
                 model.model,
-                prompt,
                 streamer,
             ),
         )

tests/conftest.py

@@ -289,16 +289,13 @@ async def set_keep_alive(self, name: str, ttl: float | None) -> None:
 
 
 def fake_run_generation(
-    manager: Any,
     request_options: Any,
     input_ids: Any,
     attention_mask: Any,
-    tokenizer: Any,
     model: Any,
-    prompt_text: str,
     streamer: Any,
 ) -> dict[str, Any]:
-    del manager, request_options, input_ids, attention_mask, tokenizer, prompt_text
+    del request_options, input_ids, attention_mask
     for chunk in model.outputs:
         streamer.feed(chunk)
     streamer.end()