Rewritten Memory Manager for ostris/ai-toolkit Double-buffered async CUDA · bf16 precision · Sub-linear rank scaling
Chernihiv, Ukraine 🇺🇦
⚠️ REQUIRES TWO FILES
Results below are only achievable when both files are replaced:
One without the other will not produce these results.
What This Solves
Standard ai-toolkit layer offloading is sequential :
GPU computes layer N ████████░░░░░░░░
Transfer weights N+1 ████████
GPU computes layer N+1 ████████
↑ GPU idle ↑At rank 1024: weight matrices = hundreds of MB per layer.
Sequential transfer = 179 sec/iter . Full training = 100+ hours .
Viking Engine makes transfer invisible :
GPU computes layer N ████████
Transfer weights N+1 ████████ ← parallel CUDA stream
GPU computes layer N+1 ████████Transfer disappears from the profiler. Completely.
⚡ Benchmark & Performance Verification (FLUX.2-Dev / RTX 4090)
LoRA Configuration Speed (s/it) VRAM Memory Status Rank 128 (Optimized) 6.70s / 6.50s 24 GB (Zero OOM / Stable) Rank 512 (Deep Gesture) 8.97s 24 GB (Double Buffered) Rank 1024 (Extreme) 22.45s 24 GB (Full 8-bit Stack Forced) Rank 1280 (Extreme) 65.80s 24 GB (Full 8-bit Stack Forced)
Benchmark - Flux2-dev, RTX 4090, Rank 128
10 3 8 First Benchmark — Flux2-dev (32B params) · RTX 4090
Rank Trainable Params Before After Speedup 16 97,517,568 baseline 5.92 s/it — 32 ~200,000,000 ~12 s/it 6.57 s/it 1.8× 512 3,120,562,176 ~179 s/it ~14 s/it 12× 1024 6,241,124,352 OOM / 179 s/it ~31 s/it ∞ / 5.8×
6.24 billion trainable parameters on a single RTX 4090. Zero OOM. Zero crashes. 19.5% of entire Flux2 model trained simultaneously.
Sub-linear scaling
Parameters ×32 → Speed ×2.4 only (rank 16 → 512)
Parameters ×64 → Speed ×5.2 only (rank 16 → 1024)The engine becomes more efficient as rank increases.
Real log — rank 1024 stabilization:
step 1 → 136.94 s/it (cold start)
step 10 → 40.51 s/it (pipeline filling)
step 30 → 33.39 s/it
step 50 → 31.96 s/it
step 100 → 31.21 s/it ← stable ✓Iteration profile at rank 32:
backward: 3.85s ✓
predict_unet: 2.01s ✓
optimizer_step: 0.08s ✓
transfer: 0.00s ← hidden inside compute ✓Architecture
Two engines in one file
Engine A — Direct path (rank ≤ 32):
1 # Overhead from Events > transfer time at small ranks
2 # Direct non-blocking wins
3 w = self . m . weight . to ( device , non_blocking = True )
4 return F . linear ( x , _dequant ( w , dtype ) , b )
Engine B — Double-buffered async (rank ≥ 512):
1
2 The code is closed , but it continues to work and fly at supersonic speeds
3
bf16 precision patch — one line, 2.7× speedup at rank 512:
1 # Viking method — before network.apply_to()
2 # todo switch everything to proper mixed precision like this
3 self . network . force_to ( self . device_torch , dtype = torch . bfloat16 )
Forces LoRA matrices float32 → bfloat16. Weight size halved.
Additional optimizations:
CPU pinned memory — GPU DMA reads directly from DRAM, no cache copySmart text encoder orchestration — Mistral-24B loads, encodes, unloads before training startssm_89 support — correct FMA values for RTX 4090 (Ada Lovelace)High-priority CUDA streams — compute stream never waits in queue
Works With All ai-toolkit Models
Flux2 (32B) was the test model — heaviest available, worst-case benchmark.
✅ FLUX.1 / FLUX.2
✅ Stable Diffusion 1.x / 2.x / 3 / 3.5
✅ SDXL
✅ Video models (Wan, HunyuanVideo, etc.)
✅ Any future model in ai-toolkit
Installation
1 git clone https://github.com/ostris/ai-toolkit
2 cd ai-toolkit
3 pip install -r requirements.txt
4
5 # Backup originals
6 cp toolkit/manager_modules.py toolkit/manager_modules_BACKUP.py
7 cp jobs/process/BaseSDTrainProcess.py jobs/process/BaseSDTrainProcess_BACKUP.py
8
9 # Place Viking Engine files from this repo
In BaseSDTrainProcess.py find line ~1778, add before network.apply_to():
self.network.force_to(self.device_torch, dtype=torch.bfloat16)Recommended config (rank 32 — balanced):
1 network :
2 type : lora
3 linear : 32
4 linear_alpha : 64
5 conv : 32
6 conv_alpha : 64
7 lokr_full_rank : true
8 lokr_factor : -1
Expected on RTX 4090 + Flux2: 5.92 – 6.57 sec/iter
Server Scalability
1 # Consumer (this repo):
2 w = weight_cpu . to ( device , non_blocking = True )
3
4 # NVLink cluster — same pattern:
5 w = weight_gpu0 . to ( device_1 , non_blocking = True )
Same double-buffering. Same CUDA Streams. Same Events.
Proof of Quality
Models trained with Viking Engine at rank 1024 are on CivitAI.
🔗
Orakul Studio — CivitAI Context
ostris — author of ai-toolkit used by thousands —
requested this code for integration. The ticket is open.
The # todo switch everything to proper mixed precision like this comment
existed in the original source. We read it and went deeper.
Built in Chernihiv, Ukraine. Basement. Artillery fire. RTX 4090.Architecture matters more than hardware.
The smell of the iron is stable. 🦊
GitHub · CivitAI · Chernihiv, Ukraine 🇺🇦 · 2026 
