Real-Time WebGPU Inference for MIDI Autocomplete with ONNX Runtime

Musicians feel latency. If you press a key on a MIDI keyboard and the sound or accompaniment arrives more than 10 milliseconds later, the disconnect disrupts your playing. When building an interactive AI system that autocompletes piano performances, sending MIDI events to a cloud server is out of the question. Even on a fast fiber connection, round-trip times hover around 30 to 80 milliseconds. Add queue wait times and inference latency, and the system becomes unusable for live play.

Running the neural network directly in the web browser solves this. By using WebGPU and ONNX Runtime Web, we can execute a 125-million parameter causal transformer model on local hardware. While running massive models locally usually requires specialized techniques like AirLLM for 70B parameter inference, browser-based execution requires a much smaller footprint to fit within standard browser memory limits and capture stylistic patterns in keyboard performances.

Tokenizing MIDI for Causal Transformers

Transformers process discrete tokens. To feed musical notes into a model, we must convert raw MIDI messages into a structured vocabulary. A naive approach might convert MIDI to a piano-roll matrix, but this wastes memory and loses timing precision.

Instead, we use an event-based representation. We define a vocabulary consisting of four main event types:

1. Note-On: Triggers a pitch (128 possible values, from C-1 to G9).
2. Note-Off: Releases a pitch (128 possible values).
3. Time-Shift: Signals time progression. We bin time steps into 100 intervals ranging from 10 milliseconds to 1 second.
4. Velocity: Sets the force of the strike. We group the 127 standard MIDI velocities into 32 discrete bins.

This yields a total vocabulary size of 388 tokens. Here is a TypeScript class that handles the serialization of MIDI events into token sequences:

type MidiEvent = 
  | { type: 'note_on'; pitch: number; velocity: number }
  | { type: 'note_off'; pitch: number }
  | { type: 'time_shift'; ms: number };
 
class MidiTokenizer {
  private vocabSize = 388;
  private timeBins = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 150, 200, 250, 300, 400, 500, 600, 700, 800, 900, 1000];
 
  encode(events: MidiEvent[]): Uint32Array {
    const tokens: number[] = [];
    for (const event of events) {
      if (event.type === 'note_on') {
        tokens.push(event.pitch); // 0-127
      } else if (event.type === 'note_off') {
        tokens.push(128 + event.pitch); // 128-255
      } else if (event.type === 'time_shift') {
        const binIndex = this.findClosestBin(event.ms);
        tokens.push(256 + binIndex); // 256-276
      }
    }
    return new Uint32Array(tokens);
  }
 
  private findClosestBin(ms: number): number {
    return this.timeBins.reduce((prev, curr, idx) => 
      Math.abs(curr - ms) < Math.abs(this.timeBins[prev] - ms) ? idx : prev, 0
    );
  }
}

This compact vocabulary keeps the model's embedding layer small. A 125M parameter model configured with an embedding dimension of 768 and 12 attention layers requires roughly 250 megabytes of memory when stored in float16 precision.

Setting Up ONNX Runtime Web with WebGPU

To run the model locally, we use ONNX Runtime Web (onnxruntime-web). We target the WebGPU execution provider, which bypasses the CPU overhead of WebGL and allows direct access to compute pipelines on modern graphics cards.

First, configure the library to load the WebGPU execution assembly files from a local directory or CDN. This setup prevents blocking the main browser thread:

import * as ort from 'onnxruntime-web/webgpu';
 
ort.env.wasm.wasmPaths = 'https://cdnjs.cloudflare.com/ajax/libs/onnxruntime-web/1.17.1/';

Next, initialize the inference session. We pass options to optimize memory allocations and enable FP16 precision:

async function loadMidiModel(modelUrl) {
  const options = {
    executionProviders: ['webgpu'],
    preferredOutputLocations: {},
  };
 
  // Pre-allocate execution provider options
  options.preferredOutputLocations = {
    'logits': 'gpu-buffer'
  };
 
  const session = await ort.InferenceSession.create(modelUrl, options);
  return session;
}

Requesting the output tensor to remain on the GPU buffer prevents the browser from copying large chunks of data back to system memory. This step minimizes latency during the autoregressive loop.

Optimizing the Autoregressive Loop with KV Caching

In standard transformer inference, computing predictions for token t requires processing all tokens from index 0 to t-1. As the sequence grows, the computational cost increases quadratically. This overhead mirrors the scaling challenges found in multi-agent systems, where context window replication drives token consumption. To maintain sub-20ms latency, we must implement Key-Value (KV) caching.

The KV cache stores the computed key and value projections from the self-attention layers of previous tokens. When generating the next token, we only pass the single new token to the model along with the cached states. The model processes the single token and appends its key and value states to the existing cache.

To support this, the ONNX model must be exported with explicit inputs and outputs for the cache tensors of each attention layer. A 12-layer model requires 24 cache inputs (one key and one value tensor per layer) and 24 corresponding outputs.

Here is how to run the autoregressive generation loop using KV caches in JavaScript:

async function generateNextToken(session, newToken, kvCache) {
  const feeds = {};
  
  // Input token shape is [batch_size, sequence_length] -> [1, 1]
  feeds['input_ids'] = new ort.Tensor('int64', new BigInt64Array([BigInt(newToken)]), [1, 1]);
  
  // Map existing KV cache tensors to inputs
  for (let i = 0; i < 12; i++) {
    feeds[`past_key_values.${i}.key`] = kvCache.keys[i];
    feeds[`past_key_values.${i}.value`] = kvCache.values[i];
  }
 
  // Run inference
  const results = await session.run(feeds);
 
  // Extract the new KV cache for the next step
  const nextKvCache = { keys: [], values: [] };
  for (let i = 0; i < 12; i++) {
    nextKvCache.keys.push(results[`present_key_values.${i}.key`]);
    nextKvCache.values.push(results[`present_key_values.${i}.value`]);
  }
 
  // Sample from the output logits
  const logits = results['logits'];
  const sampledToken = sampleLogits(logits);
 
  return { token: sampledToken, cache: nextKvCache };
}

Using this approach, execution time remains flat regardless of sequence length. On an Apple M1 GPU, generating a single token with the 125M parameter model takes roughly 12 milliseconds using the KV cache, compared to over 80 milliseconds when recalculating the full sequence at length 512.

Managing GPU Memory and Preventing Leaks

WebGPU relies on explicit memory management. In the browser environment, JavaScript relies on garbage collection, which is not designed to handle large GPU buffers. If you instantiate new tensors in every generation step, you will trigger memory churn. This leads to frame drops and audio stuttering.

To prevent this, we reuse tensors. We pre-allocate a static memory block for the KV cache that matches the maximum sequence length, such as 1024 tokens. Instead of creating new tensors at every step, we update parts of the pre-allocated buffers.

class TensorPool {
  constructor(device, size) {
    this.device = device;
    this.size = size;
    this.buffer = device.createBuffer({
      size: size,
      usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST | GPUBufferUsage.STORAGE,
    });
  }
 
  write(offset, data) {
    this.device.queue.writeBuffer(this.buffer, offset, data);
  }
}

By keeping the buffers on the GPU and using WebGPU command encoders to copy data between them, we avoid transferring data to the CPU. The CPU only receives the sampled token ID, which is a single integer.

Interfacing with Web MIDI and Audio Scheduling

Once the model generates a token, we must schedule it. The standard JavaScript event loop is too loose for musical timing. Functions like setTimeout and requestAnimationFrame can be delayed by UI rendering or layout calculations.

We solve this by combining the Web MIDI API with the Web Audio API clock. The Web Audio context clock runs on a high-priority system thread, providing a reliable timestamp reference.

First, initialize MIDI access and listen to input ports:

let midiOutput = null;
 
async function setupMidi() {
  const access = await navigator.requestMIDIAccess();
  const outputs = Array.from(access.outputs.values());
  if (outputs.length > 0) {
    midiOutput = outputs[0]; // Send to the first available output device
  }
}

When scheduling a generated note, we read the current time from the AudioContext and calculate the offset. We translate the generated time_shift tokens into absolute timestamps for the MIDI output port:

const audioCtx = new AudioContext();
let scheduledTime = audioCtx.currentTime;
 
function playNote(pitch, velocity, delayMs) {
  if (!midiOutput) return;
 
  const noteOnMessage = [0x90, pitch, velocity];
  const noteOffMessage = [0x80, pitch, 0];
 
  // Convert delay to milliseconds relative to the performance clock
  const targetTime = performance.now() + delayMs;
 
  midiOutput.send(noteOnMessage, targetTime);
  midiOutput.send(noteOffMessage, targetTime + 250); // Hold note for 250ms
}

Using the system timestamp in the send method schedules the MIDI message at the hardware driver level. This bypasses any main-thread congestion that might occur when the browser renders UI updates.

Benchmarks and Quantization Trade-offs

We evaluated the performance of this approach across three hardware configurations using a 125M parameter model. The benchmarks measure the time required to generate one token in a loop with a sequence length of 512.

FP16 precision reduces latency by more than half on devices that support it. Most modern mobile devices and laptops support WebGPU FP16 extensions. For older hardware, fall back to WebGL or WebAssembly (WASM) execution providers, though these will require reducing the model size to 35M parameters to keep latency under the 30ms threshold.

Using INT8 quantization reduces the model file size to roughly 125MB, but it does not always speed up execution on WebGPU. Because current WebGPU runtimes lack optimized integer matrix multiplication kernels for all hardware configurations, FP16 remains the most reliable option for low-latency generation.