klein.cu is a CUDA/GPU port of klein.c (the CPU-only FLUX.2 Klein text-to-image generator by Camenduru). It preserves the original program structure, model architecture, and inference flow while moving all computationally heavy work onto the GPU.
klein.cu implements the complete FLUX.2 Klein inference pipeline as CUDA kernels:
- Qwen3 text encoder (36-layer transformer, 32/8-head GQA, RoPE, SwiGLU) — outputs concatenated layers 9/18/27 → [512, 7680]
- FLUX rectified-flow transformer (5 joint-attention double blocks + 20 single blocks, AdaLN, 2D RoPE) — Euler denoising
- FLUX VAE decoder (im2col + GEMM convolutions, group norm, mid self-attention, 2× nearest upsample)
| Component | CPU (klein.c) | GPU (klein.cu) |
|---|---|---|
| Matrix multiply | cblas_sgemm |
BF16 tensor-core WMMA GEMM (mma.sync, 32×64 tile, 8 warps) with F32 accumulation; plain BF16→F32 FMA fallback on pre-Ampere |
| Attention | Full [heads, seq, seq] score matrices |
Head-by-head, query-tiled (128 queries/tile) with bounded scratch |
| Weights | BF16/F16 → F32 on host (AVX512) | Stored as BF16 in device memory (direct load, no conversion); activations stay F32 |
| Latents | Host RAM, copied per step | Device-resident through the whole denoise loop |
| Conv2d | im2col + BLAS on host | im2col + BF16 tensor-core GEMM on device |
| Tokenizer | Host (BPE is sequential) | Host (unchanged) |
| PNG/BMP | Host I/O | Host I/O (unchanged) |
The original low-RAM sequential strategy is preserved:
Load Encoder → Encode Text → Free Encoder
Load Transformer → Denoise → Free Transformer
Load VAE → Decode → Free VAE
All weights are stored as BF16 in device memory (the model is distributed in BF16, so it is loaded directly without any dtype conversion), and activations/working buffers as float32; only the final image is copied back to the host. Norm/biases/embeddings are read as BF16 and converted to F32 on the fly inside the kernels.
Instead of materializing [heads, seq, seq] score matrices (which would be ~1.5 GB for the FLUX transformer at 112×112), attention_d processes one head at a time with a [128, kv_seq] score tile. This keeps VRAM bounded for large images.
klein.cu/
├── klein_cuda.h # Shared header: types, API, CUDA_CHECK/LAUNCH macros
├── klein_cuda_kernels.cu # Core kernels: GEMM, attention, norms, RoPE, RNG, ...
├── klein_cuda_loader.cu # Safetensors reader (BF16 direct-to-device) + tokenizer
├── klein_cuda_encoder.cu # Qwen3 text encoder
├── klein_cuda_transformer.cu # FLUX rectified-flow transformer
├── klein_cuda_vae.cu # FLUX VAE decoder
├── klein_cuda_main.cu # Entry point + PNG/BMP saving
├── CMakeLists.txt # CMake build configuration
└── README.md # This file
- NVIDIA GPU with CUDA compute capability 7.0+ (16+ GB VRAM recommended for the full model). BF16 tensor cores (WMMA) are used on Ampere and newer (SM 8.0+); older GPUs fall back to a plain BF16→F32 FMA GEMM.
- CUDA Toolkit 11.0+ (tested with 12.x)
- CMake 3.18+
- A C++17 compiler (MSVC or MinGW on Windows, GCC/Clang on Linux)
mkdir build
cd build
cmake .. -DCMAKE_CUDA_ARCHITECTURES=89 # 89 = RTX 40xx; adjust for your GPU
cmake --build . --config ReleaseCommon architecture values: 75 (Turing), 80/86 (Ampere), 89 (Ada), 90 (Hopper/Blackwell).
klein_cuda.exe <model_dir> [prompt] [-s steps] [-S seed] [-W width] [-H height]
Arguments:
model_dir— Path to the FLUX.2 Klein model directoryprompt— Text description (default: "a red apple")-s steps— Denoising steps (default: 1)-S seed— Random seed (default: 42)-W width— Output width (default: 64)-H height— Output height (default: 64)
Example:
klein_cuda.exe C:/models/flux-klein "a beautiful sunset over ocean" -s 4 -S 123 -W 512 -H 512Simply run klein_cuda.exe without arguments to launch the graphical interface:
klein_cuda.exe
The GUI provides:
- Text prompt input
- Model folder selection (with browse button)
- Width/Height/Seed/Steps configuration
- Generate button
- Status display with inference time
- Generated image preview
Outputs are written to output.png and output.bmp in the current directory.
klein.cu includes detailed timing for each pipeline stage:
================================================================================
PERFORMANCE TIMINGS
================================================================================
Encoder Loading: 4.55 seconds
Transformer Load: 4.40 seconds
VAE Loading: 0.13 seconds
---------------------------------------------------------------------------
Text Encoding: 0.79 seconds
Denoising: 8.22 seconds
VAE Decoding: 0.68 seconds
---------------------------------------------------------------------------
TOTAL INFERENCE: 18.77 seconds
================================================================================
(Measured on an RTX 3090, 512×512 output, 4 denoising steps.)
| Output size | Denoising | VAE Decoding | Total |
|---|---|---|---|
| 64×64 | 2.74 s | 0.31 s | 12.86 s |
| 512×512 | 8.22 s | 0.68 s | 18.77 s |
model_dir/
├── tokenizer/tokenizer.json
├── text_encoder/
│ ├── model-00001-of-00002.safetensors
│ └── model-00002-of-00002.safetensors
├── transformer/diffusion_pytorch_model.safetensors
└── vae/diffusion_pytorch_model.safetensors
- RNG: The CPU version uses a single serial xoshiro256** stream for the initial latent. The GPU version gives each element its own stream seeded from
(seed, index). The distribution is identical, but the exact per-seed image differs from the CPU build — this is inherent to any parallel RNG re-implementation. - BF16 weights: The model is distributed in BF16 and is kept as BF16 in device memory (no conversion). GEMMs use WMMA tensor cores (
mma.sync) with F32 accumulation; activations, norms, and biases are converted to F32 on the fly. This matches the precision of the original FP32 pipeline within BF16 rounding. - Numerics:
--use_fast_mathis enabled for performance; results may differ slightly from the CPU build in the last ULP. - Win32 GUI: Native window with controls
MIT License — same as the original klein.c project.