尝鲜：Llama 3-8B-Instruct 在昇腾 NPU 上的 SGLang 推理实践

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! NPU memory allocated.")
except Exception as e:
    print(f"❌ Engine initialization failed: {e}")
    raise

为了后面批量测试方便，封装了一个简单的推理函数：

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])

性能基准

下面分别测了吞吐量、时延、显存占用，以及不同 batch size 下的表现。测试代码基于 PyTorch 和 torch_npu，直接调用模型 generate。

吞吐量

单位时间处理 token 数，看 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")

跑下来吞吐挺高，做大批量离线生成没什么压力。

推理时延

端到端延迟和单 token 延迟，关系到在线服务的响应体验。

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

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

inputs = tokenizer("Hello, explain NPU.", return_tensors="pt").to("npu")

# Warmup
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
print(f"E2E Latency: {latency_ms:.2f} ms")

input_len = inputs["input_ids"].shape[-1]
output_len = output.shape[-1]
gen_token_count = output_len - input_len
print(f"Per-Token Latency: {latency_ms/gen_token_count:.2f} ms/token")

时延很短，NPU 做实时推理也跟得上。

显存占用

用 npu-smi 和 PyTorch 接口看 HBM 使用情况。

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())

显存占用在合理范围内，8B 模型大概占十几 GB。

批量性能分析

不同 batch size 下的延迟和吞吐变化，最能看出 NPU 的并行能力。

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):
        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 延迟还下降了，这也是 NPU 的优势。

压力测试

为了验证稳定性，用长序列和高并发跑了一下。

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()

batch_sizes = [1, 2, 4, 8, 16]
seq_lengths = [64, 128, 256]
num_iters = 10
prompt = "Describe the architecture and optimization of Ascend NPU."

def stress_test(batch_size, seq_len):
    texts = [prompt] * batch_size
    inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True).to("npu")
    
    for _ in range(3):
        model.generate(**inputs, max_new_tokens=32)
    
    total_tokens = 0
    total_latency = 0.0
    
    for _ in range(num_iters):
        start = time.time()
        output = model.generate(**inputs, max_new_tokens=seq_len)
        end = time.time()
        
        gen_tokens = (output.shape[-1] - inputs["input_ids"].shape[-1]) * batch_size
        total_tokens += gen_tokens
        total_latency += (end - start)
    
    avg_latency = total_latency / num_iters
    avg_throughput = total_tokens / total_latency
    return avg_latency, avg_throughput

print("Batch, SeqLen, AvgLatency(s), AvgThroughput(tokens/s)")
for seq_len in seq_lengths:
    for bs in batch_sizes:
        avg_lat, avg_th = stress_test(bs, seq_len)
        print(f"{bs}, {seq_len}, {avg_lat:.3f}, {avg_th:.2f}")

测试结果：

可以看到，即使 sequence length 到了 256，batch 16 的时候吞吐量还是能稳定在 1688 tokens/s，没有明显波动。如果业务场景里并发量不高，单条请求的延迟也能保持在毫秒级，相当友好。

结语

Llama 3-8B-Instruct 在昇腾 NPU 上配合 SGLang 推理，整体表现符合预期。无论是吞吐、延迟还是显存开销，都没有明显短板。搭建流程也比较顺畅，从环境准备到跑通测试没花太多时间。如果后续有更大规模的模型部署需求，这套组合的扩展性也值得进一步尝试。