TileRT 将 LLM 算子分解为 tile 级任务,在 8×B200 上动态调度计算-IO-通信重叠,实现 batch=1 超低延迟推理。DeepSeek-V3.2 达 600 tok/s,GLM-5 达 500 tok/s,MTP 下 590 tok/s。已在 Z.ai 生产部署。
Q1 痛点: 传统推理框架(vLLM / SGLang / TRT-LLM)面向高吞吐 batch serving 优化:连续 batching、PagedAttention、大 batch 填满 GPU。但实时交互场景(高频交易、AI 编程辅助、长 agent 链)需要极低的单请求延迟(TPOT),batch=1 时 GPU 利用率低、kernel launch 开销占比大、计算与通信串行等待。百亿参数 MoE 模型(DeepSeek-V3.2: 671B,GLM-5: ~数百B)在 batch=1 下的解码延迟是瓶颈。
Q2 方法: Tile-level runtime engine——
RMSNormProjxWqkvia 融合了 RMSNorm + QKV 投影 + FP8 权重转换)。每个 fused op 进一步分解为细粒度 tile 任务。prepare_money() 捕获整个 forward pass 为 CUDA graph,dsa_show_hands() 单次调用 replay,消除逐 kernel launch 开销。_swizzle_mma_16x32 按 PTX ISA 的 MMA 寄存器布局(fig.88)预排权重,绕过常规 CUDA 抽象直接映射到 tensor core 指令。DsaTempVarIdx)为所有中间结果提供固定索引布局,避免动态内存分配。核心技术壁垒: tile-level 算子分解 + 跨设备动态调度。将 GEMM、Norm、量化等操作在 tile 粒度拆解后,由运行时(而非静态图编译器)动态调度到多 GPU 的计算单元上,实现计算-通信的细粒度重叠。这要求对 PTX ISA 寄存器布局、NVLink 通信原语、CUDA graph 捕获/回放机制有深度掌控,且核心调度逻辑封装在闭源 C++ 库中——Python 层仅为编排壳。
Q3 结果: 8×NVIDIA B200 单节点:
系统范围:
inject_cache() 支持 P/D 分离(prefill 可委托外部系统如 SGLang)num_devices=8传统推理框架逐 kernel 串行启动(RMSNorm → GEMM → AllReduce → ...),每步等待前一步完成。TileRT 将每个 kernel 拆为 tile 任务,由 C++ runtime 在多 GPU 上交错调度:
这种调度由 C++ 运行时动态完成,Python 层不可见——所有 tile 调度逻辑封装在 torch.ops.tilert.dsa_show_hands() 和 dsa_show_hands_prepare_money() 中。
无形式化作者证明 — 仅实证。
TileRT 是工程实现而非学术论文,无 throughput/latency 解析模型、无推导公式。代码结构揭示的设计选择:
ModelArgs.max_batch_size = 1 和 FlashSparseMLA 的 assert batch == 1 表明这不是临时限制而是根本性架构决策——所有 tile 调度、CUDA graph 捕获、temp var 布局都为单请求优化。RMSNormProjxWqkvia(1095 LOC)是最复杂的融合单元,包含 FP8/BF16 权重转换逻辑。_swizzle_mma_16x32 直接按 PTX ISA fig.88 的寄存器布局重排权重,说明开发团队不信任编译器自动优化,选择手动控制到 MMA 指令的寄存器分配层面。DsaTempVarIdx 固定索引 + validate_temp_vars_layout() 运行时校验,等价于静态 tensor 寄存器文件。放弃动态内存管理的灵活性,换取零分配开销。prepare_money() 一次性捕获全部 61 层 forward,update_sampling_config() 需完全 teardown + re-capture,说明 graph 粒度选择优先延迟而非灵活性。torch.ops.tilert.* 中,Python 只做编排。所有 tilert_forward() 最终 dispatch 到 C++ 自定义算子——社区可看到接口定义和参考实现,但无法审计核心调度器。| 版本 | 日期 | 关键改进 | DeepSeek-V3.2 | GLM-5-FP8 | MTP |
|---|---|---|---|---|---|
| v0.1.0-alpha.1 | 2025-11-20 | 初始发布 | baseline | — | — |
| v0.1.1 | 2025-12-23 | tile 调度优化 | 3–4× over baseline | — | — |
| v0.1.2-alpha.1 | 2026-01-26 | MTP 支持 | — | — | 590 tok/s (mtp=3) |
| v0.1.3 | 2026-02-14 | GLM-5 支持 | 600 tok/s | 500 tok/s | — |
| v0.1.4-dev | 2026-05-22 | 生产化 (Z.ai) | in production | in production | — |
MTPPreprocessLayer(embedding lookup + hidden state merge)+ MTP(MoeBlock + HeadProj)| 项 | 要求 |
|---|---|
| GPU | 8× NVIDIA B200(硬编码) |
| CUDA | 12.8 或 12.9 |
| Python | 3.11–3.12 |
| OS | Linux x86_64 |
| 最大序列长度 | 160K tokens |
| batch size | 1(硬编码) |
| 参数 | DeepSeek-V3.2 | GLM-5 |
|---|---|---|
| 层数 | 61 (3 dense + 58 MoE) | 78 |
| 隐藏维度 | 7168 | 6144 |
| 注意力头数 | 128 | 64 |
| 路由专家数 | 256 | 64 |
| 激活专家数 | 8 | — |
| KV LoRA rank | 512 | — |
| Q LoRA rank | 1536 | — |
| 精度 | FP8 / BF16 | FP8 |
| index heads | 64, dim=128, topk=2048 | — |
| 场景 | TileRT | 传统框架 (vLLM/SGLang) | 原因 |
|---|---|---|---|
| 单请求低延迟 | 强 (600 tok/s) | 弱 (batch=1 利用率低) | tile 级调度最大化单请求硬件利用率 |
| 高并发吞吐 | 不适用 (batch=1) | 强 | TileRT 不支持 batching |
| 长 prompt prefill | 委托外部系统 | 原生支持 | TileRT 通过 inject_cache() 实现 P/D 分离 |
| MoE 模型 | 强 (expert routing fused) | 一般 | 专家选择+gate+SiLU 融合为单 kernel |
| # | 主张 | 支撑证据 | 承上启下 |
|---|---|---|---|
| 1 | batch=1 解码延迟是实时 AI 的核心瓶颈 | README 列举 HFT、interactive AI、long-running agents、AI-assisted coding 四类场景均需 ms 级 TPOT | 定义优化目标:TPOT 而非 throughput |
| 2 | 逐 kernel 串行执行在 batch=1 下浪费硬件资源 | 每个 kernel 启动有固定开销;Norm/GEMM/AllReduce 间存在等待气泡;batch=1 compute intensity 低 | 指出现有方案的结构性缺陷 |
| 3 | 算子融合 + tile 级分解消除 kernel 间气泡 | ops/ 目录 20+ fused ops:RMSNormProjxWqkvia 将 Norm+Proj+Quant 三类操作融合;每个 fused op 进一步分解为 tile 任务 | 核心方案——从 kernel-level 下沉到 tile-level |
| 4 | 动态 tile 调度实现跨设备计算-通信重叠 | C++ runtime dsa_show_hands() 在 8 GPU 上交错调度 tile;所有 fused op 带 AllReduce 后缀(UnProjOAllReduce, ExpertDownAllReduce)表明通信融入计算管线 | tile 调度器是 latency reduction 的核心杠杆 |
| 5 | CUDA graph 整体捕获消除 launch 开销 | prepare_money() 一次捕获 61 层 forward;forward() 单次 dsa_show_hands() replay | 与 tile 调度互补——tile 管运行效率,graph 管启动效率 |
| 6 | PTX 级权重 swizzle 消除微架构浪费 | _swizzle_mma_16x32 按 MMA 寄存器布局预排权重,省去运行时 register rearrangement | 优化栈从算法层贯穿到微架构层 |
| 7 | P/D 分离 API 使方案可融入生产 | inject_cache(layer_caches) + set_cur_pos() 接受外部 prefill 系统的 KV cache | 从单机极致优化扩展到分布式部署 |
python/models/base.py #TileRTModule: 顶层抽象类(ABC),定义 golden_forward() / tilert_forward() 双路径。enable_tilert(True) 切换到 C++ 后端路径并触发权重格式转换。所有 fused op 均派生自此类。SerializableTileRTModule: 可组合的 op 流水线容器。exec_seq: list[TileRTModule] 按 register_op() 顺序存储子操作。init_tilert_weights(state_dict) 遍历 exec_seq 初始化权重;init_tilert_vars(batch_size, seq_len) 分配临时变量。python/models/deepseek_v3_2/modules/dsa.py #Dsa(SerializableTileRTModule) 构建完整的 DeepSeek-V3.2 层栈:
MlpBlock(dense FFN),调用 RMSNormUpGateSiLU + DownAllReduceMoeBlock(MLA + MoE),每层包含完整的 attention + expert routing 管线RMSNormHeadProj(final logit projection)register_op() 以 prefix=f"layer_{layer_idx}_" 命名,确保权重加载时的 state_dict key 匹配。remove_selected=True 在构造后释放已注册到 C++ 的权重张量,节省 Python 侧内存。
python/models/deepseek_v3_2/modules/mla.py #DeepSeek-V3.2 的 Multi-head Latent Attention 被分解为 8 个串行 fused op,每个 op 的输入/输出通过 DsaTempVarIdx 固定索引传递:
| # | Op 类 | 文件 | 功能 |
|---|---|---|---|
| 1 | RMSNormProjxWqkvia | ops/rmsnorm_projx_wqkvia.py (1095 LOC) | RMSNorm + QKV 联合投影 + FP8 权重解量化 |
| 2 | LayerNormRoPERotate | ops/layernorm_rope_rotate.py | LayerNorm + RoPE 旋转位置编码 |
| 3 | RmsnormProjqWqib | ops/rmsnorm_projq_wqib.py | RMSNorm + Q 投影 + index bias |
| 4 | ProjxWis | ops/projx_wis.py | index score projection(稀疏注意力用) |
| 5 | ProjqWqb | ops/projq_wqb.py | Q bias projection |
| 6 | KVRMSNorm | ops/rmsnorm_kv.py | KV 缓存归一化 |
| 7 | ProjoWKVb | ops/projo_wkvb.py | 输出投影 + KV bias |
| 8 | UnProjOAllReduce | ops/unproj_o_allreduce.py (428 LOC) | 反投影 + 8-GPU AllReduce |
每层维护 3 个 KV cache: ki_cache (compressed key, shape [seqlen, 128])、kv_cache (compressed KV, [seqlen, 512])、pe_cache (position encoding, [seqlen, 64])。
ops/rmsnorm_projx_wqkvia.py (1095 LOC) #融合 RMSNorm + QKV projection + FP8/BF16 权重解量化,是单个文件中最大的 op 实现。包含:
_swizzle_mma_16x32 / _swizzle_qmma_16x32: 按 PTX ISA fig.88 的 MMA 寄存器布局对权重做 swizzle。这两个函数是 tile-level 方案与微架构对接的关键——将权重在 Python 侧预排为 MMA 指令所需的寄存器布局,消除运行时 rearrangement。compute_kernel_type 在 "fp8" 和 "bf16" 间切换tilert_forward() 最终 dispatch 到 torch.ops.tilert.* C++ 算子ops/expert_sel_up_gate_silu.py (729 LOC) #融合 expert selection + gate/up projection + SiLU 激活。从 256 个路由专家中选择 top-8(n_activated_experts=8),单 kernel 完成选择→投影→激活。配套的 expert_down_allreduce.py(403 LOC)完成 expert 下投影 + 8-GPU AllReduce。MoE 层还包含 n_shared_experts=1 个共享专家,与 routed expert 结果相加。
ops/flash_sparse_mla.py (265 LOC) + ops/sparse_index.py (125 LOC) #flash_sparse_mla() 实现 Flash Sparse MLA:
index_n_heads=64),128 维 index key,通过 sparse_index 计算 top-2048 KV 位置(index_topk=2048)sparse_index 执行 FP8 $q \times k \to \text{ReLU} \to \text{scale} \to \text{reduce\_sum}$,输出 index score 用于 top-k 选择dsa_flash_sparse_mla_combine_16h(16 头),GLM-5 用 glm5_flash_sparse_mla_combine_8h(8 头)这套稀疏注意力机制通过 self_attn.indexer.wk 权重实现,是独立于 DeepSeek 原始论文的 TileRT 特有优化——在 batch=1 decode 时减少 KV cache 读取量。
modules/end2end.py (520 LOC) #ShowHandsDSALayer 是 Python 层与 C++ runtime 的桥接核心:
_init_weights(): 8 个 Python 线程并行加载 8 GPU 权重。每个线程构建 Dsa + optional MTP,初始化权重后调用 generate_params_with_continuous_storage() 将中间存储合并为连续内存。prepare_money(): 在每个 device 上调用 C++ dsa_show_hands_prepare_money(params, intermediates, caches, profile_logs, max_seq_len, with_mtp, is_glm5) 捕获 CUDA graph。forward(): 调用 dsa_show_hands(token_id.cpu(), active_mtp, is_glm5) 回放 CUDA graph,从每个 GPU 收集 DeviceResult。update_sampling_config(): 变更 temperature/top_p/top_k 时完全 teardown + re-capture 整个 CUDA graph。"Show Hands" 扑克术语贯穿 API:prepare_money(准备资金 = 捕获 graph)、show_hands(梭哈 = 执行 forward)、go_home(回家 = 清理)。
generator.py (550 LOC) #DSAv32Generator 提供两条生成路径:
_generate_without_mtp(): 标准自回归,逐 token 调用 decode_layer.forward(),通过 DsaTempVarIdx.TOKEN_OUT 读取 next token_generate_with_mtp(): MTP 投机解码,每步尝试接受多 token,accepted_counts 记录每步接受数P/D 分离 API:
inject_cache(layer_caches): 接受外部 prefill 系统的 (ki, kv, pe) 三元组 per layer(0 到 n_layers-1),各张量为 BF16set_cur_pos(cur_pos): 设置 C++ 后端的 RoPE 当前位置inject_last_hidden_state() 注入最后隐藏状态refs/kernel.py (354 LOC) #混合使用两种 DSL 的 golden 验证内核:
| 函数 | DSL | 功能 |
|---|---|---|
act_quant() | TileLang @tilelang.jit | per-block FP8 量化 + 动态 scaling |
fp8_gemm() | TileLang @tilelang.jit | FP8 矩阵乘 + block-scaled accumulation + L2 cache swizzle |
fp8_index() | TileLang @tilelang.jit | FP8 index score: $q \times k \to \text{ReLU} \to \text{scale} \to \text{reduce\_sum}$ |
weight_dequant() | Triton @triton.jit | FP8→BF16 block-wise 反量化 |
TileLang 是 tile-ai 自研 DSL,与 TileRT 的 TileLang/TileScale 生态共享。参考内核用于 golden_forward() 路径做正确性验证。
models/preprocess/weight_converter.py (697 LOC) #将 HuggingFace 格式权重转换为 TileRT 的分片格式:
--model_type deepseek-v32 和 --model_type glm-5temp_var_indices.py #DsaTempVarIdx 枚举定义 51 个 named indices,与 C++ DsaTempVars 一一对应:
TOKEN_OUT、HIDDEN、Q、K、V 等核心中间张量均有固定索引validate_temp_vars_layout() 在初始化时校验 Python 与 C++ 索引一致性models/glm_5/ #GLM5Generator(592 LOC): 与 DSAv32Generator 结构类似,额外支持 thinking modeModelArgsGLM5: 78 层, dim=6144, 64 heads, 64 routed expertsShowHandsDSALayer 运行时,通过 is_glm5 flag 分派到不同 C++ kernelupdate_sampling_config() 在变更 temperature/top_p/top_k 时必须完全 teardown 并重新 capture 整个 forward pass 的 CUDA graph,因为 sampling 参数被 bake 到 graph 指令中。生产中应尽量避免频繁切换 sampling 参数。self_attn.indexer.wk。TileRT 是高度 hardware-proximal 的框架(PTX 级 swizzle、tile-level 调度、CUDA graph 整体捕获),其复杂性暗示硬件层面的潜在改进:
_swizzle_mma_16x32 的存在说明 MMA 指令的 operand layout 对软件不友好,硬件若提供可配置 swizzle 模式可消除这层预处理开销