Folks with smaller VRAM but enough RAM - Token per second almost doubled for Q3 (unsloth)

#40
by nizoai - opened

https://github.com/nizoai/NTransformer

Try this. The test was done on RTX pro 5000 blackwell + 190 GB ram with Ryzen 9000

image

I thought ik_llama.cpp already did this with --prefetch-experts

always nice to have options tho

Hi! I'm one of those people with a limited hardware budget for inference, and I tried to squeeze as much out of my setup as i could without losing model quality. I'm still working on it and profiling the CPU path. Here's what i managed to do.

Hardware

  • GPU: 1× NVIDIA RTX PRO 6000 Blackwell Workstation, 96 GB (sm_120), driver 595.84
  • CPU: 2× AMD EPYC 9754
  • RAM: 1.5 TB DDR5-4800, 24×64 GB, 12 channels per socket
  • Memory bandwidth: Stream 344 GB/s per socket, 685 GB/s for both

Software

  • Linux 6.8
  • torch 2.9.1+cu128, CUDA 12.8
  • kt-kernel 0.7.0.post1 (build AVX512-BF16)
  • sglang-kt fork (SGLang + KTransformers)

Model

GLM-5.3-Flash 320B-A18B, original FP8 (306 GB), 1M context

Results

Decode

Batch Original (ms/step) Original (tok/s) Forked (ms/step) Forked (tok/s) Speedup
bs1 64.7 15.5 39.7 25.2 1.63×
bs2 — 22.3 — 32.5 1.46×
bs4 — 32.8 — 47.1 1.44×

Prefill (bs1)

Prompt Original (s) Original (tok/s) Forked (s) Forked (tok/s) Speedup
8k 15.0 546* 9.3 881* 1.61×
32k 48.7 673 28.5 1150 1.71×
65k 103.4 634* 64.4 1018* 1.61×

The base prompt in Claude Code with a few plugins and MCP servers lands right around 30k, so the first TTFB at session start dropped by almost half.

What I changed in the engine (no changes to model, only engine)

  1. Getting it to run on sm_120 (Blackwell): torch fallback for the DSA indexer top-k, DeepGEMM JIT disabled, kt-kernel built for the AVX512-BF16 path (no AMX on EPYC).
  2. Prefix cache for the hybrid model: checkpoints of the KDA/mamba recurrent states between turns (extra_buffer, track interval 256, a pool of 192 state slots). TTFT of the next turn stays around 10 s and does not grow with the conversation length; a checkpoint miss would cost +103 s per turn (the whole session re-prefilled).
  3. Fused Triton mHC kernels (hc_pre / hc_post) on the decode path: 58.6 → 46.8 ms per step.
  4. Further decode fusions: fused RMSNorm, fused gate top-k, KDA glue kernels merged (6 → 2 launches per layer), shared expert overlapped with the CPU experts, dead-mask skip in the KT wrapper — another −2.5 ms.
  5. Tiered CUDA-graph capture by context length (131k / 262k / 524k / 1M tiers): −6.2 ms.
  6. KT CPU-expert stream collapsed onto the main stream (no cross-stream handshakes per layer): −1.8 ms.
  7. Sparse-MLA prefill loop rewritten without host synchronizations: 256-row blocks instead of 8-row chunks with 3 host syncs each (135k syncs per 32k prompt before), plus three GPU expert slots with a 3-layer lookahead on the layerwise GPU prefill route, so the PCIe bus keeps streaming experts while the GPU computes the DSA layers.

Hi! I'm one of those people with a limited hardware budget for inference, and I tried to squeeze as much out of my setup as i could without losing model quality. I'm still working on it and profiling the CPU path. Here's what i managed to do.

Hardware

  • GPU: 1× NVIDIA RTX PRO 6000 Blackwell Workstation, 96 GB (sm_120), driver 595.84
  • CPU: 2× AMD EPYC 9754
  • RAM: 1.5 TB DDR5-4800, 24×64 GB, 12 channels per socket
  • Memory bandwidth: Stream 344 GB/s per socket, 685 GB/s for both

Software

  • Linux 6.8
  • torch 2.9.1+cu128, CUDA 12.8
  • kt-kernel 0.7.0.post1 (build AVX512-BF16)
  • sglang-kt fork (SGLang + KTransformers)

Model

GLM-5.3-Flash 320B-A18B, original FP8 (306 GB), 1M context

Results

Decode

Batch Original (ms/step) Original (tok/s) Forked (ms/step) Forked (tok/s) Speedup
bs1 64.7 15.5 39.7 25.2 1.63×
bs2 — 22.3 — 32.5 1.46×
bs4 — 32.8 — 47.1 1.44×

Prefill (bs1)

Prompt Original (s) Original (tok/s) Forked (s) Forked (tok/s) Speedup
8k 15.0 546* 9.3 881* 1.61×
32k 48.7 673 28.5 1150 1.71×
65k 103.4 634* 64.4 1018* 1.61×

The base prompt in Claude Code with a few plugins and MCP servers lands right around 30k, so the first TTFB at session start dropped by almost half.

What I changed in the engine (no changes to model, only engine)

  1. Getting it to run on sm_120 (Blackwell): torch fallback for the DSA indexer top-k, DeepGEMM JIT disabled, kt-kernel built for the AVX512-BF16 path (no AMX on EPYC).
  2. Prefix cache for the hybrid model: checkpoints of the KDA/mamba recurrent states between turns (extra_buffer, track interval 256, a pool of 192 state slots). TTFT of the next turn stays around 10 s and does not grow with the conversation length; a checkpoint miss would cost +103 s per turn (the whole session re-prefilled).
  3. Fused Triton mHC kernels (hc_pre / hc_post) on the decode path: 58.6 → 46.8 ms per step.
  4. Further decode fusions: fused RMSNorm, fused gate top-k, KDA glue kernels merged (6 → 2 launches per layer), shared expert overlapped with the CPU experts, dead-mask skip in the KT wrapper — another −2.5 ms.
  5. Tiered CUDA-graph capture by context length (131k / 262k / 524k / 1M tiers): −6.2 ms.
  6. KT CPU-expert stream collapsed onto the main stream (no cross-stream handshakes per layer): −1.8 ms.
  7. Sparse-MLA prefill loop rewritten without host synchronizations: 256-row blocks instead of 8-row chunks with 3 host syncs each (135k syncs per 32k prompt before), plus three GPU expert slots with a 3-layer lookahead on the layerwise GPU prefill route, so the PCIe bus keeps streaming experts while the GPU computes the DSA layers.

Hey! im actually planning to buy almost your exact same hardware setup in the future

was wondering if you tried ornith 1.5 397b I wonder how that model performs

if you havent- no big deal I was just wondering

Hey! im actually planning to buy almost your exact same hardware setup in the future

was wondering if you tried ornith 1.5 397b I wonder how that model performs

if you havent- no big deal I was just wondering

Haven't tried that model. Right now I'm running benchmarks for the kt-kernel optimizations, the GPU should free up by morning, and I'll give it a shot then.

Update: another 18–29 % off prefill

Follow-up to my earlier post on squeezing GLM-5.3-Flash 320B-A18B onto one RTX PRO 6000 Blackwell
96 GB + 2× EPYC 9754 + 1.5 TB DDR5. Same hardware, same model, same FP8 weights, no quantization
changes. Driver is now 595.91.07 (the only hardware-side change since last time).

Decode

Batch Baseline (ms/step) Baseline (tok/s)
bs1 64.7 15.5
bs2 — 22.3
bs4 — 32.8

Prefill (bs1)

Prompt Baseline (s) Baseline (tok/s)
8k 15.0 546
32k 48.7 673
65k 103.4 634

Where it is now

Decode (bs1, measured immediately after a cold prefill at that context)

Context Last post (ms/step) Now (ms/step) Now (tok/s) vs baseline
8k 39.7 33.4 29.9 1.94×
32k — 33.4 29.9 —
60k — 33.4 29.9 —
250k — 34.3 29.2 —
512k — 37.7 26.6 —

Prefill (bs1, cold, chunked_prefill_size 16384)

Prompt Baseline (s) Last post (s) Now (s) Now (tok/s) vs last post vs baseline
2k — 9.00 6.43 319 −28.6 % —
8k 15.0 9.25 6.57 1219 −29.0 % 2.28×
32k 48.7 28.40 22.59 1451 −20.5 % 2.16×
60k 103.4 † 58.80 48.26 1273 −17.9 % 2.14× †
250k — 369.66 332.35 752 −10.1 % —
512k — 1169.64 1118.64 469 −4.4 % —

† The baseline row is 65,536 tokens and the new one is 61,440, so that speedup is about 6 % flattered.
Every other row compares equal lengths.

The benefit shrinks with context because attention grows quadratically and neither of the new changes
touches attention. The useful range for an agentic coding session is the first three rows.

Still open

The biggest remaining item is the DSA indexer's eager scorer. On this card the Triton path is gated to
compute capability 8.6/8.9 and we're 12.0, so prefill always takes the eager loop, and one constant
(_DEFAULT_QUERY_CHUNK = 32) drives roughly 90 % of those 786,712 launches. Raising it would remove
about 665k launches and 3–4.5 s, but it changes the bmm batch size and therefore possibly the cuBLAS
algorithm, so it has to pass a selection-fingerprint and logit diff before it can ship. That's next.

Sign up or log in to comment