Llama 3-8B-Instruct 在昇腾 NPU 上基于 SGLang 的性能实测

接下来配置 SGLang Engine 模式，以便在脚本中直接调用昇腾 NPU 执行推理。这里指定后端为 ascend，并使用 FP16 精度以节省显存。

# sglang_engine_setup.py
import os
import sglang as sgl

os.environ['MAX_JOBS'] = '1'
os.environ['SGLANG_TARGET_BACKEND'] = 'ascend'
MODEL_PATH = os.path.expanduser("~/models/Llama-3-8B")

print("Initializing SGLang Engine (Backend: Ascend)...")
try:
    engine = sgl.Engine(
        model_path=MODEL_PATH,
        tp_size=1,
        trust_remote_code=True,
        backend="ascend",
        dtype="float16"
    )
    print("✅ Engine initialized successfully!")
except Exception as e:
    print(f"❌ Engine initialization failed: {e}")
    raise

为了方便批量测试，我们可以封装一个推理函数：

# inference_function.py
BATCH_SIZE = 4
MAX_NEW_TOKENS = 50

def run_inference(prompts):
    outputs = []
    for prompt in prompts:
        out = engine.generate(prompt, max_new_tokens=MAX_NEW_TOKENS)
        outputs.append(out)
    return outputs

test_prompts = ["Hello world!"] * BATCH_SIZE
sample_output = run_inference(test_prompts)
print("Sample output:", sample_output[0])

性能基准测试

吞吐量与延迟

吞吐量衡量单位时间内处理的 token 或样本数，是评估推理性能的核心指标。常见的有 tokens/sec 和 samples/sec。吞吐量越高，越适合多用户并发或大批量离线生成场景。

测试代码如下，注意在 Ascend 设备上需设置 device_map="npu"：

import torch
import torch_npu
from transformers import AutoTokenizer, AutoModelForCausalLM
import time

model_name = "/path/to/your/model"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name, 
    torch_dtype=torch.float16, 
    device_map="npu"
)
model.eval()

prompt = "Describe the architecture of Ascend NPU."
inputs = tokenizer(prompt, return_tensors="pt").to("npu")

# Warmup
for _ in range(5):
    model.generate(**inputs, max_new_tokens=32)

num_iters = 20
total_tokens = 0
start = time.time()

for _ in range(num_iters):
    out = model.generate(**inputs, max_new_tokens=128)
    gen_tokens = out.shape[-1] - inputs["input_ids"].shape[-1]
    total_tokens += gen_tokens

end = time.time()
throughput = total_tokens / (end - start)
print(f"Throughput: {throughput:.2f} tokens/sec")

时延（Latency）则关注单个请求的响应速度，包括 E2E Latency（端到端）和 Per-token Latency（单 token 平均时间）。低时延对于在线交互场景至关重要。

# 单次请求时延测试
inputs = tokenizer("Hello, explain NPU.", return_tensors="pt").to("npu")
for _ in range(5):
    model.generate(**inputs, max_new_tokens=16)

start = time.time()
output = model.generate(**inputs, max_new_tokens=64)
end = time.time()
latency_ms = (end - start) * 1000

input_len = inputs["input_ids"].shape[-1]
output_len = output.shape[-1]
gen_token_count = output_len - input_len
per_token_latency = latency_ms / gen_token_count

print(f"E2E Latency: {latency_ms:.2f} ms")
print(f"Per-Token Latency: {per_token_latency:.2f} ms/token")

显存占用

显存是限制大模型部署的关键资源。Ascend 提供 npu-smi 实时查看 HBM 使用情况，PyTorch 层面也可统计。

import torch_npu
allocated = torch_npu.memory.npu_memory_reserved()
cached = torch_npu.memory.npu_memory_allocated()
print(f"Reserved HBM: {allocated/1024/1024:.2f} MB")
print(f"Allocated HBM: {cached/1024/1024:.2f} MB")

# 或者使用系统命令
import subprocess
out = subprocess.check_output("npu-smi info", shell=True)
print(out.decode())

批量压力测试

批量吞吐量和时延能反映 NPU 在并发场景下的并行利用率。通过调整 batch size，可以判断最佳配置。

import torch
import torch_npu
import time
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "/path/to/your/model"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name, 
    torch_dtype=torch.float16, 
    device_map="npu"
)
model.eval()

def measure(bs=1, seq=128):
    text = "Ascend NPU performance test. " * (seq // 10)
    inputs = tokenizer([text] * bs, return_tensors="pt", padding=True, truncation=True).to("npu")
    
    for _ in range(3): # warmup
        model.generate(**inputs, max_new_tokens=32)
    
    start = time.time()
    out = model.generate(**inputs, max_new_tokens=seq)
    end = time.time()
    
    input_len = inputs["input_ids"].shape[-1]
    output_len = out.shape[-1]
    gen_tokens = (output_len - input_len) * bs
    latency = end - start
    throughput = gen_tokens / latency
    return latency, throughput, gen_tokens

print("batch_size, seq_len, latency(s), throughput(tokens/s)")
for bs in [1, 2, 4, 8, 16]:
    lat, th, tk = measure(bs=bs, seq=128)
    print(f"{bs}, 128, {lat:.3f}, {th:.2f}")

从测试结果来看，随着 batch size 增大，总吞吐量显著提升。虽然总延迟略有增加，但每个 token 的平均延迟下降，体现了 Ascend NPU 在大批量并发推理中的强大并行计算能力。

总结

本次实测表明，Llama 3-8B-Instruct 在 Ascend NPU 配合 SGLang 框架下，能够展现出高吞吐量、低延迟和稳定的性能。即使在大批量、高并发和长序列生成的复杂场景中，NPU 也能保持高效的资源利用率。开发者可以直接在云端或本地环境中快速完成模型加载、推理和性能调优，无需过度担心环境配置的复杂性。