nx_embedding.nx
buildroot/runtime/nx_embedding.nx
about
nx_embedding.nx -- token-ID -> embedding-vector lookup.
Closes the INPUT half of the inference pipeline. With this brick
shipped, the substrate can do:
token_ids -> nx_embedding_lookup -> [n_tok, hidden_dim] tensor
(the input to layer 0)
Composes only existing primitives:
NxTensor (L1 container)
NxQuantBlock / Q8 / Q4_K (L2 quantized weight containers)
LoopVerdict (control)
===== Embedding table layout ====================================
An embedding table is a [vocab_size, hidden_dim] matrix. Token ID
i selects row i. Two storage paths:
Dense: *NxTensor with shape [vocab_size, hidden_dim], i64
Q10 values. Simple memcpy of row i to out[t, :].
Quantized: *NxQuantBlock (q4_0) or *NxQuantBlockQ8 (q8_0) or
*NxQuantQ4K (q4_K). Row-by-row layout: vocab_size
blocks of (hidden_dim / BLOCK_SIZE) sub-blocks each.
Lookup dequantizes the requested row to Q10.
VRAM impact: storing a 50k x 4096 embedding table dense in Q10 is
50000 * 4096 * 8 = 1.6 GB. Quantized at q4_0 is ~150 MB (10.67x
compression). Per nx_quant_policy, embeddings stay at f16 by
default (small + sensitive), but for memory-constrained deployment
a quantized embedding table is a real ~1.5 GB save.
===== Out-of-vocab handling =====================================
Returns NX_EMB_ERR_OOV if any token ID >= vocab_size. Caller
MUST validate token IDs upstream (or accept the error verdict
and fall back). No silent UNK substitution -- explicit verdict.
genealogy_id: mikolov_2013_word2vec + bengio_2003_neural_lm +
rumelhart_hinton_williams_1986_distributed_repr
dependencies 6 imports · 3 importers
imports: nx_syscalls.nxnx_tier.nxnx_loop.nxnx_tensor.nxnx_quant_block.nxnx_quant_block_q8.nx
imported by: nx_actor_role_llm.nxnx_llm_run.nxnx_llm_run_v2.nx
call flow from main pre-order; caps 40 nodes / depth 6 declared; ↻ = already shown
structs
| none |
consts
| 55 | const NX_MAGIC_3000: i64 = 3000 |
| 56 | const NX_MAGIC_3031: i64 = 3031 |
| 57 | const NX_MAGIC_7000: i64 = 7000 |
| 58 | const NX_MAGIC_7015: i64 = 7015 |
| 62 | const NX_EMB_OK: nx_int = 0 |
| 63 | const NX_EMB_ERR_BAD_DIMS: nx_int = 1 |
| 64 | const NX_EMB_ERR_OOV: nx_int = 2 // token >= vocab_size |
| 65 | const NX_EMB_ERR_BAD_TID: nx_int = 3 // token < 0 |
| 66 | const NX_EMB_ERR_BAD_TABLE: nx_int = 4 |
| 67 | const NX_EMB_N_VERDICTS: nx_int = 5 |
functions
| 69 | func nx_emb_verdict_is_valid(v: nx_int) -> nx_int called by 1: main |
| 83 | func nx_embedding_lookup(table: *NxTensor, token_ids: *i64, n_tokens: nx_int, called by 4: _lr_step_embedmainnx_llm_generate_onenx_llm_generate_one_v2 calls 1: nx_t_is_contiguous |
| 139 | func nx_embedding_lookup_q8(table: *NxQuantBlockQ8, vocab: nx_int, hidden: nx_int, called by 1: main |
| 199 | func main() -> i64 |