Gpu tiledcosinetopk - CyrilB1531/lodestar GitHub Wiki

Development build. This page describes main, not a released package. The latest published Lodestar.Gpu is 0.2.0 — read its documentation.

Home › Gpu › GPU kernels

TiledCosineTopK

Sweeps a resident matrix with a batch of queries: cosine similarity, then top-k.

public sealed class TiledCosineTopK

Example — the shape a caller writes.

using Lodestar.Embeddings.Search;
using Lodestar.Gpu.Compute;

using var context = GpuContext.Create(preferCpu: true);
float[] rows = [1f, 0f, 0f, 1f, 1f, 1f];
using var matrix = DeviceEmbeddingMatrix.Upload(context, rows, count: 3, dimension: 2);
var kernel = new TiledCosineTopK(context);

IReadOnlyList<IReadOnlyList<SearchResult>> hits = kernel.Search(matrix, [1f, 0f], 1, 2);

int best = hits[0][0].Index;  // => 0
int howMany = hits[0].Count;  // => 2

Members — one page each.

Member What it does
TiledCosineTopK.Search The best k rows for each query in a batch

Remarks — two kernels on one stream. The first tiles the query vector through shared memory so a group reads it once rather than once per thread; the second keeps the selection on the accelerator instead of copying every score back. Each thread keeps a heap of the best rows it reads, and the group merges the threads' sorted lists, so a query's scores are read once however large k is — the repeated parallel argmax this replaced read them once per hit (#1214).

Hits are SearchResult, the record EmbeddingIndex.Search returns, so a caller moving between the two paths changes no type.

Ties break by row index ascending, matching EmbeddingIndex.Search. That is the one thing the CPU sort and the GPU reduction have to agree on by construction rather than by luck, and a corpus of identical rows is what the tests use to prove it.

Build this once. Loading compiles, and a first launch on a freshly loaded kernel measures ILGPU's compiler — which is why the package's performance page measures its gate with the warm-up excluded.

The group size is clamped to what the accelerator allows and rounded down to a power of two: ILGPU's CPU accelerator caps a group dimension at 16, not 256, so a kernel hard-coded to 256 threads runs on a GPU and refuses to launch exactly where correctness is tested.

Applies to — net10.0, netstandard2.1.

See also — DeviceEmbeddingMatrix, SearchResult, the namespace index.