• Stefy Lanza (nextime / spora )'s avatar
    embeddings(bge-m3): vectorize sparse/colbert aggregation on GPU · 59f6ebe0
    Stefy Lanza (nextime / spora ) authored
    The per-token Python loops that built the sparse {token_id: weight} dict and
    the colbert vector lists held the GIL between forward passes, starving the GPU
    (steady ~10-36% util with brief bursts). Replace them with GPU tensor ops:
    
    - sparse: relu(sparse_linear·h) with invalid positions (specials/padding)
      zeroed, then max-pooled per token id via a single scatter_reduce(amax) into a
      (B, vocab) matrix; per-row nonzero extraction is O(nnz), not O(L).
    - colbert: one C-level tolist per row over its valid tokens (vs per-token).
    - valid mask computed on-device via torch.isin.
    
    Output byte-matches the previous per-token logic (dense/sparse/colbert within
    1e-4). Benefit grows with batch size, so larger client batches translate to GPU
    work instead of Python-loop time.
    Co-Authored-By: 's avatarClaude Opus 4.8 <noreply@anthropic.com>
    Claude-Session: https://claude.ai/code/session_01Mw2KQiswmD69T45fTfjKwW
    59f6ebe0
embeddings.py 85.4 KB