static_embeddings 0.1.3 → 0.1.4

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 7ec20ad222235414c2e6f081a2cea5f932a7e3004cb78dcca5552c20208bd6a8
-  data.tar.gz: 80553b44c62b46625562accee74ac9bfe689400f05258e551dc9603105f64f7a
+  metadata.gz: f72b45c2d6fa5310e3b6d5a9b0a86ab0430fc4af71f6811f9aa02655de6272a8
+  data.tar.gz: 51f0bb96a3b0adc5f2799e1df425cf1747ca521d92ae90afba59e0014f1d4106
 SHA512:
-  metadata.gz: 7fc51c583e8d0f05e5e6b7181767450e4fc93601180a638444cb715be0321d993454003d93bd9bd9f7154e5d678b57dc0fdeba7d5f86c81154f0b61886898101
-  data.tar.gz: 1739214fe214cbb86a81e643878a371547b36365d408f2e81e7e95c6c0f0b66c2e214c839f41744347c5015c82832f5a34410e585eb504dab4fe2d3a106ff459
+  metadata.gz: 53630b742cd2723656b5d826a82f57042dcaab1e3259758aa6a71eac20ba387a23b3fb3905f0f36f3ce641b122229f366c9af1f029205b44cd82110f538a49cb
+  data.tar.gz: 199cd93e9898b7457d5e19b9cf1857b4d78ffc46777ca26778b094fa70ee02559dc90733e4e815eba74297ef9c066940792cef09e022801ef48c704e868fa8db

data/CHANGELOG.md

@@ -1,5 +1,90 @@
 # Changelog
 
+## 0.1.4 (unreleased)
+
+Hot-path work in the C runtime. Token ids, the pooling contract and f16
+rounding are unchanged: the ASCII parity sweep and the differential fuzz
+digest against `StaticEmbeddings::Reference` still match. The
+`model2vec.StaticModel` oracle was re-run for 0.1.4 against the same
+`potion-retrieval-32m` snapshot and `.semb` as 0.1.3: 31/31 id rows, 31/31
+vectors, `min_cosine=0.9999999999989528`, `max_abs_all=2.980232238769531e-07`
+— the same printed numbers as 0.1.3. L2 still accumulates sum-of-squares in
+double; 0.1.4 does that with SIMD pairwise adds.
+
+Measured on an M1 Pro, `potion-retrieval-32m`, `samples/run_all.sh`
+`DURATION=25`, production build (no alloc-stats). The 0.1.3 comparison run
+had `STATIC_EMBEDDINGS_ALLOC_STATS=1`, which the docs cost at about 3% on
+`embed` and 13% on `tokenize`; differences smaller than that, or inside the
+usual 5–10% laptop spread, are not claimed.
+
+### Changed
+
+- **L2 normalisation is a SIMD double reduction plus one scale into `out`.**
+  It used to `scale_copy` with `1.0` and then walk the vector twice in scalar
+  double. On this model `dim=512` and short in-vocabulary English, that L2
+  slot was about 19% of an ASCII `embed_batch` sample tree. The fused path
+  has a NEON/`__aarch64__` kernel and an SSE2 kernel; the scalar tail is
+  unchanged. ASCII batch went 307k → 338k texts/s on the M1 Pro (+10%). The
+  SSE2 path is compiled, not timed.
+
+- **The AArch64 f16 top-k kernel decodes 16 halves per iteration, not 8.**
+  The x86 F16C kernel is untouched. Same 50 000-row, `k=10`, ASCII matrix:
+
+  ```text
+                       f32        f16
+  x86-64, f16c       3.00 ms    1.65 ms    f16 1.8x faster   (0.1.2, F16C)
+  M1 Pro, neon-fp16  2.36 ms    1.81 ms    f16 1.3x faster   (0.1.4)
+  ```
+
+  The M1 Pro row used to read 2.29 / 2.88 ms, f16 1.3x *slower*. `docs/PERFORMANCE.md`
+  and the README follow the new numbers. On the lookup-table fallback f16 is
+  still usually slower than f32.
+
+- **Scratch buffers are reused per OS thread, and freed when the thread
+  exits.** `embed`, `tokenize`, `embed_batch` and `embed_token_ids` used to
+  `malloc`/`free` a scratch set on every call. A first cut of 0.1.4 used
+  `__thread` storage with no destructor, which leaked the heap buffers inside
+  the slot — bounded for a Puma pool, unbounded on the fiber-scheduler path
+  that `rb_thread_create`s one OS thread per large call. The slot is now a
+  `pthread_key` / `FlsAlloc` value whose destructor frees it. On release the
+  slot is trimmed back to the reserve sizes, so one `max_tokens: false`
+  document does not pin megabytes on that thread until process exit.
+  `tokenize` and the small `embed` path release the slot through `rb_ensure`,
+  so `rb_ary_push` raising cannot leave `in_use` stuck. `memory_smoke` calls
+  acquire/release so ASan/valgrind actually see this code.
+
+- **ASCII WordPiece copies an all-ASCII word as bytes** instead of
+  encoding UTF-8 through the codepoint buffer, then hashes and splits on the
+  trie with identity offsets. Output is the same ids. On the OOV tokenize
+  probes this is inside the run-to-run floor.
+
+- **Latin-1 (`U+0080..U+00FF`) lower/NFD/Mn/P/C/Zs tables are built at
+  load.** Codepoints below 256 skip the binary search into the mmap'd maps.
+  ASCII (`< 0x80`) still uses the existing 128-entry class table. Not visible
+  as texts/s on the English batch corpus.
+
+- **Vocabulary compare is an inline unaligned equality test** instead of
+  libc `memcmp`. Trie child lookup special-cases one- and two-edge nodes
+  before the linear/binary search. Same ids; no measured throughput claim.
+
+Two things were measured and not shipped. Unrolling `add_row` to 32 floats
+and stretching the prefetch distance to 8 cost about 13% on
+`random_pooling_hot_path` — that loop is already memory-bound on random 2 KiB
+row gathers. Walking the ASCII trie without `cps2` (`start += matched_len`)
+cost about 38% on the OOV tokenize probes.
+
+### Added
+
+- **Batch samples print `mean_tokens_per_text`, `unk_ratio` and
+  `tokens_per_sec`.** 307k vs 103k texts/s is not readable without token
+  density: on this corpus ASCII is 31 tokens/text and 0 unk, hashes 73 and 0,
+  unicode 68 and 4.4% unk.
+
+- **`test/scratch_tls_test.rb`.** When the instrumented allocator is loaded:
+  scratch bytes plateau across eight generations of helper threads, and a
+  large `max_tokens: false` embed must not leave the calling thread's slot
+  pinned. The `alloc_stats` CI job is what actually runs it.
+
 ## 0.1.3 (unreleased)
 
 Hardening and tooling, mostly borrowed from the sibling `tg_geometry` gem after

data/README.md

@@ -132,10 +132,11 @@ in float32. For a 512-dimensional model one vector goes from 2048 to 1024 bytes.
 It is accepted by `embed`, `embed_batch`, `embed_with_stats`, `embed_token_ids`,
 `embed_token_ids_with_stats`, `pack`, `unpack`, and the top-k helpers.
 
-Choose `f16` for the RAM and storage it saves, not for speed. It halves the
-bytes a top-k scan streams, but every row still has to be decoded before
-scoring, and which effect wins is a property of the machine — this repository's
-own samples show `f16` winning on x86-64 F16C and losing on M1 Pro. See
+Choose `f16` for the RAM and storage it saves; speed is a property of the
+machine and of the decode kernel. It halves the bytes a top-k scan streams, but
+every row still has to be decoded before scoring — this repository's own
+samples show `f16` winning on x86-64 F16C and, as of 0.1.4, on M1 Pro neon-fp16
+as well. On the lookup-table fallback it usually loses. See
 `docs/PERFORMANCE.md` before assuming either.
 
 `StaticEmbeddings.simd_backend` reports the live kernel: `"neon-fp16"`,
@@ -323,7 +324,7 @@ cc -O2 -std=c99 -Wall -Wextra -Iext/static_embeddings \
   ext/static_embeddings/se_unicode.c \
   ext/static_embeddings/se_tokenizer.c \
   ext/static_embeddings/se_embed.c \
-  -lm -o tmp/memory_smoke
+  -lm -pthread -o tmp/memory_smoke
 ./tmp/memory_smoke tmp/test-tiny.semb
 ```
 

data/docs/ARCHITECTURE.md

@@ -174,10 +174,13 @@ storage/transport choice, not a different model. `embed_array` and
 
 The f32 compute path uses small explicit SIMD kernels where they keep the format
 simple: row accumulation and scaling use NEON/SSE when the compiler target
-exposes them, with scalar fallback everywhere else. `dot_top_k`/`cosine_top_k`
+exposes them, L2 sum-of-squares uses SIMD double (NEON on AArch64, SSE2 on
+x86) so the inverse stays a `1/sqrt` of a double accumulation, and there is a
+scalar fallback everywhere else. `dot_top_k`/`cosine_top_k`
 share the SIMD dot kernel for `format: :f32`; `cosine_top_k` additionally
 accumulates each row's sum of squares so it can divide by the true norms.
-`format: :f16` top-k decodes half components on the fly. What is still
+`format: :f16` top-k decodes half components on the fly (16-wide on AArch64
+NEON-FP16, 16-wide on x86 F16C). What is still
 deliberately absent is a dependency on BLAS: pooling is gathering random
 embedding rows, not a dense matrix multiply.
 
@@ -293,8 +296,13 @@ cadence in bytes so a long ASCII run is no less interruptible. The pooling loop
 checks every 256 rows, and the top-k scan every 1024 rows. `unblock_cancel` sets that flag when Ruby
 interrupts a GVL-free region.
 
-No global mutable state exists in C. The model is immutable after load and
-scratch buffers are per call.
+No global mutable state exists in C. The model is immutable after load.
+Scratch buffers are reused per OS thread via `pthread_key` / `FlsAlloc`, and
+the destructor frees them when the thread exits — including the short-lived
+threads the fiber scheduler path creates per large call. On release the slot
+is trimmed back to the reserve sizes, so a single unlimited-`max_tokens`
+document does not pin its working set until the thread dies. A nested call on
+the same thread allocates a one-off heap scratch that is freed on release.
 
 ## Where this fits in a RAG pipeline
 

data/docs/LIMITATIONS.md

@@ -57,8 +57,10 @@ the entry says whether that is a decision or just unfinished work.
 ## Storage formats
 
 - **`f32` and `f16` only.** No int8, no binary quantisation.
-- **`f16` is a storage trade-off, not a speed feature.** See
-  `docs/PERFORMANCE.md`; it wins on some CPUs and loses on others.
+- **`f16` is a storage encoding, not a different model.** Top-k over an f16
+  matrix can be faster or slower than f32 depending on the decode kernel; see
+  `docs/PERFORMANCE.md`. `embed_batch(format: :f16)` still encodes the blob
+  with the scalar converter.
 - **`f16` rounding is half-up, not ties-to-even.** Blobs written by this gem can
   differ from NumPy or PyTorch by one ULP on exact halfway values.
 

data/docs/MODEL_AUDIT.md

@@ -12,26 +12,28 @@ runtime half of it even though the file bytes are untouched.
 |---|---|---|
 | 0.1.1 | `parity OK`, recorded below | superseded |
 | 0.1.2 | **not re-run** | superseded before release |
-| 0.1.3 | `parity OK`, recorded below | current |
+| 0.1.3 | `parity OK` | tokenizer / pooling contract; same numbers as 0.1.4 |
+| 0.1.4 | `parity OK`, recorded below | current |
 
 0.1.2 changed `is_control()`, which changes token ids for any input containing
 `U+007F`. `StaticEmbeddings::Reference` cannot settle whether the new behaviour
 matches HuggingFace, because it is an implementation twin of the C runtime
 written in this repository. The 0.1.3 release candidate was checked against a
 fresh upstream `model2vec.StaticModel` oracle that includes DEL, control
-characters, Unicode, OOV, long-word and truncation rows.
+characters, Unicode, OOV, long-word and truncation rows. 0.1.4 re-ran that
+oracle against the SIMD L2 runtime on the same snapshot and `.semb`.
 
 Source model:
 
 - Hugging Face repository: `minishlab/potion-retrieval-32M`
 - Hugging Face snapshot: `6fc8051fab2a1e0ee76689cf08c853792ac285e7`
 - Oracle implementation: `model2vec.StaticModel.from_pretrained`
-- Python package: `model2vec 0.9.0`
+- Python package: `model2vec 0.9.0` (`tokenizers 0.23.1`, `numpy 2.5.2`)
 - Oracle file: `tmp/model2vec_oracle.json`
 - Oracle rows in this recorded run: `31`
 - Oracle dimension: `512`
 - Oracle max length: `512`
-- Runtime at time of this record: `static_embeddings 0.1.3`
+- Runtime at time of this record: `static_embeddings 0.1.4`
 
 Converted `.semb`:
 
@@ -44,7 +46,8 @@ Converted `.semb`:
 Parity command:
 
 ```bash
-python tools/model2vec_oracle.py minishlab/potion-retrieval-32M \
+python tools/model2vec_oracle.py \
+  ~/.cache/huggingface/hub/models--minishlab--potion-retrieval-32M/snapshots/6fc8051fab2a1e0ee76689cf08c853792ac285e7 \
   --out tmp/model2vec_oracle.json
 
 bundle exec rake parity \
@@ -52,7 +55,7 @@ bundle exec rake parity \
   ORACLE=tmp/model2vec_oracle.json
 ```
 
-Parity result:
+Parity result (0.1.4, 2026-08-31):
 
 ```text
 rows=31
@@ -64,6 +67,11 @@ vector_failures=[]
 parity OK
 ```
 
+The printed `min_cosine` and `max_abs_all` are identical to the 0.1.3 record.
+Worst vector row is still idx=29 (`"word" + 400 spaces`, 80 800 bytes) at
+`max_abs=2.9802322e-07`. SIMD pairwise double L2 did not move the oracle
+agreement on this corpus.
+
 Decision:
 
 - Tokenization parity: pass
@@ -75,6 +83,6 @@ Decision:
 - Long input / truncation behavior: pass
 - DEL / control-character behavior: pass
 
-Accepted for runtime and benchmark use **under 0.1.3**. The `.semb` file is
-unchanged and its SHA256 still matches; the runtime half of the record was
-refreshed after the DEL fix.
+Accepted for runtime and benchmark use **under 0.1.4**. The `.semb` file is
+unchanged and its SHA256 still matches. The runtime half of the record was
+refreshed after the 0.1.4 L2 kernel change.

data/docs/PERFORMANCE.md

@@ -22,33 +22,42 @@ ruby tools/benchmark.rb [model.semb]
 
 ## Where the time actually goes
 
-Tokenization dominates; pooling is a handful of row reads and adds. That is why
-the tokenizer is the file to profile, and why SIMD in the pooling loop is not
-where the wins are — that loop is memory-bound on random row lookups and the
-compiler already vectorises it.
-
-Out-of-vocabulary text is the main cliff. In-vocabulary words hit the
-vocabulary hash directly; everything else falls through to trie-driven subword
-splitting, which on a synthetic all-OOV corpus measured several times the
-per-text cost. Non-English input on an English model pays this on top of the
-`[UNK]` quality problem.
+Which file to profile depends on the model and the text. On a short-dimension
+demo model, or on synthetic all-OOV input, tokenization dominates. On
+`potion-retrieval-32m` (`dim=512`) with in-vocabulary English, pooling is the
+larger share: each token is a random 2 KiB row add, and before 0.1.4 the scalar
+L2 pass showed up as about 19% of an ASCII `embed_batch` sample tree. SIMD in
+that L2 slot was a real win; another unroll of the add loop was not — it is
+memory-bound on those gathers, and the compiler already vectorises it.
+
+Out-of-vocabulary text is the remaining tokenizer cliff. In-vocabulary words
+hit the vocabulary hash directly; everything else falls through to trie-driven
+subword splitting, which on a synthetic all-OOV corpus measured several times
+the per-text cost. Non-English input on an English model pays this on top of
+the `[UNK]` quality problem. Quote `mean_tokens_per_text` and `unk_ratio` from
+the batch samples before comparing texts/s across modes: hashes are slower
+than ASCII on this corpus mostly because they emit ~73 tokens/text against
+~31, not only because they miss the hash.
 
 ## `f16` is a storage trade-off
 
-`format: :f16` halves the bytes a top-k scan streams and doubles the decode
-work. Which one wins is a property of the machine, and this repository's own
-sample runs disagree with each other — same corpus, same `k`, native decode
-kernel in both cases:
+`format: :f16` halves the bytes a top-k scan streams and costs a decode per
+row. Which one wins is a property of the machine and of the decode kernel, and
+this repository's own sample runs disagree with each other — same corpus, same
+`k`, native decode kernel in both cases:
 
 ```text
                      f32        f16
-x86-64, f16c       3.00 ms    1.65 ms    f16 1.8x faster
-M1 Pro, neon-fp16  2.29 ms    2.88 ms    f16 1.3x slower
+x86-64, f16c       3.00 ms    1.65 ms    f16 1.8x faster   (F16C kernel, unchanged)
+M1 Pro, neon-fp16  2.36 ms    1.81 ms    f16 1.3x faster   (0.1.4, 16-wide NEON)
 ```
 
-Neither ratio transfers. Confirm `StaticEmbeddings.simd_backend`, then measure
-on the hardware you will run on. On the lookup-table fallback `f16` is usually
-slower than `f32`.
+The M1 Pro row in 0.1.3 was 2.29 / 2.88 ms, f16 1.3x *slower*, with an 8-wide
+decode. Neither ratio transfers. Confirm `StaticEmbeddings.simd_backend`, then
+measure on the hardware you will run on. On the lookup-table fallback `f16` is
+usually slower than `f32`. `embed_batch(format: :f16)` still encodes the
+returned blob with the scalar half converter, so a batch that is not a top-k
+scan can still be slower than `f32` on the same machine.
 
 ## C allocation counters
 

data/ext/static_embeddings/se_embed.c

@@ -8,6 +8,7 @@
 #define SE_HAVE_NEON 1
 #elif defined(__SSE__)
 #include <xmmintrin.h>
+#include <emmintrin.h>
 #define SE_HAVE_SSE 1
 #endif
 
@@ -71,15 +72,49 @@ static void scale_copy(float *out, const float *acc, uint32_t dim, float inv) {
         out[j] = acc[j] * inv;
 }
 
-void se_l2_normalize(float *vec, uint32_t dim) {
+static double sum_sq_d(const float *vec, uint32_t dim) {
+    uint32_t j = 0;
     double sum = 0.0;
-    for (uint32_t j = 0; j < dim; j++)
+#if defined(SE_HAVE_NEON) && defined(__aarch64__)
+    float64x2_t s0 = vdupq_n_f64(0.0);
+    float64x2_t s1 = vdupq_n_f64(0.0);
+    for (; j + 7 < dim; j += 8) {
+        float32x4_t a = vld1q_f32(vec + j);
+        float32x4_t b = vld1q_f32(vec + j + 4);
+        float64x2_t a0 = vcvt_f64_f32(vget_low_f32(a));
+        float64x2_t a1 = vcvt_f64_f32(vget_high_f32(a));
+        float64x2_t b0 = vcvt_f64_f32(vget_low_f32(b));
+        float64x2_t b1 = vcvt_f64_f32(vget_high_f32(b));
+        s0 = vaddq_f64(s0, vmulq_f64(a0, a0));
+        s1 = vaddq_f64(s1, vmulq_f64(a1, a1));
+        s0 = vaddq_f64(s0, vmulq_f64(b0, b0));
+        s1 = vaddq_f64(s1, vmulq_f64(b1, b1));
+    }
+    sum = vaddvq_f64(vaddq_f64(s0, s1));
+#elif defined(SE_HAVE_SSE)
+    __m128d s0 = _mm_setzero_pd();
+    __m128d s1 = _mm_setzero_pd();
+    for (; j + 3 < dim; j += 4) {
+        __m128 v = _mm_loadu_ps(vec + j);
+        __m128d lo = _mm_cvtps_pd(v);
+        __m128d hi = _mm_cvtps_pd(_mm_movehl_ps(v, v));
+        s0 = _mm_add_pd(s0, _mm_mul_pd(lo, lo));
+        s1 = _mm_add_pd(s1, _mm_mul_pd(hi, hi));
+    }
+    double tmp[2];
+    _mm_storeu_pd(tmp, _mm_add_pd(s0, s1));
+    sum = tmp[0] + tmp[1];
+#endif
+    for (; j < dim; j++)
         sum += (double)vec[j] * (double)vec[j];
+    return sum;
+}
+
+void se_l2_normalize(float *vec, uint32_t dim) {
+    double sum = sum_sq_d(vec, dim);
     if (sum <= 0.0)
         return;
-    float inv = (float)(1.0 / sqrt(sum));
-    for (uint32_t j = 0; j < dim; j++)
-        vec[j] *= inv;
+    scale_copy(vec, vec, dim, (float)(1.0 / sqrt(sum)));
 }
 
 static se_status_t embed_ids_core(const se_model_t *model, se_scratch_t *sc, const uint32_t *ids,
@@ -128,8 +163,12 @@ static se_status_t embed_ids_core(const se_model_t *model, se_scratch_t *sc, con
     }
 
     if (model->meta.normalization_type == SE_NORMALIZATION_L2) {
-        scale_copy(out, acc, dim, 1.0f);
-        se_l2_normalize(out, dim);
+        double sum = sum_sq_d(acc, dim);
+        if (sum <= 0.0) {
+            memset(out, 0, (size_t)dim * sizeof(float));
+            return SE_OK;
+        }
+        scale_copy(out, acc, dim, (float)(1.0 / sqrt(sum)));
     } else {
         const float inv = 1.0f / (float)used;
         scale_copy(out, acc, dim, inv);

data/ext/static_embeddings/se_f16.c

@@ -221,17 +221,31 @@ static float dot_product_f16_neon(const float *q, const uint8_t *row, size_t dim
     size_t j = 0;
     float32x4_t a0 = vdupq_n_f32(0.0f);
     float32x4_t a1 = vdupq_n_f32(0.0f);
+    float32x4_t a2 = vdupq_n_f32(0.0f);
+    float32x4_t a3 = vdupq_n_f32(0.0f);
+    for (; j + 15 < dim; j += 16) {
+        float16x8_t h0 =
+            vreinterpretq_f16_u16(vld1q_u16((const uint16_t *)(const void *)(row + j * 2)));
+        float16x8_t h1 =
+            vreinterpretq_f16_u16(vld1q_u16((const uint16_t *)(const void *)(row + (j + 8) * 2)));
+        float32x4_t r0 = vcvt_f32_f16(vget_low_f16(h0));
+        float32x4_t r1 = vcvt_f32_f16(vget_high_f16(h0));
+        float32x4_t r2 = vcvt_f32_f16(vget_low_f16(h1));
+        float32x4_t r3 = vcvt_f32_f16(vget_high_f16(h1));
+        a0 = vmlaq_f32(a0, vld1q_f32(q + j), r0);
+        a1 = vmlaq_f32(a1, vld1q_f32(q + j + 4), r1);
+        a2 = vmlaq_f32(a2, vld1q_f32(q + j + 8), r2);
+        a3 = vmlaq_f32(a3, vld1q_f32(q + j + 12), r3);
+    }
     for (; j + 7 < dim; j += 8) {
         float16x4_t h0 =
             vreinterpret_f16_u16(vld1_u16((const uint16_t *)(const void *)(row + j * 2)));
         float16x4_t h1 =
             vreinterpret_f16_u16(vld1_u16((const uint16_t *)(const void *)(row + (j + 4) * 2)));
-        float32x4_t r0 = vcvt_f32_f16(h0);
-        float32x4_t r1 = vcvt_f32_f16(h1);
-        a0 = vmlaq_f32(a0, vld1q_f32(q + j), r0);
-        a1 = vmlaq_f32(a1, vld1q_f32(q + j + 4), r1);
+        a0 = vmlaq_f32(a0, vld1q_f32(q + j), vcvt_f32_f16(h0));
+        a1 = vmlaq_f32(a1, vld1q_f32(q + j + 4), vcvt_f32_f16(h1));
     }
-    float32x4_t sumv = vaddq_f32(a0, a1);
+    float32x4_t sumv = vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3));
     float dot = vaddvq_f32(sumv);
     for (; j < dim; j++)
         dot += q[j] * se_f16_bits_to_float(read_u16le(row + j * 2));
@@ -243,8 +257,30 @@ static float dot_and_row_sq_f16_neon(const float *q, const uint8_t *row, size_t
     size_t j = 0;
     float32x4_t d0 = vdupq_n_f32(0.0f);
     float32x4_t d1 = vdupq_n_f32(0.0f);
+    float32x4_t d2 = vdupq_n_f32(0.0f);
+    float32x4_t d3 = vdupq_n_f32(0.0f);
     float32x4_t s0 = vdupq_n_f32(0.0f);
     float32x4_t s1 = vdupq_n_f32(0.0f);
+    float32x4_t s2 = vdupq_n_f32(0.0f);
+    float32x4_t s3 = vdupq_n_f32(0.0f);
+    for (; j + 15 < dim; j += 16) {
+        float16x8_t h0 =
+            vreinterpretq_f16_u16(vld1q_u16((const uint16_t *)(const void *)(row + j * 2)));
+        float16x8_t h1 =
+            vreinterpretq_f16_u16(vld1q_u16((const uint16_t *)(const void *)(row + (j + 8) * 2)));
+        float32x4_t r0 = vcvt_f32_f16(vget_low_f16(h0));
+        float32x4_t r1 = vcvt_f32_f16(vget_high_f16(h0));
+        float32x4_t r2 = vcvt_f32_f16(vget_low_f16(h1));
+        float32x4_t r3 = vcvt_f32_f16(vget_high_f16(h1));
+        d0 = vmlaq_f32(d0, vld1q_f32(q + j), r0);
+        d1 = vmlaq_f32(d1, vld1q_f32(q + j + 4), r1);
+        d2 = vmlaq_f32(d2, vld1q_f32(q + j + 8), r2);
+        d3 = vmlaq_f32(d3, vld1q_f32(q + j + 12), r3);
+        s0 = vmlaq_f32(s0, r0, r0);
+        s1 = vmlaq_f32(s1, r1, r1);
+        s2 = vmlaq_f32(s2, r2, r2);
+        s3 = vmlaq_f32(s3, r3, r3);
+    }
     for (; j + 7 < dim; j += 8) {
         float16x4_t h0 =
             vreinterpret_f16_u16(vld1_u16((const uint16_t *)(const void *)(row + j * 2)));
@@ -257,8 +293,8 @@ static float dot_and_row_sq_f16_neon(const float *q, const uint8_t *row, size_t
         s0 = vmlaq_f32(s0, r0, r0);
         s1 = vmlaq_f32(s1, r1, r1);
     }
-    float dot = vaddvq_f32(vaddq_f32(d0, d1));
-    float row_sq = vaddvq_f32(vaddq_f32(s0, s1));
+    float dot = vaddvq_f32(vaddq_f32(vaddq_f32(d0, d1), vaddq_f32(d2, d3)));
+    float row_sq = vaddvq_f32(vaddq_f32(vaddq_f32(s0, s1), vaddq_f32(s2, s3)));
     for (; j < dim; j++) {
         float r = se_f16_bits_to_float(read_u16le(row + j * 2));
         dot += q[j] * r;

data/ext/static_embeddings/se_format.c

@@ -61,6 +61,34 @@ static uint32_t hash_bytes_continue(uint32_t h, const uint8_t *data, size_t len)
     return h;
 }
 
+static int se_memeq(const uint8_t *a, const uint8_t *b, size_t n) {
+    while (n >= 8) {
+        uint64_t ua, ub;
+        memcpy(&ua, a, 8);
+        memcpy(&ub, b, 8);
+        if (ua != ub)
+            return 0;
+        a += 8;
+        b += 8;
+        n -= 8;
+    }
+    if (n >= 4) {
+        uint32_t ua, ub;
+        memcpy(&ua, a, 4);
+        memcpy(&ub, b, 4);
+        if (ua != ub)
+            return 0;
+        a += 4;
+        b += 4;
+        n -= 4;
+    }
+    while (n--) {
+        if (*a++ != *b++)
+            return 0;
+    }
+    return 1;
+}
+
 int se_vocab_lookup_piece(const se_model_t *model, const uint8_t *prefix, size_t prefix_len,
                           const uint8_t *bytes, size_t len, uint32_t *id_out) {
     if (prefix_len > UINT32_MAX || len > UINT32_MAX || prefix_len > UINT32_MAX - len)
@@ -79,9 +107,9 @@ int se_vocab_lookup_piece(const se_model_t *model, const uint8_t *prefix, size_t
         if (slot->token_id == SE_SLOT_EMPTY)
             return 0;
         if (slot->hash == h && slot->str_len == total_len) {
-            const char *token = model->vocab_strings + slot->str_off;
-            if ((prefix_len == 0 || memcmp(token, prefix, prefix_len) == 0) &&
-                (len == 0 || memcmp(token + prefix_len, bytes, len) == 0)) {
+            const uint8_t *token = (const uint8_t *)(model->vocab_strings + slot->str_off);
+            if ((prefix_len == 0 || se_memeq(token, prefix, prefix_len)) &&
+                (len == 0 || se_memeq(token + prefix_len, bytes, len))) {
                 *id_out = slot->token_id;
                 return 1;
             }
@@ -103,6 +131,17 @@ static inline uint32_t trie_find_child(const se_trie_t *trie, const se_trie_node
     if (count == 0)
         return SE_SLOT_EMPTY;
 
+    if (count == 1)
+        return trie->edges[start].byte == byte ? trie->edges[start].child : SE_SLOT_EMPTY;
+
+    if (count == 2) {
+        if (trie->edges[start].byte == byte)
+            return trie->edges[start].child;
+        if (trie->edges[start + 1u].byte == byte)
+            return trie->edges[start + 1u].child;
+        return SE_SLOT_EMPTY;
+    }
+
     if (count <= 8) {
         for (uint32_t i = 0; i < count; i++) {
             const se_trie_edge_t *edge = &trie->edges[start + i];
@@ -342,6 +381,37 @@ static se_status_t validate_norm_ranges(const se_range_t *ranges, uint32_t count
     return SE_OK;
 }
 
+static void fill_map256(const se_map_entry_t **out, const se_map_entry_t *entries, uint32_t count) {
+    memset(out, 0, 256 * sizeof(*out));
+    for (uint32_t i = 0; i < count; i++) {
+        if (entries[i].cp < 256)
+            out[entries[i].cp] = &entries[i];
+    }
+}
+
+static void fill_range256(uint8_t *out, const se_range_t *ranges, uint32_t count) {
+    memset(out, 0, 256);
+    for (uint32_t i = 0; i < count; i++) {
+        uint32_t lo = ranges[i].lo;
+        uint32_t hi = ranges[i].hi;
+        if (lo > 255)
+            continue;
+        if (hi > 255)
+            hi = 255;
+        for (uint32_t cp = lo; cp <= hi; cp++)
+            out[cp] = 1;
+    }
+}
+
+static void prepare_norm_fast_tables(se_model_t *model) {
+    fill_map256(model->norm.lower256, model->norm.lower, model->norm.lower_count);
+    fill_map256(model->norm.nfd256, model->norm.nfd, model->norm.nfd_count);
+    fill_range256(model->norm.mn256, model->norm.mn, model->norm.mn_count);
+    fill_range256(model->norm.punct256, model->norm.punct, model->norm.punct_count);
+    fill_range256(model->norm.control256, model->norm.control, model->norm.control_count);
+    fill_range256(model->norm.whitespace256, model->norm.whitespace, model->norm.whitespace_count);
+}
+
 static se_status_t parse_norm_tables(se_model_t *model, const uint8_t *base, struct section sec,
                                      se_error_t *err) {
     memset(&model->norm, 0, sizeof(model->norm));
@@ -417,8 +487,13 @@ static se_status_t parse_norm_tables(se_model_t *model, const uint8_t *base, str
     rc = validate_norm_ranges(model->norm.control, model->norm.control_count, "control", err);
     if (rc != SE_OK)
         return rc;
-    return validate_norm_ranges(model->norm.whitespace, model->norm.whitespace_count, "whitespace",
-                                err);
+    rc = validate_norm_ranges(model->norm.whitespace, model->norm.whitespace_count, "whitespace",
+                              err);
+    if (rc != SE_OK)
+        return rc;
+
+    prepare_norm_fast_tables(model);
+    return SE_OK;
 }
 
 static se_status_t validate_vocab_hash(const se_model_t *model, struct section vocab_strings,

data/ext/static_embeddings/se_internal.h

@@ -136,6 +136,12 @@ typedef struct {
     uint32_t control_count;
     const se_range_t *whitespace;
     uint32_t whitespace_count;
+    const se_map_entry_t *lower256[256];
+    const se_map_entry_t *nfd256[256];
+    uint8_t mn256[256];
+    uint8_t punct256[256];
+    uint8_t control256[256];
+    uint8_t whitespace256[256];
 } se_norm_tables_t;
 
 typedef struct {
@@ -382,6 +388,9 @@ static inline int se_is_ascii_boundary(uint32_t cp) {
 void se_scratch_init(se_scratch_t *s);
 void se_scratch_free(se_scratch_t *s);
 int se_scratch_reserve(se_scratch_t *s, uint32_t dim);
+se_scratch_t *se_scratch_acquire(uint32_t dim);
+void se_scratch_release(se_scratch_t *s);
+void se_scratch_drop_thread(void);
 
 size_t se_prefix_boundary_len(const se_model_t *model, const uint8_t *input, size_t input_len,
                               size_t target, size_t backscan);

data/ext/static_embeddings/se_tokenizer.c

@@ -4,6 +4,18 @@
 #include <stdlib.h>
 #include <string.h>
 
+#ifndef _WIN32
+#include <pthread.h>
+#else
+#ifndef WIN32_LEAN_AND_MEAN
+#define WIN32_LEAN_AND_MEAN
+#endif
+#ifndef _WIN32_WINNT
+#define _WIN32_WINNT 0x0600
+#endif
+#include <windows.h>
+#endif
+
 #define SE_CANCEL_CHECK_MASK 0x3ffu
 
 void se_scratch_init(se_scratch_t *s) {
@@ -75,14 +87,18 @@ static int push_id(se_scratch_t *sc, size_t *n_ids, uint32_t id) {
     return 1;
 }
 
+#define SE_SCRATCH_CPS_KEEP   256
+#define SE_SCRATCH_IDS_KEEP   512
+#define SE_SCRATCH_BYTES_KEEP 1024
+
 int se_scratch_reserve(se_scratch_t *s, uint32_t dim) {
-    if (!grow_u32(&s->cps, &s->cps_cap, 256))
+    if (!grow_u32(&s->cps, &s->cps_cap, SE_SCRATCH_CPS_KEEP))
         return 0;
-    if (!grow_u32(&s->cps2, &s->cps2_cap, 256))
+    if (!grow_u32(&s->cps2, &s->cps2_cap, SE_SCRATCH_CPS_KEEP))
         return 0;
-    if (!grow_u32(&s->ids, &s->ids_cap, 256))
+    if (!grow_u32(&s->ids, &s->ids_cap, SE_SCRATCH_IDS_KEEP))
         return 0;
-    if (!grow_bytes(&s->bytes, &s->bytes_cap, 256))
+    if (!grow_bytes(&s->bytes, &s->bytes_cap, SE_SCRATCH_BYTES_KEEP))
         return 0;
     if (!grow_float(&s->acc, &s->acc_cap, dim))
         return 0;
@@ -90,21 +106,223 @@ int se_scratch_reserve(se_scratch_t *s, uint32_t dim) {
     return 1;
 }
 
+static int shrink_u32(uint32_t **buf, size_t *cap, size_t keep) {
+    size_t bytes = 0;
+    void *p;
+
+    if (*cap <= keep)
+        return 1;
+    if (!*buf) {
+        *cap = 0;
+        return 1;
+    }
+    if (!se_array_bytes(keep, sizeof(uint32_t), &bytes))
+        return 0;
+    p = se_realloc(SE_ALLOC_SCRATCH, *buf, bytes);
+    if (!p)
+        return 0;
+    *buf = (uint32_t *)p;
+    *cap = keep;
+    return 1;
+}
+
+static int shrink_bytes(uint8_t **buf, size_t *cap, size_t keep) {
+    size_t bytes = 0;
+    void *p;
+
+    if (*cap <= keep)
+        return 1;
+    if (!*buf) {
+        *cap = 0;
+        return 1;
+    }
+    if (!se_array_bytes(keep, sizeof(uint8_t), &bytes))
+        return 0;
+    p = se_realloc(SE_ALLOC_SCRATCH, *buf, bytes);
+    if (!p)
+        return 0;
+    *buf = (uint8_t *)p;
+    *cap = keep;
+    return 1;
+}
+
+static void se_scratch_trim(se_scratch_t *s) {
+    (void)shrink_u32(&s->cps, &s->cps_cap, SE_SCRATCH_CPS_KEEP);
+    (void)shrink_u32(&s->cps2, &s->cps2_cap, SE_SCRATCH_CPS_KEEP);
+    (void)shrink_u32(&s->ids, &s->ids_cap, SE_SCRATCH_IDS_KEEP);
+    (void)shrink_bytes(&s->bytes, &s->bytes_cap, SE_SCRATCH_BYTES_KEEP);
+}
+
+typedef struct {
+    se_scratch_t scratch;
+    int in_use;
+} se_tls_scratch_t;
+
+static void se_scratch_tls_dtor(void *p) {
+    se_tls_scratch_t *tls = (se_tls_scratch_t *)p;
+    if (!tls)
+        return;
+    se_scratch_free(&tls->scratch);
+    se_free(tls);
+}
+
+#ifndef _WIN32
+static pthread_key_t se_scratch_key;
+static pthread_once_t se_scratch_once = PTHREAD_ONCE_INIT;
+
+static void se_scratch_key_init(void) {
+    (void)pthread_key_create(&se_scratch_key, se_scratch_tls_dtor);
+}
+
+static se_tls_scratch_t *se_scratch_tls_get(void) {
+    (void)pthread_once(&se_scratch_once, se_scratch_key_init);
+    return (se_tls_scratch_t *)pthread_getspecific(se_scratch_key);
+}
+
+static se_tls_scratch_t *se_scratch_tls_slot(void) {
+    se_tls_scratch_t *tls = se_scratch_tls_get();
+    if (tls)
+        return tls;
+
+    tls = (se_tls_scratch_t *)se_malloc(SE_ALLOC_SCRATCH, sizeof(*tls));
+    if (!tls)
+        return NULL;
+    memset(tls, 0, sizeof(*tls));
+    se_scratch_init(&tls->scratch);
+    if (pthread_setspecific(se_scratch_key, tls) != 0) {
+        se_free(tls);
+        return NULL;
+    }
+    return tls;
+}
+
+static void se_scratch_tls_clear(void) {
+    se_tls_scratch_t *tls = se_scratch_tls_get();
+    if (!tls)
+        return;
+    (void)pthread_setspecific(se_scratch_key, NULL);
+    se_scratch_tls_dtor(tls);
+}
+#else
+static DWORD se_fls_index = FLS_OUT_OF_INDEXES;
+static INIT_ONCE se_fls_once = INIT_ONCE_STATIC_INIT;
+
+static VOID WINAPI se_scratch_fls_dtor(PVOID p) {
+    se_scratch_tls_dtor(p);
+}
+
+static BOOL CALLBACK se_scratch_fls_init(PINIT_ONCE once, PVOID param, PVOID *ctx) {
+    (void)once;
+    (void)param;
+    (void)ctx;
+    se_fls_index = FlsAlloc(se_scratch_fls_dtor);
+    return se_fls_index != FLS_OUT_OF_INDEXES;
+}
+
+static se_tls_scratch_t *se_scratch_tls_get(void) {
+    if (se_fls_index == FLS_OUT_OF_INDEXES)
+        return NULL;
+    return (se_tls_scratch_t *)FlsGetValue(se_fls_index);
+}
+
+static se_tls_scratch_t *se_scratch_tls_slot(void) {
+    if (!InitOnceExecuteOnce(&se_fls_once, se_scratch_fls_init, NULL, NULL) ||
+        se_fls_index == FLS_OUT_OF_INDEXES)
+        return NULL;
+
+    se_tls_scratch_t *tls = se_scratch_tls_get();
+    if (tls)
+        return tls;
+
+    tls = (se_tls_scratch_t *)se_malloc(SE_ALLOC_SCRATCH, sizeof(*tls));
+    if (!tls)
+        return NULL;
+    memset(tls, 0, sizeof(*tls));
+    se_scratch_init(&tls->scratch);
+    if (!FlsSetValue(se_fls_index, tls)) {
+        se_free(tls);
+        return NULL;
+    }
+    return tls;
+}
+
+static void se_scratch_tls_clear(void) {
+    se_tls_scratch_t *tls = se_scratch_tls_get();
+    if (!tls)
+        return;
+    (void)FlsSetValue(se_fls_index, NULL);
+    se_scratch_tls_dtor(tls);
+}
+#endif
+
+se_scratch_t *se_scratch_acquire(uint32_t dim) {
+    se_tls_scratch_t *tls = se_scratch_tls_slot();
+    if (!tls)
+        return NULL;
+
+    if (!tls->in_use) {
+        if (!se_scratch_reserve(&tls->scratch, dim))
+            return NULL;
+        tls->in_use = 1;
+        return &tls->scratch;
+    }
+
+    se_scratch_t *heap = (se_scratch_t *)se_malloc(SE_ALLOC_SCRATCH, sizeof(*heap));
+    if (!heap)
+        return NULL;
+    se_scratch_init(heap);
+    if (!se_scratch_reserve(heap, dim)) {
+        se_scratch_free(heap);
+        se_free(heap);
+        return NULL;
+    }
+    return heap;
+}
+
+void se_scratch_release(se_scratch_t *s) {
+    se_tls_scratch_t *tls;
+
+    if (!s)
+        return;
+
+    tls = se_scratch_tls_get();
+    if (tls && s == &tls->scratch) {
+        tls->in_use = 0;
+        se_scratch_trim(&tls->scratch);
+        return;
+    }
+    se_scratch_free(s);
+    se_free(s);
+}
+
+void se_scratch_drop_thread(void) {
+    se_tls_scratch_t *tls = se_scratch_tls_get();
+    if (tls && tls->in_use)
+        tls->in_use = 0;
+    se_scratch_tls_clear();
+}
+
 static int is_whitespace(const se_model_t *m, uint32_t cp) {
     if (cp < 0x80)
         return se_is_ascii_whitespace(cp);
+    if (cp < 256)
+        return m->norm.whitespace256[cp];
     return se_range_contains(m->norm.whitespace, m->norm.whitespace_count, cp);
 }
 
 static int is_control(const se_model_t *m, uint32_t cp) {
     if (cp < 0x80)
         return (cp < 0x20 || cp == 0x7f) && cp != '\t' && cp != '\n' && cp != '\r';
+    if (cp < 256)
+        return m->norm.control256[cp];
     return se_range_contains(m->norm.control, m->norm.control_count, cp);
 }
 
 static int is_punct(const se_model_t *m, uint32_t cp) {
     if (cp < 0x80)
         return se_is_ascii_punct(cp);
+    if (cp < 256)
+        return m->norm.punct256[cp];
     return se_range_contains(m->norm.punct, m->norm.punct_count, cp);
 }
 
@@ -154,12 +372,18 @@ static int decode_one(const uint8_t *src, size_t len, size_t *i, uint32_t *cp_ou
 static int normalization_stable(const se_model_t *m, uint32_t cp) {
     if (cp < 0x80)
         return 1;
-    if (m->meta.do_lower_case && se_map_lookup(m->norm.lower, m->norm.lower_count, cp))
-        return 0;
+    if (m->meta.do_lower_case) {
+        const se_map_entry_t *lower =
+            cp < 256 ? m->norm.lower256[cp] : se_map_lookup(m->norm.lower, m->norm.lower_count, cp);
+        if (lower)
+            return 0;
+    }
     if (m->meta.strip_accents) {
-        if (se_map_lookup(m->norm.nfd, m->norm.nfd_count, cp))
+        const se_map_entry_t *nfd =
+            cp < 256 ? m->norm.nfd256[cp] : se_map_lookup(m->norm.nfd, m->norm.nfd_count, cp);
+        if (nfd)
             return 0;
-        if (se_range_contains(m->norm.mn, m->norm.mn_count, cp))
+        if (cp < 256 ? m->norm.mn256[cp] : se_range_contains(m->norm.mn, m->norm.mn_count, cp))
             return 0;
     }
     return 1;
@@ -233,6 +457,7 @@ typedef struct {
     size_t n_ids;
     size_t n_unk;
     size_t segment_len;
+    int segment_ascii;
     se_token_stats_t *stats;
     se_error_t *err;
     volatile sig_atomic_t *cancelled;
@@ -248,6 +473,8 @@ static se_status_t oom(token_state_t *st, const char *where) {
 }
 
 static int append_segment_cp(token_state_t *st, uint32_t cp) {
+    if (cp >= 0x80)
+        st->segment_ascii = 0;
     if (!grow_u32(&st->scratch->cps, &st->scratch->cps_cap, st->segment_len + 1))
         return 0;
     st->scratch->cps[st->segment_len++] = cp;
@@ -272,7 +499,8 @@ static int cap_after_append(token_state_t *st) {
 typedef enum { WORDPIECE_OK = 1, WORDPIECE_OOM = 0, WORDPIECE_INVALID = -1 } wordpiece_status_t;
 
 static wordpiece_status_t wordpiece(const se_model_t *m, se_scratch_t *sc, const uint32_t *word,
-                                    size_t word_len, size_t *n_ids, size_t *n_unk) {
+                                    size_t word_len, size_t *n_ids, size_t *n_unk,
+                                    int known_ascii) {
     const se_meta_t *meta = &m->meta;
 
     if (word_len == 0)
@@ -285,30 +513,49 @@ static wordpiece_status_t wordpiece(const se_model_t *m, se_scratch_t *sc, const
         return 1;
     }
 
-    size_t bytes_need = 0;
-    if (!se_checked_mul_size(word_len, 4, &bytes_need) ||
-        !se_checked_add_size(bytes_need, 4, &bytes_need))
-        return 0;
-    if (!grow_bytes(&sc->bytes, &sc->bytes_cap, bytes_need))
-        return 0;
-    if (!grow_u32(&sc->cps2, &sc->cps2_cap, word_len + 1))
-        return 0;
+    int ascii_identity = known_ascii;
+    if (!ascii_identity) {
+        ascii_identity = 1;
+        for (size_t k = 0; k < word_len; k++) {
+            if (word[k] >= 0x80) {
+                ascii_identity = 0;
+                break;
+            }
+        }
+    }
 
     size_t blen = 0;
-    for (size_t k = 0; k < word_len; k++) {
-        if (blen > UINT32_MAX)
+    if (ascii_identity) {
+        if (!grow_bytes(&sc->bytes, &sc->bytes_cap, word_len))
+            return 0;
+        if (!grow_u32(&sc->cps2, &sc->cps2_cap, word_len + 1))
             return 0;
-        sc->cps2[k] = (uint32_t)blen;
-        uint32_t cp = word[k];
-        if (cp < 0x80) {
-            sc->bytes[blen++] = (uint8_t)cp;
-        } else {
-            blen += se_utf8_encode(cp, sc->bytes + blen);
+        for (size_t k = 0; k < word_len; k++) {
+            sc->bytes[k] = (uint8_t)word[k];
+            sc->cps2[k] = (uint32_t)k;
         }
+        blen = word_len;
+        sc->cps2[word_len] = (uint32_t)word_len;
+    } else {
+        size_t bytes_need = 0;
+        if (!se_checked_mul_size(word_len, 4, &bytes_need) ||
+            !se_checked_add_size(bytes_need, 4, &bytes_need))
+            return 0;
+        if (!grow_bytes(&sc->bytes, &sc->bytes_cap, bytes_need))
+            return 0;
+        if (!grow_u32(&sc->cps2, &sc->cps2_cap, word_len + 1))
+            return 0;
+
+        for (size_t k = 0; k < word_len; k++) {
+            if (blen > UINT32_MAX)
+                return 0;
+            sc->cps2[k] = (uint32_t)blen;
+            blen += se_utf8_encode(word[k], sc->bytes + blen);
+        }
+        if (blen > UINT32_MAX)
+            return 0;
+        sc->cps2[word_len] = (uint32_t)blen;
     }
-    if (blen > UINT32_MAX)
-        return 0;
-    sc->cps2[word_len] = (uint32_t)blen;
 
     if (word_len <= meta->max_token_chars) {
         uint32_t exact_id = 0;
@@ -358,7 +605,8 @@ static wordpiece_status_t wordpiece(const se_model_t *m, se_scratch_t *sc, const
 static se_status_t append_wordpiece(token_state_t *st, const uint32_t *word, size_t word_len,
                                     int *stop) {
     wordpiece_status_t wp =
-        wordpiece(st->model, st->scratch, word, word_len, &st->n_ids, &st->n_unk);
+        wordpiece(st->model, st->scratch, word, word_len, &st->n_ids, &st->n_unk,
+                  word_len == 1 ? (word[0] < 0x80) : st->segment_ascii);
     if (wp == WORDPIECE_OOM)
         return oom(st, "tokenizing");
     if (wp == WORDPIECE_INVALID) {
@@ -375,6 +623,7 @@ static se_status_t flush_segment(token_state_t *st, int *stop) {
         return SE_OK;
     se_status_t rc = append_wordpiece(st, st->scratch->cps, st->segment_len, stop);
     st->segment_len = 0;
+    st->segment_ascii = 1;
     return rc;
 }
 
@@ -394,13 +643,30 @@ static se_status_t feed_token_cp(token_state_t *st, uint32_t cp, int *stop) {
     return SE_OK;
 }
 
+static const se_map_entry_t *lower_entry(const se_model_t *m, uint32_t cp) {
+    if (cp < 256)
+        return m->norm.lower256[cp];
+    return se_map_lookup(m->norm.lower, m->norm.lower_count, cp);
+}
+
+static const se_map_entry_t *nfd_entry(const se_model_t *m, uint32_t cp) {
+    if (cp < 256)
+        return m->norm.nfd256[cp];
+    return se_map_lookup(m->norm.nfd, m->norm.nfd_count, cp);
+}
+
+static int is_mn(const se_model_t *m, uint32_t cp) {
+    if (cp < 256)
+        return m->norm.mn256[cp];
+    return se_range_contains(m->norm.mn, m->norm.mn_count, cp);
+}
+
 static se_status_t emit_lowered(token_state_t *st, uint32_t cp, int *stop) {
     if (st->model->meta.do_lower_case) {
         if (cp >= 'A' && cp <= 'Z')
             cp += 32;
         else if (cp >= 0x80) {
-            const se_map_entry_t *e =
-                se_map_lookup(st->model->norm.lower, st->model->norm.lower_count, cp);
+            const se_map_entry_t *e = lower_entry(st->model, cp);
             if (e) {
                 for (uint32_t k = 0; k < e->len; k++) {
                     se_status_t rc = feed_token_cp(st, e->out[k], stop);
@@ -418,16 +684,16 @@ static se_status_t emit_stripped(token_state_t *st, uint32_t cp, int *stop) {
     if (!st->model->meta.strip_accents || cp < 0x80)
         return emit_lowered(st, cp, stop);
 
-    const se_map_entry_t *e = se_map_lookup(st->model->norm.nfd, st->model->norm.nfd_count, cp);
+    const se_map_entry_t *e = nfd_entry(st->model, cp);
     if (!e) {
-        if (se_range_contains(st->model->norm.mn, st->model->norm.mn_count, cp))
+        if (is_mn(st->model, cp))
             return SE_OK;
         return emit_lowered(st, cp, stop);
     }
 
     for (uint32_t k = 0; k < e->len; k++) {
         uint32_t d = e->out[k];
-        if (se_range_contains(st->model->norm.mn, st->model->norm.mn_count, d))
+        if (is_mn(st->model, d))
             continue;
         se_status_t rc = emit_lowered(st, d, stop);
         if (rc != SE_OK || *stop)
@@ -557,6 +823,7 @@ se_status_t se_tokenize(const se_model_t *model, se_scratch_t *sc, const uint8_t
     st.stats = stats;
     st.err = err;
     st.cancelled = cancelled;
+    st.segment_ascii = 1;
 
     size_t i = 0;
     size_t iterations = 0;

data/ext/static_embeddings/static_embeddings.c

@@ -274,10 +274,14 @@ static void batch_worker_run(batch_job_t *job, se_scratch_t *scratch) {
 
 static void *batch_execute(void *arg) {
     batch_job_t *job = (batch_job_t *)arg;
-    se_scratch_t scratch;
-    se_scratch_init(&scratch);
-    batch_worker_run(job, &scratch);
-    se_scratch_free(&scratch);
+    se_scratch_t *scratch = se_scratch_acquire(job->model->meta.dim);
+    if (!scratch) {
+        se_error_set(&job->error, SE_ERR_OOM, "out of memory while sizing scratch buffers");
+        job->failed = 1;
+        return NULL;
+    }
+    batch_worker_run(job, scratch);
+    se_scratch_release(scratch);
     return NULL;
 }
 
@@ -928,6 +932,32 @@ static VALUE embed_one_via_batch(VALUE self, VALUE text, VALUE max_tokens_opt,
     return embed_texts_internal(self, text, 0, 1, max_tokens_opt, format, validation, stats);
 }
 
+typedef struct {
+    se_scratch_t *scratch;
+    const se_model_t *model;
+    const uint8_t *input;
+    size_t input_len;
+    uint32_t max_tokens;
+    float *out;
+    se_token_stats_t *stats;
+    se_error_t err;
+    se_status_t rc;
+} embed_one_scratch_job_t;
+
+static VALUE embed_one_scratch_body(VALUE arg) {
+    embed_one_scratch_job_t *job = (embed_one_scratch_job_t *)(uintptr_t)arg;
+    job->rc = se_embed_one(job->model, job->scratch, job->input, job->input_len, job->max_tokens,
+                           job->out, job->stats, &job->err, NULL);
+    return Qnil;
+}
+
+static VALUE embed_one_scratch_ensure(VALUE arg) {
+    embed_one_scratch_job_t *job = (embed_one_scratch_job_t *)(uintptr_t)arg;
+    se_scratch_release(job->scratch);
+    job->scratch = NULL;
+    return Qnil;
+}
+
 static VALUE embed_one_value(VALUE self, VALUE text, VALUE max_tokens_opt,
                              se_vector_format_t format, se_encoding_validation_t validation,
                              se_token_stats_t *stats) {
@@ -947,25 +977,28 @@ static VALUE embed_one_value(VALUE self, VALUE text, VALUE max_tokens_opt,
     rb_enc_associate(result, binary_encoding);
     float *out = (float *)RSTRING_PTR(result);
 
-    se_scratch_t scratch;
-    se_scratch_init(&scratch);
-    if (!se_scratch_reserve(&scratch, dim)) {
-        se_scratch_free(&scratch);
+    se_token_stats_t local_stats;
+    embed_one_scratch_job_t job;
+    memset(&job, 0, sizeof(job));
+    job.model = &w->model;
+    job.input = (const uint8_t *)RSTRING_PTR(text);
+    job.input_len = (size_t)RSTRING_LEN(text);
+    job.max_tokens = resolve_max_tokens(&w->model, max_tokens_opt);
+    job.out = out;
+    job.stats = stats ? stats : &local_stats;
+    se_error_clear(&job.err);
+
+    job.scratch = se_scratch_acquire(dim);
+    if (!job.scratch)
         rb_raise(rb_eNoMemError, "out of memory");
-    }
 
-    se_error_t err;
-    se_error_clear(&err);
-    se_token_stats_t local_stats;
-    se_status_t rc =
-        se_embed_one(&w->model, &scratch, (const uint8_t *)RSTRING_PTR(text),
-                     (size_t)RSTRING_LEN(text), resolve_max_tokens(&w->model, max_tokens_opt), out,
-                     stats ? stats : &local_stats, &err, NULL);
-    se_scratch_free(&scratch);
+    rb_ensure(embed_one_scratch_body, (VALUE)(uintptr_t)&job, embed_one_scratch_ensure,
+              (VALUE)(uintptr_t)&job);
     RB_GC_GUARD(text);
+    RB_GC_GUARD(result);
 
-    if (rc != SE_OK)
-        raise_se(&err);
+    if (job.rc != SE_OK)
+        raise_se(&job.err);
 
     return result;
 }
@@ -1001,6 +1034,38 @@ static VALUE model_embed_with_stats(int argc, VALUE *argv, VALUE self) {
     return hash;
 }
 
+typedef struct {
+    se_scratch_t *scratch;
+    const se_model_t *model;
+    const uint8_t *input;
+    size_t input_len;
+    uint32_t max_tokens;
+    se_token_stats_t stats;
+    se_error_t err;
+    se_status_t rc;
+    VALUE ids;
+} tokenize_scratch_job_t;
+
+static VALUE tokenize_scratch_body(VALUE arg) {
+    tokenize_scratch_job_t *job = (tokenize_scratch_job_t *)(uintptr_t)arg;
+    job->rc = se_tokenize(job->model, job->scratch, job->input, job->input_len, job->max_tokens,
+                          &job->stats, &job->err, NULL);
+    if (job->rc != SE_OK)
+        return Qnil;
+
+    job->ids = rb_ary_new_capa((long)job->stats.token_count);
+    for (uint32_t i = 0; i < job->stats.token_count; i++)
+        rb_ary_push(job->ids, UINT2NUM(job->scratch->ids[i]));
+    return Qnil;
+}
+
+static VALUE tokenize_scratch_ensure(VALUE arg) {
+    tokenize_scratch_job_t *job = (tokenize_scratch_job_t *)(uintptr_t)arg;
+    se_scratch_release(job->scratch);
+    job->scratch = NULL;
+    return Qnil;
+}
+
 static VALUE model_tokenize(int argc, VALUE *argv, VALUE self) {
     VALUE text, opts;
     rb_scan_args(argc, argv, "1:", &text, &opts);
@@ -1011,32 +1076,27 @@ static VALUE model_tokenize(int argc, VALUE *argv, VALUE self) {
     check_text_encoding_mode(
         text, -1, resolve_encoding_validation(lookup_option(opts, id_validate_encoding)));
 
-    uint32_t max_tokens = resolve_max_tokens(&w->model, lookup_option(opts, id_max_tokens));
-
-    se_scratch_t scratch;
-    se_scratch_init(&scratch);
-    if (!se_scratch_reserve(&scratch, w->model.meta.dim)) {
-        se_scratch_free(&scratch);
+    tokenize_scratch_job_t job;
+    memset(&job, 0, sizeof(job));
+    job.model = &w->model;
+    job.input = (const uint8_t *)RSTRING_PTR(text);
+    job.input_len = (size_t)RSTRING_LEN(text);
+    job.max_tokens = resolve_max_tokens(&w->model, lookup_option(opts, id_max_tokens));
+    job.ids = Qnil;
+    se_error_clear(&job.err);
+
+    job.scratch = se_scratch_acquire(w->model.meta.dim);
+    if (!job.scratch)
         rb_raise(rb_eNoMemError, "out of memory");
-    }
 
-    se_token_stats_t stats;
-    se_error_t err;
-    se_error_clear(&err);
-    se_status_t rc = se_tokenize(&w->model, &scratch, (const uint8_t *)RSTRING_PTR(text),
-                                 (size_t)RSTRING_LEN(text), max_tokens, &stats, &err, NULL);
-    if (rc != SE_OK) {
-        se_scratch_free(&scratch);
-        raise_se(&err);
-    }
+    rb_ensure(tokenize_scratch_body, (VALUE)(uintptr_t)&job, tokenize_scratch_ensure,
+              (VALUE)(uintptr_t)&job);
+    RB_GC_GUARD(text);
 
-    VALUE ids = rb_ary_new_capa((long)stats.token_count);
-    for (uint32_t i = 0; i < stats.token_count; i++)
-        rb_ary_push(ids, UINT2NUM(scratch.ids[i]));
+    if (job.rc != SE_OK)
+        raise_se(&job.err);
 
-    se_scratch_free(&scratch);
-    RB_GC_GUARD(text);
-    return ids;
+    return job.ids;
 }
 
 typedef struct {
@@ -1077,17 +1137,15 @@ static VALUE num2ull_at_value(VALUE arg) {
 
 static void *ids_execute(void *arg) {
     ids_job_t *job = (ids_job_t *)arg;
-    se_scratch_t scratch;
-    se_scratch_init(&scratch);
-    if (!se_scratch_reserve(&scratch, job->model->meta.dim)) {
+    se_scratch_t *scratch = se_scratch_acquire(job->model->meta.dim);
+    if (!scratch) {
         se_error_set(&job->error, SE_ERR_OOM, "out of memory while sizing scratch buffers");
-        se_scratch_free(&scratch);
         return NULL;
     }
     se_error_clear(&job->error);
-    se_embed_ids(job->model, &scratch, job->ids, job->n_ids, job->out, &job->stats, &job->error,
+    se_embed_ids(job->model, scratch, job->ids, job->n_ids, job->out, &job->stats, &job->error,
                  &job->cancelled);
-    se_scratch_free(&scratch);
+    se_scratch_release(scratch);
     return NULL;
 }
 

data/lib/static_embeddings/version.rb

@@ -1,3 +1,3 @@
 module StaticEmbeddings
-  VERSION = "0.1.3"
+  VERSION = "0.1.4"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: static_embeddings
 version: !ruby/object:Gem::Version
-  version: 0.1.3
+  version: 0.1.4
 platform: ruby
 authors:
 - Roman Haydarov