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 is a handle: every copy is the same engine, and its last owner gives everything back — the weights, the kit, the CUDA context. embed_close is an optional early release.

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
}

: EmbedImpl

: EmbedImpl {
    GpuKit kit
    EmbedCfg cfg
    GkBuf wemb
    GkBuf pemb
    GkBuf temb
    GkBuf elnw
    GkBuf elnb
    ( Vec EmbedLayer ) layers
    Unigram tok
    b has_tok
    b ok
    // the weights as a lease: `loaded` says whether kit + weights exist
    // right now; `dir`/`gpu` are what a reload opens. The tokenizer and
    // the config stay for the engine's life — a request tokenizes on its
    // own fiber while the weights may be gone.
    b loaded
    String dir
    i gpu
    // the weight arena: device memory carved from 128 MB chunks (~20
    // allocations for BGE-M3's 389 tensors, and ~20 frees on an unload)
    ( Vec i ) arena
    ( Vec i ) arenasz
    ( Vec GpuBuffer ) arenao  // the chunks' owners: clearing this frees them
    i arena_cur
    i arena_off
    i arena_cap
    ( Vec GpuCopy ) up_q  // every tensor's upload, sent as ONE batch
}

: Embed { s ctl }

An Embed is a handle on its state in an rcbox (stdlib/core/rcbox.nu): every copy is the same engine, and the last owner releases it.

@ Embed_drop sink Embed h → v

@ 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_unload Embed e__h → v

Give the device back: every weight (the arena chunks), the kit — its buffer pool, its kernels, the CUDA context. What stays is the engine: config, tokenizer, the model dir for the reload. Idempotent.

@ embed_close sink Embed e → v

Early release (optional): the last owner of the engine gives back the weights, the kit and the CUDA context.

: 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.

---

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 // the served Embed handle's ctl word, lent by embed_serve's caller (0 = not serving)

: ~ s g_em_token ``

: ~ s g_em_name ``

: ~ i g_em_reqs 0

: ~ i g_em_unload_ms 0 // 0 = the weights stay for the process's life

── the weights as a lease: --unload-after ─────────────────────────

The model thread is the only thread that touches the device, so it is the one that lets the weights go and brings them back: a job arriving at an unloaded engine reloads it first (embed_reload — the tokenizer and config never left, so the request was already tokenized on its own fiber); a wake with no job and the idle clock past the limit unloads (embed_unload — the arena, the pool, the kernels, the CUDA context). The wakes come from a ticker thread that broadcasts the request cond five times a second while the flag is on; without it the model thread sleeps until a job and the limit could never fire.

: ~ i g_em_idle_since 0 // monotonic_ns at the end of the last job

: ~ i g_em_loads 0 // 1 is the load before the port opened

: ~ i g_em_unloads 0

: ~ i g_em_load_ms 0 // the last load's wall time

: EmJobImpl

: EmJobImpl {
    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 holds one word, so the queue is two pointers, and its lock and two conditions sit in an EmSync block the server allocates once and keeps for the process (the model thread and the ticker outlive any one scope, as the handlers do). The submitting fiber owns the job (an EmJob handle) and keeps it until it has seen done; the queue and the model thread only borrow its address meanwhile.

: EmJob { s ctl }

@ EmJob_drop sink EmJob h → v

: ~ i g_q_head 0

: ~ i g_q_tail 0

: ~ b g_q_stop F

: EmSync

: EmSync {
    Mutex m  // guards the queue, g_q_stop, the idle clock and g_em_reqs
    Cond req  // a job was queued, the server is stopping, or a ticker wake
    Cond done  // a job was finished
}

: ~ i g_q_sync 0 // *EmSync as an address (0 = never served)

---

main.nu

packages/embed — the embedding server. CLI:

embed serve <model-dir> [--addr HOST:PORT] [--token T] [--maxseq N] [--unload-after S] [--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