Skip to main content

Module embedding

Module embedding 

Source
Expand description

Embedding lookup kernel.

Matches embedding-lookup-v1.yaml. output[i] = W[token_ids[i]] — table lookup with bounds checking.

Each function provides one of three backends:

  • fn embedding_scalar(...) – Pure Rust scalar reference (ground truth)
  • unsafe fn embedding_avx2(...) – AVX2 SIMD implementation
  • fn embedding_ptx() -> &'static str – PTX assembly source string

Functions§

embedding_avx2
AVX2 embedding lookup – delegates to scalar (memory-bound, no compute).
embedding_ptx
PTX assembly for embedding lookup.
embedding_scalar
Embedding table lookup (scalar reference).