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.
This commit is contained in:
0z5a
2026-09-02 15:00:16 +08:00
parent 800981d85a
commit 01bcd0d105
5 changed files with 117 additions and 20 deletions
+1 -1
View File
@@ -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]]]
}
+4
View File
@@ -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