flashinfer.attention.prims_ts.batch_prefill

flashinfer.attention.prims_ts.batch_prefill(q: Tensor, k: Tensor, v: Tensor, *, qo_indptr: Tensor | None = None, kv_indptr: Tensor | None = None, mask_type: Literal['dense', 'causal', 'variable_window'] = 'dense', window_left: int = -1, variable_window_token_starts: Tensor | None = None, variable_window_token_ends: Tensor | None = None, sm_scale: float | None = None, output_scale: float = 1.0, out_dtype: dtype | None = None, out: Tensor | None = None) → Tensor

Run one-shot fixed or packed-ragged task-scheduled context attention.

Fixed tensors use [B, S, H, D] storage. Providing both cumulative int32 offset tensors selects packed [total_tokens, H, D] storage. D may be 128 or 256. Causal masking is bottom-right aligned. window_left=-1 disables the left window; a positive value selects the private head-paired GQA policy and retains at most window_left + 1 keys at each causal row, including when S_q < S_kv.

Parameters:

  • q (torch.Tensor) – Fixed or packed query, key, and value tensors.

  • k (torch.Tensor) – Fixed or packed query, key, and value tensors.

  • v (torch.Tensor) – Fixed or packed query, key, and value tensors.

  • qo_indptr (torch.Tensor, optional) – Cumulative query and K/V offsets for packed-ragged input.

  • kv_indptr (torch.Tensor, optional) – Cumulative query and K/V offsets for packed-ragged input.

  • mask_type ({"dense", "causal", "variable_window"}) – Attention mask mode. variable_window is supported only for fixed-shape inputs.

  • window_left (int) – Left sliding-window extent, or -1 to disable the window.

  • variable_window_token_starts (torch.Tensor, optional) – Inclusive per-query K bounds required for variable_window. Both must be CUDA int32 tensors shaped [B, Sq] and satisfy 0 <= starts[b, q] <= ends[b, q] < Sk.

  • variable_window_token_ends (torch.Tensor, optional) – Inclusive per-query K bounds required for variable_window. Both must be CUDA int32 tensors shaped [B, Sq] and satisfy 0 <= starts[b, q] <= ends[b, q] < Sk.

  • sm_scale (float, optional) – Softmax scale; defaults to the inverse square root of head size.

  • output_scale (float) – Scale applied to the attention output.

  • out_dtype (torch.dtype, optional) – Requested output dtype.

  • out (torch.Tensor, optional) – Caller-owned output tensor.