perf: batch ragged prefill requests

- Pack prompts with a shared prefix start and attention backend into one forward.

- Select per-request final logits from cumulative query lengths.

- Cover ragged tokens, logprobs, scheduling, and documentation.

astrai/inference/runtime/executor.py

@@ -307,30 +307,39 @@ class Executor:
     def execute_prefill(
         self,
         tasks: List[Task],
-        prompt_len: int,
         start_pos: int = 0,
         return_logprobs: bool = False,
     ):
-        if start_pos >= prompt_len:
-            return []
-
         tasks = sorted(tasks, key=lambda t: t.task_id)
         batch_sz = len(tasks)
+        prompt_lens = [len(t.prompt_ids) for t in tasks]
+        if any(start_pos >= prompt_len for prompt_len in prompt_lens):
+            raise ValueError("prefill start_pos must precede every prompt end")
+        q_lens = [prompt_len - start_pos for prompt_len in prompt_lens]
 
         input_ids = torch.tensor(
-            [token for t in tasks for token in t.prompt_ids[start_pos:prompt_len]],
+            [token for t in tasks for token in t.prompt_ids[start_pos:]],
             dtype=torch.long,
             device=self.device,
         )
 
         task_ids = [t.task_id for t in tasks]
-        position_ids = torch.arange(
-            start_pos, prompt_len, dtype=torch.long, device=self.device
-        ).repeat(batch_sz)
+        position_ids = torch.cat(
+            [
+                torch.arange(
+                    start_pos, prompt_len, dtype=torch.long, device=self.device
+                )
+                for prompt_len in prompt_lens
+            ]
+        )
 
         with (
             torch.inference_mode(),
-            timed(f"execute_prefill b={batch_sz} prompt_len={prompt_len}", logger),
+            timed(
+                f"execute_prefill b={batch_sz} tokens={sum(q_lens)} "
+                f"q_len={min(q_lens)}..{max(q_lens)}",
+                logger,
+            ),
         ):
             outputs = self.model(
                 input_ids,
@@ -342,10 +351,10 @@ class Executor:
                 ),
                 fwd="prefill",
             )
-            q_len = prompt_len - start_pos
-            logits = outputs["logits"][
-                torch.arange(1, batch_sz + 1, device=self.device) * q_len - 1
-            ]
+            last_token_indices = (
+                torch.tensor(q_lens, dtype=torch.long, device=self.device).cumsum(0) - 1
+            )
+            logits = outputs["logits"][last_token_indices]
 
         step_out, _ = self._sample_logits(logits, tasks, return_logprobs)
         return tasks, step_out

astrai/inference/scheduler.py

@@ -177,16 +177,14 @@ class InferenceScheduler:
             for t in to_prefill:
                 t.input_tokens = len(t.prompt_ids)
 
-            groups: Dict[Tuple[int, int, Optional[AttentionBackend]], List[Task]] = {}
+            groups: Dict[Tuple[int, Optional[AttentionBackend]], List[Task]] = {}
             for t in to_prefill:
                 start_pos = min(
                     self._task_cache.task_cached(t.task_id), len(t.prompt_ids) - 1
                 )
-                groups.setdefault((len(t.prompt_ids), start_pos, t.backend), []).append(
-                    t
-                )
+                groups.setdefault((start_pos, t.backend), []).append(t)
 
-            for (prompt_len, start_pos, _), group in groups.items():
+            for (start_pos, _), group in groups.items():
                 backend = group[0].backend
                 backend_context = (
                     attn_backend(backend) if backend is not None else nullcontext()
@@ -196,7 +194,7 @@ class InferenceScheduler:
                     self._metrics.record([t.task_id for t in group], "prefill"),
                 ):
                     prefilled, step_out = self._executor.execute_prefill(
-                        group, prompt_len, start_pos, return_logprobs=return_logprobs
+                        group, start_pos=start_pos, return_logprobs=return_logprobs
                     )
 
                 for t, out in zip(prefilled, step_out):