howard-hou

/

EmbeddingRWKV

like 6

Hi, and thanks for releasing EmbeddingRWKV — the state-centric / reusable-state
angle is really nice.

I wanted to run the text embedder on CPU-only machines, but the reference
rwkv_emb is CUDA-only: it JIT-compiles the .cu kernel at import and hard-codes
device="cuda" / torch.half, so it won't run without a GPU.

So I reimplemented the text embedding path in pure PyTorch (fp32, CPU), with
no CUDA dependency. I kept your RWKV-7 layer classes as-is (they're already pure
torch); the only new piece is a CPU WKV-7 operator, which I verified is
numerically equivalent to the reference kernel:

kernel:     w = exp(-0.6065 * sigmoid(u))      # 0.6065 = exp(-0.5)
pure path:  w = exp(-exp(-softplus(-u) - 0.5))
            exp(-softplus(-u) - 0.5) == exp(-0.5) * sigmoid(u)   → identical decay

It loads only the text tower + the retrieval head (the vit.* / proj.* vision
weights are dropped, so no SigLIP download), auto-detects the architecture from
the state dict (works for the 0.1B / 0.4B / 1.4B checkpoints), and exposes a
small encode(texts, kind="query"|"passage") API with the retrieval instruction
applied on the query side.

Repo (Apache-2.0, attribution + NOTICE included): "https://codeberg.org/scarletwolf_ai/embedding-rwkv-cpu"

from embedding_rwkv_cpu import EmbeddingRWKV
m = EmbeddingRWKV("rwkv0b1-emb-curriculum.pth")
q = m.encode(["how to sort a list in python"], kind="query")
d = m.encode(["Use sorted() or list.sort() to order a list."], kind="passage")
print((d @ q[0]).tolist())

One CPU note that might interest others: the recurrence is memory-bandwidth
bound on the [B,H,N,N] state, not compute bound — so a small batch (≈4) that
keeps the state in cache is actually the fastest; larger batches spill to RAM and
slow down.

Happy to open a PR (here on HF, or wherever you prefer) if you'd like to link or
include it. Thanks again for the model!

PS : small data point, not a benchmark: on a FR→EN code-documentation retrieval

• task (French queries, English docs, ~56 queries), the 0.1B text embedder scored
• notably below same-size multilingual transformer encoders. The MTEB-English
• number doesn't seem to transfer to this cross-lingual/code setting — might be
• useful signal for a future multilingual/code training mix.