flashinfer.sdsa.quantize_write_index_qk_nvfp4

flashinfer.sdsa.quantize_write_index_qk_nvfp4(key: Tensor, slots: Tensor, query: Tensor, query_rows: Tensor, *, key_cache: Tensor, key_scales: Tensor, packed_query: Tensor, query_scales: Tensor) → None

Quantize indexer keys and queries to NVFP4 in one launch: the keys into a paged key cache and the queries into rows, the layouts indexer_scores() reads.

Each key and query gets one E4M3 scale per 16 values: their amax over 6. There is no per-tensor global scale, so queries and keys are quantized the same way. Write all keys of a cache at the same scale: do not scale them by a factor that changes between calls, or their scores stop being comparable. Values above 2688 in magnitude saturate.

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:

  • key (torch.Tensor) – BF16 or FP16 [tokens, 128].

  • slots (torch.Tensor) – int64 [tokens]: each key’s cache slot s. Negative slots are skipped.

  • query (torch.Tensor) – BF16 or FP16 [queries, 128].

  • query_rows (torch.Tensor) – int64 [queries]: each query’s row r. Negative rows are skipped.

  • key_cache (torch.Tensor) – uint8 [pages, 64, 64], output: a key’s 64 packed E2M1 bytes go to key_cache[s // 64, s % 64].

  • key_scales (torch.Tensor) – float8_e4m3fn [pages, 64, 8], output: a key’s 8 scales go to key_scales[s // 64, s % 64].

  • packed_query (torch.Tensor) – uint8 [rows, 64], output: a query’s 64 packed E2M1 bytes go to row r.

  • query_scales (torch.Tensor) – float8_e4m3fn [rows, 8], output: a query’s 8 scales go to row r.