部署

tok = AutoTokenizer.from_pretrained(model_dir)
model = AutoModelForCausalLM.from_pretrained(model_dir, dtype=torch.float32).to(dev).eval()
streamer = TextStreamer(tok, skip_prompt=True, skip_special_tokens=True)  # 每凑成一个词就打印并 flush
for line in sys.stdin:
  if not line.strip():
    continue
  msgs = [{"role": "user", "content": line.strip()}]
  ids = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt", return_dict=True)["input_ids"].to(dev)
  with torch.inference_mode():
    model.generate(ids, max_new_tokens=64, do_sample=False, streamer=streamer)

const EOS: u32 = 2; // <|im_end|>
const MAX_NEW: usize = 64;

/// Greedy generation; `emit` gets each newly complete piece of text.
fn generate(model: &mut ModelWeights, tok: &Tokenizer, prompt: &[u32], max_new: usize, dev: &Device, mut emit: impl FnMut(&str) -> Result<()>) -> Result<usize> {
  let mut ids: Vec<u32> = Vec::new();
  let mut printed = 0; // bytes of the decoded answer already emitted
  let mut logits = model.forward(&Tensor::new(prompt, dev)?.unsqueeze(0)?, 0)?.squeeze(0)?;
  while ids.len() < max_new {
    let next = logits.argmax(0)?.to_scalar::<u32>()?;
    if next == EOS {
      break;
    }
    ids.push(next);
    // decode the whole answer and emit only the new suffix, once it no longer ends in a partial character
    let text = tok.decode(&ids, true).map_err(|e| anyhow!(e))?;
    if text.len() > printed && !text.ends_with('\u{fffd}') {
      emit(&text[printed..])?;
      printed = text.len();
    }
    logits = model.forward(&Tensor::new(&[next], dev)?.unsqueeze(0)?, prompt.len() + ids.len() - 1)?.squeeze(0)?;
  }
  Ok(ids.len())
}

fn main() -> Result<()> {
  let t0 = Instant::now();
  let args: Vec<String> = std::env::args().collect();
  let dev = match args.get(1).map(String::as_str) {
    Some("cpu") => Device::Cpu,
    Some("metal") => Device::new_metal(0)?,
    _ => bail!("usage: perf-serve <cpu|metal> [model_dir]"),
  };
  let (mut model, tok) = load(args.get(2), &dev)?;
  let t_load = t0.elapsed();
  // warm-up: a 2-token prefill and one decode step touch every kernel a request uses
  let warm = tok.encode(chat_prompt("hi").as_str(), false).map_err(|e| anyhow!(e))?.get_ids().to_vec();
  generate(&mut model, &tok, &warm[..2], 2, &dev, |_| Ok(()))?;
  eprintln!("ready: load {:.1} ms, warm-up done at {:.1} ms", t_load.as_secs_f64() * 1e3, t0.elapsed().as_secs_f64() * 1e3);

  let stdout = std::io::stdout();
  for line in std::io::stdin().lock().lines() {
    let line = line?;
    if line.trim().is_empty() {
      continue;
    }
    let t = Instant::now();
    let prompt = tok.encode(chat_prompt(line.trim()).as_str(), false).map_err(|e| anyhow!(e))?.get_ids().to_vec();
    let mut first: Option<f64> = None;
    let mut out = stdout.lock();
    let n = generate(&mut model, &tok, &prompt, MAX_NEW, &dev, |piece| {
      first.get_or_insert(t.elapsed().as_secs_f64() * 1e3);
      out.write_all(piece.as_bytes())?;
      Ok(out.flush()?)
    })?;
    writeln!(out)?;
    out.flush()?;
    eprintln!("request: {} prompt tokens, {n} new tokens, first piece {:.1} ms, total {:.1} ms", prompt.len(), first.unwrap_or(0.0), t.elapsed().as_secs_f64() * 1e3);
  }
  Ok(())
}

$ printf 'What is the capital of France?\nName three primary colors.\n' | perf-serve cpu
The capital of France is Paris.
The primary colors are:

1. Red: A warm, deep red hue that is often associated with fire, blood, and earth.
2. Blue: A deep, deep blue hue that is often associated with water, sky, and sky.
3. Yellow: A bright, warm yellow hue that is

ready: load 79.2 ms, warm-up done at 92.9 ms
request: 37 prompt tokens, 7 new tokens, first piece 52.8 ms, total 88.4 ms
request: 35 prompt tokens, 64 new tokens, first piece 48.1 ms, total 377.9 ms

   10.7 ms  'The'
   13.3 ms  ' primary'
   15.9 ms  ' colors'
   18.4 ms  ' are'
   20.9 ms  ':'
   23.5 ms  '\n\n'
...
  164.0 ms  ' that'
  166.5 ms  ' is'
  167.5 ms  '\n'
65 chunks; stderr: request: 35 prompt tokens, 64 new tokens, first piece 10.7 ms, total 167.5 ms

[profile.release]
strip = true          # 去掉符号表
# lto = "fat"         # 跨 crate 的链接时优化
# codegen-units = 1   # 每个 crate 只切一个代码生成单元
# panic = "abort"     # panic 时直接退出，不展开栈
# opt-level = "s"     # 按体积优化

perf-serve:
	/System/Library/Frameworks/Metal.framework/Versions/A/Metal (compatibility version 1.0.0, current version 382.5.3)
	/usr/lib/libSystem.B.dylib (compatibility version 1.0.0, current version 1359.0.0)
	/System/Library/Frameworks/Foundation.framework/Versions/C/Foundation (compatibility version 300.0.0, current version 5027.0.69)
	/System/Library/Frameworks/CoreFoundation.framework/Versions/A/CoreFoundation (compatibility version 150.0.0, current version 5027.0.69)
	/usr/lib/libobjc.A.dylib (compatibility version 1.0.0, current version 228.0.0)
	/usr/lib/libiconv.2.dylib (compatibility version 7.0.0, current version 7.0.0)
	/System/Library/Frameworks/Accelerate.framework/Versions/A/Accelerate (compatibility version 1.0.0, current version 4.0.0)

error: no matching package named `candle-kernels` found
location searched: crates.io index
required by package `candle-core v0.11.0`

error: failed to run custom build command for `onig_sys v69.9.3`
...
  error occurred in cc-rs: failed to find tool "aarch64-linux-gnu-gcc": No such file or directory (os error 2)

// 放在旁边，读文件
let mut r = std::io::BufReader::new(std::fs::File::open(format!("{dir}/SmolLM2-135M-Instruct-Q8_0.gguf"))?);
let content = gguf_file::Content::read(&mut r)?;
let model = ModelWeights::from_gguf(content, &mut r, dev)?;

// 放在旁边，mmap（PERF_MMAP=1）
let map = unsafe { memmap2::Mmap::map(&file)? };
let mut r = std::io::Cursor::new(&map[..]);
let content = gguf_file::Content::read(&mut r)?;
let model = ModelWeights::from_gguf(content, &mut r, dev)?;

// 编进二进制（--features embed）
static GGUF: &[u8] = include_bytes!("SmolLM2-135M-Instruct-Q8_0.gguf");
static TOKENIZER: &[u8] = include_bytes!("tokenizer.json");

let mut r = std::io::Cursor::new(GGUF);
let content = gguf_file::Content::read(&mut r)?;
let model = ModelWeights::from_gguf(content, &mut r, dev)?;
let tok = Tokenizer::from_bytes(TOKENIZER).map_err(|e| anyhow!(e))?;