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:
@@ -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]]]
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user