| --- |
| 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) |
|
|
| [](https://huggingface.co/spaces/dejanseo/fanout-diffusion) |
| [](https://opensource.org/licenses/Apache-2.0) |
| [](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 |
| ``` |
| |
