跳到主要内容

从框架到 Kernel:在 MUSA 上用 Torch Profiler、MSYS 与 MCU 分级调优

从框架到 Kernel:MUSA 分级调优

本文不是一篇只展示最终数字的“跑分记录”,而是一套可以复制执行的调试流程。我们会在 MTT S5000 上跑通一个 Hugging Face BERT 模型,依次用 Torch Profiler、Moore Perf System(msys)和 Moore Perf Compute(mcu)定位不同层级的问题,再实施 batching 与 FP16 优化,最后用无 Profiler 基准和数值误差验证收益。

先看结果

google/bert_uncased_L-2_H-128_A-2seq_len=128、每轮处理 64 条样本的条件下:

  • 逐样本串行 FP32:450.34 samples/s
  • Batch FP32:28065.69 samples/s
  • Batch FP16:28607.44 samples/s
  • 最终相对基线提升:63.52x
  • Batch 场景峰值显存:64.86 MiB → 33.05 MiB
  • FP32/FP16 输出余弦相似度:0.99999913

更重要的是,三个工具给出了能互相验证的证据链:

  1. Torch Profiler 发现 aten::linear 和 attention 被重复调用 64 倍。
  2. MSYS 发现大量短 Kernel、Launch Gap,并给出真实热点 Kernel 名。
  3. MCU 发现基线热点只有 1 个 block,而优化后有 256 个 blocks。

因此,本次性能问题的根因不是“某个 GEMM Kernel 写坏了”,而是工作粒度太小,硬件没有被铺开。

三层性能调试漏斗

1. 环境与前提

本文使用的真实环境:

你需要满足:

  • 机器能访问 MUSA 设备;
  • 已安装 torch、torch_musa、msys、mcu;
  • 至少预留 1 GB 磁盘空间保存 Trace 和报告。MSYS JSON/REP 可能达到数百 MB。

检查版本和设备:

python - <<'PY'
import platform
import torch
import torch_musa

print("python:", platform.python_version())
print("torch:", torch.__version__)
print("torch_musa:", getattr(torch_musa, "__version__", "unknown"))
print("musa_available:", torch.musa.is_available())
print("device_count:", torch.musa.device_count())
print("device_name:", torch.musa.get_device_name(0))
print("device_properties:", torch.musa.get_device_properties(0))
PY

msys --version
mcu --version
mcu --list-sections

本文环境中 mcu --list-sections 返回三个默认可用 Section:

SpeedOfLight
LaunchStats
MemoryWorkloadAnalysis

安装 Transformers:

python -m pip install transformers

本文使用的模型为 google/bert_uncased_L-2_H-128_A-2(Tiny BERT,2 层、Hidden Size 128)。官方下载页:https://huggingface.co/google/bert_uncased_L-2_H-128_A-2 。若 Hugging Face 直连不稳定,可使用镜像 https://hf-mirror.com/google/bert_uncased_L-2_H-128_A-2 ,或设置:

export HF_ENDPOINT=https://hf-mirror.com
export HF_HUB_DISABLE_XET=1

下载完成后,正式测试切换为离线模式,避免网络抖动污染结果:

export HF_HUB_OFFLINE=1

2. 准备可重复的推理脚本

创建实验目录:

mkdir -p ~/musa_profile_lab/results
cd ~/musa_profile_lab

保存下面的脚本为 bert_infer.py

import argparse
import json
import os
import time
from pathlib import Path

os.environ.setdefault("HF_HUB_DISABLE_XET", "1")

import torch
import torch_musa # noqa: F401
from transformers import AutoModel

MODES = {
"serial-fp32": (False, torch.float32),
"batch-fp32": (True, torch.float32),
"batch-fp16": (True, torch.float16),
}


def parse_args():
p = argparse.ArgumentParser()
p.add_argument(
"--model",
default="google/bert_uncased_L-2_H-128_A-2",
)
p.add_argument("--mode", choices=MODES, default="serial-fp32")
p.add_argument("--batch-size", type=int, default=64)
p.add_argument("--seq-len", type=int, default=128)
p.add_argument("--warmup", type=int, default=3)
p.add_argument("--iterations", type=int, default=10)
p.add_argument("--torch-profile", action="store_true")
p.add_argument("--trace", default="")
p.add_argument("--seed", type=int, default=2026)
p.add_argument("--pre-delay", type=float, default=0.0)
return p.parse_args()


def make_inputs(model, batch_size, seq_len, device, seed):
gen = torch.Generator(device="cpu").manual_seed(seed)
vocab = int(model.config.vocab_size)
input_ids = torch.randint(
100,
vocab,
(batch_size, seq_len),
generator=gen,
dtype=torch.long,
)
attention_mask = torch.ones(
(batch_size, seq_len),
dtype=torch.long,
)
return input_ids.to(device), attention_mask.to(device)


def forward_workload(
model,
input_ids,
attention_mask,
batched,
):
if batched:
with torch.profiler.record_function(
"batched_bert_forward"
):
return model(
input_ids=input_ids,
attention_mask=attention_mask,
).last_hidden_state


pieces = []

with torch.profiler.record_function(
"serial_bert_forward"
):
for i in range(input_ids.shape[0]):
pieces.append(
model(
input_ids=input_ids[i:i + 1],
attention_mask=attention_mask[i:i + 1],
).last_hidden_state
)
return torch.cat(pieces, dim=0)


def main():
args = parse_args()
batched, dtype = MODES[args.mode]
device = torch.device("musa:0")

model = AutoModel.from_pretrained(
args.model,
local_files_only=False,
)
model.eval().to(device)
if dtype == torch.float16:
model.half()

input_ids, attention_mask = make_inputs(
model,
args.batch_size,
args.seq_len,
device,
args.seed,
)

with torch.inference_mode():
for _ in range(args.warmup):
output = forward_workload(
model,
input_ids,
attention_mask,
batched,
)
torch.musa.synchronize()

if args.pre_delay:
time.sleep(args.pre_delay)

profile = None
if args.torch_profile:
activities = [
torch.profiler.ProfilerActivity.CPU,
torch.profiler.ProfilerActivity.MUSA,
]
profile = torch.profiler.profile(
activities=activities,
record_shapes=True,
profile_memory=True,
with_stack=False,
)
profile.__enter__()

start = time.perf_counter()
for _ in range(args.iterations):
output = forward_workload(
model,
input_ids,
attention_mask,
batched,
)
torch.musa.synchronize()
elapsed = time.perf_counter() - start

if profile is not None:
profile.__exit__(None, None, None)
trace = args.trace or (
f"results/{args.mode}.json"
)
Path(trace).parent.mkdir(
parents=True,
exist_ok=True,
)
profile.export_chrome_trace(trace)
table = profile.key_averages(
group_by_input_shape=True
).table(
sort_by="self_device_time_total",
row_limit=30,
)
print("=== TORCH_PROFILER_TOP ===")
print(table)

checksum = float(output.float().mean().cpu())
result = {
"model": args.model,
"mode": args.mode,
"dtype": str(dtype),
"batch_size": args.batch_size,
"seq_len": args.seq_len,
"iterations": args.iterations,
"total_s": elapsed,
"batch_latency_ms": (
elapsed * 1000 / args.iterations
),
"sample_latency_ms": (
elapsed
* 1000
/ (args.iterations * args.batch_size)
),
"samples_per_s": (
args.iterations
* args.batch_size
/ elapsed
),
"checksum": checksum,
"max_memory_mb": (
torch.musa.max_memory_allocated()
/ 1024**2
),
}
print(
"RESULT_JSON="
+ json.dumps(result, sort_keys=True)
)


if __name__ == "__main__":
main()

这个脚本有几个刻意的设计:

  • 三种模式共用相同输入和随机种子;
  • eval()inference_mode() 排除训练开销;
  • warmup 后才计时;
  • 计时边界显式同步 MUSA;
  • record_function 给 Torch Trace 添加高层语义;
  • --pre-delay 用于配合 MSYS 跳过加载和 warmup。

先做一个小规模 Smoke Test:

python bert_infer.py \
--mode batch-fp32 \
--batch-size 4 \
--seq-len 32 \
--warmup 1 \
--iterations 1

只要输出包含 RESULT_JSON= 且 checksum 不是 NaN,就说明推理链路已经跑通。

3. 先跑无 Profiler 基准

Profiler 会引入额外开销,所以任何优化对比都必须先用无 Profiler 基准建立事实。

export HF_HUB_OFFLINE=1
cd ~/musa_profile_lab

python bert_infer.py \
--mode serial-fp32 \
--batch-size 64 \
--seq-len 128 \
--warmup 3 \
--iterations 10

python bert_infer.py \
--mode batch-fp32 \
--batch-size 64 \
--seq-len 128 \
--warmup 3 \
--iterations 10

python bert_infer.py \
--mode batch-fp16 \
--batch-size 64 \
--seq-len 128 \
--warmup 3 \
--iterations 10

实测:

  • 串行 FP32:142.116 ms / 64 samples450.34 samples/s
  • Batch FP32:2.280 ms / 64 samples28065.69 samples/s
  • Batch FP16:2.237 ms / 64 samples28607.44 samples/s

Batch FP32 相对串行提升 62.32x。FP16 在 Batch FP32 基础上只再提升 1.93%,但将峰值显存从 64.86 MiB 降到 33.05 MiB

无 Profiler 基准:Batching 是主要收益来源

4. 第一级:Torch Profiler 找框架级重复

4.1 为什么先用 Torch Profiler

Torch Profiler 的成本最低,而且最接近模型代码。它适合回答:

  • 哪些 Operator 被调用;
  • 每个 Operator 被调用多少次;
  • 输入 Shape 是否符合预期;
  • 是否存在 Python 循环、重复分配或同步;
  • 优化后计算图是否真的变小。

采集三个模式:

python bert_infer.py \
--mode serial-fp32 \
--batch-size 64 \
--seq-len 128 \
--warmup 2 \
--iterations 1 \
--torch-profile \
--trace results/torch_serial_fp32.json

python bert_infer.py \
--mode batch-fp32 \
--batch-size 64 \
--seq-len 128 \
--warmup 2 \
--iterations 1 \
--torch-profile \
--trace results/torch_batch_fp32.json

python bert_infer.py \
--mode batch-fp16 \
--batch-size 64 \
--seq-len 128 \
--warmup 2 \
--iterations 1 \
--torch-profile \
--trace results/torch_batch_fp16.json

将 JSON 拷贝到本机后,可以用 Chrome Trace Viewer 或 Perfetto 打开:

chrome://tracing
https://ui.perfetto.dev/

4.2 如何解读

单个 64-sample 工作单元中:

  • musaLaunchKernel2768 → 45 → 24
  • aten::linear832 → 13 → 13
  • scaled_dot_product_attention128 → 2 → 2
  • aten::empty2880 → 47 → 39

linear 和 attention 恰好减少 64 倍,说明根因就是逐样本循环。

Torch Profiler:逐样本循环将算子调用放大约 64 倍

4.3 本层能做出的决策

看到这种调用结构后,不应该立刻去调 GEMM Tile,而应该先消除 Python 循环。

本层优化收益:

  • FP32 Launch 数减少 98.37%
  • FP16 Launch 数减少 99.13%
  • 端到端吞吐提升 62.32x

4.4 当前版本的边界

本文的 torch/torch_musa 组合虽然提供:

torch.profiler.ProfilerActivity.MUSA

但 Trace 主要包含 privateuse1_runtime 的 Launch API,没有完整设备 Kernel Duration,也没有可直接排序的 self_musa_time_total

这不是继续猜测的理由,而是切换到 MSYS 的信号。

5. 第二级:MSYS 看系统时间线

5.1 MSYS 解决什么问题

MSYS 类似 Nsight Systems,用来回答:

  • CPU Runtime API 与设备 Kernel 如何排列;
  • GPU 是否持续有工作;
  • 是否存在大段 Gap;
  • 是否有大量很短的 Kernel;
  • 哪个 Kernel 占总设备时间最多;
  • 是否存在不必要的 memcpy 或 synchronize。

5.2 采集稳态区间

脚本先 warmup,再通过 --pre-delay 4 等待;MSYS 用 --delay 3 启动采集。这样模型加载和稳态区间更容易分离。

基线:

msys profile \
-t musa \
--delay 3 \
-o results/msys_serial \
/usr/bin/python bert_infer.py \
--mode serial-fp32 \
--batch-size 64 \
--seq-len 128 \
--warmup 2 \
--iterations 20 \
--pre-delay 4

优化后:

msys profile \
-t musa \
--delay 3 \
-o results/msys_batch \
/usr/bin/python bert_infer.py \
--mode batch-fp16 \
--batch-size 64 \
--seq-len 128 \
--warmup 2 \
--iterations 1200 \
--pre-delay 4

生成统计:

msys stats \
-r musa_gpu_kern_sum,musa_api_sum,musa_gpu_sum,musa_kern_exec_sum,musa_api_gpu_sum \
-f json \
-o results/serial_stats \
results/msys_serial.msys-rep

msys analyze \
-f json \
-o results/serial_analysis \
results/msys_serial.msys-rep

msys export \
-t json \
-f true \
-o results/msys_serial_trace.json \
results/msys_serial.msys-rep

对 optimized report 重复同样操作。

5.3 MSYS 实测发现

串行 FP32 稳态:

  • 每个 64-sample 工作单元约 3088 个硬件事件;
  • GPU In-use 约 17.47%
  • 首要热点是 musa_asm_sgemm_nt_256_128x128_epilogue_robust64
  • 该 Kernel 占 Kernel 时间 61.76%,平均 25.142 us
  • CopyLastContiguousKernel7.34%

Batch FP16 稳态:

  • 每个工作单元约 38 个硬件事件,减少约 81.3x
  • GPU In-use 约 15.47%
  • 热点切换为 TCE FP16 GEMM,平均 14.776 us
  • 出现 musa_asm_htce_flash_atten_512_256x128x128_persistence

In-use 没有变成 100% 并不代表优化失败。这个模型只有 2 层、Hidden Size 128,本身太小。优化后的价值是用相近的几毫秒处理更多请求。

5.4 两个容易踩的坑

第一个坑:在某些环境中启用 OS Runtime Trace:

msys profile -t musa,osrt ...

会导致 Python 在 init_fs_encoding 阶段失败:

FileExistsError: [Errno 17] File exists: '/usr/lib/python3.10'

-t musa 的最小测试正常,所以这是 OSRT 注入与当前环境/Python 的兼容性问题,不是 Python 安装损坏。本次直接使用 -t musa 完成分析。

第二个坑:msys profile 会读取 stdin 以响应 q。不要把多个 Profile 命令通过同一个 stdin 连续喂给它,否则后续 Shell 内容可能被 Profiler 吃掉。

6. 第三级:MCU 看单个热点 Kernel

6.1 为什么必须最后用 MCU

MCU 会 Replay Application 多次收集 Counter。它的开销最高,因此必须先从 MSYS 获得真实热点名,再用精确 -k-c 1 采集。

不要用 MCU 下打印的模型 latency 做端到端对比,因为 Replay 会将一次推理放大很多倍。

基线热点:

mcu \
-f \
-o results/mcu_serial \
-k musa_asm_sgemm_nt_256_128x128_epilogue_robust64 \
-c 1 \
--sections SpeedOfLight,LaunchStats,MemoryWorkloadAnalysis \
/usr/bin/python bert_infer.py \
--mode serial-fp32 \
--batch-size 64 \
--seq-len 128 \
--warmup 1 \
--iterations 1

优化后热点:

mcu \
-f \
-o results/mcu_batch \
-k musa_asm_hhhhssgemm_nt_tce_128_64x64B128_epilogue_stage4 \
-c 1 \
--sections SpeedOfLight,LaunchStats,MemoryWorkloadAnalysis \
/usr/bin/python bert_infer.py \
--mode batch-fp16 \
--batch-size 64 \
--seq-len 128 \
--warmup 1 \
--iterations 1

两次采集都执行了 8 个 Replay Pass,并分别生成 789 行 Counter Data:

results/mcu_serial.mcu-rep
results/mcu_batch.mcu-rep

6.2 LaunchStats 给出的决定性证据

基线:

Grid  = (1, 1, 1)
Block = (256, 1, 1)

优化后:

Grid  = (2, 128, 1)  # 256 blocks
Block = (128, 1, 1)

6.3 本层应该做什么决策

这组数据说明:

  • 问题不是 SGEMM 的局部指令效率;
  • 只有 1 个 block 时,手写 Kernel 也无法弥补 55 个空闲 MP;
  • 正确方向是扩大工作粒度;
  • Batch 后切换到 TCE FP16 GEMM 和 Flash Attention,说明优化路径已经生效。

MCU 的收益不是再凑一个端到端百分比,而是验证优化机制,并阻止错误的 Kernel 微优化。

7. 继续做 Batch Sweep

对于 Tiny Model,固定 Batch 64 不一定是吞吐上限。继续测试:

for b in 1 4 16 64 128 256; do
python bert_infer.py \
--mode batch-fp16 \
--batch-size "$b" \
--seq-len 128 \
--warmup 5 \
--iterations 30
done

实测:

  • Batch 1:470.20 samples/s24.19 MiB
  • Batch 4:1740.95 samples/s24.19 MiB
  • Batch 16:7007.44 samples/s24.19 MiB
  • Batch 64:27229.32 samples/s33.05 MiB
  • Batch 128:56227.06 samples/s57.18 MiB
  • Batch 256:113368.17 samples/s104.88 MiB

Batch Latency 始终约为 2.13–2.35 ms,吞吐近似线性增长。

FP16 Batch Sweep:吞吐近似线性增长,显存成本后段上升

生产系统不应盲目固定 Batch 256,而应根据两个约束动态选择:

  1. 延迟 SLO:请求最多允许等待多久;
  2. 显存预算:并发模型、KV Cache 和其他服务还需要多少空间。

比较实用的做法是动态 Batching + 长度分桶:

  • 在一个短时间窗口中聚合请求;
  • 按 Sequence Length 分 Bucket;
  • 达到 Batch 上限或等待超时就立即发车;
  • 避免长短请求混在一起产生过多 Padding。

8. 验证 FP16 数值误差

性能优化不能只看速度。使用相同权重、相同输入分别执行 FP32 和 FP16:

import json
import torch
import torch_musa
from transformers import AutoModel

model_id = "google/bert_uncased_L-2_H-128_A-2"
device = torch.device("musa:0")

model = AutoModel.from_pretrained(model_id)
model.eval().to(device)

g = torch.Generator(device="cpu").manual_seed(2026)
ids = torch.randint(
100,
model.config.vocab_size,
(64, 128),
generator=g,
dtype=torch.long,
).to(device)
mask = torch.ones_like(ids)

with torch.inference_mode():
fp32 = model(
input_ids=ids,
attention_mask=mask,
).last_hidden_state.float().cpu()

model.half()
fp16 = model(
input_ids=ids,
attention_mask=mask,
).last_hidden_state.float().cpu()

diff = (fp32 - fp16).abs()
cosine = torch.nn.functional.cosine_similarity(
fp32.flatten().double(),
fp16.flatten().double(),
dim=0,
)

print(json.dumps({
"cosine_similarity": float(cosine),
"mean_abs_error": float(diff.mean()),
"max_abs_error": float(diff.max()),
}, indent=2))

本文结果:

cosine_similarity = 0.999999131985
mean_abs_error = 0.00111281115
max_abs_error = 0.0157820880

这个误差对本实验可接受,但真实业务仍要使用任务指标验证,例如准确率、召回率、生成质量或回归集通过率。

9. 一套可以复用的分级调试流程

Step 0:无 Profiler 基准

  • 固定输入、Seed、Batch 和 Sequence Length;
  • Warmup 后再计时;
  • 计时边界同步设备;
  • 排除模型下载和加载;
  • 保存 Checksum 或业务输出。

Step 1:Torch Profiler

先看 Operator、Shape、调用次数、Python Annotation 和内存分配。

如果调用数随样本数线性膨胀,优先修:

  • Batching
  • Cache
  • Python 循环
  • 同步
  • 重复分配

Step 2:MSYS

查看:

  • GPU Gap
  • 短 Kernel 数量
  • Memcpy/Sync
  • 热点 Kernel 汇总
  • 优化前后事件数量

最终拿到真实 Kernel 名。

Step 3:MCU

精确指定热点 Kernel(-k);限制采集数量(-c 1);采集 SpeedOfLight、LaunchStats、MemoryWorkloadAnalysis 三个 Section。

只有以下情况才值得考虑自定义 Kernel:

  • Grid 已足够大;
  • 热点占比高且稳定;
  • MCU 明确显示 Occupancy、带宽、Bank Conflict 或 Tile 存在瓶颈;
  • 框架级重复和系统级 Gap 已经消除。

Step 4:回到无 Profiler 基准

重新比较:

  • 吞吐
  • P50/P95/P99 延迟
  • 峰值显存
  • 正确性
  • 业务指标

10. 常见问题

为什么 FP16 只提升 1.93%?

因为这是一个 2 层、Hidden Size 128 的 Tiny BERT。Batch 后一次推理只有约 2.3 ms,Launch 和框架开销占比较高。更大的模型通常会体现更明显的 FP16 算力收益。

为什么 GPU In-use 只有约 15–17%?

In-use 低不等于吞吐低。该模型工作量太小,优化后用相近时间处理了几十倍样本。对在线服务来说,吞吐、尾延迟和成本比“看起来满载”更重要。

为什么不用 MCU 直接分析整个模型?

MCU 会 Replay 并采集大量 Counter。直接抓整个模型会很慢、报告很大,也难以解读。正确方式是先用 MSYS 找热点,再精确抓一个 Kernel。

Profiler 下的延迟为什么比普通 Benchmark 大?

Profiler 会记录 Shape、内存、Runtime Event,MCU 还会 Replay。Profiler 的时间用于定位,不用于最终性能结论。

如何打开报告?

  • Torch Trace:Perfetto 或 Chrome Trace Viewer
  • MSYS:msys export -t json 后用 Trace Viewer
  • MCU:mcu-ui results/mcu_batch.mcu-rep

在无桌面环境中,可以通过 X11/VNC 转发 MCU UI;也可以读取 .mcu-rep 中的 Kernel Metadata。

11. 产物清单

本文实验目录:

~/musa_profile_lab/

主要产物:

bert_infer.py
validate_outputs.py
summarize_results.py
TUTORIAL.md
results/summary.json
results/torch_*.json
results/msys_*.msys-rep
results/msys_*_trace.json
results/*_stats_*.json
results/*_analysis_*.json
results/mcu_serial.mcu-rep
results/mcu_batch.mcu-rep
results/validation.json
results/batch_sweep.txt

结语

这次优化最重要的收获不是 63.52x 这个数字,而是确定了每个工具应该在什么时候使用:

  • Torch Profiler 找框架级重复;
  • MSYS 找系统级 Gap 和真实热点;
  • MCU 验证单 Kernel 是否匹配硬件。

正确的性能调试不是一开始就钻进最底层,而是用证据逐层缩小问题:

在本文案例中,框架层 batching 解决了最大问题;FP16 带来更合适的 TCE/Flash Attention 路径和显存收益;MCU 则证明了底层机制,并避免把时间浪费在无法填满硬件的单 Block SGEMM 上。