OrakulStorm/ai-toolkit-bonememory

⚡ Viking Engine — BoneMemory

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:

File	Purpose
toolkit/manager_modules.pyd	⚡ Async double-buffer engine
jobs/process/BaseSDTrainProcess.py	🎯 bf16 precision patch

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

orakul_report_folder_logs

orakul_report.txt

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.

Speed ​​demonstration: (Rank 512)

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.
Higher rank = longer GPU compute = more time to hide transfer latency.

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):

# Overhead from Events > transfer time at small ranks
# Direct non-blocking wins
w = self.m.weight.to(device, non_blocking=True)
return F.linear(x, _dequant(w, dtype), b)

Engine B — Double-buffered async (rank ≥ 512):

The code is closed, but it continues to work and fly at supersonic speeds

bf16 precision patch — one line, 2.7× speedup at rank 512:

# Viking method — before network.apply_to()
# todo switch everything to proper mixed precision like this
self.network.force_to(self.device_torch, dtype=torch.bfloat16)

Forces LoRA matrices float32 → bfloat16. Weight size halved.
DMA transfer halved. Overlap efficiency increased.

Additional optimizations:

• CPU pinned memory — GPU DMA reads directly from DRAM, no cache copy
• Smart text encoder orchestration — Mistral-24B loads, encodes, unloads before training starts
• sm_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

git clone https://github.com/ostris/ai-toolkit
cd ai-toolkit
pip install -r requirements.txt

# Backup originals
cp toolkit/manager_modules.py toolkit/manager_modules_BACKUP.py
cp jobs/process/BaseSDTrainProcess.py jobs/process/BaseSDTrainProcess_BACKUP.py

# 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):

network:
  type: lora
  linear: 32
  linear_alpha: 64
  conv: 32
  conv_alpha: 64
  lokr_full_rank: true
  lokr_factor: -1

Expected on RTX 4090 + Flux2: 5.92 – 6.57 sec/iter

---

Server Scalability

# Consumer (this repo):
w = weight_cpu.to(device, non_blocking=True)

# NVLink cluster — same pattern:
w = weight_gpu0.to(device_1, non_blocking=True)

Same double-buffering. Same CUDA Streams. Same Events.
Maps directly to tensor parallelism on H100/A100 clusters.

---

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