Skip to content

[FEA] IVF-PQ transform: faster PQ-encode kernels (SIMT / tensor-core backends) #3009

Open

Description

@jamxia155

Problem
In ivf_pq::transform, PQ encoding of residuals is a significant part of the cost at large scale.

Solution
Add alternative encode backends for PER_SUBSPACE codebooks, selectable via a parameter:
baseline (current), simt, tensor (TF32/FP16 tensor cores), tensor3 (3-pass split, ~exact).

Measured (microbench, us/row): baseline 0.370, simt 0.143, tensor 0.064, tensor3 0.0715.
Accuracy vs exact FP64: tensor ~0.08% of codes differ; tensor3 ~exact. Best 1B config uses tensor3.

Open questions for maintainers

  • API shape: a field in the transform/index params vs an automatic choice (default stays baseline or tensor3?).
  • Is a ~0.08% code mismatch acceptable for tensor, or should only the ~exact variant ship?
  • PER_CLUSTER codebooks are not supported by the new backends (falls back to baseline).

Prototype has an env var (LANCE_CUVS_PQ_ENCODE_BACKEND); upstream would replace it with a parameter.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions