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
+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