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.