dejanseo/fanout-diffusion

Fanout Diffusion Retriever (540k)

A sub-millisecond, hardware-accelerated continuous consistency retrieval engine that expands broad search queries into coherent 10-query fan-out sets in dense embedding space ($D = 768$) in a single forward pass.

👉 Try the Live Web Playground on Hugging Face Spaces

Built with 1-bit Tensor Core PTX MMA kernels, 1-step consistency distillation, and INT4 outer quantization, delivering:

• 0.329 ms single-query latency (RTX 4090, 1000x faster than legacy multi-step EDM)
• 417,035 QPS native throughput (4,170,350 fanout vectors/sec in standalone C++)
• 1.64 MB model checkpoint (INT4 outer + 1-bit weights)
• 97.5% prompt alignment parity with ground truth targets across all 55,819 queries in the 540k dataset
• Zero-dependency deployment via native C++ executable, ONNX Runtime, and Triton Inference Server

---

Performance & Architecture Comparison

Model Architecture	Precision	Forward Steps	Single-Query Latency	Batch 256 Latency	Throughput (QPS)	Parity vs GT	Checkpoint Size
Legacy EDM Baseline	FP32	16 (Heun ODE)	330.0 ms	~4,200 ms	~61 QPS	100.0%	165.2 MB
B1 Consistency (PyTorch)	1-bit + FP32	1 (Single-step)	0.812 ms	3.25 ms	~78,769 QPS	97.5%	7.35 MB
B1 Consistency QAT (INT4)	1-bit + INT4	1 (Single-step)	0.515 ms	1.82 ms	~140,650 QPS	97.5%	1.64 MB
ONNX Runtime (TensorRT/CUDA)	1-bit graph	1 (Single-step)	0.650 ms	1.95 ms	~131,280 QPS	97.5%	30.64 MB
Native C++ Engine (mma.sync)	1-bit PTX	1 (Single-step)	0.329 ms	0.614 ms	417,035 QPS	97.5%	7.35 MB

---

Semantic Quality Evaluation (540k Dataset)

Evaluated across all 55,819 test queries (558,190 fanout vectors) in data/diffusion_dataset_540k.pt:

Evaluation Metric	Legacy EDM (16 steps)	1-Step Consistency Champion	Relative Parity
Prompt Alignment (Cosine Sim)	0.7001	0.6829	97.5%
Pairwise Diversity	0.3500	0.3282	93.8%
Target Manifold MSE	0.0000	0.000934	High fidelity
Subquery Coverage ($\ge 0.60$)	71.4%	69.02%	96.7%
Inference Time (Full 55,819 queries)	4.8 hours	2.64 seconds	6,500x speedup

---

1-Bit Architectural Frontier Sweep (Ablation Study)

An empirical frontier sweep (F1–F7) was executed across backbone depth, direct 768-dim embedding projections, wide MLPs, Gram matrix diversity loss, and multi-step consistency leaps:

Run	Configuration	Val Loss	Alignment	Diversity	Latency	Outcome
Champion B1 (Production)	$512\text{d}, 1024\text{ MLP}, \text{Gram } 0.5$	0.00093	0.683	0.320	0.200 ms	Locked Production Baseline; crispest semantic specialization.
Frontier F1	$4\text{L Deep}, 512\text{d}, 1024\text{ MLP}$	0.00179	0.505	0.181	7.063 ms	Increased depth slowed inference 1.7x without 1-step benefit.
Frontier F2	$2\text{L}, 768\text{d Native (No Bottleneck)}$	0.00196	0.540	0.184	3.988 ms	Bypassing projection bottlenecks boosted alignment (+0.035).
Frontier F3	$2\text{L}, 512\text{d}, 2048\text{ Wide MLP}$	0.00188	0.500	0.217	4.039 ms	Wider MLPs increased slot separation (+18% diversity).
Frontier F4	$2\text{L}, 512\text{d}, \text{Gram } 1.5$	0.00231	0.420	0.336	4.235 ms	High diversity, but excessive repulsion degraded alignment.
Frontier F5	$2\text{L}, 512\text{d}, 2\text{-Step Consistency Leap}$	0.00073	0.708	0.039	0.350 ms	Highest numerical alignment, but diversity collapsed across slots.
Frontier F6	$2\text{L}, 768\text{d}, 2048\text{ MLP}, \text{Gram } 1.5$	0.00221	0.479	0.275	4.225 ms	Compounding F2+F3+F4 caused semantic drift from 1.5 Gram repulsion.
Frontier F7	$2\text{L}, 768\text{d}, 2048\text{ MLP}, \text{Gram } 0.3$	0.00154	0.600	0.121	5.554 ms	Calibrated Gram restored alignment to 0.600, but slots clustered.

---

Repository Files

• checkpoints/champion_b1_consistency_1step_qat.pt (1.64 MB): The champion INT4 QAT 1-step consistency checkpoint.
• checkpoints/champion_b1_consistency_1step.onnx (30.64 MB): Standalone ONNX export compatible with ONNX Runtime & Triton.
• models/fanout_1bit_weights.bin (7.35 MB): Raw packed binary weights for native C++ mma.sync execution.
• bin/fanout_engine.exe (0.48 MB): Standalone native Windows C++ binary with CUDA PTX MMA kernel.
• fanout_deployment_bundle.zip (26.26 MB): Self-contained production deployment bundle.
• data/taxonomy_embeddings.pt (17.14 MB): Pre-computed Google Product Taxonomy embeddings ($D = 768$).
• data/diffusion_dataset_540k.pt (1.8 GB): 55,819 query sets (558,190 query-fanout pairs) embedded via google/embeddinggemma-300m.
• diffusion_540k_best.pt (165.2 MB): Legacy 16-step EDM FP32 checkpoint for baseline comparison.
• load_model.py: Universal loader supporting ONNX, PyTorch 1-step, and legacy EDM.

---

Quickstart

1. ONNX Runtime (Zero Build Dependencies)

pip install huggingface_hub onnxruntime-gpu sentence-transformers

import torch
from sentence_transformers import SentenceTransformer
from load_model import load_fanout_onnx, sample_consistency_onnx

# 1. Load ONNX model directly from Hugging Face Hub
session = load_fanout_onnx(repo_id="dejanseo/fanout-diffusion")

# 2. Encode broad query using Google embeddinggemma-300m
embedder = SentenceTransformer("google/embeddinggemma-300m", device="cuda")
query_vector = embedder.encode(["running shoes and athletic sneakers"], convert_to_tensor=True, normalize_embeddings=True)

# 3. Generate 10 continuous fan-out vectors in a single pass (sub-millisecond)
fanout_vectors = sample_consistency_onnx(session, query_vector)
# fanout_vectors shape: [1, 10, 768]

2. Standalone C++ Engine (fanout_engine.exe)

# Run 10,000 queries through native PTX Tensor Core kernel
fanout_engine.exe --model models/fanout_1bit_weights.bin --batch 256 --warmup 200 --benchmark 1000

Output:

Batch 256 Latency: 0.614 ms | Throughput: 417,035 queries/sec (4,170,350 vectors/sec)
Single-Query Latency: 0.329 ms

3. Triton Inference Server

The repository includes a production-ready Triton configuration:

docker run --gpus all --rm -p 8000:8000 -p 8001:8001 -p 8002:8002 \
  -v $(pwd)/triton_model_repository:/models \
  nvcr.io/nvidia/tritonserver:24.08-py3 \
  tritonserver --model-repository=/models