Dart Selector Kernels#

void trt_edgellm::kernel::computeDartRowNorms(
rt::Tensor const &embeds,
rt::Tensor &norms,
cudaStream_t stream
)#

Per-row L1 and L2 norms over the hidden dimension.

Computes, for every row of a [numRows, hiddenSize] FP16 matrix, the L1 norm (used for DART pivot scoring) and the L2 norm (used as the cosine-similarity denominator).

Parameters:

  • embeds – [in] Row-major FP16 matrix [numRows, hiddenSize] (a [1, S, H] tensor viewed as [S, H])

  • norms – [out] FP32 output [numRows, 2]; column 0 = L1, column 1 = L2

  • stream – [in] CUDA stream for execution

Throws:

std::runtime_error – on shape/dtype mismatch

void trt_edgellm::kernel::computeDartPivotDots(
rt::Tensor const &embeds,
int32_t const *pivotIndices,
int32_t numPivots,
rt::Tensor &dots,
cudaStream_t stream
)#

Dot products of a small set of pivot rows against all rows.

dots[p][r] = sum_h embeds[pivotIndices[p]][h] * embeds[r][h], accumulated in FP32. Cosine similarity is then dots[p][r] / (L2[pivot_p] * L2[r]) on the host.

Parameters:

  • embeds – [in] Row-major FP16 matrix [numRows, hiddenSize]

  • pivotIndices – [in] Device INT32 array of row indices, length numPivots (<= kDartMaxPivots)

  • numPivots – [in] Number of pivot rows

  • dots – [out] FP32 output [numPivots, numRows]

  • stream – [in] CUDA stream for execution

Throws:

std::runtime_error – on shape/dtype mismatch or numPivots out of range