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)
工作流视频——点图播放
工作流是这样:
- 导入文档,生成 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:
演示视频——点图播放
我们的应用:
- 在 30 毫秒以内查完 3600 万以上向量。
- 1 秒以内生成回答。
搞定!
我们刚搭出了最快的 RAG 栈:用 BQ 做高效检索,用超快的 serverless 部署来跑这套 AI 工作流。
代码在这里:https://github.com/patchy631/ai-engineering-hub/tree/main/fastest-rag-milvus-groq
工作流再放一遍方便对照 👇
需要提一句:binary quantization 让向量搜索这一层效率高得离谱。但生产环境里,检索很少只是一次向量查询。
真实场景的 agent 会同时从 Slack、GitHub、Jira、数据库、文档拉 context。这意味着鉴权、同步、query 路由、权限、rerank 全都跟检索速度同等重要。
所以把 BQ 看成检索基础设施里关键的一环,不是检索的全部。向量层越快、越轻量,你就越有余力投到让检索可靠的其他事上。
最近一篇推文里我系统讲过整套检索基础设施长什么样,拆了 Google、Microsoft 这些公司在生产环境里到底怎么给 agent 喂 context。
往下读:
_[作者引用了 @avichawla 自己关于检索基础设施的另一条推文——见原 X 串]
收工!
如果你喜欢这篇教程:
来这里找我 → @_avichawla
每天我都会分享 DS、ML、LLM、RAG 相关的教程和洞见。