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_encode e text out ) → b out = the embedding vector ( embed_dim e ) ( embed_close e )

Inference is per text (batch 1). The sequence is padded to a quantised length and the padding masked out of attention and pooling (see __em_bucket) — the numbers are the numbers of the unpadded run, and what the quantisation buys is a forward that stops compiling kernels and allocating device memory after the first few requests.

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_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

@ 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_ids_norm * Embed e ( Vec i ) ids ( Vec f ) out b normalize → b

Forward over an already-tokenized id sequence.

The CUDA context is thread-local, and this may be called from a server's worker rather than from whatever thread opened the model, so the device is bound to the caller first. It is a no-op on the CPU backend and a pointer store on CUDA.

---

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. 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
    ( Vec f ) out
    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