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
    b ok
    // the weights as a lease: `loaded` says whether kit + weights exist
    // right now; `dir`/`gpu` are what a reload opens.
    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).
    // Each weight is a view of its chunk and holds it: the chunk goes back
    // when the last weight in it does.
    ( Vec GkBuf ) arena
    i arena_cur  // the chunk being carved (its index in arena; -1 none)
    i arena_off
    i arena_cap
    // every tensor's upload, sent as ONE batch: its device view, and its
    // tensor's index in the file
    ( Vec GkBuf ) up_dst
    ( Vec i ) up_src
}

: EmbedStat

: EmbedStat {
    b ok
    b loaded
    String backend
    String device
    i kernels
}

What the server reports while a forward runs: kept beside the engine, so reading it never waits for the device.

: Embed

: Embed {
    Unigram tok
    ( Locked EmbedCfg ) cfg
    ( Locked EmbedStat ) st
    ( Locked EmbedImpl ) l
}

An Embed is a handle: every copy is the same engine, and the last owner releases it. The tokenizer and the config stay for the engine's life — a request tokenizes on its own fiber, under no lock the forward holds — while the engine itself (the kit, the weights) is Locked: one forward, one load or one unload at a time.

@ Embed_share Embed h → Embed

@ 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_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_unload Embed 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_reload Embed h → !v String

Bring the weights back after embed_unload. F carries the loader's message; the engine stays usable for tokenizing either way.

@ embed_loaded Embed h → b

@ embed_set_pooling Embed h i mode b normalize → v

Pooling override: mode EM_POOL_CLS | EM_POOL_MEAN, normalize on/off. The engine's copy follows, after any forward that is running.

@ embed_set_maxseq Embed h i n → v

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

@ embed_dim Embed h → i

@ embed_ok Embed h → b

@ embed_backend Embed h → String

@ embed_device_name Embed h → String

@ embed_kernel_count Embed h → i

Kernels the engine's kit has compiled so far (0 while unloaded) — a server that stopped meeting new shapes stops growing this.

@ embed_maxseq Embed h → i

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

@ embed_tokenize Embed h 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 h 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 h 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 h ( Vec i ) ids ( Vec f ) out → b

@ embed_encode_batch Embed h ( 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 h ( 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 globals set at start: the queue (a Channel) and the server's state (a Locked), reached from every fiber and thread in safe code.

API

: EmServe

: EmServe {
    ? Embed model  // one owner of the served model, while serving
    i reqs
    // ── 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 tick with the idle clock past the
    // limit unloads (embed_unload — the arena, the pool, the kernels, the
    // CUDA context). The ticks come from a ticker thread, five a second
    // while the limit is on; without them the model thread sleeps until
    // a job and the limit could never fire.
    i unload_ms  // 0 = the weights stay for the process's life
    i idle_since  // monotonic_ns at the end of the last job
    i loads  // 1 is the load before the port opened
    i unloads
    i load_ms  // the last load's wall time
    b serving
}

The server's state: the served model while serving, the request count and the weights' lease clock — a Locked global set at start, so the handlers, the model thread and the ticker all reach it in safe code.

: ( Locked EmServe ) g_em ( __em_serve_blank )

: EmCfg { String token String name }

The bearer token (empty = open) and the served model's name, set when serving starts and read by every request as copies. (String globals holding copies of the caller's strings could be replaced under a reader in another thread; the compiler refuses that in safe code.)

: ( Locked EmCfg ) g_em_cfg ( __em_cfg0 )

: i EM_JOB 0

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

One job per waiting request on a channel the model thread drains; the job carries the channel its answer goes back on, and the submitting fiber waits on that. A tick (the unload timer) and a stop travel the same queue, so the model thread waits on one thing.

: i EM_TICK 1

: i EM_STOP 2

: EmJob

: EmJob {
    i kind
    ( Vec i ) ids  // flat tokens for the whole request
    ( Vec i ) offs  // B+1 offsets into ids
    i nout  // B*dim floats come back
    b normalize
    ( Channel ( Vec f ) ) reply  // the embeddings, or an empty Vec on failure
}

: ( Channel EmJob ) g_q ( __em_queue_new )

@ embed_serve Embed e s name s host i port s token i unload_s → 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] [--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