[Web] WebAssembly: MatMulNBits dequantizes its whole weight matrix to float32 on every call (1 GB scratch for an LLM output layer; WebAssembly memory never shrinks; slow decode) #33795

Describe the issue

In onnxruntime-web's WebAssembly build, every MatMulNBits call dequantizes its whole weight matrix to float32 into a scratch buffer of K × N × 4 bytes, then runs an fp32 GEMM. This happens on every call, including decode steps with M = 1.

MLAS has no QNBit GEMM kernels for WebAssembly: platform.cpp sets QNBitGemmDispatch only for AVX2/AVX512, NEON, RVV and LASX. So MatMulNBits::PrePack returns early (!MlasIsQNBitGemmAvailable(...)), and Compute always falls back to ComputeBUnpacked, which starts with:

// TODO(fajin): move B dequant to prepack
auto tmp_b_data_ptr = IAllocator::MakeUniquePtr<float>(allocator, SafeInt<size_t>(K_) * N_, true);

In a browser that has two costs.

Memory. The scratch buffer is as large as the biggest weight matrix expanded to float32, and WebAssembly memory never shrinks, so the page keeps that peak for good. For a 4-bit LLM output layer (2,048 × 128,000, 131 MB of 4-bit weights) the buffer is 1,048 MB. With onnx-community/LFM3-1.7B-ONNX (model_q4.onnx, 1,297 MB of weights), WebAssembly memory goes from 1,363 MB after loading to 2,366 MB on the first prompt. iOS gives a Safari tab 2,048 MB (Memory Limits: active 2048 in the device log), and Safari terminated the tab's web process as the first reply started.

Speed. Each token dequantizes every 4-bit weight in the model to float32 again, about 7 GB of float32 writes per token for a 1.7B model. LFM3 on the CPU writes at 3.8 tokens/s on an M3 Ultra (16 threads, Chrome) and 0.3 tokens/s on an iPhone 16 Pro Max (Safari). For comparison, llama.cpp compiled to WebAssembly (through wllama) writes 77 tokens/s on the same Mac with LFM3 at a similar 4-bit size (Q4_K_M).

A workaround we use: splitting each large MatMulNBits along N into ≤ 64 MB-as-float32 chunks (Slices of the same weights plus a Concat), before creating the session. That brings LFM3's peak WebAssembly memory to 1,644 MB, with identical greedy output, and it now runs on the iPhone. Speed is unchanged.

Suggestions:

1. Bound the scratch buffer in ComputeBUnpacked: dequantize and multiply B in column tiles that reuse one small buffer, instead of materializing all of K × N. Peak memory becomes independent of N, and for M = 1 the tile stays in cache.
2. Add a WASM SIMD (simd128) QNBit kernel to MLAS, at least a fused dequantize-and-dot GEMV for M = 1, so decode never dequantizes the whole matrix.
3. Or move B dequantization into PrePack as the TODO says. A float32 copy costs 8× the 4-bit bytes, though, so (1) or (2) fit browsers better.

To reproduce

1. Make a model with one MatMulNBits node the shape of an LLM output layer (Python, pip install onnx numpy):

# make_model.py: one MatMulNBits node, K = 2048, N = 128000 (or argv[1]), 4-bit, block 32.
import sys
import numpy as np
import onnx
from onnx import TensorProto, helper

N = int(sys.argv[1]) if len(sys.argv) > 1 else 128000
K, block = 2048, 32
blocks = K // block
rng = np.random.default_rng(0)
B = rng.integers(0, 256, size=(N, blocks, block // 2), dtype=np.uint8)
scales = (rng.random(N * blocks, dtype=np.float32) * 0.01).astype(np.float32)
node = helper.make_node(
    "MatMulNBits", ["A", "B", "scales"], ["Y"], domain="com.microsoft",
    K=K, N=N, bits=4, block_size=block, accuracy_level=4,
)
graph = helper.make_graph(
    [node], "matmulnbits",
    [helper.make_tensor_value_info("A", TensorProto.FLOAT, [1, "M", K])],
    [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, "M", N])],
    [helper.make_tensor("B", TensorProto.UINT8, B.shape, B.tobytes(), raw=True),
     helper.make_tensor("scales", TensorProto.FLOAT, scales.shape, scales.tobytes(), raw=True)],
)
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("com.microsoft", 1)])
model.ir_version = 10
onnx.save(model, sys.argv[2] if len(sys.argv) > 2 else "matmulnbits.onnx")

python make_model.py 128000 matmulnbits.onnx
python make_model.py 8000 matmulnbits-8000.onnx

2. Serve this page next to the models, then open index.html?model=matmulnbits.onnx (add &threads=4 on a cross-origin-isolated page):

<!doctype html>
<meta charset="utf-8" />
<pre id="out">running…</pre>
<script type="module">
  // Track WebAssembly memory before onnxruntime loads.
  const memories = new Set()
  const Memory = WebAssembly.Memory
  WebAssembly.Memory = new Proxy(Memory, {
    construct(target, args) {
      const memory = Reflect.construct(target, args)
      memories.add(memory)
      return memory
    }
  })
  for (const name of ["instantiate", "instantiateStreaming"]) {
    const original = WebAssembly[name]
    WebAssembly[name] = async (...args) => {
      const result = await original(...args)
      const exports = (result.instance ?? result).exports ?? {}
      for (const value of Object.values(exports))
        if (value instanceof Memory) memories.add(value)
      return result
    }
  }
  const mb = (bytes) => Math.round(bytes / 1e6)
  const wasmMB = () => mb([...memories].reduce((sum, m) => sum + m.buffer.byteLength, 0))
  const log = (line) => (document.getElementById("out").textContent += `\n${line}`)

  const params = new URLSearchParams(location.search)
  const file = params.get("model") ?? "matmulnbits.onnx"
  const ort = await import(
    "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.31.1/dist/ort.wasm.bundle.min.mjs"
  )
  ort.env.wasm.numThreads = Number(params.get("threads") ?? 1)
  const model = new Uint8Array(await (await fetch(file)).arrayBuffer())
  const session = await ort.InferenceSession.create(model, { executionProviders: ["wasm"] })
  log(`${file}, ${ort.env.wasm.numThreads} thread(s)`)
  log(`after session creation: WebAssembly memory ${wasmMB()} MB`)
  const K = 2048
  for (const M of [1, 1, 16]) {
    const A = new ort.Tensor("float32", new Float32Array(M * K).fill(0.01), [1, M, K])
    const t0 = performance.now()
    await session.run({ A })
    log(`after a run with M=${M}: ${(performance.now() - t0).toFixed(0)} ms, WebAssembly memory ${wasmMB()} MB`)
  }
</script>

Actual (Chrome 152, macOS, M3 Ultra; the timings come from a background tab, so only compare them with each other):

Expected: a decode step (M = 1) shouldn't need a float32 copy of the whole weight matrix: memory near the 4-bit weights, and time near reading them once.

Urgency

4-bit LLMs that fit an iPhone's memory crash Safari on the CPU path, and decode on WebAssembly is 20× slower than llama.cpp's WebAssembly build.

Platform

Web Browser

OS Version

macOS 27.2 (Chrome 152.0.7977.167); iOS 27.2 (Safari, iPhone 16 Pro Max)

ONNX Runtime Installation

Released Package

ONNX Runtime Version or Commit ID

onnxruntime-web 1.31.1 (ort.wasm.bundle.min.mjs)

ONNX Runtime API

JavaScript

Architecture

ARM64

Execution Provider

Default CPU

Model File

Generated by the script above; also https://huggingface.co/onnx-community/LFM3-1.7B-ONNX (onnx/model_q4.onnx)

Is this a quantized model?

Yes