Skip to content

ARCHITECTURE.md

Latest commit

 

History

History
161 lines (116 loc) · 7.98 KB

ARCHITECTURE.md

File metadata and controls

161 lines (116 loc) · 7.98 KB
Raw

Architecture

This document describes the implementation of the native inference engine. If you only want to build and run the service, see README.md.

Model structure

The service loads the OpenAI Privacy Filter checkpoint (openai/privacy-filter) directly from safetensors and runs it on Burn 0.21's Flex CPU backend. The architecture matches the HuggingFace token-classification layout:

  1. Token embedding [vocab_size, hidden_size]
  2. num_hidden_layers transformer layers, each:
    • RMSNorm
    • GQA attention with attention sinks and a bidirectional sliding-window mask
    • RMSNorm
    • Sparse MoE MLP
  3. Final RMSNorm
  4. Score head: Linear(hidden_size, num_labels) — 33 classes for this checkpoint.

Attention

Each layer uses grouped-query attention (num_attention_heads Q heads, num_key_value_heads K/V heads, head_dim per head). Keys and values are repeated to match the query head count inside repeat_kv.

Rotary position embeddings are YaRN-scaled with NTK-by-parts blending between interpolation and extrapolation frequencies. Tables are precomputed at startup up to n_ctx.

The sliding-window mask is a full banded [seq_len, seq_len] matrix: cells where |i - j| > sliding_window are set to -1e9 before softmax. In addition, every layer learns a per-head attention sink — an additive logit applied to a dedicated sink column concatenated to the attention scores.

MoE MLP

The MLP is a sparse Mixture-of-Experts block:

  • A router scores each token against all experts.
  • Per-token top-k experts are selected and weighted by a local softmax.
  • Tokens are bucketed by routed expert. Each non-empty bucket is processed in parallel via rayon.
  • Expert inner loops use the gemm crate's hand-tuned AVX2 / AVX-512 / NEON microkernels. The bias vector is fused into the GEMM via beta = 1.0.
  • The activation function is (up + 1) · gate · σ(α · gate) with α = 1.702 and clamp limits ±7.0.

Variable-length batches

forward_batch operates on a flat [total_tokens, hidden_size] tensor. Attention scores, softmax, and attn·V are computed per-sequence at exact length — there is no padding mask and no wasted attention compute on padded positions. Projections and MoE dispatch remain batched across the whole merge.

Windowing and long-text handling

Long texts are split into contiguous non-overlapping chunks of at most n_ctx tokens (ids.chunks(n_ctx)). This matches the reference implementation exactly: example_to_windows in the Python code hardcodes stride = window_size, and the eval runner even panics with "overlapping windows are not supported" if a token were ever seen twice.

The reference codebase contains log-sum-exp aggregation infrastructure (logprob_logsumexp, counts) that looks like it averages overlapping window predictions, but in practice every token always has count == 1. PR #27 in the reference repo added a fast-path that skips the per-token aggregation loop entirely for this common case.

Numerical fidelity

The implementation is tested against the Python transformers reference:

  • Tokenization parity: Rust tokenizer IDs match Python exactly.
  • Argmax label parity: Per-token argmax labels match the reference forward.
  • Logit parity: Full-model logits are within tolerance of the reference. The observed worst-case relative error across the parity fixtures is ≈ 4 × 10⁻⁴.
  • Batched vs single-sequence: forward_batch and forward produce identical per-sequence results up to ≈ 2 × 10⁻⁵ (difference is only matmul reduction reordering at rank-4).

Weights ship as bf16 on disk and are widened to f32 in memory, which roughly doubles weight size (~2.7 GB → ~5.4 GB).

Request coalescer

By default the engine is wrapped in a coalescing layer that merges concurrent anonymize requests into a single batched forward pass:

  • When a request arrives and the engine is idle, the worker waits up to 200 ms for additional requests before submitting the batch.
  • If the previous batch was a singleton, the wait is skipped (assumes sequential workload).
  • A soft token cap prevents batches from growing without bound.
  • The work queue is bounded; when full the service returns HTTP 503.

You can disable the coalescer by setting PF_COALESCE_DISABLE=true.

Full configuration reference

These environment variables are exposed for tuning. Most users should not need to change them.

Variable Default Description
PF_MODEL_DIR ./model Directory for model files
PF_BIND 0.0.0.0:10123 HTTP listen address
PF_DECODE_MODE viterbi viterbi (constrained CRF decode) or argmax (independent per-token)
PF_MAX_BATCH_TOKENS 4 * n_ctx Max tokens packed into one batched forward pass
PF_BUCKET_RATIO (off) Length-bucketing guard for the packer. Set to > 1.0 to enable. With the CPU varlen attention path this is usually unnecessary.
PF_COALESCE_WAIT_MS 200 Idle wait when coalescer is free
PF_COALESCE_QUEUE 256 Bounded work-queue capacity for the coalescer
PF_COALESCE_DISABLE false Skip the coalescer wrapper

Tests and benchmarks

Unit tests

Primitive numerical tests (RMSNorm, SwiGLU, GEMM wrapper, repeat_kv, sub-batch planning) — no weights required:

cargo test --release --lib

Integration tests

Gated on PF_TEST_MODEL_DIR (default ./model). Missing weights cause the tests to print a skip note and pass.

# Python parity (tokenization, argmax labels, raw logits, span extraction)
PF_TEST_MODEL_DIR=./model cargo test --release --test python_parity -- --nocapture

# Batched-vs-single parity
PF_TEST_MODEL_DIR=./model cargo test --release --test batched_parity -- --nocapture

# Coalescer correctness
PF_TEST_MODEL_DIR=./model cargo test --release --test coalesce -- --nocapture

Regenerating the Python reference fixture

tests/fixtures/reference.safetensors holds token IDs, intermediate activations, and full-model logits dumped from the Python transformers reference. Regenerate after a model or fixture change:

# The script declares its own dependencies via PEP 723 inline metadata,
# so uv installs them automatically in an ephemeral environment:
uv run tests/scripts/dump_reference.py ./model tests/fixtures/reference.safetensors

# Or run it directly — the shebang calls uv for you:
./tests/scripts/dump_reference.py ./model tests/fixtures/reference.safetensors

If you prefer to manage your own Python environment, the script needs transformers>=5.6, torch, safetensors, and numpy installed.

NixOS users: nix develop includes uv and sets LD_LIBRARY_PATH with the system libraries that prebuilt PyPI wheels link against, so the commands above work without any extra setup.

Criterion benchmarks

PF_TEST_MODEL_DIR=./model cargo bench

Benchmark groups cover batch-of-16 similar texts, batch-of-16 varied texts, sequential, and concurrent access patterns, with and without the coalescer.