Repository navigation
Description
Is it related to a problem?
ivf_pq::transform (and extend) always applies rotation_matrix via GEMM, even when force_random_rotation=false
and dim == rot_dim, i.e. when the rotation is exactly the identity. In that case the GEMM computes x * I = x: it changes
nothing and is pure overhead. Skipping it whenever the identity is known is a free win: bit-identical output, no accuracy
or memory trade-off, and strictly less work. The only cost is a small cached value on the index (see Solution); no
serialization or public API change is needed.
How large the win is
The skip is worthwhile at any size; the analysis below only quantifies it and shows why it grows with dimension.
Per row, with D = dim = rot_dim, K = n_lists, pq_dim sub-vectors of pq_len = D / pq_dim, and 2^b codewords:
| Step | FLOPs/row | Scaling in D |
|---|---|---|
| Rotation GEMM | 2 D^2 | quadratic |
| Coarse assignment, exact | 2 K D | linear |
| Coarse assignment, approximate (grouped, M ~ sqrt(K) groups, p probes) | ~2 D sqrt(K) (1 + p) | linear |
| PQ encode | 2 pq_dim 2^b pq_len = 2^(b+1) D | linear |
Every other step is linear in D, so the rotation's share rises with D. It exceeds the PQ-encode cost once D > 2^b,
exceeds exact assignment once D > K, and exceeds approximate assignment once D > ~sqrt(K) (1 + p). Modern embedding
models (1-4K dims) are already past the first and, with approximate assignment, the last of these thresholds, and
dimensions keep growing. The rotation matrix is also D^2 floats, so it is a growing memory-traffic cost as well.
Note that the GEMM may run at a lower effective throughput than the assignment kernels (e.g. plain FP32 vs TF32 tensor cores),
so its share of time can exceed its share of FLOPs.
Solution
Have the index lazily determine, once, whether its rotation matrix is exactly the identity, and cache the result in a
non-serialized member of index_impl (next to the existing lazily initialized, never-serialized members such as
centers_int8_ / rotation_matrix_int8_; e.g. a mutable std::optional<bool> rotation_is_identity_). The first
transform/extend call checks the D x D matrix on the device once; later calls read the cached value. When it is set,
transform/extend bypass the GEMM and use the input directly as the rotated data. Output is bit-identical.
Why check the matrix rather than record what build was asked to do: the identity is a property of the matrix, and
the build-from-components overloads accept a caller-supplied rotation_matrix, so force_random_rotation and
dim == rot_dim alone do not tell whether it is the identity. A check of the actual matrix is correct for a freshly
built index, a component-built index and a deserialized index alike, so there is no serialization format change and
older index files benefit without versioning.
Points to settle in review:
- Invalidation: the cache must be reset wherever the matrix can be written, and
view_implholds a view onto
caller-owned memory whose contents can change. A conservative start is to cache only forowning_impland re-check
views on each call. - The lazy fill happens in a
constmethod, so concurrent first calls on one index would race. The existing*_int8_
caches share this; astd::once_flagwould make it safe if preferred.
Example measurements (anecdotal, one setup: H100, CUDA 13.4, 1M rows, D = 2048, K = 32768, 256 sub-vectors, 8-bit codes):
- Exact (unfused) assignment: rotation is ~6% of FLOPs; skipping it cut transform_s from 0.986 to 0.85 (-14%).
- Grouped approximate assignment (see [FEA] Faster coarse-list assignment in IVF-PQ transform for large n_lists (grouped-GEMM assigner) #3007): rotation is ~52% of FLOPs; skipping it cut transform_s from
~0.63 to ~0.42 (-32%; two runs each, run-to-run spread ~10%).
The saving per row is the same cost in both cases (~0.14-0.20 s per 1M rows here); it is a larger fraction when the
rest of the transform is cheaper or D is larger.
Alternatives
Callers could avoid rotation themselves, but that needs a different API path. The change is also beneficial to stock
cuvsIvfPqTransform with no API change.
Additional context
We have a working patch and will open a PR. The prototype checks the matrix on every call and gates the skip on an env
var for A/B; the upstream version would be always-on with the cached member (no knob, no per-call check).
Metadata
Metadata
Assignees
Labels
Type
Projects
- StatusShow more project fieldsNeeds Triage
Activity