Optimizing Distilled SLM Models for Browser-Local Rust and WASM Inference in TinyRustLM

Full report

On this page

Executive summary

TinyRustLM’s central engineering problem is not just “make a small model run in the browser.” It is to keep the model useful while staying inside browser memory ceilings, WebAssembly’s feature fragmentation, slow startup paths, and a strict local-only privacy posture. The current TinyRustLM implementation already exposes several of the hard limits: model loading currently performs a full host-to-WASM copy, rejects transfers above 128 MiB, materializes runtime-owned tensor storage after the copy, preallocates a full-context KV cache, serializes all public operations through a single global mutex, and executes prompt prefill token-by-token instead of through a batched path. MiRust’s own documentation also states that no comparable hardware/browser throughput dataset exists yet for the current snapshot, so production targets must be treated as launch goals rather than measured claims.

The most important architectural recommendation is to optimize for browser reality first, not for maximum raw model quality. For TinyRustLM, that means a decoder-only, deep-thin distilled family with tied embeddings and grouped-query attention, shipped primarily in q4 and q8 weights-only formats, with a worker-owned runtime, a chunked loader ABI, a paged KV cache, and dual kernel paths for simd128 and scalar fallback. This matches what recent small-model work has found useful for on-device models: MobileLLM reports gains from deep-thin architectures, embedding sharing, and grouped-query attention at sub-billion scale, while SmolLM2’s 135M and 360M configurations already use GQA-style layouts with fewer KV heads than attention heads.

A second critical conclusion is that browser-local privacy has to be proved operationally, not merely asserted in product copy. TinyRustLM’s site says the prompt “never leaves your device” and that the browser uses the local .slm artifact directly, but production support still needs smoke tests that intercept all network traffic, all WebSocket attempts, and all service-worker-mediated requests while generation runs. Playwright’s request routing and event APIs, together with service-worker blocking, are sufficient to turn this into a release gate.

The core difficulty in the browser runtime

WebAssembly in the browser is a constrained CPU target, not a normal native process. Rust’s wasm32-unknown-unknown target supports core and alloc, but major std features are missing or degraded, and std::thread::spawn panics. Rust’s default enabled WebAssembly feature set also evolves over time, which means TinyRustLM cannot assume simd128, threads, or other newer proposals exist everywhere; if a binary contains unsupported instructions, the engine cannot execute it. Rust explicitly recommends conditional compilation for WebAssembly target features such as simd128, rather than assuming runtime dispatch like on x86.

Browser concurrency is similarly conditional. WebAssembly threads rely on SharedArrayBuffer, and broader access to SharedArrayBuffer, precise timers, and memory-measurement APIs requires a cross-origin-isolated page via COOP and COEP headers. Without cross-origin isolation, timer precision is coarsened and thread-based runtime designs are unavailable or brittle across browsers. That means TinyRustLM must be designed to run acceptably in a single-thread or limited-thread mode, with threads treated as an optimization tier rather than a baseline requirement.

The memory story is even more restrictive than the compute story. TinyRustLM’s current loader contract allocates WASM memory, copies the entire artifact into linear memory, and then materializes runtime-owned tensor storage. MiRust documents that peak loading memory can exceed artifact size because the browser may hold the fetched ArrayBuffer, the copied WASM bytes, and the decoded tensor storage at the same time. For larger models, this transient duplication is often more dangerous than steady-state decode memory.

KV cache growth is the other structural difficulty. TinyRustLM’s current KV layout is contiguous Vec<f32> storage indexed as [layer][position][kv_head][head_dim], with memory given by 2 × layers × max_context × kv_heads × head_dim × 4. Because the current implementation allocates for the full declared context up front, long contexts increase memory before any useful work is done, and long prompts make both prefill and attention history reads more expensive. MiRust also notes that the current forward path rejects head_count != kv_head_count, so although the container can describe GQA or MQA, the present runtime cannot exploit them yet.

Model architecture recommendations for WASM

For TinyRustLM’s browser target, the preferred architecture family is a distilled, decoder-only transformer with deep-thin scaling, tied embeddings, RoPE, SwiGLU-style feed-forward blocks, and GQA or MQA from the start. MobileLLM’s findings support deep-thin networks, embedding sharing, and grouped-query attention for sub-billion on-device models, and SmolLM2’s released 135M and 360M configs reflect the same direction. The official SmolLM2 135M config uses 30 layers, hidden size 576, 9 attention heads, 3 KV heads, 8,192 max positions, tied embeddings, and vocab size 49,152. The 360M config uses 32 layers, hidden size 960, 15 attention heads, 5 KV heads, and the same 8,192-position context class.

That architecture choice matters directly for browser memory. Using TinyRustLM’s documented KV formula with SmolLM2-style shapes, a plain f32 KV cache for the 135M model is about 90 MiB at 2,048 tokens, 180 MiB at 4,096, and 360 MiB at 8,192. For the 360M model, the same arithmetic is about 160 MiB at 2,048, 320 MiB at 4,096, and 640 MiB at 8,192. Those numbers are before extra scratch, logits, host-side duplication during load, and browser overhead. In other words, the advertised maximum context is not the right default context for browser-local execution.

For shipping profiles, the practical recommendation is to separate three browser tiers. The universal tier should be smaller distilled models below roughly 100M–135M parameters in q4, running with a default 2K context and a scalar fallback. The standard tier should be 135M q4 with a worker runtime and optional simd128, defaulting to 2K or 4K context depending on detected memory and isolation state. The high-quality tier should be 360M q4 only when TinyRustLM has four things simultaneously: streaming load, worker isolation, GQA execution, and at least paged or quantized KV. Without those, 360M is too close to browser memory cliffs on long prompts. This recommendation is an engineering inference from TinyRustLM’s current loader behavior, SmolLM2’s configs, and the documented steady-state KV formula.

A further recommendation is to make GQA support a first-class runtime milestone. SmolLM2 already uses fewer KV heads than attention heads, and MobileLLM identifies grouped-query attention as part of the on-device design space. TinyRustLM’s current restriction that head_count == kv_head_count is therefore not just a missing feature; it blocks one of the most important memory levers available for browser inference.

Matrix layout and quantized kernel design

TinyRustLM’s current storage model is already pointed in the right direction for browser inference. MiRust describes runtime-owned storage variants as plain Vec<f32> for f32, Vec<i8> plus row scales for q8, and packed Vec<u8> plus block scales and block size for q4. It also documents “selective dequantization”: matrix operations dispatch directly by storage type, while small vectors or embedding rows can be copied into reusable f32 scratch, avoiding accidental full-model expansion. That is exactly the right baseline for a browser runtime where memory traffic is usually more expensive than arithmetic.

The next step is to make the tensor layout explicitly kernel-oriented. TinyRustLM’s SLM1 format already gives a 108-byte fixed header, 64-byte tensor entries, and a tensor-data section aligned to 64 bytes. Although the current runtime copies payloads into Rust-owned vectors rather than borrowing aligned file pointers, that alignment is still valuable: it supports predictable chunk boundaries for streaming loaders and makes it reasonable to repack weights once at load time into microkernel-specific layouts.

For decode, TinyRustLM should optimize the batch-1 case first. MiRust’s performance notes say matrix-vector operations dominate weight reads in decode, and in the current implementation every generated token performs the full stack of projections, attention work, and feed-forward operations. That strongly suggests a kernel strategy centered on GEMV-like microkernels, not generic GEMM abstractions. The right packed layout for weights is therefore output-channel-major and tile-friendly: each destination-row tile should store its quantized weight block(s) and scale data contiguously, with scales brought close to the payload they decode. This minimizes pointer chasing in WASM linear memory and makes the kernel’s access pattern regular for both SIMD and scalar paths. The exact tile shape is an implementation choice, but it should be selected to keep one activation strip hot while streaming weight rows sequentially.

For q8, the preferred browser path is a weights-only kernel: activations remain in f32 scratch, while the kernel streams i8 weights and row scales and accumulates to f32. For q4, the kernel should unpack nibble pairs into signed integers in registers, apply the block scale, and accumulate directly without constructing large temporary dequantized buffers. The important design rule is that q4 and q8 kernels must be direct kernels, not “dequantize matrix then multiply” routines, because MiRust’s own memory-bandwidth discussion shows that weight reads dominate and direct quantized kernels are the only way to reduce that traffic materially.

SIMD support should be treated as a family of binaries, not a single runtime branch. Rust documents simd128 as a target feature, and also warns that WebAssembly binaries must only contain instructions the engine understands. In practice, TinyRustLM should ship at least two builds: a scalar MVP-compatible build and a simd128 build. A third experimental build can target relaxed SIMD once the product deliberately narrows browser coverage. This is also where XNNPACK is instructive as a reference point: its README explicitly positions WebAssembly MVP, WebAssembly SIMD, and WebAssembly Relaxed SIMD as distinct supported architectures for optimized inference primitives.

The scalar fallback path is not optional quality debt; it is part of the correctness contract. TinyRustLM’s deterministic sampling notes already warn that reproducibility depends on floating-point ordering, quantized kernels, softmax behavior, and tie handling. Therefore, the scalar and SIMD kernels should share identical scale interpretation, rounding rules, and reduction order as much as practical, and greedy smoke tests should run against both binaries on fixed prompts.

KV-cache paging, context, and generation behavior

TinyRustLM’s current contiguous full-capacity f32 KV cache is simple but expensive. It is simple because indexing is trivial and commit discipline is clean: MiRust documents that K and V for a token are written layer by layer and only exposed with commit_len(position + 1) after the full layer loop completes. It is expensive because the cache is reserved for the full declared context at model load, and because long prompts multiply both serial prefill cost and history reads.

The right next design is a paged KV cache. PagedAttention shows why: paging reduces KV-cache waste and makes memory management more flexible for long-running decode workloads. Browser-local TinyRustLM does not need the whole vLLM serving model, but it does need the same core idea: fixed-size pages, per-layer page tables, and an append-only write path that preserves the current “commit after complete token” discipline. For TinyRustLM, page sizes of 16 to 32 tokens are a sensible starting point because they keep metadata modest while limiting over-allocation at short contexts.

Compression should be staged. KVQuant shows that sub-4-bit KV compression can be accurate, especially with per-channel key quantization and pre-RoPE key quantization, but the more recent system-oriented work is a useful warning: SAW-INT4 argues that only a narrow subset of 4-bit designs survives real serving constraints such as paged layouts, regular memory access, and fused attention execution, and a broader practical study finds that throughput gains from compression can diminish inside production-style attention stacks, while response length can sometimes increase and cancel end-to-end speedups. For TinyRustLM, that means phase one should be paged f32 or paged int8 KV; phase two can add int4 KV only behind a quality gate with measured latency, output-length, and regression fixtures.

Context defaults should be lower than architectural maxima. SmolLM2’s 8K-capable configs are useful, but browser-local defaults should be chosen from memory arithmetic, not model cards. A safe product stance is 2K default context for all universal-browser profiles, 4K as an opt-in for 135M on cross-origin-isolated desktops, and 360M at 4K only when paged or quantized KV is enabled. An 8K mode should be marked experimental until TinyRustLM has streaming load, GQA execution, and production-quality cache paging.

Prompt prefill deserves separate treatment from decode. MiRust documents that prefill is currently fully serial and that there is no batched prompt path. Because prefill time scales with prompt length and runs before the first generated token, adding batched prefill is one of the highest-return optimizations after worker migration and vectorization. In practical terms, TinyRustLM should have two performance counters in every benchmark and every diagnostic dump: prefill tokens per second and decode tokens per second. Collapsing them into one “tokens/s” number hides the exact browser bottleneck you need to fix.

On generation controls, TinyRustLM should preserve its current deterministic default of temperature=0, top_k=1, top_p=1, and fixed seed for smoke tests. For stochastic decoding, the current XorShift64 design with fixed candidate buffers is a good browser choice because it avoids vocabulary-sized heap churn per token. Repetition penalty should be implemented exactly and cheaply: Hugging Face’s reference processor applies the penalty once per seen token, multiplying negative logits and dividing positive logits, with optional prompt_ignore_length to exclude the prompt from the penalty set. EOS handling should support multiple eos_token_id values, retain TinyRustLM’s existing “remove terminal EOS from prompt before prefill” behavior, and add a forced-EOS-at-max-length safeguard for clean termination. After applying logits processors, TinyRustLM should renormalize scores, because Hugging Face explicitly recommends renormalize_logits=True when processors may break normalization.

Loading, validation, and startup path

For the WASM binary itself, WebAssembly.instantiateStreaming() is the right baseline because MDN describes it as the most efficient way to load and instantiate Wasm, avoiding the older ArrayBuffer step for the module. That optimization matters for TinyRustLM’s code artifact, but it does not solve model loading. The .slm file is not a Wasm module, and TinyRustLM’s current ABI still forces an all-at-once copy into WASM linear memory.

So TinyRustLM needs a streaming .slm loader, not a fake “memory-mapped” loader. In today’s browser model, the closest equivalents are Blob.stream() for chunked reads, FileReaderSync inside workers for synchronous file slices, and OPFS FileSystemSyncAccessHandle inside dedicated workers for higher-performance local file access. MDN specifically notes that OPFS is private to the page’s origin, optimized for performance, supports in-place access, and that FileSystemSyncAccessHandle is only available inside dedicated workers. That is enough to build a browser-native replacement for the current full-file transfer ABI: a chunk handle, incremental parser state, progressive SHA-256, and transactional activation only after the final verified chunk lands.

Startup validation has to be split into cheap and expensive phases. MiRust documents the current admission sequence as: parse fixed header and checksum, parse tensor directory, parse tokenizer section, verify tensor hashes and required shapes, decode or copy tensors into typed storage, resolve indices, allocate KV/scratch/logits, and only then install the model into the runtime. It also documents all-or-nothing commit and rollback semantics. That is a good transactional integrity model, but it is a bad latency model if expensive decode happens before cheap rejection paths are exhausted.

The recommendation is a two-phase verifier. Phase one should be outside the hot path and as cheap as possible: verify a canonical SHA-256 plus signed manifest, then parse only enough header and directory data to reject incompatible artifacts quickly. MiRust’s own checksum page says the current checksum is useful for accidental corruption but does not prove publisher identity or authenticity, and explicitly recommends a SHA-256 plus signed canonical manifest binding the artifact to format version, dimensions, tokenizer hash, lineage, and compatibility. Phase two can then perform heavier tensor validation and optional kernel-specific repacking, ideally into OPFS so subsequent runs avoid repeating the whole cost.

Diagnostics, privacy proof, benchmarks, and acceptance criteria

TinyRustLM already has the beginnings of a useful diagnostic surface. MiRust lists fields such as model-loaded state, last error, load time, prompt and generated token counts, token speed, peak scratch, KV length, quantization mode, tokenizer output, logits summary, selected token, top-k summary, and sampling configuration. It also warns that the current JSON schema is informal and that tokenizer IDs or logits summaries may themselves reveal prompt characteristics, so production telemetry should default to no transmission and clearly separate operator diagnostics from user-safe exports.

For production support, that diagnostic schema should be expanded and versioned. The minimum additional fields should be artifact SHA-256, manifest-signature verdict, browser and engine version, crossOriginIsolated state, worker mode, thread mode, SIMD mode, loader bytes copied, loader bytes streamed, total model heap estimate, current and reserved KV bytes, page count if paged KV is enabled, prefill tok/s, decode tok/s, time-to-first-token, and whether any network request occurred after generation began. Timing should use performance.now() because it is monotonic and high resolution, and memory should opportunistically use measureUserAgentSpecificMemory() when the browser exposes it.

Privacy smoke tests should be a release gate. In CI, launch Chromium with Playwright, create a browser context with serviceWorkers: 'block', route or observe all requests at the browser-context level, and assert that after the initial page assets finish loading there are no POST requests, no model-inference API calls, no unexpected cross-origin fetches, and no WebSocket traffic during generation. Record any remaining same-origin asset fetches explicitly. Also assert that the page remains functional when disconnected from the network after the .slm artifact is present in OPFS, which demonstrates that inference itself is local. This turns TinyRustLM’s “prompt never leaves your device” promise into an executable test rather than a marketing statement.

Because MiRust states that there is no published comparable tokens-per-second dataset for the current snapshot, the following are recommended launch targets, not measured TinyRustLM baselines. They should be evaluated on at least two reference classes: a cross-origin-isolated desktop browser with simd128, and a compatibility browser path without threads and with scalar fallback.

For models smaller than 135M, the launch target should be a q4 artifact that is usable across the widest set of browsers: ready-to-generate in at most 1.5 seconds from warm local cache and 4 seconds from a cold local file on the reference desktop; prefill at 12 tok/s or better on scalar and 25 tok/s or better on SIMD; decode at 7 tok/s or better on scalar and 15 tok/s or better on SIMD; and live memory under 220 MiB at a 2K context. This tier is the safest default for strict privacy-first deployments. The recommendation is grounded in the existing TinyRustLM q4 storage model and the browser memory duplication hazards documented in the current loader.

For a 135M-class model, the practical target is q4 first and q8 second. Under TinyRustLM-style storage, a 135M model is roughly 80.5 MiB in q4 and about 129.6 MiB in q8 before extra container metadata, while raw f32 would be about 515 MiB. The launch target should be at most 2.5 seconds ready-to-generate from warm cache and 6 seconds from cold local file on the reference desktop; prefill at 8 tok/s scalar and 20 tok/s SIMD; decode at 4 tok/s scalar and 10 tok/s SIMD; and memory under 350 MiB at 2K context, under 450 MiB at 4K context with paged or compressed KV.

For a 360M-class model, q4 should be considered the primary browser format. Under the same extrapolation, q4 is roughly 214.6 MiB and q8 about 345.6 MiB before runtime overhead, while f32 would be about 1.37 GiB. The launch target should be at most 5 seconds from warm cache and 12 seconds from cold local file on the reference desktop; prefill at 3 tok/s scalar and 8 tok/s SIMD; decode at 1.5 tok/s scalar and 4 tok/s SIMD; and memory under 650 MiB at 2K context, under 850 MiB at 4K context with paged or int8 KV enabled. An 8K mode should not pass acceptance until TinyRustLM proves stable behavior under long-context regressions, because the uncompressed f32 KV arithmetic alone reaches about 640 MiB at 8K.

Acceptance criteria should therefore be explicit. A release should not ship unless all of the following are true on the supported matrix: deterministic greedy output matches fixtures for both scalar and SIMD builds; model activation is transactional with no partially loaded runtime state after failure; no remote inference traffic is observed in smoke tests; prefill and decode counters are reported separately; the runtime stays worker-owned so the main thread remains responsive; and memory usage remains within the product’s published envelope for the default context. Where stochastic decoding is offered, TinyRustLM should document reproducibility scope honestly: the current MiRust guidance is correct that matching seed and settings are necessary but not sufficient if floating-point ordering, softmax, kernel implementations, or tie handling diverge.

Open questions and limitations

Some important implementation details remain intentionally framed here as recommendations rather than observations. The MiRust documentation is highly useful for the current TinyRustLM shape and constraints, but it also makes clear that some pages are research or implementation evidence rather than claims of production support. In particular, there is no current public TinyRustLM benchmark corpus for comparable browser hardware, and there is no current public proof that GQA execution, paged KV, or chunked .slm loading already exist in the released runtime.

The most consequential unresolved engineering choice is how aggressive TinyRustLM should be with KV compression in production. The literature now shows both promise and risk: quantization can reduce memory sharply, but practical serving studies also show that some methods lose their advantage when integrated into real attention stacks or shift output lengths in unhelpful ways. The safest roadmap is still clear, though: first add GQA execution and paged KV, then add int8 KV, then consider int4 only with benchmark-backed acceptance.