From 01bcd0d105465a0648a50d930351286e00afa614 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 2 Sep 2026 12:10:16 +0800 Subject: [PATCH] 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 | 35 ++++++----- astrai/inference/scheduler.py | 10 ++-- docs/developer/architecture.md | 2 +- docs/guides/inference.md | 4 ++ tests/inference/test_scheduler.py | 86 ++++++++++++++++++++++++++++ 5 files changed, 117 insertions(+), 20 deletions(-) diff --git a/astrai/inference/runtime/executor.py b/astrai/inference/runtime/executor.py index 612a26a..8d21c95 100644 --- a/astrai/inference/runtime/executor.py +++ b/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 diff --git a/astrai/inference/scheduler.py b/astrai/inference/scheduler.py index f568fe2..234e307 100644 --- a/astrai/inference/scheduler.py +++ b/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): diff --git a/docs/developer/architecture.md b/docs/developer/architecture.md index 7cf96d8..791d331 100644 --- a/docs/developer/architecture.md +++ b/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]]] } diff --git a/docs/guides/inference.md b/docs/guides/inference.md index 1add499..56c9b5e 100644 --- a/docs/guides/inference.md +++ b/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 diff --git a/tests/inference/test_scheduler.py b/tests/inference/test_scheduler.py index 4e65cbd..2576527 100644 --- a/tests/inference/test_scheduler.py +++ b/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: