pymllm Runtime Design

总览

pymllm 是 mllm 的 Python serving runtime。它不是传统意义上的 mllm C++ Backend，而是一套围绕 PyTorch/CUDA 生态构建的在线推理服务运行时。当前实现面向 Jetson Orin 等边缘 GPU 设备，重点支持 Qwen3、Qwen3-VL 和 Qwen3.5 系列模型。

它的设计参考了 SGLang serving runtime 的核心分层，但进行了明显收缩：当前主路径以 单机单 GPU 为目标，优先保证在 Jetson 上可运行、可调试、可扩展，而不是覆盖大规模 分布式 serving 的全部复杂度。

pymllm runtime architecture

Figure 1: pymllm runtime architecture.

整体分层

从开发者视角看，pymllm 可以分为五层：

  1. 服务入口层：FastAPI HTTP server，提供 OpenAI-compatible API 和原生 /generate API。

  2. 配置层：ServerConfig、ModelConfig、QuantizationConfig 统一解析 模型路径、dtype、调度参数、缓存参数、量化参数和加速开关。

  3. 控制面：Engine 启动 tokenizer、scheduler、detokenizer 子进程，并在主进程中 维护 request/response 状态。

  4. 数据面：scheduler 持有 GPU-owning ModelRunnerProcess，负责 batch 构造、 KV cache 分配、prefix cache 命中、forward 和 sampling。

  5. 加速层：FlashInfer、CUDA Graph、Triton、CUTLASS 和 mllm-kernel 提供 attention、 quantization、GEMM 和缓存写入等高频算子。

进程拓扑

Engine 在启动时创建三个子进程，并在主进程中保留 request/response 管理逻辑：

Main Process
  ├── FastAPI Server
  ├── Engine
  └── RequestResponseProcess
         │
         │ ZMQ
         ▼
TokenizerProcess
         │
         │ ZMQ or shared queue
         ▼
SchedulerProcess
  └── ModelRunnerProcess  (in-process, owns GPU resources)
         │
         │ ZMQ
         ▼
DetokenizerProcess
         │
         │ ZMQ
         ▼
RequestResponseProcess

这个拓扑的核心取舍是：GPU 资源由 scheduler 进程内的 ModelRunnerProcess 直接持有。 这样 scheduler 可以在同一进程中完成调度、KV cache 资源释放、prefix cache 更新和模型 forward，避免再引入 model worker 进程之间的 GPU 资源同步。

请求生命周期

一次 chat completion 请求的典型路径如下：

  1. HTTP server 接收请求并转换为 GenerateReqInput。

  2. RequestResponseProcess 为请求分配 request id，并把请求送入 tokenizer。

  3. TokenizerProcess 调用 tokenizer / processor，生成 TokenizedGenerateReqInput。

  4. SchedulerProcess 接收 tokenized request，创建 Req，放入等待队列。

  5. scheduler 根据 token budget、running request 数量和 prefill/decode 状态构造 ScheduleBatch。

  6. ModelRunnerProcess 为 batch 分配 request slot 和 KV slot，执行 prefix matching。

  7. ModelRunner 构造 ForwardBatch，初始化 attention backend metadata，调用模型 forward，并对 logits 做 sampling。

  8. scheduler 更新每个 Req 的输出 token、finished reason 和 timing 字段。

  9. DetokenizerProcess 将 token id 转回文本。

  10. HTTP server 以普通 JSON 或 SSE streaming 形式返回结果。

控制面：Engine 与配置

pymllm.configs.server_config.ServerConfig 是服务运行时的主配置对象。它覆盖：

  • 模型和 tokenizer：model_path、tokenizer_path、load_format、dtype。

  • HTTP server：host、port、api_key、served_model_name。

  • 调度与内存：max_running_requests、max_total_tokens、max_prefill_tokens、 mem_fraction_static。

  • 加速后端：attention_backend、gdn_decode_backend、disable_cuda_graph、 enable_torch_compile。

  • IPC 与多模态传输：enable_shared_queue、tensor_transport_mode、 cuda_ipc_pool_size_mb。

  • 观测与调试：log_level、decode_log_interval。

Engine 启动前会加载 HuggingFace config，解析 EOS token、默认输出长度和 dtype，并确保 model/tokenizer 路径可用。启动后，Engine 会监控子进程健康状态；任一核心子进程异常退出， 服务会被标记为 unhealthy。

调度器

SchedulerProcess 是 pymllm 的中心调度组件。它负责：

  • 接收 tokenized requests。

  • 将输入请求转换为内部 Req 状态。

  • 根据 prefill/decode 状态构造 ScheduleBatch。

  • 控制 max_running_requests、max_total_tokens、max_prefill_tokens 等资源约束。

  • 在请求结束或中止时释放 request slot 和 KV slot。

  • 将 decode token 发送给 detokenizer。

当前调度策略以 FCFS 和单 GPU 资源约束为主。max_prefill_tokens 用于限制一轮调度 可接纳的 prefill token 数；长 prompt 的运行时 chunked prefill 切分仍待后续接入。

ModelRunner

ModelRunner 是真正执行模型 forward 的组件。它在初始化阶段完成：

  1. 设置 CUDA device 和默认 dtype。

  2. 加载模型类和 safetensors 权重。

  3. 解析模型 metadata，例如 layer 数、head 数、head dim、context length。

  4. 初始化 request-to-token pool、token-to-KV pool 和 KV allocator。

  5. 初始化 attention backend。

  6. 预热 cuBLAS。

  7. 按配置捕获 decode CUDA Graph。

forward 阶段分为 extend 和 decode 两类：

  • extend / prefill：处理 prompt token，写入 KV cache，并返回每个请求最后一个 token 的 logits。

  • decode：每个请求生成一个新 token，复用已有 KV cache 和 attention metadata。

KV cache 与 prefix cache

pymllm.mem_cache.memory_pool 中的 KV 管理采用三层结构：

ReqToTokenPool
    maps (request slot, position) -> kv index

TokenToKVPoolAllocator
    manages free integer KV slots

KVPool
    stores per-layer K/V tensors on GPU

TokenToKVPoolAllocator 使用 free-list 管理 KV slot，并通过批量释放接口降低大量请求结束或 prefix cache eviction 时的开销。KVPool 在条件满足时会调用 mllm-kernel 的 store_cache JIT kernel 写入 K/V；否则回退到 PyTorch indexing。

Prefix cache 当前有三种实现：

  • RadixCache：标准 radix-tree prefix cache。

  • ChunkCache：关闭 radix cache 时使用的简单缓存路径。

  • MambaRadixCache：为包含 GDN / Mamba-like 状态的 hybrid 模型预留的状态缓存路径。

当启用 RadixCache 时，extend batch 会先执行 prefix matching。命中的 prefix token 不再 重复计算，但对应 radix tree 节点会被 lock，直到请求结束或资源释放时再 unlock。

IPC 与多模态数据传输

普通控制消息通过 ZMQ 传输。多模态请求中的大 tensor 可以走 shared queue fast path， 由 enable_shared_queue 和 tensor_transport_mode 控制。

tensor_transport_mode 支持三种模式：

模式

行为

适用场景

default

GPU tensor 先拷到 CPU，再放入 POSIX shared memory。

最稳妥，调试优先。

cuda_ipc

GPU tensor 通过 CUDA IPC handle 跨进程共享。

避免 GPU->CPU 拷贝，但长服务中可能有 PyTorch IPC 生命周期问题。

cuda_ipc_pool

使用预分配 GPU workspace，发送方回收 chunk。

面向生产服务的推荐 GPU tensor 传输方式。

与 mllm C++ Backend 的关系

pymllm 和 cpu_backend、qnn_backend、ascend_backend 的层级不同：

  • C++ Backend 接入的是 mllm C++ 的 Tensor、Op、Module、Dispatcher 和设备 allocator。

  • pymllm 接入的是 Python/PyTorch serving pipeline，主要服务于在线推理、模型加载、 KV cache、调度和 CUDA kernel 集成。

  • mllm-kernel 是两者可以共享思想的低层 kernel 工具包，但当前 pymllm 更直接依赖 其中的 Python JIT CUDA kernel。