Add bounded ZeroGPU BF16 CUDA control

README.md

@@ -1,13 +1,18 @@
 ---
-title: Embedding Quantization Cuda Control
-emoji: 🏃
-colorFrom: gray
-colorTo: green
+title: Embedding Quantization CUDA Control
+emoji: 🔬
+colorFrom: blue
+colorTo: yellow
 sdk: gradio
-sdk_version: 6.20.0
-python_version: '3.12'
+sdk_version: 6.5.1
 app_file: app.py
 pinned: false
+license: apache-2.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Embedding Quantization CUDA Control
+
+A bounded ZeroGPU control lane for the matched local MLX embedding
+quantization experiment. The first milestone reproduces the frozen
+Qwen3-Embedding-0.6B BF16 vectors on CUDA and compares them with the saved MLX
+BF16 vectors. CUDA results do not reproduce MLX/Metal performance.

app.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+import json
+import os
+import time
+from pathlib import Path
+
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+import torch.nn.functional as F
+from transformers import AutoModel, AutoTokenizer
+
+
+MODEL_PATH = Path("/models/qwen3-embedding-0.6b")
+DATA_ROOT = Path(os.environ.get("REPRO_DATA_ROOT", "/data"))
+INPUT_PATH = DATA_ROOT / "inputs/retrieval_pairs.json"
+MLX_REFERENCE = DATA_ROOT / "local-reference/qwen3-embedding-0.6b/bf16.npz"
+OUTPUT_DIR = DATA_ROOT / "cloud-results/qwen3-embedding-0.6b"
+TASK = (
+    "Given a natural-language search query, retrieve the single passage "
+    "that best answers it"
+)
+
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, padding_side="left")
+model = AutoModel.from_pretrained(
+    MODEL_PATH,
+    torch_dtype=torch.bfloat16,
+    trust_remote_code=True,
+).to("cuda").eval()
+
+
+def detailed_instruction(query: str) -> str:
+    return f"Instruct: {TASK}\nQuery:{query}"
+
+
+def encode(text: str) -> np.ndarray:
+    batch = tokenizer(text, return_tensors="pt", truncation=True, max_length=32768)
+    batch = {key: value.to("cuda") for key, value in batch.items()}
+    with torch.inference_mode():
+        output = model(**batch)
+        vector = F.normalize(output.last_hidden_state[:, -1, :].float(), p=2, dim=-1)
+    return vector[0].cpu().numpy()
+
+
+def retrieval_metrics(queries: np.ndarray, documents: np.ndarray) -> tuple[dict, np.ndarray, np.ndarray]:
+    scores = queries @ documents.T
+    order = np.argsort(-scores, axis=1)
+    ranks = np.array([
+        int(np.where(order[index] == index)[0][0]) + 1
+        for index in range(len(queries))
+    ])
+    positive = np.diag(scores)
+    negative = scores.copy()
+    np.fill_diagonal(negative, -np.inf)
+    margins = positive - negative.max(axis=1)
+    return {
+        "pair_count": len(queries),
+        "top1": float(np.mean(ranks == 1)),
+        "recall_at_5": float(np.mean(ranks <= 5)),
+        "mrr": float(np.mean(1.0 / ranks)),
+        "mean_margin": float(margins.mean()),
+        "minimum_margin": float(margins.min()),
+        "mean_rank": float(ranks.mean()),
+        "worst_rank": int(ranks.max()),
+    }, scores, ranks
+
+
+@spaces.GPU(duration=300)
+def run_bf16_control() -> dict:
+    pairs = json.loads(INPUT_PATH.read_text())
+    query_texts = [detailed_instruction(item["query"]) for item in pairs]
+    document_texts = [item["document"] for item in pairs]
+
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats()
+        torch.cuda.synchronize()
+    started = time.perf_counter()
+    queries = np.stack([encode(text) for text in query_texts])
+    documents = np.stack([encode(text) for text in document_texts])
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+    elapsed = time.perf_counter() - started
+    metrics, scores, ranks = retrieval_metrics(queries, documents)
+    metrics.update({
+        "lane": "cuda-zerogpu-bf16",
+        "model": "Qwen/Qwen3-Embedding-0.6B",
+        "source_revision": "97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3",
+        "elapsed_seconds": elapsed,
+        "texts_per_second": len(query_texts + document_texts) / elapsed,
+        "torch_version": torch.__version__,
+        "cuda_device": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
+        "cuda_peak_bytes": torch.cuda.max_memory_allocated() if torch.cuda.is_available() else None,
+    })
+
+    comparison = None
+    if MLX_REFERENCE.exists():
+        reference = np.load(MLX_REFERENCE)
+        ref_all = np.concatenate([reference["queries"], reference["documents"]])
+        cuda_all = np.concatenate([queries, documents])
+        aligned = np.sum(ref_all * cuda_all, axis=1)
+        score_delta = scores - reference["scores"]
+        comparison = {
+            "mean_aligned_cosine_cuda_vs_mlx_bf16": float(aligned.mean()),
+            "minimum_aligned_cosine_cuda_vs_mlx_bf16": float(aligned.min()),
+            "score_rmse_cuda_vs_mlx_bf16": float(np.sqrt(np.mean(score_delta ** 2))),
+            "queries_with_rank_change": int(np.count_nonzero(ranks - reference["ranks"])),
+        }
+
+    result = {"metrics": metrics, "mlx_bf16_comparison": comparison}
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(
+        OUTPUT_DIR / "cuda-bf16.npz",
+        queries=queries,
+        documents=documents,
+        scores=scores,
+        ranks=ranks,
+        metrics=np.array(json.dumps(metrics)),
+    )
+    (OUTPUT_DIR / "cuda-bf16.json").write_text(json.dumps(result, indent=2) + "\n")
+    return result
+
+
+with gr.Blocks() as demo:
+    gr.Markdown(
+        "# Embedding Quantization CUDA Control\n"
+        "Runs one bounded 48-text BF16 control and writes the result to the attached private bucket. "
+        "The requested GPU duration is capped at five minutes."
+    )
+    run_button = gr.Button("Run 0.6B BF16 CUDA control", variant="primary")
+    output = gr.JSON(label="Result")
+    run_button.click(
+        fn=run_bf16_control,
+        outputs=output,
+        concurrency_limit=1,
+        api_name="run_bf16_control",
+    )
+
+
+if __name__ == "__main__":
+    demo.queue(default_concurrency_limit=1).launch()

requirements.txt

@@ -0,0 +1,4 @@
+transformers>=4.51,<5
+numpy>=2,<3
+spaces>=0.45
+torch>=2.8