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(
+        position_ids = torch.cat(
-            start_pos, prompt_len, dtype=torch.long, device=self.device
+            [
-        ).repeat(batch_sz)
+                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
+            last_token_indices = (
-            logits = outputs["logits"][
+                torch.tensor(q_lens, dtype=torch.long, device=self.device).cumsum(0) - 1
-                torch.arange(1, batch_sz + 1, device=self.device) * q_len - 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(
+                groups.setdefault((start_pos, t.backend), []).append(t)
-                    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):

docs/developer/architecture.md

@@ -822,7 +822,7 @@ classDiagram
             +InferenceWorkspace _workspace
             +Optional[str] device
             +Optional[torch.dtype] dtype
-            +execute_prefill(tasks, prompt_len, start_pos=0)
+            +execute_prefill(tasks, start_pos=0)
             +execute_decode(tasks, return_logprobs=False) Union[List[int], List[Tuple[int, float]]]
         }
 

docs/guides/inference.md

@@ -105,6 +105,10 @@ handle the call, inference raises an error rather than silently switching.
 
 `CudaBackend` prefill path: writes K/V, then calls `attn_paged_prefill` — a ragged-batch (paged) prefill kernel that reads K/V directly from the flat pool via `req_to_token`, addressing each request's `q_len`/`kv_len` through `qo_indptr` and `kv_indptr`. No explicit K/V gather needed.
 
+The scheduler packs requests with the same prefix-cache start position and
+attention backend into one prefill forward even when their prompt lengths differ.
+Requests with different prefix hit lengths remain separate batches.
+
 Fallback: when `CudaBackend` cannot handle an input (wrong dtype or head_dim), `FlashAttnBackend` is tried next (if installed), then `TorchNativeBackend`.
 
 This fallback is performed by the public `attention(...)` policy entry point

tests/inference/test_scheduler.py

@@ -142,6 +142,67 @@ def test_step_splits_decode_batch_by_request_backend():
     ]
 
 
+def test_step_batches_ragged_prefill_with_shared_cache_start():
+    scheduler = object.__new__(InferenceScheduler)
+    scheduler._cache = SimpleNamespace(page_size=64)
+    scheduler._task_cache = MagicMock()
+    scheduler._task_cache.task_cached.return_value = 0
+    scheduler._metrics = MetricsCollector()
+    scheduler._executor = MagicMock()
+
+    short = Task("short", [1, 2, 3])
+    long = Task("long", [4, 5, 6, 7, 8])
+    for task in (short, long):
+        scheduler._metrics.register(task.task_id)
+
+    scheduler._executor.execute_prefill.return_value = (
+        [long, short],
+        [11, 12],
+    )
+
+    produced, aborted = scheduler._step([short, long])
+
+    assert aborted == []
+    assert produced == [long, short]
+    scheduler._executor.execute_prefill.assert_called_once_with(
+        [short, long], start_pos=0, return_logprobs=False
+    )
+    assert long.output_ids == [11]
+    assert short.output_ids == [12]
+
+
+def test_execute_prefill_packs_ragged_prompts_and_selects_last_logits():
+    executor = object.__new__(Executor)
+    executor.device = torch.device("cpu")
+    executor.task_cache = MagicMock()
+    executor.task_cache.bind.return_value = MagicMock()
+    executor._workspace = MagicMock()
+    all_logits = torch.arange(42, dtype=torch.float32).reshape(6, 7)
+    executor.model = MagicMock(return_value={"logits": all_logits})
+    executor._sample_logits = MagicMock(
+        return_value=([101, 102], torch.tensor([101, 102]))
+    )
+
+    task_b = Task("b", [20, 21, 22, 23, 24])
+    task_a = Task("a", [10, 11, 12])
+
+    tasks, output = executor.execute_prefill([task_b, task_a], start_pos=1)
+
+    assert tasks == [task_a, task_b]
+    assert output == [101, 102]
+    model_args, model_kwargs = executor.model.call_args
+    assert model_args[0].tolist() == [11, 12, 21, 22, 23, 24]
+    assert model_kwargs["position_ids"].tolist() == [1, 2, 1, 2, 3, 4]
+    executor.task_cache.bind.assert_called_once_with(
+        ["a", "b"], executor._workspace, start_pos=1
+    )
+    sample_args, sample_kwargs = executor._sample_logits.call_args
+    torch.testing.assert_close(sample_args[0], all_logits[[1, 5]])
+    assert sample_args[1] == [task_a, task_b]
+    assert sample_args[2] is False
+    assert sample_kwargs == {}
+
+
 def test_scheduler_concurrent_add_remove_task(mock_model_and_tokenizer):
     """Test concurrent add and remove task operations."""
     mock_model, mock_tokenizer = mock_model_and_tokenizer
@@ -418,6 +479,31 @@ def test_run_batch_return_logprobs_aligned(device):
         scheduler.stop()
 
 
+def test_ragged_prefill_matches_sequential_greedy_tokens_and_logprobs(device):
+    scheduler, _tok, _model = _make_real_scheduler(device)
+    prompts = [
+        [10, 20, 30],
+        [5, 6, 7, 8],
+        [40, 41, 42, 43, 44],
+    ]
+    try:
+        ragged = scheduler.run_batch(
+            prompts, max_tokens=1, temperature=0, return_logprobs=True
+        )
+        sequential = [
+            scheduler.run_batch(
+                [prompt], max_tokens=1, temperature=0, return_logprobs=True
+            )[0]
+            for prompt in prompts
+        ]
+
+        assert [result[0] for result in ragged] == [result[0] for result in sequential]
+        for ragged_result, sequential_result in zip(ragged, sequential):
+            assert ragged_result[1] == pytest.approx(sequential_result[1], abs=1e-6)
+    finally:
+        scheduler.stop()
+
+
 def test_run_batch_respects_max_tokens(device):
     scheduler, _tok, _model = _make_real_scheduler(device)
     try: