Add bounded larger-family CUDA quantizer controls

app.py

@@ -371,6 +371,114 @@ def run_family_bf16_control(family: str) -> dict:
         }
 
 
+def _run_family_quantized_control(family: str, variant: str) -> dict:
+    if family not in BF16_FAMILIES:
+        raise ValueError(f"unsupported family: {family}")
+    if variant not in QUANTIZERS:
+        raise ValueError(f"unsupported quantizer: {variant}")
+
+    spec = BF16_FAMILIES[family]
+    config = quantization_config(variant)
+    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]
+    family_tokenizer = AutoTokenizer.from_pretrained(
+        spec["path"], padding_side="left", trust_remote_code=True
+    )
+
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
+    torch.cuda.synchronize()
+    allocation_before = int(torch.cuda.memory_allocated())
+    load_started = time.perf_counter()
+    family_model = AutoModel.from_pretrained(
+        spec["path"],
+        quantization_config=config,
+        device_map={"": 0},
+        trust_remote_code=True,
+    ).eval()
+    torch.cuda.synchronize()
+    load_seconds = time.perf_counter() - load_started
+    allocation_after_load = int(torch.cuda.memory_allocated())
+
+    encode_started = time.perf_counter()
+    queries = np.stack([
+        encode_with(family_model, text, family_tokenizer) for text in query_texts
+    ])
+    documents = np.stack([
+        encode_with(family_model, text, family_tokenizer) for text in document_texts
+    ])
+    torch.cuda.synchronize()
+    encode_seconds = time.perf_counter() - encode_started
+    metrics, scores, ranks = retrieval_metrics(queries, documents)
+    metrics.update({
+        "lane": f"cuda-zerogpu-{variant}",
+        "family": family,
+        "model": spec["source"],
+        "source_revision": spec["revision"],
+        "quantizer": variant,
+        "quantization_config": config.to_dict(),
+        "load_seconds": load_seconds,
+        "encode_seconds": encode_seconds,
+        "texts_per_second": len(query_texts + document_texts) / encode_seconds,
+        "torch_version": torch.__version__,
+        "cuda_device": torch.cuda.get_device_name(0),
+        "cuda_allocation_before_load": allocation_before,
+        "cuda_allocation_after_load": allocation_after_load,
+        "cuda_incremental_model_allocation": allocation_after_load - allocation_before,
+        "cuda_peak_bytes": int(torch.cuda.max_memory_allocated()),
+    })
+    output_dir = DATA_ROOT / "cloud-results" / family
+    result = {
+        "metrics": metrics,
+        "cuda_bf16_comparison": compare_vectors(
+            queries,
+            documents,
+            scores,
+            ranks,
+            output_dir / "cuda-bf16.npz",
+            "vs_cuda_bf16",
+        ),
+        "mlx_bf16_comparison": compare_vectors(
+            queries,
+            documents,
+            scores,
+            ranks,
+            DATA_ROOT / "local-results/full-q4-q6-q8" / family / "quality/bf16.npz",
+            "vs_mlx_bf16",
+        ),
+    }
+    output_dir.mkdir(parents=True, exist_ok=True)
+    stem = f"cuda-{variant}"
+    np.savez_compressed(
+        output_dir / f"{stem}.npz",
+        queries=queries,
+        documents=documents,
+        scores=scores,
+        ranks=ranks,
+        metrics=np.array(json.dumps(metrics)),
+    )
+    (output_dir / f"{stem}.json").write_text(json.dumps(result, indent=2) + "\n")
+
+    del family_model, family_tokenizer
+    gc.collect()
+    torch.cuda.empty_cache()
+    return result
+
+
+@spaces.GPU(duration=300)
+def run_family_quantized_control(family: str, variant: str) -> dict:
+    try:
+        return _run_family_quantized_control(family, variant)
+    except Exception as error:
+        return {
+            "diagnostic_error": type(error).__name__,
+            "diagnostic_message": str(error),
+            "diagnostic_traceback": traceback.format_exc(),
+        }
+
+
 with gr.Blocks() as demo:
     gr.Markdown(
         "# Embedding Quantization CUDA Control\n"
@@ -416,6 +524,24 @@ with gr.Blocks() as demo:
         concurrency_limit=1,
         api_name="run_family_bf16_control",
     )
+    gr.Markdown("## Larger-family CUDA-native quantizer controls")
+    quant_family = gr.Dropdown(
+        choices=list(BF16_FAMILIES),
+        value="qwen3-embedding-8b",
+        label="Model family",
+    )
+    family_quantizer = gr.Dropdown(
+        choices=list(QUANTIZERS), value="bnb-int8", label="Quantizer"
+    )
+    family_quant_button = gr.Button("Run bounded family quantizer control")
+    family_quant_output = gr.JSON(label="Family quantized result")
+    family_quant_button.click(
+        fn=run_family_quantized_control,
+        inputs=[quant_family, family_quantizer],
+        outputs=family_quant_output,
+        concurrency_limit=1,
+        api_name="run_family_quantized_control",
+    )
 
 
 if __name__ == "__main__":