flashinfer.sdsa.gather_key_sample

flashinfer.sdsa.gather_key_sample(key_cache: Tensor, block_table: Tensor, offsets: Tensor, sample_ends: Tensor, *, sample_keys: Tensor, samples_per_run: int, pages_per_run: int, key_scales: Tensor | None = None, sample_scales: Tensor | None = None) → None

Copy a strided random sample of each request’s paged keys into dense sample rows, whose scores give thresholds for indexer_threshold_select().

Viewed as 64-key pages (sample_keys.view(-1, 64, d)), the sample can be scored by indexer_scores(). It runs as one launch without host synchronization, so it can be captured in a CUDA graph.

Parameters:

  • key_cache (torch.Tensor) – As in indexer_scores().

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

  • offsets (torch.Tensor) – uint8 [capacity]: sample slot s takes the key at offsets[s] (below 64 * pages_per_run) into run s // samples_per_run.

  • sample_ends (torch.Tensor) – int32 [requests]: each request’s number of samples. Later slots are left as they are.

  • sample_keys (torch.Tensor) – BF16 [requests, capacity, 128] or NVFP4 uint8 [..., 64], output, capacity a multiple of 128.

  • samples_per_run (int) – Samples per run of pages_per_run pages.

  • pages_per_run (int) – Pages per run.

  • key_scales (torch.Tensor, optional) – NVFP4 key scales, as in indexer_scores().

  • sample_scales (torch.Tensor, optional) – [requests, capacity, 8], output: the sampled keys’ NVFP4 scales.