flashinfer.token_sparse.group_union_masks

flashinfer.token_sparse.group_union_masks(positions: Tensor, query_rows: Tensor, group_unions: Tensor, requests: Tensor, block_table: Tensor, *, union: Tensor, union_sizes: Tensor, mask: Tensor, num_keys: int) → None

Build the inputs of masked dense attention from key positions: each group of consecutive tokens’ union of keys, and each token’s bitmask over it.

Neighboring tokens, such as a request’s prefill tokens, select many of the same keys, so a group of them can attend to their union as one varlen batch entry, each token masked to its own keys. The tokens form len(requests) equal groups. It runs as one launch without host synchronization, so it can be captured in a CUDA graph (run it once before capture to compile it).

Parameters:

  • positions (torch.Tensor) – int32 [tokens, K], as in token_sparse_kv_indices().

  • query_rows (torch.Tensor) – int32 [tokens]: each token’s row of mask, -1 for none.

  • group_unions (torch.Tensor) – int32 [groups]: each group’s row of union, -1 for none.

  • requests (torch.Tensor) – int32 [groups]: each group’s request, a row of block_table.

  • block_table (torch.Tensor) – int32 [requests, blocks]: each request’s pages of 64 keys.

  • union (torch.Tensor) – int32, output: a row gets its group’s union of U keys, as cache slots in ascending key order.

  • union_sizes (torch.Tensor) – int32, output: U for each union row.

  • mask (torch.Tensor) – int32, output: bit j of word w of a token’s row is set when the token selected union column 32w + j. Only words w < 4 ceil(U / 128) are written.

  • num_keys (int) – The number of keys: positions at or past it count as -1. The block table must cover it, and shared memory grows with it.