---
language:
- en
tags:
- diffusion
- consistency-models
- tensor-core
- 1-bit
- int4
- retrieval
- query-expansion
- onnx
- triton
- embeddinggemma
license: apache-2.0
spaces:
- dejanseo/fanout-diffusion
---

# Fanout Diffusion Retriever (540k)

[![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Live%20Demo-blue)](https://huggingface.co/spaces/dejanseo/fanout-diffusion)
[![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-green.svg)](https://opensource.org/licenses/Apache-2.0)
[![ONNX](https://img.shields.io/badge/ONNX-Universal%20Export-brightgreen)](https://huggingface.co/dejanseo/fanout-diffusion/blob/main/checkpoints/champion_b1_consistency_1step.onnx)

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](https://huggingface.co/spaces/dejanseo/fanout-diffusion)**

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)

```bash
pip install huggingface_hub onnxruntime-gpu sentence-transformers
```

```python
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`)

```bash
# 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:
```bash
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
```