与 Python 生态集成实战

实战：构建高性能 BPE Tokenizer

为什么 Tokenizer 是性能瓶颈？

在大型语言模型（LLM）的服务端，Tokenizer（分词器）是一个被严重忽视的性能瓶颈。当处理长文档或高并发请求时，Tokenization 可能占据总推理时间的 10-30%。

以 GPT-2 的 BPE（Byte-Pair Encoding）算法为例，Python 实现处理 100万 token 约需要 2-3 秒，而 Mojo 实现可以将这个过程压缩到 50ms 以内，加速约 40-60 倍。

BPE 核心算法概念

Mojo Tokenizer 核心实现

# fast_tokenizer.mojo
from algorithm import parallelize
from collections import Dict
from sys.info import num_physical_cores

struct BPETokenizer:
    var vocab: Dict[String, Int]       # token → id
    var merges: List[Tuple[String, String]]  # 合并规则列表
    var vocab_size: Int

    fn __init__(inout self, vocab_path: String):
        self.vocab = Dict[String, Int]()
        self.merges = List[Tuple[String, String]]()
        self.vocab_size = 0
        self._load_vocab(vocab_path)

    fn encode(self, text: String) -> List[Int]:
        """将文本转换为 token ID 列表。"""
        # 1. 初始化：每个 UTF-8 字节作为初始 token
        var tokens = List[String]()
        for i in range(len(text)):
            tokens.append(String(text[i]))

        # 2. 按优先级应用合并规则
        for merge in self.merges:
            var left = merge[].get[0, String]()
            var right = merge[].get[1, String]()
            var merged = left + right
            tokens = self._apply_merge(tokens, left, right, merged)

        # 3. 转换为 ID
        var ids = List[Int]()
        for t in tokens:
            if t[] in self.vocab:
                ids.append(self.vocab[t[]])
            else:
                ids.append(0)  # unk_token
        return ids

    fn encode_batch(self, texts: List[String]) -> List[List[Int]]:
        """并行批量编码。"""
        var results = List[List[Int]](len(texts))

        @parameter
        fn encode_one(i: Int):
            results[i] = self.encode(texts[i])

        # 并行处理批次中的每个文本
        parallelize[encode_one](len(texts), num_physical_cores())
        return results

    fn _apply_merge(
        self,
        tokens: List[String],
        left: String,
        right: String,
        merged: String
    ) -> List[String]:
        var new_tokens = List[String]()
        var i = 0
        while i < len(tokens):
            if i < len(tokens) - 1 and tokens[i] == left and tokens[i+1] == right:
                new_tokens.append(merged)
                i += 2
            else:
                new_tokens.append(tokens[i])
                i += 1
        return new_tokens

    fn _load_vocab(inout self, path: String):
        # 实际实现中从 JSON 文件加载词汇表
        print("从", path, "加载词汇表")

性能对比

调用 Hugging Face Transformers

from python import Python
from max.engine import InferenceSession
from max.tensor import Tensor, TensorSpec

def huggingface_with_max():
    var transformers = Python.import_module("transformers")
    var torch = Python.import_module("torch")

    # 1. 用 Hugging Face 加载 tokenizer（Python 侧）
    var tokenizer = transformers.AutoTokenizer.from_pretrained(
        "bert-base-uncased"
    )

    # 2. Tokenize 输入文本
    var texts = ["Hello, how are you?", "Mojo is amazing!"]
    var encoding = tokenizer(
        texts,
        padding=True,
        truncation=True,
        max_length=128,
        return_tensors="np"  # 返回 numpy 数组
    )

    # 3. 导出模型为 ONNX（一次性操作）
    # 已有 .onnx 文件则跳过此步
    var model_hf = transformers.AutoModel.from_pretrained("bert-base-uncased")
    torch.onnx.export(
        model_hf,
        (encoding["input_ids"], encoding["attention_mask"]),
        "bert.onnx",
        opset_version=14
    )

    # 4. 用 MAX Engine 加载并推理（Mojo 侧，高性能）
    var session = InferenceSession()
    var model = session.load("bert.onnx")

    var outputs = model.execute(
        "input_ids", encoding["input_ids"],
        "attention_mask", encoding["attention_mask"]
    )

    print("BERT 输出形状：", outputs.get[DType.float32]("last_hidden_state").shape())
    # [2, 128, 768] — batch=2, seq_len=128, hidden_dim=768

def main():
    huggingface_with_max()

混合 Python + Mojo 推理流水线

架构设计：职责分离

# pipeline.mojo — 混合推理流水线
from python import Python
from max.engine import InferenceSession
from max.tensor import Tensor
from algorithm import vectorize
from math import exp
from time import now

fn softmax_simd(
    logits: UnsafePointer[Float32],
    n: Int
) -> List[Float32]:
    """SIMD 加速的 Softmax。"""
    # 找最大值（数值稳定性）
    var max_val: Float32 = logits[0]
    for i in range(1, n):
        if logits[i] > max_val:
            max_val = logits[i]

    # 计算 exp(x - max) 并求和
    var probs = List[Float32](n)
    var total: Float32 = 0.0
    for i in range(n):
        probs[i] = exp(logits[i] - max_val)
        total += probs[i]

    # 归一化
    for i in range(n):
        probs[i] /= total
    return probs

def classify_image(image_path: String) -> String:
    # Python 侧：图像加载和预处理
    var PIL = Python.import_module("PIL.Image")
    var np = Python.import_module("numpy")
    var json = Python.import_module("json")

    var img = PIL.open(image_path).resize((224, 224))
    var arr = np.array(img, dtype="float32") / 255.0
    arr = np.expand_dims(np.transpose(arr, (2,0,1)), 0)

    # Mojo 侧：MAX Engine 推理
    var session = InferenceSession()
    var model = session.load("resnet50.onnx")
    var t0 = now()
    var outputs = model.execute("input", arr)
    print("推理耗时：", (now()-t0)//1_000_000, "ms")

    # Python 侧：加载类别标签并输出
    var labels_file = open("imagenet_labels.json")
    var labels = json.load(labels_file)
    var output = outputs.get[DType.float32]("output")
    var top_class = int(np.argmax(np.array(output)))
    return String(labels[top_class])

Mojo 标准库演进路线图

当前状态（2024 年）

Mojo 仍在快速迭代中，标准库覆盖率相比 Python 还有较大差距。了解当前状态有助于做出合理的技术选型：

社区生态与局限性

当前优势

• MAX Engine：AI 推理领域已经生产就绪，有真实客户案例
• Python 兼容性：可以无缝调用整个 Python/PyTorch 生态，技术风险低
• SIMD/并行原语：在计算密集场景已经非常成熟，性能出色
• 积极的社区：Discord 社区活跃，GitHub issue 响应快速

当前局限性

适合现在就用 Mojo 的场景

• AI/ML 推理服务的性能优化（配合 MAX Engine）
• 计算密集的数值计算内核（矩阵操作、信号处理、图像处理）
• 需要 Python 生态 + 极致性能的混合系统
• 研究/实验性的 AI 系统编程探索