Runtime Inference Engine for Web Browsers Using the SmolLM2-360M-Instruct

What if your browser could run its own AI inference engine — without sending a single prompt to a backend?

Read through a lightweight, fully client-side runtime inference engine architecture using SmolLM2-360M-Instruct, with ONNX Runtime Web on a dedicated Web Worker architecture for local LLM capabilities directly into the browser.

Built for speed, privacy, and a seamless developer experience, it features:

Local AI Inference: Run SmolLM2‑360M‑Instruct entirely on the user's device.

WebGPU Acceleration: Hardware-accelerated inference with automatic CPU fallback for broader browser compatibility.

Background Execution: Offload model loading, tokenization, and generation to a Web Worker, keeping the main thread responsive.

Persistent Model Caching: Use IndexedDB to avoid repeatedly downloading model assets.

Optimized Token Generation: Implement a KV-cache-based autoregressive generation loop for efficient token-by-token inference.

Editor-Style Autocomplete: Generate contextual suggestions, press Tab to accept, or Esc to dismiss.

It keeps inference entirely off the main thread and uses a KV‑cache generation loop for fast token‑by‑token output.

---

1. Project structure & dependencies

ai-autocomplete/
├── index.html
├── main.js
├── workers/
│   └── inference.worker.js
├── models/
│   └── SmolLM2-360M-Instruct/          # served statically
│       ├── onnx/
│       │   └── model_q4f16.onnx        # 270 MB
│       ├── tokenizer.json
│       └── tokenizer_config.json
└── package.json

npm install onnxruntime-web @huggingface/transformers

@huggingface/transformers is used only for tokenization/decoding; the actual ONNX inference is still raw onnxruntime-web. This avoids writing a BPE tokenizer by hand while keeping full control over the ONNX session and generation loop.

---

2. Web Worker – model loading, caching, and generation

workers/inference.worker.js

import * as ort from "onnxruntime-web/webgpu";
import { AutoTokenizer } from "@huggingface/transformers";

// ------------------------------------------------------------------
// Configuration
// ------------------------------------------------------------------
const MODEL_BASE = "/models/SmolLM2-360M-Instruct";
const MODEL_URL = `${MODEL_BASE}/onnx/model_q4f16.onnx`;
const TOKENIZER_URL = `${MODEL_BASE}/tokenizer.json`;
const TOKENIZER_CONFIG_URL = `${MODEL_BASE}/tokenizer_config.json`;

const DB_NAME = "smollm2-cache";
const DB_VERSION = 1;
const STORE_NAME = "models";
const MODEL_KEY = "smollm2-360m-instruct-q4f16";

// ------------------------------------------------------------------
// State
// ------------------------------------------------------------------
let session = null;
let tokenizer = null;
let initializing = null;

// ------------------------------------------------------------------
// IndexedDB helpers (model binary is large → IndexedDB is more
// reliable than Cache API for this use case)
// ------------------------------------------------------------------
function openDB() {
  return new Promise((resolve, reject) => {
    const request = indexedDB.open(DB_NAME, DB_VERSION);
    request.onupgradeneeded = () => {
      const db = request.result;
      if (!db.objectStoreNames.contains(STORE_NAME)) {
        db.createObjectStore(STORE_NAME);
      }
    };
    request.onsuccess = () => resolve(request.result);
    request.onerror = () => reject(request.error);
  });
}

async function getModelFromCache() {
  const db = await openDB();
  return new Promise((resolve, reject) => {
    const tx = db.transaction(STORE_NAME, "readonly");
    const store = tx.objectStore(STORE_NAME);
    const req = store.get(MODEL_KEY);
    req.onsuccess = () => resolve(req.result ?? null);
    req.onerror = () => reject(req.error);
  });
}

async function saveModelToCache(arrayBuffer) {
  const db = await openDB();
  return new Promise((resolve, reject) => {
    const tx = db.transaction(STORE_NAME, "readwrite");
    const store = tx.objectStore(STORE_NAME);
    const req = store.put(arrayBuffer, MODEL_KEY);
    req.onsuccess = () => resolve();
    req.onerror = () => reject(req.error);
  });
}

// ------------------------------------------------------------------
// Model & tokenizer loading
// ------------------------------------------------------------------
async function loadModelBytes() {
  const cached = await getModelFromCache();
  if (cached) {
    console.log("[worker] model loaded from IndexedDB");
    return cached;
  }

  console.log("[worker] downloading model…");
  const response = await fetch(MODEL_URL);
  if (!response.ok) throw new Error(`Model download failed: ${response.status}`);

  const buffer = await response.arrayBuffer();
  await saveModelToCache(buffer);
  console.log("[worker] model cached in IndexedDB");
  return buffer;
}

async function loadTokenizer() {
  // Transformers.js can load the tokenizer directly from URL.
  tokenizer = await AutoTokenizer.from_pretrained(MODEL_BASE, {
    // Avoid downloading the full model again – only tokenizer files.
    // The library will fetch tokenizer.json / tokenizer_config.json.
    // If your server structure differs, pass explicit URLs.
  });
  console.log("[worker] tokenizer ready");
}

// ------------------------------------------------------------------
// Session creation – WebGPU first, WASM fallback
// ------------------------------------------------------------------
async function createSession(modelBytes) {
  // q4f16 is optimised for WebGPU; on WASM it may not initialise
  // reliably. We still try WASM as a last resort.
  const tryWebGPU = async () => {
    if (!navigator.gpu) throw new Error("WebGPU not available");
    const adapter = await navigator.gpu.requestAdapter();
    if (!adapter) throw new Error("No WebGPU adapter");

    return ort.InferenceSession.create(modelBytes, {
      executionProviders: ["webgpu"],
      // graphOptimizationLevel: "all", // optional
    });
  };

  const tryWasm = async () => {
    return ort.InferenceSession.create(modelBytes, {
      executionProviders: ["wasm"],
    });
  };

  try {
    console.log("[worker] trying WebGPU…");
    return await tryWebGPU();
  } catch (e) {
    console.warn("[worker] WebGPU failed:", e.message);
    console.log("[worker] falling back to WASM…");
    return await tryWasm();
  }
}

async function initialize() {
  if (session) return session;
  if (initializing) return initializing;

  initializing = (async () => {
    await loadTokenizer();
    const modelBytes = await loadModelBytes();
    session = await createSession(modelBytes);
    console.log("[worker] session ready");
    return session;
  })();

  try {
    return await initializing;
  } finally {
    initializing = null;
  }
}

// ------------------------------------------------------------------
// Generation loop with KV cache
// ------------------------------------------------------------------
/**
 * Greedy generation using past_key_values for fast decoding.
 * Returns only the newly generated text (the completion).
 */
async function generateCompletion(promptText, maxNewTokens = 20) {
  const model = await initialize();

  // SmolLM2-Instruct uses a ChatML-like template.
  // For autocomplete we wrap the existing text in a simple user turn
  // and let the model continue the assistant turn.
  const messages = [
    {
      role: "user",
      content:
        "Continue the following text naturally, without repeating it:\n\n" +
        promptText,
    },
  ];

  // Apply chat template and tokenize.
  const chatPrompt = tokenizer.apply_chat_template(messages, {
    tokenize: false,
    add_generation_prompt: true, // primes the assistant turn
  });

  const inputs = await tokenizer(chatPrompt, { return_tensors: "np" }); // returns input_ids, attention_mask

  let inputIds = inputs.input_ids; // Int32Array or BigInt64Array
  let attentionMask = inputs.attention_mask;

  // ONNX expects int64.
  inputIds = BigInt64Array.from(inputIds, BigInt);
  attentionMask = BigInt64Array.from(attentionMask, BigInt);

  const inputIdsTensor = new ort.Tensor("int64", inputIds, [1, inputIds.length]);
  const attentionMaskTensor = new ort.Tensor("int64", attentionMask, [
    1,
    attentionMask.length,
  ]);

  // Discover past‑key‑value input names from the session.
  const inputNames = model.inputNames;
  const pastKeyNames = inputNames.filter((n) => n.includes("past_key_values"));
  const pastKeyValues = {};

  // Initial run: no past.
  let feeds = {
    input_ids: inputIdsTensor,
    attention_mask: attentionMaskTensor,
  };

  let outputs = await model.run(feeds);

  // Collect the first generated token.
  const generatedIds = [];
  const logits = outputs.logits; // shape [1, seq_len, vocab_size]

  let nextTokenId = argmax(logits, logits.dims[1] - 1);
  generatedIds.push(nextTokenId);

  // Prepare past_key_values for the next step.
  for (const name of pastKeyNames) {
    if (outputs[name]) pastKeyValues[name] = outputs[name];
  }

  // Autoregressive loop.
  for (let step = 1; step < maxNewTokens; step++) {
    // Stop on EOS (SmolLM2 EOS token id is typically 2, but check
    // tokenizer.eos_token_id if needed).
    if (nextTokenId === tokenizer.eos_token_id) break;

    const currentInput = BigInt64Array.from([BigInt(nextTokenId)]);
    const currentAttention = BigInt64Array.from([1n]);

    feeds = {
      input_ids: new ort.Tensor("int64", currentInput, [1, 1]),
      attention_mask: new ort.Tensor("int64", currentAttention, [1, 1]),
      ...pastKeyValues, // feed previous KV cache
    };

    outputs = await model.run(feeds);

    const stepLogits = outputs.logits;
    nextTokenId = argmax(stepLogits, 0); // seq_len = 1
    generatedIds.push(nextTokenId);

    // Update KV cache with new outputs.
    for (const name of pastKeyNames) {
      if (outputs[name]) pastKeyValues[name] = outputs[name];
    }
  }

  // Decode only the newly generated tokens.
  const completion = tokenizer.decode(generatedIds, {
    skip_special_tokens: true,
  });

  return completion;
}

/** Greedy argmax over the last vocabulary dimension. */
function argmax(tensor, seqIndex = 0) {
  const data = tensor.data;
  const vocabSize = tensor.dims[tensor.dims.length - 1];
  const offset = seqIndex * vocabSize;
  let bestId = 0;
  let bestVal = -Infinity;
  for (let i = 0; i < vocabSize; i++) {
    const v = data[offset + i];
    if (v > bestVal) {
      bestVal = v;
      bestId = i;
    }
  }
  return bestId;
}

// ------------------------------------------------------------------
// Message handling
// ------------------------------------------------------------------
self.onmessage = async ({ data }) => {
  const { type, text, requestId } = data;

  try {
    if (type === "init") {
      await initialize();
      self.postMessage({ type: "ready" });
      return;
    }

    if (type === "suggest") {
      const completion = await generateCompletion(text, 20);

      self.postMessage({
        type: "suggestions",
        requestId,
        suggestions: completion ? [completion] : [],
      });
    }
  } catch (error) {
    console.error("[worker] error:", error);
    self.postMessage({
      type: "error",
      requestId,
      message: error.message,
    });
  }
};

Key points in the worker

• IndexedDB stores the raw ArrayBuffer once; subsequent page loads skip the 260 MB download.
• WebGPU detection happens inside the worker; navigator.gpu is available in Workers. If it fails, WASM is used.
• KV cache reuses past_key_values across steps, which is essential for fast token generation.
• Chat template follows SmolLM2‑Instruct’s expected format.
• The generation loop is greedy (fastest). You can later add top‑k/top‑p sampling if needed.

---

3. Main thread – editor, debounce, and keyboard shortcuts

main.js

const aiWorker = new Worker(
  new URL("./workers/inference.worker.js", import.meta.url),
  { type: "module" }
);

let requestId = 0;
let debounceTimer;
let activeRequestId = null;
let currentSuggestion = "";

const editor = document.querySelector("#editor");
const suggestionsBox = document.querySelector("#suggestions");

// ------------------------------------------------------------------
// Initialize model in the background.
// ------------------------------------------------------------------
aiWorker.postMessage({ type: "init" });

aiWorker.onmessage = ({ data }) => {
  switch (data.type) {
    case "ready":
      console.log("AI model initialised");
      break;

    case "suggestions": {
      // Discard stale responses.
      if (data.requestId !== activeRequestId) return;

      suggestionsBox.innerHTML = "";

      if (!data.suggestions.length) {
        currentSuggestion = "";
        return;
      }

      currentSuggestion = data.suggestions[0];

      const el = document.createElement("div");
      el.className = "suggestion";
      el.textContent = currentSuggestion;

      el.addEventListener("click", () => applySuggestion());
      suggestionsBox.appendChild(el);
      break;
    }

    case "error":
      console.error("AI inference:", data.message);
      break;
  }
};

// ------------------------------------------------------------------
// Debounce typing
// ------------------------------------------------------------------
editor.addEventListener("input", () => {
  clearTimeout(debounceTimer);

  // Hide any old suggestion while the user is typing.
  suggestionsBox.innerHTML = "";
  currentSuggestion = "";

  debounceTimer = setTimeout(() => {
    activeRequestId = ++requestId;
    aiWorker.postMessage({
      type: "suggest",
      requestId: activeRequestId,
      text: editor.value,
    });
  }, 400);
});

// ------------------------------------------------------------------
// Tab to accept, Escape to dismiss
// ------------------------------------------------------------------
editor.addEventListener("keydown", (e) => {
  if (e.key === "Tab" && currentSuggestion) {
    e.preventDefault();
    applySuggestion();
  }

  if (e.key === "Escape") {
    suggestionsBox.innerHTML = "";
    currentSuggestion = "";
  }
});

function applySuggestion() {
  if (!currentSuggestion) return;

  // Insert at the current cursor position.
  const start = editor.selectionStart;
  const end = editor.selectionEnd;
  const before = editor.value.slice(0, start);
  const after = editor.value.slice(end);

  editor.value = before + currentSuggestion + after;
  const newPos = start + currentSuggestion.length;
  editor.setSelectionRange(newPos, newPos);

  suggestionsBox.innerHTML = "";
  currentSuggestion = "";
  editor.focus();
}

---

4. HTML & CSS – clean inline suggestion UI

index.html

<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8" />
  <title>AI Autocomplete – SmolLM2</title>
  <style>
    body {
      font-family: system-ui, sans-serif;
      max-width: 800px;
      margin: 2rem auto;
      padding: 0 1rem;
    }

    #editor {
      width: 100%;
      min-height: 200px;
      padding: 0.75rem;
      font-size: 1rem;
      font-family: inherit;
      border: 1px solid #ccc;
      border-radius: 6px;
      resize: vertical;
    }

    #suggestions {
      margin-top: 0.5rem;
    }

    .suggestion {
      display: inline-block;
      padding: 0.35rem 0.6rem;
      background: #f0f0f0;
      border: 1px solid #ddd;
      border-radius: 4px;
      cursor: pointer;
      font-size: 0.95rem;
      color: #333;
      white-space: pre-wrap;
    }

    .suggestion:hover {
      background: #e8e8e8;
    }

    .hint {
      margin-top: 0.4rem;
      font-size: 0.8rem;
      color: #777;
    }
  </style>
</head>
<body>
  <h1>AI Autocomplete</h1>

  <textarea
    id="editor"
    placeholder="Start writing…"
    rows="10"
  ></textarea>

  <div id="suggestions"></div>
  <div class="hint">Press <strong>Tab</strong> to accept · <strong>Esc</strong> to dismiss</div>

  <script type="module" src="/main.js"></script>
</body>
</html>

The Tab/Esc interaction mirrors common editor behaviour.

---

5. Why this design meets your requirements

Requirement	How it is handled
SmolLM2‑360M‑Instruct	Model file model_q4f16.onnx is loaded from the specified path. Tokenizer is loaded from the same base directory.
Download once, store in IndexedDB	getModelFromCache / saveModelToCache use IndexedDB, which is more reliable than the Cache API for large binaries.
Instant start (after first load)	Model bytes are read from IndexedDB synchronously in the worker; no network request on repeat visits. KV cache keeps generation fast.
Fast token generation	Autoregressive loop reuses past_key_values; greedy decoding is the fastest sampling strategy.
WebGPU / CPU fallback	navigator.gpu is checked inside the worker; WebGPU is tried first, then WASM. Note: q4f16 is optimised for WebGPU; if WASM initialisation fails, the worker surfaces an error so you can show a user‑friendly message or load a q8 model instead.
Background inference	Everything (download, tokenization, session creation, generation) runs inside inference.worker.js; the main thread only sends text and renders suggestions.
Editor‑style suggestions	Debounced input, stale‑request filtering via requestId, inline suggestion box, Tab to apply, Esc to cancel.

---

6. Production notes & tuning

• Q4F16 + WASM caveat – model_q4f16.onnx is specifically optimised for WebGPU. On browsers without WebGPU the WASM fallback may stall during initialisation. If you need guaranteed CPU support, serve model_q8.onnx (or model_int8.onnx) and use it when navigator.gpu is absent.
• Tokenizer files – @huggingface/transformers will fetch tokenizer.json and tokenizer_config.json from the model directory. Make sure those files are served alongside the ONNX model.
• Max new tokens – 20 tokens is usually enough for a sentence‑level completion. Increase it carefully; longer generations block the worker for longer.
• Stale requests – The worker returns the requestId it received; the main thread compares it with activeRequestId and ignores outdated suggestions.
• Memory – After generation the KV cache is discarded. If you want to keep the session warm for the next keystroke, the ONNX session itself stays loaded in the worker, which is the main performance win.

This gives you a private, offline‑capable autocomplete engine that runs entirely in the browser, uses the GPU when available, and falls back to the CPU without blocking the editor.