过去6-8个月，我一直在努力转向ML系统和AI基础设施领域。以下是我最喜欢的一些工作——

1. 一个小型Python推理引擎，主要作为我实现主要推理技术的试验台，在RTX 4060 Ti上启用前缀缓存后，吞吐量约600 tokens/s - https://github.com/JINO-ROHIT/tachyon…
2. 写了一篇博客，使用CuTe DSL在Ada上性能超越cuBLAS - https://jino-rohit.github.io/blogs/08_cute_dsl_matmul.html…
3. 一篇关于NCCL集体通信的详细文章 - https://jino-rohit.github.io/blogs/11_collective_communication.html…
4. 系列博客文章，介绍torch compile的工作原理和内部机制 - https://jino-rohit.github.io/blogs/
5. 我在这里维护大部分ML系统相关工作的笔记和实验 - https://github.com/JINO-ROHIT/ml-systems-notes…
6. 在SGLang和llm-compressor（vLLM）中的开源工作。

我正努力成为一名更强大的ML系统和推理工程师。接下来几个月我应该专注于提升什么？

JINO-ROHIT/tachyon

来源：https://github.com/JINO-ROHIT/tachyon

tachyon

一个可在消费级硬件上运行的LLM推理引擎。

1. 每个请求拥有独立状态，你可以按请求级别选择启用或禁用KV缓存。
2. 支持连续批处理：解码步骤被批处理以充分利用GPU，请求一旦从活动批次中完成即被回收，允许新请求进入批次。
3. 预填充步骤使用较小的批次大小进行批处理，以避免内存溢出。
4. 向量化操作以提高批处理性能。
5. 前缀缓存：当存在共享系统提示或少量示例等时，节省昂贵的预填充计算。
6. 仅需3行代码即可调用引擎并执行推理！

用法

1. 首先下载Llama 1B的权重。

hf_hub_download(
    repo_id=f"meta-llama/Llama-3.2-1B-Instruct",
    filename="model.safetensors",
    local_dir=f"Llama-3.2-1B-Instruct"
)

2. 调用引擎。

from tachyon.engine.llm import Engine
engine = Engine("meta-llama/Llama-3.2-1B-Instruct")
print(engine.generate_text("Explain AGI"))  # 单个请求

# 多个请求
outputs = engine.generate_text([
    "Explain AGI",
    "What is vLLM?",
    "Tell me about SGLang"
])

for o in outputs:
    print(o)

OpenAI兼容服务器

你可以启动服务器，然后使用openai库调用模型：

from openai import OpenAI

client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="anything"
)

resp = client.chat.completions.create(
    model="llama-3.2-1b-instruct",
    messages=[{"role": "user", "content": "Hello"}],
    stream=False,
)

3. 基准测试脚本

python3 benchmark.py

当前基准测试结果（RTX 4060 Ti 16GB）

待办事项

• 实现Llama 3系列模型。
• 将其构建为服务引擎。
• 编写基准测试脚本并检查延迟和吞吐量。
• 添加KV缓存。
• 连续批处理。
• 前缀缓存。
• 添加OpenAI兼容API服务器。
• 测试torch compile的效果（开启PR，回头再处理）。
• 分页注意力。
• 更多技术。