model.nu

packages/embed/src/model.nu — the XLM-RoBERTa-family text encoder, pure NURL on gpukit's dev-layer kernel library.

This is the model class behind BGE-M3, multilingual-e5 and friends: word/position/type embeddings + LayerNorm, N transformer blocks (bidirectional self-attention, exact-erf GELU FFN, post-LN residuals), then CLS or mean pooling and an optional L2 normalize. Every operator is a gkd_* kernel — the same dtype-generic library tensor and onnx run on — so this file contains NO kernel sources, only the wiring.

A model is a DIRECTORY (the Hugging Face layout): config.json hidden_size / num_hidden_layers / … (read here) tokenizer.json Unigram vocabulary (packages/tokenizer) model.safetensors f32 weights (packages/safetensor, mmap-backed)

( embed_open dir ) → !Embed String ( embed_open_dev dir gpu ) → !Embed String (gpu -1 = best) ( embed_encode e text out ) → b out = the embedding vector ( embed_encode_batch e ids offs out norm ) → b B texts, one call ( embed_dim e ) ( embed_close e )

Inference is BATCHED where the fused attention runs (CUDA): texts are grouped longest-first into chunks of at most EM_ROWS_BUDGET padded device rows, and a chunk is ONE forward — every kernel sees all of its sequences at once, which is what turns thirty short texts from thirty launch-bound forwards into one. Sequences are padded to a quantised length, batches to a quantised count, and every bit of padding is masked out of attention and pooling (see __em_bucket) — the numbers are the numbers of the unpadded, unbatched run, and what the quantisation buys is a forward that stops compiling kernels and allocating device memory after the first few requests. On the CPU backend (or a head width the fused kernel refuses) each text runs alone through the composed path, exactly as before.

Numerics are true float32 on the device. Verified against sentence-transformers (BGE-M3): cosine ≥ 0.9999 on a multilingual corpus (tests/embed_test.sh).

API

: i EM_POOL_CLS 0

pooling modes

: i EM_POOL_MEAN 1

: EmbedCfg

: EmbedCfg {
    i layers
    i heads
    i dim
    i ffn
    i vocab
    i maxpos  // position table rows (XLM-R: max_seq + pad offset 2)
    i padid
    f eps
    i pool
    b normalize
    i maxseq  // token cap per text (specials included)
}

: EmbedLayer

: EmbedLayer {
    GkBuf qw GkBuf qb
    GkBuf kw GkBuf kb
    GkBuf vw GkBuf vb
    GkBuf ow GkBuf ob
    GkBuf ln1w GkBuf ln1b
    GkBuf iw GkBuf ib
    GkBuf dw GkBuf db
    GkBuf ln2w GkBuf ln2b
}

: Embed

: Embed {
    * GpuKit kit
    EmbedCfg cfg
    GkBuf wemb
    GkBuf pemb
    GkBuf temb
    GkBuf elnw
    GkBuf elnb
    ( Vec EmbedLayer ) layers
    * Unigram tok
    b has_tok
    b ok
}

@ embed_open s dir → !*Embed String

Open a model directory. Pooling defaults to CLS + normalize (the BGE convention); callers can override with embed_set_pooling.

@ embed_open_dev s dir i gpu → !*Embed String

The same, with the device chosen by the caller: gpu is a CUDA device ordinal (CUDA enumeration order — fastest first by default — not nvidia-smi's PCI order), or -1 for the best device / the $NURL_GPU_DEVICE override. A named ordinal must BE a CUDA device: falling back to the CPU backend behind an explicit --gpu would be hiding exactly the mistake the flag exists to make loud.

@ embed_set_pooling * Embed e i mode b normalize → v

Pooling override: mode EM_POOL_CLS | EM_POOL_MEAN, normalize on/off. (cfg is an inline struct; a field write through two levels is not an lvalue in NURL, so the setters rebuild the struct.)

@ embed_set_maxseq * Embed e i n → v

Cap on tokens per text (specials included); clamped to the model's position table.

@ embed_dim * Embed e → i

@ embed_ok * Embed e → b

@ embed_backend * Embed e → s

@ embed_device_name * Embed e → s

@ embed_maxseq * Embed e → i

@ embed_close * Embed e → v

: i EM_ROWS_BUDGET 16384

A fused-attention chunk holds at most this many padded device rows (batch · padded length). 16384 rows of BGE-M3 activations peak around 1.5 GB of pooled device buffers — and two full-length 8192-token texts still share one forward.

: i EM_BATCH_MAX 64

… and at most this many sequences, so a flood of tiny texts still produces a forward whose per-sequence host work (padding, masks, pooling rows) stays a rounding error.

@ embed_tokenize * Embed e s text ( Vec i ) out → b

Tokenize (Unigram, <s>…</s>) with truncation to cfg.maxseq: the head of the sequence is kept and </s> re-appended, sentence-transformers style.

@ embed_encode * Embed e s text ( Vec f ) out → b

The embedding for one text, normalized per the engine's configuration. out receives embed_dim floats.

@ embed_encode_norm * Embed e s text ( Vec f ) out b normalize → b

The same, with the L2 normalize decided by the caller — a request carrying "normalize": false must not have to reconfigure the engine (and a concurrent server must not be able to observe it doing so).

@ embed_encode_ids * Embed e ( Vec i ) ids ( Vec f ) out → b

@ embed_encode_batch * Embed e ( Vec i ) ids ( Vec i ) offs ( Vec f ) out b normalize → b

Embed B already-tokenized texts in one call. ids is flat token storage — text t occupies ids[offs[t] .. offs[t+1]), offs has B+1 entries — and out must already hold B·dim floats; text t's vector is written in place at row t.

Where the fused attention runs (CUDA, a supported head width), texts are sorted longest-first and packed greedily into chunks — a chunk's padded length is its LONGEST member's bucket, so the sort is what keeps short texts from being padded out to a long stranger's length — and each chunk is one batched forward (__em_fwd_batch). Everywhere else each text runs alone on the composed path, as it always has.

The CUDA context is thread-local and this may be called from a server's worker rather than the thread that opened the model, so the device is bound to the caller first.

@ embed_encode_ids_norm * Embed e ( Vec i ) ids ( Vec f ) out b normalize → b

Forward over one already-tokenized id sequence — a batch of one.

---

serve.nu

packages/embed/src/serve.nu — the embedding model as a service.

embed serve <model-dir> [--addr 0.0.0.0:8000] [--token T] [--maxseq N]

The HTTP surface mirrors the reference FastAPI embedding service, so existing clients work unchanged:

POST /create_embedding {"text": "..." | ["...", …], "normalize": true} {"texts": ["...", …]} (same thing) → {"embeddings": [[…]], "model": "…", "dimension": N} GET /create_embedding?text=…&normalize=true (single text) GET /health {"status":"healthy", "model", "model_loaded", "device": "cuda"|"cpu", …} GET / the same (a browser poking the port should learn something, not get a 404)

Auth: no --token → open server (bind loopback!). With a token, requests must carry Authorization: Bearer <t> (or ?token=<t> for clients that cannot set headers); the compare is constant-time over the configured token.

Concurrency. One model on one device can run one forward at a time — that part is not a choice. Everything ELSE about a request is: reading the socket, parsing the JSON, tokenizing (the Unigram engine is read-only, so it is re-entrant), and serialising a few thousand floats back out, which for a batch is the larger half of the work. So the server is fiber-per-connection on the async runtime, and the forward is handed to ONE dedicated model thread over a queue — one job per REQUEST, so a request's texts reach the model together and run as a few padded batched forwards (embed_encode_batch), not text by text. A single worker used to mean a slow or idle client stalled every other client behind it; now it does not, and a batch's tokenizing and JSON overlap with another request's arithmetic.

The forward runs on a thread rather than on the requesting fiber for a hard reason, not a stylistic one: async fibers get 64 KB stacks, and NVRTC — which the first forward at a new sequence length still calls — wants far more than that. Compiling a kernel on a fiber segfaults inside libnvrtc. One model thread with an ordinary 8 MB stack also keeps every CUDA call on one thread, which is where a context wants to be.

Handlers are top-level functions over module globals, not closures — closure environments are manual in NURL and a server's handlers live for the process (same idiom as whisper/nurllama).

API

: ~ i g_em 0 // *Embed as an address (0 = not serving)

: ~ s g_em_token ``

: ~ s g_em_name ``

: ~ i g_em_reqs 0

: EmJob

: EmJob {
    i next
    ( Vec i ) ids  // flat tokens for the whole request
    ( Vec i ) offs  // B+1 offsets into ids
    ( Vec f ) out  // B*dim floats, written in place
    b normalize
    b done
    b ok
}

── The model queue ───────────────────────────────────────────────────

One job per waiting request, linked through the jobs themselves — a module global cannot hold a Vec (or a Mutex) built by a call, so the queue is two pointers and the synchronisation primitives are kept as the two words of each one's Cell. The submitting fiber owns the job and frees it once it has seen done; the model thread only fills it in.

: ~ i g_q_head 0

: ~ i g_q_tail 0

: ~ b g_q_stop F

: ~ i g_q_m_ptr 0

: ~ i g_q_m_bytes 0

: ~ i g_q_req_ptr 0

: ~ i g_q_req_bytes 0

: ~ i g_q_done_ptr 0

: ~ i g_q_done_bytes 0

@ embed_serve * Embed e s name s host i port s token → i

Serve e (borrowed for the server's lifetime). Blocks until stopped.

---

main.nu

packages/embed — the embedding server. CLI:

embed serve <model-dir> [--addr HOST:PORT] [--token T] [--maxseq N] [--pool cls|mean] [--no-normalize] embed text <model-dir> <text…> one-shot: print the vector (CSV)

A model directory is the Hugging Face layout: config.json + tokenizer.json + model.safetensors (f32). BGE-M3 works as-is; any XLM-RoBERTa-family encoder does (multilingual-e5: --pool mean). The model argument can be a local directory OR a Hugging Face ref (embed serve BAAI/bge-m3): a ref is fetched into the shared ~/.nurl cache via the hub package and the downloaded directory is used.

API

@ main → i