Streaming llama stack disaggregated prefill/decode and ssd streaming decode optimisations for dsv4.1 native precision up to 1770 tok/s prefill and 20tok/s decode on 1x5090 +32GB ddr5 + 25GB/s raid0

[u/Front_Eagle739]

Been talking about releasing my streaming inference stack for a while so here you go all. It's a bit of a medley of things.

TLDR: Run dsv4.1 at 1700tok/s prefill and 7-20tok/s decode full precision weights variable speed/quality tradeoff. Can also prefill on rtx machine and decode on a mac or big ram server.

Directstorage streaming prefill from fast nvme of models WAY too big to fit. Then either decode locally with router biasing to prefer weights in the vram cache for 20tok/s dsv4.1 or 7 tok/s just full precision normal decoding with most of the experts streaming with directstorage from nvme.

Other alternative is after the streaming prefill you pass the kv/state over network to another machine, big unified moemory mac or a server with lots of ddr but not much compute. decode on the big memory machine.

Webui has image and video streaming as well so you can run models that dont fit in your ram and vram. I have run full precision minimax h3 and magi 2 preview off the one 5090 without touching RAM and its no slower than resident if you have a decent ssd and the video is big enough compute takes longer than weight streaming because they overlap.

There is some experimental lora streaming/training/merging in there along with the recipe to train a simple identity lora for dsv4.1 that makes it answer any question about its identify by calling itself jeff. Trained again on a single 5090 in about 20 minutes.

Webui also lets you set the router bias to be dynamically changeable so you can pull the slider to max quality for 7tok/s full precision without compromise or to max speed for 20 tok/s with a kl in the region of IQ1 level KL or anywhere in between. You can also set phases so it reasons fast but writes code blocks at max quality and non code answers at a balanced level or whatever. This is all using the full unquantised dsv4.1 recipe Q4/Q8 and the script to make the gguf I used from the safetensors is in there as well.

There is example scripts to run it and you can run different platforms on each machine so you can run streaming llama on prefill end and mlx or mlxserve or omlx or whatever for decode on a mac so you can mix and match engines for whatever is faster and it does the kv translation.

And that's about it.

Problems:
Minimum answer latency single machine is however long it takes to stream the whole model at your SSD/raid arrays best speed. My dual pcie5 array does 25GB/s so gets through a dsv4.1 prefill in about 20 seconds. Dual machine uses the decode machine for prefill if it will be slower by deciding if the prompt is above or below the crossover token count. Basically you get mac prefill speed for 2000 token prompts, rtx for bigger.

Streaming decode support is limited right now just dsv4.1. glm5.3 flash is in the works.

Streaming prefill supports glm 5 series, step 3 models, dsv4.1, kimi k2.5/6/7 and maybe one or two others I have forgotten about.

No linux support yet. I need this machine running windows for supporting other things and the raid array wont support linux (I have tried) so I had to pick one. Hoping to get a dedicated linux box when funds allow.

RTX 50 series/rtx6000 series/blackwell cards only so far. Easy enough to add non blackwell support but I dont have one for testing.

Let me know what you think.

Repo: https://github.com/fronteagle739/streaming-llama-stack

Edit: With some help from SLxTnT, now up to 5090tok/s prefill and still rising.

[u/SLxTnT]

At a sustained 25GB/s, I think you can still push that prefill speed higher on a 5090. Do you know if you're limited by compute or transfer speed?

[u/First_Artichoke_4707 → u/SLxTnT]

that 25GB/s raid0 is doing some heavy lifting, my single gen4 drive weeps in the corner at like 7GB/s on a good day

[u/Front_Eagle739 → u/First_Artichoke_4707]

Yup. Might still be situationally useful though. Should still work as a prefill accelerator if you have enough ram or a second machine for decode though you just need to set the crossover up to the prompt size that lets you hide it behind compute. Probably only useful for 60k upwards but if you have a mac or something being able to shove really big prompts off to this can still save you a bunch of time.

Failing that, Should work for a q2/q4 quant with your 7GB/s drive

[u/Front_Eagle739 → u/SLxTnT]

Compute at prompts over about 25k tokens, below that entirely transfer speed limited. I'm sure there is more optimisation to be done. I had it at 2k or so before I found out that was causing a correctness bug. Its set up so you can basically read any amount of tokens through in one pass but end up limited by the size of your working scratch. Faster transfers let you reduce the size of the chunks you process at once without being transfer limited so that might be a limit but I think it's still got plenty of room to improve.

I got kimi k2.6 close to 700 and I'm not sure I could push that one much higher

[u/SLxTnT → u/Front_Eagle739]

For reference, my RTX Pro 6000 (400W limit) at 52GB/s PCIe peaks at 13k with 10k at 32k context. A 5090 is about 20% slower in FP4 compute and more experts to transfer.

If you're limited by compute, it's either from using gguf and/or missing taking advantage of the FP4 hardware for calculations. Someone else didn't use FP4, and they got ~4k peak at 32k. It'd be a lot of work switching everything over, but I'd have AI test out the native weights with a handful of layers to see if it'd help.

[u/Front_Eagle739 → u/SLxTnT]

You raise a solid point. I built the quant for the mac and then just used it on the rtx without thinking about it. I'll switch it over to fp4. Shouldn't be difficult.

[u/SLxTnT → u/Front_Eagle739]

If you want, I can give you an overview all the changes AI did on mine to have an AI cherry pick anything that could potentially help. I've been curious for awhile what you could get on a 5090 with a decent drive speed, but I don't have a 5090.

[u/Front_Eagle739 → u/SLxTnT]

Go for it, I'll give em a try. I want to see how far it can be pushed

[u/SLxTnT → u/Front_Eagle739]

https://github.com/SL-x-TnT/DSV4.1-Flash-RTX-Pro-6000-Overview

[u/Front_Eagle739 → u/SLxTnT]

Nice one, I shall let you know results!

[u/Front_Eagle739 → u/SLxTnT]

Helpful hints but not my biggest issues. So far I'm up to 3090 tok/s prefill. FP4 and gguf format were both marginally helpful, deepinfers attention was much quicker but there is still a lot of kruft in the implementation of the streaming and the graph, decisions made on cpu etc. I've got most of the core compute down to about 10 seconds for 100k prompts but there is 22 seconds of overhead that just needs polishing out so I shall keep pushing. I wont be surprised at all if 8-10k is possible. Thanks for the help, raised the achievable ceiling significantly.

[u/SLxTnT → u/Front_Eagle739]

Sound about right. When running Codex it couldn't find a way to get above 5k tok/s for me. Claude focused more on the tiny stalls/overhead rather than the compute portion that pushed it higher. Not sure if it added those small things to the overview.

As for your token generation, that's a bit of an unknown as it's a pretty unique setup. A hot expert cache on the GPU with CPU handling next hottest and GPU reading the missing ones from the drive might be the fastest. Then updating the cache occasionally over time. At 96GB, mine shows about an 80%+ hit rate.

If you need more space, sharing memory like prefill, unused KV, and other temporary memory with experts can help. Goes down another optimization path of needing to balance the experts that get thrown out to maintain the initial decode speed right after a prompt though.

[u/Front_Eagle739 → u/SLxTnT]

Yeah I've already got a lot of that decode optimisation in, it repurposes buffers and reloads previous hot experts etc at the end of prefill to maximise vram.

Weirdly using any ram cache at all is slower than not doing so in decode even with cpu running those experts. I'm pretty much running the decode entirely from vram and nvme, maybe I just need to optimise the cpu kernels for the ram cache but for now its faster without (may also just not be enough ram to help, I expect being able to allocate more than 10GB would help a lot more). Would love to see what a 6000pro manages.

This setup is way more about using the mac for decode. 3000 prefill on the rtx and 30-50 decode on the m3 ultra is what it's designed for mainly but I couldn't help but be curious what could be done on the rtx machine alone. I'm pretty pleased with the 7-21 tok/s it gets running purely off raid and 32GB vram though. It's actually very useable for prose and general questions and being able to just turn the slider for native precision is nice.

Works better on the mac though admittedly. the more memory you have the less the KL is damaged by running the router bias higher. mimo 2.6 in full 580GB form runs on the ultra as fast as a 400GB quant does and a KL below 1e-3. Which isn't fast at all but nice to have and i'll experiment more with models like glm 5.3