Laya-TS
A port of Convai Innovations Laya System One Model
Laya-TS๐
Porting Laya to TypeScript: Bringing โSystem 1โ Decision Models to WebGPU and NodeJS.๐
The Motivation & Goals๐
In modern AI architectures, generative LLMs (System 2) are frequently overused for tasks that only require reflexive judgmentโsuch as triage, ticket routing, moderation, and intent classification. Forcing these fast decisions through autoregressive text generation introduces hundreds of milliseconds of latency, excessive inference costs, risk of hallucinations, and uncalibrated confidence scores.
To solve this, Nandakishor Mukkunnoth and ConvAI Innovations introduced Laya: a non-autoregressive โSystem 1โ decision engine. By evaluating arbitrary inputs against typed decision primitives (choice, score, and noul) in a single forward pass Laya returns mathematically calibrated probabilities trained via RLCD (Reinforcement Learning for Calibrated Decisions) without generating a single token of text.
However, Laya was strictly tied into to Python and PyTorch. For my ideas, This created a deployment barrier: you needed much heavier compute to run these applications in lower-powered places.
My Goal:๐
Re-engineer the entire Laya inference pipeline into a cross-platform, zero-Python TypeScript stack, bringing that super-fact calibrated decision-making directly into web browsers via WebGPU and into lean Node.js backends, but while attempting to maintain byte-exact numerical parity with the PyTorch reference implementation.
Key Achievements๐
1. Zero-Python TypeScript Core๐
I ported the entire downstream inference pipeline into a runtime-agnostic package (inference-core):
- Re-implemented sequence construction (buildSequence.ts) and exact [MASK] marker token alignment for ModernBERT and mmBERT backbones.
- Ported the Act Head linear projections and feature extraction (feats.ts, actHead.ts) to calculate entropy and margin statistics outside the neural graph.
- Added post-hoc temperature scaling (calibrate.ts) and typed answer rendering (outputShaping.ts)
Client-Side WebGPU Inference in the Browser๐
With browser-runtime, models run locally on consumer hardware using onnx-runtime-web:
- Zero Network Latency & Total Privacy: Decisions happen directly in the browser; sensitive user text and tickets never leave the device.
- Hardware Acceleration & Graceful Fallback: Automatically leverages WebGPU on supported devices with a seamless WebAssembly fallback and built-in client memory budgeting (memoryBudget.ts) to prevent mobile tab crashes.
- Interactive Demo: Shipped a live, browser-based WebGPU evaluation app (browser-parity-check).
3. Lean Server Runtime & Microservice๐
For backend workloads, server-runtime packages onnxruntime-node into a lightweight engine:Eliminates the need for PyTorch, CUDA toolkits, or Python runtime dependencies on the host (verified in bare node:22-slim containers).
- Shipped a drop-in HTTP microservice (systemone-server) implementing the official POST /v1/systemone protocol with CLI and environment variable configuration.
4. End-to-End ONNX Tooling & Mathematical Parity๐
To guarantee fidelity:
- Built an automated offline export pipeline (export) using TorchDynamo opset 18, dynamic shapes, and FP16/INT8 quantization (quantize.py).
- Designed an automated contract parity test suite (contract-parity-check) and Playwright headless browser suite (browser-parity-check)
- validating that TypeScript outputs match PyTorch ground truth down to machine precision.