Extracting maximum efficiency on older gpus from Mojo kernel and max serve

First of all thanks to the Modular team for creating Mojo and MAX.

I wanted to share one experiment we have been doing at https://cogitx.ai/ (we are early stage AI startup) because I think it shows something interesting about what is still possible even on an older GPU like the NVIDIA T4.

Our original problem was not really single user inference. For agentic workloads we need multiple persistent users, growing context, paged KV cache and good decode throughput at the same time.

With llama.cpp we could run GGUF models very nicely for one or two requests, but concurrency was where things started becoming difficult for our workload. vLLM with gguf is python wrapper so you are throwing every efficiency out of the box.

So we started going lower level with Mojo and MAX and tuning the path specifically for SM75.

The model here is Qwen3-1.7B on a single NVIDIA T4 16GB.

Our first optimized path reached this:

Batch v5 step v7 step v8 step v8 aggregate
1 13.69 ms 13.96 ms 8.03 ms 124.5 tok/s
2 28.02 ms 14.43 ms 8.17 ms 244.9 tok/s
4 29.29 ms 14.71 ms 8.43 ms 474.4 tok/s
8 31.55 ms 15.37 ms 9.57 ms 835.9 tok/s
10 32.58 ms 15.74 ms 10.73 ms 931.9 tok/s

The thing I found very interesting here is the scaling.

At B=1 we were doing around 124 tok/s and at B=10 we were still giving around 93 tok/s/user, or ~932 tok/s aggregate.

This came after progressively replacing/tuning different pieces rather than treating inference as one GEMM benchmark.

We worked on:

  • custom paged decode attention
  • SM75/Turing specific HMMA decode linears
  • quantized projection paths
  • Q4_0 dequantization
  • RMSNorm / residual paths
  • RoPE
  • KV writes and paged KV layout
  • warp-level reductions and tiling
  • FlashDecode/FlashAttention-style attention paths
  • batching and scheduling around multiple active sequences

The first change was mainly around paged decode attention. Then , we moved the decode linear path to our SM75 HMMA, which brought B=10 from ~635 tok/s to ~932 tok/s.

After that we moved the model path back toward the format we actually care about in deployment: GGUF Q4 weights, instead of stopping at the W8A16 experiment.

That gave us our current short-context peak:

10 concurrent requests
Qwen3-1.7B Q4 GGUF
NVIDIA T4 / SM75
256-token context
~1337 generated tokens/sec aggregate

This was probably the number that surprised me the most.

The T4 is a 2018 Turing GPU with roughly 320 GB/s theoretical memory bandwidth. Our independent bandwidth probe on the same GPU was around 311.6 GB/s, so once we started looking at the decoder from a bytes moved perspective, a lot of the optimisation decisions became much clearer.

But obviously 1,337 tok/s at 256 context is not the full story.

Our actual target is an agentic workload, something more like:

~3K system + tools + policies
        ↓
user turn / tool calls
        ↓
delta prefill
        ↓
decode
        ↓
reuse paged KV
        ↓
next turn
        ↓
context keeps growing

So we also measured how the same B=10 path behaves as KV grows.

Context Aggregate decode
256 1,337 tok/s
~1K 1,015 tok/s
~2K 840 tok/s
~3K 699 tok/s
~4K 613 tok/s

And in one of our more realistic agent style runs at roughly 3.1K active context, we were around 633 tok/s aggregate for the 10 users.

For me this result is actually more important than the 1.3K headline.

Once context starts growing, the problem very quickly stops being only about making the projection kernel faster. Paged KV access, attention traffic, prefix reuse, delta prefill, scheduling and interference between prefill/decode start dominating the design.

We built a separate stateful harness for exactly this reason. It models multiple persistent users with a large shared system or tool prompt, incremental user turns, reused paged KV state and eventually context compaction.

That is where MAX has been particularly interesting for us.

Instead of having to choose between a high level serving runtime and low level kernel control, we can keep MAX for the graph/runtime/cache/scheduling side and selectively own the kernels where the hardware or quantization format needs something different.

The path we are experimenting with now is roughly:

GGUF Q4 weights
      ↓
custom SM75 quantized kernels
      ↓
paged / quantized KV
      ↓
custom decode attention
      ↓
MAX graph + runtime
      ↓
continuous multi-user serving

But getting from ~307 tok/s at B=10 in the early path to ~932 tok/s, and eventually ~1.34K tok/s on the GGUF-native short-context path, on the same T4 has been a very interesting exercise.

More than anything, Mojo has made it possible for us to work at the level of warps, lanes, tiles, registers and memory movement without having to build the entire serving system in CUDA/C++.

Happy to share more of the kernel breakdown and profiling numbers if useful.

2 Likes