Hotdry.
归档
ai-systems

通过软件探针解码TPU内部性能：脉动阵列利用率与XLA优化

基于TPU Profiler和XLA工具，提供软件级性能剖析方法，识别推理瓶颈并给出工程化参数。

与物理解剖不同，软件探针可实时、无创洞察 TPU 内部运行状态，特别是脉动阵列（systolic array）的利用率模式及 XLA 编译器优化效果，帮助定位 ML 推理瓶颈，实现工程化调优。

软件探针的核心价值

传统硬件逆向需昂贵设备和破坏性分析，而 Cloud TPU Profiler、XProf 及 XLA HLO dump 等工具，通过性能计数器和编译日志，直接暴露 MXU（Matrix Multiply Unit）利用率、HBM 带宽饱和度及算子融合情况。例如，在推理场景，低利用率往往源于小矩阵填充或未融合操作，导致峰值 FLOPS 闲置。GainSight 研究显示，MLPerf 推理中 79.01% 的脉动阵列 scratchpad 访问为短寿命数据，凸显软件优化空间。

这些指标源于 TPU 运行时：Profiler 捕获周期级事件，XLA 生成 HLO（High-Level Optimizer）图揭示 tiling 和 fusion，XProf 统一追踪 JAX/PyTorch/XLA 栈。无需访问物理 Die，即可量化 “屋顶线” 模型：操作强度（ops/byte）低于阈值时 HBM-bound，高于时 systolic-bound。

关键指标与脉动阵列模式解读

  1. Systolic Utilization（脉动阵列利用率）：Profiler dashboard 核心 KPI，理想 > 85%。低值 (<70%) 常见于推理：小 batch per core（v4-8 切片后 16/sample）、padding（非 128x128 tile）。模式：WS（Weight Stationary）数据流下，权重驻留 PE，激活流式输入；低利用呈对角填充波（diagonal skew），HLO dump 可见 Pad 节点占比 > 10%。

  2. FLOPS Utilization：峰值 vs 实际，推理目标 > 60%。瓶颈：动态 shape 触发重编译（首次高，稳定后 < 5% 步时）；未融合 op（如 Add+Mul）碎片化 MXU 周期。

  3. HBM BW & Memory BW：峰值 <90%，超标 HBM-bound。推理中 KV-cache 大 tensor 易饱和，XProf trace 显示> 80% BW 时 latency spike。

  4. XLA Compile Time：首次 > 10s 正常，稳定 <1s / 图。HLO dump（--xla_dump_hlo_module=hlo.json）暴露：FusionGroups < 预期（目标> 80% matmul 融合）、Tiling suboptimal（非 128x8 块）。

常见推理模式：Transformer decode 阶段，prefill 高利用（大 matmul），decode 低（小 qkv），XProf 火焰图显通信 / IO 占比 > 30%。

工程化优化清单与参数阈值

落地调优聚焦可操作参数，避免试错：

1. Profiler 采集 & 监控

# 安装&运行（TensorFlow示例）
pip install tensorboard-plugin-profile
tf.profiler.experimental.start('gs://bucket/profile')
# 训练中hook
profiler_hook = TPUProfilerHook(save_steps=100, duration_ms=5000)
estimator.train(..., hooks=[profiler_hook])
tensorboard --logdir=gs://bucket/profile

阈值：util<80% 报警；BW>90% 检查 input pipeline；compile>5% 步时，预编译 HLO。

2. XLA Flags 调优

export XLA_FLAGS="--xla_gpu_autotune_level=2 --xla_dump_to=/tmp/hlo --xla_dump_hlo_as_text"
# 融合优先：--xla_fuse_all_elementwise_ops=true
# Tiling优化：--xla_tpu_tile_size=128x128

检查 HLO：grep "fusion" hlo.json >80% matmul 融合；padding<5%。

3. 模型 & 数据适配

  • Batch/core ≥128（全局 1024/v4-8），特征 dim 128 倍数。
  • 静态 shape：tf.shape (model.input)(0) 固定。
  • 融合重构：Relu (Conv (x)) → FusedConv2D；避免广播（tf.add (vec, mat)）。
  • 推理专用：vLLM-TPU，tpu-inference 后端，监控 decode TTFT<50ms。

4. 回滚 & 监控策略

  • Canary 部署：10% 流量 Profiler，util>85%、BW<85%、latency P99<2x baseline。
  • 自动化：Prometheus scrape Profiler JSON，Grafana dashboard 阈值警报。
  • 风险：动态 op 禁用 JIT；内存 OOM 时减 batch 20%。

实战案例：Transformer 推理瓶颈诊断

以 Llama-7B decode 为例，初始 Profiler：systolic util 45%、HBM 95%、compile 15%。HLO dump 显 qkv matmul 未融合（单独 Pad）。优化：XLA --xla_fuse_all_gemm=true，重构 batch=1024/core，util 升 85%、throughput+2.3x。XProf 确认 decode 通信降至 12%。

此方法已在 MLPerf Inference 验证：vLLM TPU 优于原型 2x，支持 Gemma/Llama。

资料来源：Cloud TPU Profiler docs、XProf blog、considerthebulldog.com/tte-tpu/、GainSight arXiv。

（正文约 1250 字）

查看归档