怎么把 RAG 内存压到 1/32(带代码讲解)

Perplexity、Azure、HubSpot 等公司用 binary quantization(二值量化)把 RAG 内存压到 1/32,下面带代码讲清楚。


业界常用一个简单技巧,能把 RAG 内存压到大概 1/32:

  • Perplexity 在它的搜索索引里用了
  • Azure 在它的搜索流水线里用了
  • HubSpot 在它的 AI 助手里用了

为了把这套学透,我们来搭一个 RAG 系统,30 毫秒以内查 3600 万级向量。

支撑这套的技术叫 Binary Quantization(二值量化,下文 BQ)。

技术栈:

  • LlamaIndex 做编排(开源)
  • Milvus 当向量数据库(开源)
  • Kimi-K2 当 LLM,托管在 Groq 上(hosted)

workflow video preview

工作流视频——点图播放

工作流是这样:

workflow gif

  • 导入文档,生成 binary embedding。
  • 建一个 binary 向量索引,把 embedding 存进向量数据库。
  • 取与用户 query 最相似的 top-k 文档。
  • LLM 基于额外的 context 生成回答。

我们来实现一下!


1) 加载数据

我们用 LlamaIndex 的 directory reader 工具把文档读进来。它能读各种数据格式——Markdown、PDF、Word、PowerPoint、图片、音频和视频都行。

from llama_index.core import SimpleDirectoryReader

loader = SimpleDirectoryReader(
    input_dir=docs_dir,
    required_exts=[".pdf"],
    recursive=True
)

docs = loader.load_data()
documents = [doc.text for doc in docs]

这里为了讲清楚用了一个简单的 PDF 目录读取器,但实际场景里,你的导入流水线大概率会从多个数据源拉数据,格式各异,预处理步骤也不一样。

2) 生成 binary embedding

接下来生成文本 embedding(float32),再把它们转成 binary 向量——内存和存储就这么变成原来的 1/32。

import numpy as np
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

embed_model = HuggingFaceEmbedding(
    model_name="BAAI/bge-large-en-v1.5",
    trust_remote_code=True,
    cache_folder='./hf_cache'
)

for context in batch_iterate(documents, batch_size=512):
    # Generate float32 vector embeddings
    batch_embeds = embed_model.get_text_embedding_batch(context)
    # Convert float32 vectors to binary vectors
    embeds_array = np.array(batch_embeds)
    binary_embeds = np.where(embeds_array > 0, 1, 0).astype(np.uint8)
    # Convert to bytes array
    packed_embeds = np.packbits(binary_embeds, axis=1)
    byte_embeds = [vec.tobytes() for vec in packed_embeds]

    binary_embeddings.extend(byte_embeds)

这一步就是 binary quantization。

3) 向量索引

二值量化做完之后,把向量存进 Milvus 向量数据库并建索引,方便后续高效检索。

from pymilvus import MilvusClient, DataType

# Initialize client and schema
client = MilvusClient("milvus_binary_quantized.db")
schema = client.create_schema(auto_id=True, enable_dynamic_fields=True)

# Add fields to schema
schema.add_field(field_name="context", datatype=DataType.VARCHAR)
schema.add_field(field_name="binary_vector", datatype=DataType.BINARY_VECTOR)

# Create index parameters for binary vectors
index_params = client.prepare_index_params()
index_params.add_index(
    field_name="binary_vector",
    index_name="binary_vector_index",
    index_type="BIN_FLAT",  # Exact search for binary vectors
    metric_type="HAMMING"   # Hamming distance for binary vectors
)

# Create collection with schema and index
client.create_collection(
    collection_name="fastest-rag",
    schema=schema,
    index_params=index_params
)

# Insert data to index
client.insert(
    collection_name="fastest-rag",
    data=[
        {"context": context, "binary_vector": binary_embedding}
        for context, binary_embedding in zip(batch_context, binary_embeddings)
    ]
)

索引是专门优化数据检索性能的数据结构。

4) 检索

检索阶段我们做这几件事:

  • 把用户 query 做 embedding,再对它做二值量化。
  • 用汉明距离(Hamming distance)作搜索度量,比较 binary 向量。
  • 取最相似的 top 5 chunk。
  • 把这些 chunk 加进 context。

实现如下:

# Generate float32 query embedding
query_embedding = embed_model.get_query_embedding(query)
# Apply binary quantization to query
binary_query = binary_quantize(query_embedding)

# Perform similarity search using Milvus
search_results = client.search(
    collection_name="fastest-rag",
    data=[binary_query],
    anns_field="binary_vector",
    search_params={"metric_type": "HAMMING"},
    output_fields=["context"],
    limit=5  # Retrieve top 5 similar chunks
)

# Store retrieved context
full_context = []
for res in search_results:
    context = res["payload"]["context"]
    full_context.append(context)

5) 生成

接着搭生成流水线,用 Kimi-K2 instruct 模型,托管在 Groq 上跑——Groq 自己宣称最快的 AI 推理服务。

我们把 query 和检索到的 context 一起塞进一个 prompt 模板,再交给 LLM。

from llama_index.llms.groq import Groq
from llama_index.core.base.llms.types import (
    ChatMessage, MessageRole )

llm = Groq(
    model="moonshotai/kimi-k2-instruct",
    api_key=groq_api_key,
    temperature=0.5,
    max_tokens=1000
)

prompt_template = (
    "Context information is below.\n"
    "---------------------\n"
    "CONTEXT: {context}\n"
    "---------------------\n"
    "Given the context information above think step by step "
    "to answer the user's query in a crisp and concise manner. "
    "In case you don't know the answer say 'I don't know!'.\n"
    "QUERY: {query}\n"
    "ANSWER: "
)

query = "Provide concise breakdown of the document"

prompt = prompt_template.format(context=full_context, query=query)
user_msg = ChatMessage(role=MessageRole.USER, content=prompt)

# Stream response from LLM
streaming_response = llm.stream_complete(user_msg.content)

最后用 Streamlit 包一层界面!

为了测规模和推理速度,我们在 PubMed 数据集(3600 万以上向量)上跑了一遍。

来段 demo:

demo video preview

演示视频——点图播放

我们的应用:

  • 在 30 毫秒以内查完 3600 万以上向量。
  • 1 秒以内生成回答。

搞定!

我们刚搭出了最快的 RAG 栈:用 BQ 做高效检索,用超快的 serverless 部署来跑这套 AI 工作流。

代码在这里:https://github.com/patchy631/ai-engineering-hub/tree/main/fastest-rag-milvus-groq

工作流再放一遍方便对照 👇

workflow gif

需要提一句:binary quantization 让向量搜索这一层效率高得离谱。但生产环境里,检索很少只是一次向量查询。

真实场景的 agent 会同时从 Slack、GitHub、Jira、数据库、文档拉 context。这意味着鉴权、同步、query 路由、权限、rerank 全都跟检索速度同等重要。

所以把 BQ 看成检索基础设施里关键的一环,不是检索的全部。向量层越快、越轻量,你就越有余力投到让检索可靠的其他事上。

最近一篇推文里我系统讲过整套检索基础设施长什么样,拆了 Google、Microsoft 这些公司在生产环境里到底怎么给 agent 喂 context。

往下读:

_[作者引用了 @avichawla 自己关于检索基础设施的另一条推文——见原 X 串]


收工!

如果你喜欢这篇教程:

来这里找我 → @_avichawla

每天我都会分享 DS、ML、LLM、RAG 相关的教程和洞见。