Introducing Matmul as a Service, backed by tinygrad

A
Ali Ben
Member of Technical Staff
Blog post illustration

We’re pivoting to AI.

Introducing Matmul as a Service (MaaS), backed by tinygrad. Run any computation you want on our infra. Pay for the memory and FLOPS you use.

Why pivot from server monitoring?

Monitoring was a good business to be in. But AI was looking even better. Seeing all that capital flowing to NVIDIA and to data centers. We couldn’t just stay on the sidelines.

Running AI models is actually very simple. It’s a lot of matrix manipulations. Multiply one matrix by another (matmul), add another matrix, do that again a billion times, and you basically get AGI (or a poorly designed website, or a scam email, whatever you choose).

But how do we differentiate ourselves from all the competition?

Our beautiful API

The standard approach is to ship compiled programs to the device. That’s what CUDA does: complex kernels with indexing, flow control and memory management are sent to the GPU and run as one unit.

We wanted something easier. We only care about giving you the ability to offload two things to us: memory to store data and the actual compute. That’s really all programming is. Every programming language has:

  • A way to handle memory.
  • A way to read and write from and to that memory.
  • A way to run computations on the data stored in that memory.

So we designed our HTTP backend around these concepts:

MethodRouteUsage
POST/allocAllocate a buffer, get back an ID
PUT/buf/<buffer_id>Write data into a buffer
GET/buf/<buffer_id>Read data from a buffer
DELETE/buf/<buffer_id>Free a buffer
POST/instrExecute one instruction on the backend

So now we have a very clean API. But nobody wants to write HTTP calls by hand. How do you actually use this?

Enter tinygrad

Other neural network frameworks like PyTorch are heavy. The logic for a simple ReLU on CUDA is deeply embedded somewhere in the stack. tinygrad uses a really simple set of device-agnostic operations instead (UOps). Think RISC vs x86. High-level tensor operations are decomposed into a very small set of primitives. Every computation you write in tinygrad is transformed into a graph of these UOps.

Adding support for a new backend is just implementing these UOps. Once that’s done, the framework doesn’t care if it runs on a GPU, a TPU, or a shoebox.

The best proof of this is its impressive set of runtimes: CUDA, Metal, AMD, and more.

Their mission is to commoditize the petaflop. We intend to sell it by the unit. So we added MaaS to the list.

Working on the runtime

A new tinygrad runtime is very easy to implement. There aren’t many layers of abstraction, and half of them are completely device-agnostic. All the complex logic is handled above device-specific code.

Our runtime is ~100 lines of Python in a single file. Yet it makes use of all the speed and optimization of the tinygrad framework.

A new runtime needs four parts:

  • Allocator: Handles allocating memory, freeing memory, and transferring data in and out of the device (between host and device)
  • Renderer: Knows what the target can do, and the specifics about memory layout. Turns the graph into code.
  • Compiler: Turns the string generated by Renderer into code.
  • Program: Runs the code.

tinygrad runtimes work as follows: the codegen pipeline linearizes the graph into a list of UOps, renders it into source code (via Renderer), and compiles it (via Compiler) into a TinyELF. The already-compiled TinyELF is then passed to Program, which ships it to the device and runs it. The whole program runs as one unit on the target device: a GPU, a CPU, etc.

But we’re Matmul as a Service, not Program as a Service. We want to keep our backend minimal, so we did something different.

We keep all control flow, indexing, and layout inside Python, and only send STORE and ALU operations on the remote backend. STORE is for sending data to the remote backend. ALU is all the arithmetic operations.

Yes, it means we need to send an HTTP request for every operation, which means a couple billion HTTP requests for a 1024x1024 matmul. But who cares? Bandwidth is free.

Since we don’t need a Compiler or a Renderer, and only need the raw list of UOps, we’ll only have to Allocator, and Program.

Allocator

The Allocator handles memory: alloc/free, copyin/copyout (host <> device), and offset (a view into a subregion of a buffer, which is just bookkeeping, no data moves).

We want all memory to live on our remote backend, because it’s more convenient, and also because we can charge for memory on top of compute (more $$$). Every allocation, every byte transfer, every free goes over HTTP:

class MaasAllocator(Allocator['MaasDevice']):
  def __init__(self, dev):
    super().__init__(dev)
    self.conn = http.client.HTTPConnection(dev.host, dev.port)

  def _alloc(self, size, options):
    self.conn.request("POST", "/alloc", json.dumps({"size": size}).encode(), ...)
    return MaasBuffer(json.loads(self.conn.getresponse().read())["id"], 0)

  def _copyin(self, dest, src):
    self.conn.request("PUT", f"/buf/{dest.buf_id}", bytes(src), ...)
    self.conn.getresponse().read()

  def _copyout(self, dest, src):
    self.conn.request("GET", f"/buf/{src.buf_id}")
    dest[:] = self.conn.getresponse().read()[src.offset:src.offset+len(dest)]

  def _free(self, opaque, options):
    self.conn.request("DELETE", f"/buf/{opaque.buf_id}")
    self.conn.getresponse().read()

  def _offset(self, buf, size, offset):
    return MaasBuffer(buf.buf_id, buf.offset + offset)

Program

Normally, a runtime’s Program receives an already-compiled TinyELF and ships the compiled binary to the device. The whole kernel runs as one unit on the target, and minimizes communication overhead.

We don’t do that. We interpret the UOps list ourselves in Python, keeping all control flow, indexing, and layout local. Only STORE and ALU operations get sent to the remote backend over HTTP.

class MaasProgram(Program['MaasDevice']):
  def __init__(self, dev, obj:TinyELF):
    ...

  def __call__(self, *bufs, global_size=(1,1,1), local_size=(1,1,1), vals=(), wait=False, **kw):

    values, pbufs, pvals, masks = {}, list(bufs), list(vals), [True]
    i = 0
    while i < len(self.uops):
      u = self.uops[i]
      sv = [values[v] for v in u.src if v.dtype is not dtypes.void]

      if u.op is Ops.RANGE:
        values[u] = 0 if u not in values else values[u]+1
        if values[u] == sv[0]:           # done looping?
          del values[u]
          i = self.loop_ends[u]+1         # skip past the loop
          continue
        i += 1
        continue

      if u.op in GroupOp.ALU:
        operands = [{"buf": v[0], "off": v[1]} if isinstance(v, tuple) else _encode_value(v) for v in sv]
        resp = self._send({"op": u.op.name, "operands": operands})
        values[u] = _decode_value(json.loads(resp)["value"])   # store the result
        i += 1
        continue

      ...

Two examples of how operations are implemented inside our new tinygrad runtime

First results

Addition

We started with the simplest thing that could possibly work:

from tinygrad import Tensor

a = Tensor([1.0, 2.0, 3.0])
b = Tensor([4.0, 5.0, 6.0])
c = a + b
print(c.numpy())

As you can see, nothing in this code says where it runs (yet another nice perk of tinygrad). The run will depend on the DEV environment variable:

DEV=CPU python3 test_add.py      # runs on CPU
DEV=CUDA python3 test_add.py     # runs on NVIDIA GPU
DEV=METAL python3 test_add.py    # runs on Apple Metal
DEV=MAAS python3 test_add.py     # runs on our HTTP backend

We intercept the call to Program and print out the UOps that get sent. Here’s what the addition kernel looks like:

   0 Ops.CONST           :            dtypes.weakint    []                3
   1 Ops.CAST            :            dtypes.int        ['3']             dtypes.int
   2 Ops.PARAM           :            dtypes.float      [1]               ParamArg(0, ...)
   3 Ops.PARAM           :            dtypes.float      [1]               ParamArg(1, ...)
   4 Ops.PARAM           :            dtypes.float      [1]               ParamArg(2, ...)
   5 Ops.RANGE           : 0          dtypes.int        [1]               (0, AxisType.WEAK)
   6 Ops.INDEX           : 0          dtypes.float      [3, 5]            None
   7 Ops.LOAD            : 0          dtypes.float      [6]               None
   8 Ops.INDEX           : 0          dtypes.float      [4, 5]            None
   9 Ops.LOAD            : 0          dtypes.float      [8]               None
  10 Ops.INDEX           : 0          dtypes.float      [2, 5]            None
  11 Ops.ADD             : 0          dtypes.float      [7, 9]            None
  12 Ops.STORE           : 0          dtypes.void       [10, 11]          None
  13 Ops.END             :            dtypes.void       [12, 5]           None
  14 Ops.SINK            :            dtypes.void       [13]              KernelInfo(...

Pretty easy to wrap your head around:

  1. CONST + CAST: the constant 3 (the loop bound), cast to an int
  2. PARAM: the three tensor arguments (input a, input b, output c)
  3. RANGE: a loop from 0 to 3
  4. INDEX: compute the index into each tensor using the loop variable
  5. LOAD: carry a (buffer_id, offset) reference forward
  6. ADD: send both references to the backend, which reads and adds them (one HTTP request)
  7. STORE: write the result into c[i] (another HTTP request)
  8. END: loop back
  9. SINK: the kernel terminator

So for a 3-element addition, we send 3 adds and 3 stores. That’s 6 HTTP requests to add three numbers. The loads are just local reference passing. The actual reads happen server-side as part of each ADD. Beautiful.

And it works:

$ DEV=MAAS python3 extra/maas/test_add.py
[5. 7. 9.]

Matmul

Now for the real test. A 2×2 matmul:

from tinygrad import Tensor

a = Tensor([[1.0, 2.0],[3.0, 4.0]])
b = Tensor([[5.0, 6.0],[7.0, 8.0]])
c = a @ b
print(c.numpy())

We can do the exact same and intercept the UOps list:

   0 Ops.CONST           :            dtypes.weakint                           []                               4
   1 Ops.CAST            :            dtypes.int                               ['4']                            dtypes.int
   2 Ops.PARAM           :            dtypes.float                             [1]                              ParamArg(0, dtypes.float, device='MAAS')
   3 Ops.PARAM           :            dtypes.float                             [1]                              ParamArg(1, dtypes.float, device='MAAS')
   4 Ops.PARAM           :            dtypes.float                             [1]                              ParamArg(2, dtypes.float, device='MAAS')
   5 Ops.CONST           :            dtypes.weakint                           []                               1
   6 Ops.CAST            :            dtypes.int                               ['1']                            dtypes.int
   7 Ops.CONST           :            dtypes.weakint                           []                               2
   8 Ops.CAST            :            dtypes.int                               ['2']                            dtypes.int
   9 Ops.RANGE           : 1          dtypes.int                               [8]                              (1, AxisType.WEAK)
  10 Ops.SHL             : 1          dtypes.int                               [9, 6]                           None
  11 Ops.INDEX           : 1          dtypes.float                             [3, 10]                          None
  12 Ops.LOAD            : 1          dtypes.float                             [11]                             None
  13 Ops.MULACC          : 1          dtypes.int                               [9, 8, 6]                        None
  14 Ops.INDEX           : 1          dtypes.float                             [3, 13]                          None
  15 Ops.LOAD            : 1          dtypes.float                             [14]                             None
  16 Ops.RANGE           : 1,2        dtypes.int                               [8, 9]                           (2, AxisType.WEAK)
  17 Ops.ADD             : 1,2        dtypes.int                               [16, 8]                          None
  18 Ops.INDEX           : 1,2        dtypes.float                             [4, 17]                          None
  19 Ops.LOAD            : 1,2        dtypes.float                             [18]                             None
  20 Ops.INDEX           : 1,2        dtypes.float                             [4, 16]                          None
  21 Ops.LOAD            : 1,2        dtypes.float                             [20]                             None
  22 Ops.MULACC          : 1,2        dtypes.int                               [9, 8, 16]                       None
  23 Ops.INDEX           : 1,2        dtypes.float                             [2, 22]                          None
  24 Ops.MUL             : 1,2        dtypes.float                             [15, 19]                         None
  25 Ops.MULACC          : 1,2        dtypes.float                             [12, 21, 24]                     None
  26 Ops.STORE           : 1,2        dtypes.void                              [23, 25]                         None
  27 Ops.END             : 1          dtypes.void                              [26, 16]                         None
  28 Ops.END             :            dtypes.void                              [27, 9]                          None
  29 Ops.SINK            :            dtypes.void                              [28]                             KernelInfo(
  1. CONST + CAST: constants 4, 1, 2 (dimensions and stride)
  2. PARAM: the three matrix arguments
  3. RANGE (outer): loop over output rows
  4. SHL: shift-left, compute the offset for the row
  5. INDEX + LOAD: read elements from b
  6. RANGE (inner): loop over the reduction dimension
  7. MULACC: multiply-accumulate
  8. STORE: write the result
  9. END × 2: close both loops
$ DEV=MAAS python3 extra/maas/test_matmul.py
[[19. 22.]
 [43. 50.]]

This is how you can run your matmuls on our infrastructure. This is the future. This is matmul as a service. Yes you’ll need ~20 HTTP requests for a 2x2 matmul. But again who cares? It runs on the cloud.

Benchmarks

We needed to benchmark this. What better way than to run a real model? We decided to run Kimi K3 on MaaS.

A 1024x1024 matmul alone is over a billion multiply-accumulates, or HTTP requests to MaaS. Kimi K3 has… well, even more. Preliminary results have been encouraging. We’ll share comprehensive benchmarks in a follow-up post.

Now for the important part: how much does this cost?

Pricing

We charge per unit of compute and per unit of memory, because that’s what you use:

ResourceUnitPrice
ALU operationper HTTP request$0.0001
Memoryper byte-hour$0.00001

We believe this is competitively priced. NVIDIA charges you for the GPU whether you use it or not. AI Labs charge you per token (what is even a token?). We only charge you for what you compute. If you don’t compute anything, you don’t pay anything.

We’re also raising a seed round. We want to move our compute backend away from our VPS to some dedicated servers and start using GPUs. This way it’s even faster.


Obviously none of the above is real. Matmul as a Service is not a product. We are not raising a seed round. You cannot run Kimi K3 one HTTP request at a time.

But we do build the thing that would have monitored the servers it ran on.

Simple Observability is a platform that provides full visibility into your servers. It collects logs and metrics through a lightweight agent, supports endpoint monitoring, cron monitoring, and exposes everything through a single web interface with centrally managed configuration.

To get started, visit simpleobservability.com.

The agent is open source and available on GitHub.