TileRT: Tile-Based Runtime for Ultra-Low-Latency LLM Inference

framework tile-ai-tilert
inferencelow-latencytile-schedulingCUDA-graphMoEFP8

TileRT — L2 #

§1 TL;DR #

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 生产部署。

§2 Q1 · Q2 · Q3 #

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——

  1. 算子融合 → tile 分解: 将传统独立的 RMSNorm、Projection、Quantization、AllReduce 等操作融合为单个 tile-level fused op(如 RMSNormProjxWqkvia 融合了 RMSNorm + QKV 投影 + FP8 权重转换)。每个 fused op 进一步分解为细粒度 tile 任务。
  2. 动态调度: 运行时(C++ 后端)在多设备上动态调度 tile 任务,实现计算、IO 和跨设备通信的细粒度重叠,消除流水线气泡。
  3. CUDA graph 全链路捕获: prepare_money() 捕获整个 forward pass 为 CUDA graph,dsa_show_hands() 单次调用 replay,消除逐 kernel launch 开销。
  4. PTX ISA 级权重映射: _swizzle_mma_16x32 按 PTX ISA 的 MMA 寄存器布局(fig.88)预排权重,绕过常规 CUDA 抽象直接映射到 tensor core 指令。
  5. 静态 tensor 寄存器文件: 51 个 named temp var(DsaTempVarIdx)为所有中间结果提供固定索引布局,避免动态内存分配。
  6. 核心技术壁垒: tile-level 算子分解 + 跨设备动态调度。将 GEMM、Norm、量化等操作在 tile 粒度拆解后,由运行时(而非静态图编译器)动态调度到多 GPU 的计算单元上,实现计算-通信的细粒度重叠。这要求对 PTX ISA 寄存器布局、NVLink 通信原语、CUDA graph 捕获/回放机制有深度掌控,且核心调度逻辑封装在闭源 C++ 库中——Python 层仅为编排壳。

    Q3 结果: 8×NVIDIA B200 单节点:

    • DeepSeek-V3.2: 600 tok/s(v0.1.3)
    • GLM-5-FP8: 500 tok/s(v0.1.3)
    • MTP(mtp=3): 590 tok/s,mean accepted length 2.77(v0.1.2)
    • v0.1.1 → v0.1.3 累计 3–4× speedup over baseline
    • 已在 Z.ai 生产部署(GLM-5.1-highspeed)

    系统范围:

    • 阶段覆盖: decode 为主,通过 inject_cache() 支持 P/D 分离(prefill 可委托外部系统如 SGLang)
    • 并行维度: 8-way TP(AllReduce 在每个 fused op 内完成),不支持 PP/DP/EP
    • 部署模式: 单节点 8×B200,硬编码 num_devices=8
    • 服务模式: batch=1 only,不支持连续 batching

    §3 架构 / 方法图 #

    整体执行流 #

    sequenceDiagram participant CLI as generate.py CLI participant Gen as DSAv32Generator participant E2E as ShowHandsDSALayer participant CPP as C++ TileRT Runtime CLI->>Gen: __init__(model_args, max_new_tokens, with_mtp) Gen->>E2E: __init__() → _init_weights() E2E->>E2E: 8 threads parallel weight loading loop per GPU (0..7) E2E->>E2E: Dsa(model_args, device_id, 8) E2E->>E2E: dsa.init_tilert_weights(state_dicts) end E2E->>CPP: dsa_show_hands_prepare_money()
    [CUDA graph capture] CLI->>Gen: generate(prompt) Gen->>Gen: tokenize → prompt_tokens loop autoregressive decode Gen->>E2E: forward(token_id) E2E->>CPP: dsa_show_hands(token_id)
    [graph replay] CPP-->>E2E: DeviceResult E2E-->>Gen: next_token Gen->>Gen: check EOS / prompt mask end Gen-->>CLI: decoded text

    模型层架构 #

    graph TD subgraph Dsa["Dsa Model Stack (61 layers)"] EMB[ParallelEmbedding] --> L0[Layer 0–2: MlpBlock × 3
    dense FFN] L0 --> L3[Layer 3–60: MoeBlock × 58
    MLA + MoE per layer] L3 --> HEAD[Layer 61: RMSNormHeadProj
    final logits] end subgraph MoeDetail["MoeBlock = MLA + MoE"] MLA_IN[hidden_in] --> MLA[Mla: 8-op pipeline] MLA --> RES1[+ residual] RES1 --> ROUTE[RMSNormExpertProj
    router scores] ROUTE --> UP[ExpertSelectUpGateSiLU
    top-8 select + gate + SiLU] UP --> DOWN[ExpertDownAllReduce
    down proj + AllReduce] DOWN --> RES2[+ residual → hidden_out] end subgraph MLA8["MLA 8-Op Pipeline"] H[hidden] --> OP1[① RMSNormProjxWqkvia
    Norm+QKV+FP8] OP1 --> OP2[② LayerNormRoPERotate] OP2 --> OP3[③ RmsnormProjqWqib] OP3 --> OP4[④ ProjxWis] OP4 --> OP5[⑤ ProjqWqb] OP5 --> OP6[⑥ KVRMSNorm] OP6 --> OP7[⑦ ProjoWKVb] OP7 --> OP8[⑧ UnProjOAllReduce] end

    Tile-level 调度原理 #

    传统推理框架逐 kernel 串行启动(RMSNorm → GEMM → AllReduce → ...),每步等待前一步完成。TileRT 将每个 kernel 拆为 tile 任务,由 C++ runtime 在多 GPU 上交错调度:

    • GPU 0 的 GEMM tile 与 GPU 1 的 RMSNorm tile 并行执行
    • 某 GPU 的计算 tile 与另一 GPU 的 NVLink 通信 tile 重叠
    • AllReduce 被拆为 tile 级 send/recv,与后续层的 RMSNorm 计算重叠

    这种调度由 C++ 运行时动态完成,Python 层不可见——所有 tile 调度逻辑封装在 torch.ops.tilert.dsa_show_hands()dsa_show_hands_prepare_money() 中。

    §4 作者证明 #

    无形式化作者证明 — 仅实证。

    TileRT 是工程实现而非学术论文,无 throughput/latency 解析模型、无推导公式。代码结构揭示的设计选择:

    1. batch=1 硬约束: ModelArgs.max_batch_size = 1FlashSparseMLAassert batch == 1 表明这不是临时限制而是根本性架构决策——所有 tile 调度、CUDA graph 捕获、temp var 布局都为单请求优化。
      1. 跨算子边界融合: 不是简单的 kernel fusion(如 FlashAttention 融合 QKV attention),而是将 Norm + Projection + Quantization + Communication 四类操作跨边界融合。RMSNormProjxWqkvia(1095 LOC)是最复杂的融合单元,包含 FP8/BF16 权重转换逻辑。
        1. PTX 级性能控制: _swizzle_mma_16x32 直接按 PTX ISA fig.88 的寄存器布局重排权重,说明开发团队不信任编译器自动优化,选择手动控制到 MMA 指令的寄存器分配层面。
          1. 静态内存布局: 51 个 DsaTempVarIdx 固定索引 + validate_temp_vars_layout() 运行时校验,等价于静态 tensor 寄存器文件。放弃动态内存管理的灵活性,换取零分配开销。
            1. 整图 CUDA graph 捕获: prepare_money() 一次性捕获全部 61 层 forward,update_sampling_config() 需完全 teardown + re-capture,说明 graph 粒度选择优先延迟而非灵活性。
              1. 闭源核心 + 开源壳: 核心 tile 调度逻辑在 C++ torch.ops.tilert.* 中,Python 只做编排。所有 tilert_forward() 最终 dispatch 到 C++ 自定义算子——社区可看到接口定义和参考实现,但无法审计核心调度器。
              2. §5 实验与数据 #

                性能演进 #

                版本日期关键改进DeepSeek-V3.2GLM-5-FP8MTP
                v0.1.0-alpha.12025-11-20初始发布baseline
                v0.1.12025-12-23tile 调度优化3–4× over baseline
                v0.1.2-alpha.12026-01-26MTP 支持590 tok/s (mtp=3)
                v0.1.32026-02-14GLM-5 支持600 tok/s500 tok/s
                v0.1.4-dev2026-05-22生产化 (Z.ai)in productionin production

                MTP 投机解码效果 #

                • mtp=3 配置下 synthetic workload: 590 tok/s
                • Mean accepted length: 2.77 tokens/step(理论上限 4 = 1 base + 3 MTP heads)
                • Min accepted: 1, Max accepted: 4
                • MTP 需额外模块: MTPPreprocessLayer(embedding lookup + hidden state merge)+ MTP(MoeBlock + HeadProj)

                硬件要求与约束 #

                要求
                GPU8× NVIDIA B200(硬编码)
                CUDA12.8 或 12.9
                Python3.11–3.12
                OSLinux x86_64
                最大序列长度160K tokens
                batch size1(硬编码)

                模型规格对比 #

                参数DeepSeek-V3.2GLM-5
                层数61 (3 dense + 58 MoE)78
                隐藏维度71686144
                注意力头数12864
                路由专家数25664
                激活专家数8
                KV LoRA rank512
                Q LoRA rank1536
                精度FP8 / BF16FP8
                index heads64, dim=128, topk=2048

                Workload 适用性 #

                场景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

                §6 论证链 #

                #主张支撑证据承上启下
                1batch=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 的核心杠杆
                5CUDA graph 整体捕获消除 launch 开销prepare_money() 一次捕获 61 层 forward;forward() 单次 dsa_show_hands() replay与 tile 调度互补——tile 管运行效率,graph 管启动效率
                6PTX 级权重 swizzle 消除微架构浪费_swizzle_mma_16x32 按 MMA 寄存器布局预排权重,省去运行时 register rearrangement优化栈从算法层贯穿到微架构层
                7P/D 分离 API 使方案可融入生产inject_cache(layer_caches) + set_cur_pos() 接受外部 prefill 系统的 KV cache从单机极致优化扩展到分布式部署

                §7 实现 cross-reference #

                7.1 核心抽象层 — 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) 分配临时变量。

                7.2 模型构建 — python/models/deepseek_v3_2/modules/dsa.py #

                Dsa(SerializableTileRTModule) 构建完整的 DeepSeek-V3.2 层栈:

                • Layer 0–2: MlpBlock(dense FFN),调用 RMSNormUpGateSiLU + DownAllReduce
                • Layer 3–60: MoeBlock(MLA + MoE),每层包含完整的 attention + expert routing 管线
                • Layer 61: RMSNormHeadProj(final logit projection)

                register_op()prefix=f"layer_{layer_idx}_" 命名,确保权重加载时的 state_dict key 匹配。remove_selected=True 在构造后释放已注册到 C++ 的权重张量,节省 Python 侧内存。

                7.3 MLA 8 步流水线 — python/models/deepseek_v3_2/modules/mla.py #

                DeepSeek-V3.2 的 Multi-head Latent Attention 被分解为 8 个串行 fused op,每个 op 的输入/输出通过 DsaTempVarIdx 固定索引传递:

                #Op 类文件功能
                1RMSNormProjxWqkviaops/rmsnorm_projx_wqkvia.py (1095 LOC)RMSNorm + QKV 联合投影 + FP8 权重解量化
                2LayerNormRoPERotateops/layernorm_rope_rotate.pyLayerNorm + RoPE 旋转位置编码
                3RmsnormProjqWqibops/rmsnorm_projq_wqib.pyRMSNorm + Q 投影 + index bias
                4ProjxWisops/projx_wis.pyindex score projection(稀疏注意力用)
                5ProjqWqbops/projq_wqb.pyQ bias projection
                6KVRMSNormops/rmsnorm_kv.pyKV 缓存归一化
                7ProjoWKVbops/projo_wkvb.py输出投影 + KV bias
                8UnProjOAllReduceops/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])。

                7.4 最大 fused op — 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++ 算子

                7.5 MoE 路由 — 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 结果相加。

                7.6 稀疏注意力 — ops/flash_sparse_mla.py (265 LOC) + ops/sparse_index.py (125 LOC) #

                flash_sparse_mla() 实现 Flash Sparse MLA:

                • 约束: batch=1, seqlen≤4, split_size=64
                • 使用 64 个 index heads(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 选择
                • 根据模型分派: DeepSeek-V3.2 用 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 读取量。

                7.7 运行时编排 — 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(回家 = 清理)。

                7.8 自回归生成 — 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),各张量为 BF16
                • set_cur_pos(cur_pos): 设置 C++ 后端的 RoPE 当前位置
                • MTP 模式额外需要 inject_last_hidden_state() 注入最后隐藏状态

                7.9 参考内核 — refs/kernel.py (354 LOC) #

                混合使用两种 DSL 的 golden 验证内核:

                函数DSL功能
                act_quant()TileLang @tilelang.jitper-block FP8 量化 + 动态 scaling
                fp8_gemm()TileLang @tilelang.jitFP8 矩阵乘 + block-scaled accumulation + L2 cache swizzle
                fp8_index()TileLang @tilelang.jitFP8 index score: $q \times k \to \text{ReLU} \to \text{scale} \to \text{reduce\_sum}$
                weight_dequant()Triton @triton.jitFP8→BF16 block-wise 反量化

                TileLang 是 tile-ai 自研 DSL,与 TileRT 的 TileLang/TileScale 生态共享。参考内核用于 golden_forward() 路径做正确性验证。

                7.10 权重转换 — models/preprocess/weight_converter.py (697 LOC) #

                将 HuggingFace 格式权重转换为 TileRT 的分片格式:

                • 支持 --model_type deepseek-v32--model_type glm-5
                • 输出 8 GPU 各自的分片权重文件(对应 8-way TP)
                • 从 v0.1.3 起不再需要预转换的 HuggingFace checkpoint,直接从官方权重转换

                7.11 临时变量索引 — temp_var_indices.py #

                DsaTempVarIdx 枚举定义 51 个 named indices,与 C++ DsaTempVars 一一对应:

                • TOKEN_OUTHIDDENQKV 等核心中间张量均有固定索引
                • validate_temp_vars_layout() 在初始化时校验 Python 与 C++ 索引一致性
                • 所有 GPU 中间结果通过这些固定索引在 fused ops 间传递——等效于硬编码的 tensor 寄存器文件,零动态分配

                7.12 GLM-5 支持 — models/glm_5/ #

                • GLM5Generator(592 LOC): 与 DSAv32Generator 结构类似,额外支持 thinking mode
                • ModelArgsGLM5: 78 层, dim=6144, 64 heads, 64 routed experts
                • 复用同一套 ShowHandsDSALayer 运行时,通过 is_glm5 flag 分派到不同 C++ kernel

                关键实现细节 #

                1. CUDA graph re-capture 代价: update_sampling_config() 在变更 temperature/top_p/top_k 时必须完全 teardown 并重新 capture 整个 forward pass 的 CUDA graph,因为 sampling 参数被 bake 到 graph 指令中。生产中应尽量避免频繁切换 sampling 参数。
                  1. index_topk=2048 稀疏注意力: 每层用 64 个 index heads 计算 FP8 index score,选出 top-2048 KV 位置用于 FlashSparseMLA。相比全量 attention(160K positions),这将 KV 读取量减少约 78×,是 batch=1 低延迟的关键优化。权重来自 self_attn.indexer.wk
                  2. Software → Hardware 逆向启示 #

                    TileRT 是高度 hardware-proximal 的框架(PTX 级 swizzle、tile-level 调度、CUDA graph 整体捕获),其复杂性暗示硬件层面的潜在改进:

                    • 硬件原生 tile 调度器: 若 GPU 提供 tile-level task queue(类似 GPUOS 的 persistent kernel 但有硬件调度支持),C++ runtime 调度逻辑可大幅简化
                    • 跨 SM 通信原语: AllReduce 被融入计算 kernel;若硬件提供 SM 间直接通信通道,tile 间数据交换延迟可进一步降低
                    • 可配置 operand swizzle: _swizzle_mma_16x32 的存在说明 MMA 指令的 operand layout 对软件不友好,硬件若提供可配置 swizzle 模式可消除这层预处理开销