Pure-PyTorch CPU inference for the text embedder (no CUDA)
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 referencerwkv_emb is CUDA-only: it JIT-compiles the .cu kernel at import and hard-codesdevice="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.
Update — two important corrections to my earlier post.
1. Bug fixed in the CPU port: over-length texts collapsed to a constant embedding.
The pooling only reads the hidden state at EOS positions. My preprocessing appended the EOS before truncating to max_tokens , so any text longer than the cap lost its EOS, the pooling mask went empty, and the head mapped the zero vector to the same constant embedding for every over-length text. I measured cosine 1.0000 between two completely unrelated long documents (Python code vs. a cooking recipe); 0.159 after the fix. Fixed in the port: truncate to max_tokens - 1 first, then append EOS (commit 0098fb9).
Practical takeaway for anyone using state-pooling embedders: always make sure the EOS survives truncation. Once it does, state pooling degrades gracefully on truncated inputs , a truncated Python doc now lands near related Python content (cos 0.49) instead of in the constant-point cluster.
2. My "underperforms on cross-lingual retrieval" conclusion was largely an artifact.
Two compounding issues on my side: the EOS bug above, and a 256-token cap I had set as a CPU concession. Causal state pooling suffers far more from truncation than transformer mean-pooling , cutting a passage at 256 tokens amputates the final state. Re-run fairly at 512 tokens on my FR→EN code-documentation benchmark, the 0.1B jumped from 26/56 to 48/56 hit@3 on par with multilingual-e5-small (118M) and embeddinggemma-300m, ahead of e5-base (278M), and close behind the best in-class transformer I tested (granite-107m, 53/56). Its out-of-domain refusal calibration was actually the best of the panel.
So: MTEB-EN does transfer to multilingual retrieval better than my first post suggested , the model just needs its full context and an intact EOS. After a small domain fine-tune (Kaggle free tier), the 0.1B matched that 107M transformer on hit@1 on this benchmark; I'm putting together a full write-up of the campaign (including a scale effect on pooled evals worth knowing about) and will link it here.
Great model ,even more so now that I'm measuring it properly.
Full write-up is up: https://scarletwolf.ai/en/blog/fine-tune-embedder-rwkv.html , the complete campaign with exact numbers,
including the EOS bug details, the pool-vs-full-scale scale effect, and why we kept granite in production despite the RNN taking the out-of-domain refusal
crown.