什么时候选 candle

import time
T0 = time.perf_counter()
import sys

import torch
T_TORCH = time.perf_counter()
from transformers import AutoModelForCausalLM, AutoTokenizer
T_TF = time.perf_counter()

model_dir = sys.argv[1]  # 本地的 SmolLM2-135M-Instruct 权重
steps = int(sys.argv[2])
torch.set_num_threads(8)
tok = AutoTokenizer.from_pretrained(model_dir)
model = AutoModelForCausalLM.from_pretrained(
  model_dir, dtype=torch.float32, attn_implementation="eager"
).eval()
T_LOAD = time.perf_counter()
ids = tok("The capital of France is", return_tensors="pt").input_ids
with torch.inference_mode():
  out = model.generate(ids, max_new_tokens=steps, do_sample=False)
T_GEN = time.perf_counter()
print(repr(tok.decode(out[0, ids.shape[1]:].tolist())))
# ... 把 T0 到 T_GEN 的各段写到 stderr

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("metal") => Device::new_metal(0)?,
    _ => Device::Cpu,
  };
  let model_dir = args.get(2).context("model dir")?; // 例如 "SmolLM2-135M-Instruct"
  let steps: usize = args.get(3).context("steps")?.parse()?;

  let tok = Tokenizer::from_file(format!("{model_dir}/tokenizer.json")).map_err(|e| anyhow!(e))?;
  let cfg: LlamaConfig = serde_json::from_slice(&std::fs::read(format!("{model_dir}/config.json"))?)?;
  let cfg = cfg.into_config(false);
  // ...
  let vb = unsafe {
    VarBuilder::from_mmaped_safetensors(&[format!("{model_dir}/model.safetensors")], dtype, &dev)?
  };
  let model = Llama::load(vb, &cfg)?;
  let mut cache = Cache::new(true, DType::F32, &cfg, &dev)?;
  let t_load = t0.elapsed();
  // ...
  let ids = tok.encode("The capital of France is", false).map_err(|e| anyhow!(e))?;
  let ids = ids.get_ids();
  let input = Tensor::new(ids, &dev)?.unsqueeze(0)?;
  let logits = model.forward(&input, 0, &mut cache)?.squeeze(0)?;
  let mut next = logits.argmax(0)?.to_scalar::<u32>()?;
  let t_first = t0.elapsed();

  let mut out = vec![next];
  let mut pos = ids.len();
  for _ in 1..steps {
    let x = Tensor::new(&[next], &dev)?.unsqueeze(0)?;
    let l = model.forward(&x, pos, &mut cache)?.squeeze(0)?;
    next = l.argmax(0)?.to_scalar::<u32>()?;
    pos += 1;
    out.push(next);
  }
  println!("{:?}", tok.decode(&out, false).map_err(|e| anyhow!(e))?);
  // ... 把 t_load、t_first 和总时间写到 stderr
  Ok(())
}

' Paris. Paris is the largest city in France and the capital of the French department of the Espace'

" Paris. Paris is the largest city in France and the capital of the French department of the Espace"

x = torch.tensor([[0.1, 0.9, 0.3, 0.7], [0.5, 0.2, 0.8, 0.4]])
v, i = torch.topk(x, 2, dim=-1)
print("topk values =", v.tolist(), "indices =", i.tolist())
mask = torch.tensor([[False, True, False, True], [True, False, False, False]])
print("masked_fill =", x.masked_fill(mask, float("-inf")).tolist())
print("nonzero =", mask.nonzero().tolist())
print("einsum =", torch.einsum("i,j->ij", torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0, 5.0])).tolist())
q = k = vv = torch.arange(256, dtype=torch.float32).reshape(1, 1, 4, 64) / 256
print("sdpa =", F.scaled_dot_product_attention(q, k, vv, is_causal=True)[0, 0, -1, :4].tolist())

topk values = [[0.8999999761581421, 0.699999988079071], [0.800000011920929, 0.5]] indices = [[1, 3], [2, 0]]
masked_fill = [[0.10000000149011612, -inf, 0.30000001192092896, -inf], [-inf, 0.20000000298023224, 0.800000011920929, 0.4000000059604645]]
nonzero = [[0, 1], [0, 3], [1, 0]]
einsum = [[3.0, 4.0, 5.0], [6.0, 8.0, 10.0]]
sdpa = [0.6980971693992615, 0.7020034790039062, 0.7059096693992615, 0.7098159790039062]

use candle_core::{Device, Tensor};

fn main() -> candle_core::Result<()> {
  let dev = Device::Cpu;
  let x = Tensor::new(&[[0.1f32, 0.9, 0.3, 0.7], [0.5, 0.2, 0.8, 0.4]], &dev)?;
  let mask = Tensor::new(&[[0u8, 1, 0, 1], [1, 0, 0, 0]], &dev)?;
  let _top = x.topk(2)?;
  let _filled = x.masked_fill(&mask, f32::NEG_INFINITY)?;
  let _nz = mask.nonzero()?;
  let _outer = candle_core::einsum("i,j->ij", &[&x, &x])?;
  Ok(())
}

error[E0425]: cannot find function `einsum` in crate `candle_core`
  --> src/bin/missing.rs:10:29
   |
10 |   let _outer = candle_core::einsum("i,j->ij", &[&x, &x])?;
   |                             ^^^^^^ not found in `candle_core`

error[E0599]: no method named `topk` found for struct `candle_core::Tensor` in the current scope
   --> src/bin/missing.rs:7:16
    |
  7 |   let _top = x.topk(2)?;
    |                ^^^^ method not found in `candle_core::Tensor`
    |
...
109 |     fn topk(&self, topk: usize) -> Result<TopKOutput>;
    |        ---- the method is available for `candle_core::Tensor` here
    |
    = help: items from traits can only be used if the trait is in scope
help: trait `TopKLastDimOp` which provides `topk` is implemented but not in scope; perhaps you want to import it
    |
  1 + use candle_transformers::models::deepseek2::TopKLastDimOp;
    |

error[E0599]: no method named `masked_fill` found for struct `candle_core::Tensor` in the current scope
 --> src/bin/missing.rs:8:19
  |
8 |   let _filled = x.masked_fill(&mask, f32::NEG_INFINITY)?;
  |                   ^^^^^^^^^^^ method not found in `candle_core::Tensor`

error[E0599]: no method named `nonzero` found for struct `candle_core::Tensor` in the current scope
  --> src/bin/missing.rs:9:18
   |
 9 |   let _nz = mask.nonzero()?;
   |                  ^^^^^^^ method not found in `candle_core::Tensor`
   |
...
help: trait `NonZeroOp` which provides `nonzero` is implemented but not in scope; perhaps you want to import it
   |
 1 + use candle_transformers::models::deepseek2::NonZeroOp;
   |

// topk = 降序 argsort + narrow + gather
let idx = x.arg_sort_last_dim(false)?.narrow(D::Minus1, 0, 2)?.contiguous()?;
println!("topk values = {:?} indices = {:?}", x.gather(&idx, D::Minus1)?.to_vec2::<f32>()?, idx.to_vec2::<u32>()?);
// masked_fill = where_cond，参数顺序：mask 为真取第一个
let mask = Tensor::new(&[[0u8, 1, 0, 1], [1, 0, 0, 0]], &dev)?;
let neg = Tensor::new(f32::NEG_INFINITY, &dev)?.broadcast_as(x.shape())?;
println!("masked_fill = {:?}", mask.where_cond(&neg, &x)?.to_vec2::<f32>()?);
// einsum("i,j->ij") = 外积
let a = Tensor::new(&[1f32, 2.0], &dev)?;
let b = Tensor::new(&[3f32, 4.0, 5.0], &dev)?;
println!("einsum = {:?}", a.unsqueeze(1)?.broadcast_mul(&b.unsqueeze(0)?)?.to_vec2::<f32>()?);

topk values = [[0.9, 0.7], [0.8, 0.5]] indices = [[1, 3], [2, 0]]
masked_fill = [[0.1, -inf, 0.3, -inf], [-inf, 0.2, 0.8, 0.4]]
einsum = [[3.0, 4.0, 5.0], [6.0, 8.0, 10.0]]

let q = (Tensor::arange(0f32, 256.0, &dev)?.reshape((1, 1, 4, 64))? / 256.0)?;
match candle_nn::ops::sdpa(&q, &q, &q, None, true, 0.125, 1.0) {
  Ok(t) => println!("sdpa cpu = {:?}", t.to_vec3::<f32>()),
  Err(e) => println!("sdpa cpu: Err({e})"),
}
let out = attention(&q, &q, &q)?; // 手拼：matmul、tril2 掩码、softmax_last_dim、matmul
println!("attention = {:?}", out.squeeze(0)?.squeeze(0)?.get(3)?.narrow(0, 0, 4)?.to_vec1::<f32>()?);
let mdev = Device::new_metal(0)?;
let qm = q.to_device(&mdev)?;
let m = candle_nn::ops::sdpa(&qm, &qm, &qm, None, true, 0.125, 1.0)?;
println!("sdpa metal = {:?}", m.squeeze(0)?.squeeze(0)?.get(3)?.narrow(0, 0, 4)?.to_vec1::<f32>()?);

sdpa cpu: Err(SDPA has no cpu impl)
attention = [0.69809717, 0.7020034, 0.70590967, 0.7098159]
sdpa metal = [0.69809717, 0.7020035, 0.70590967, 0.709816]

candle-cpu load-only bf16: n=5 real min=0.07s median=0.08s peakRSS min=543.2MiB max=544.8MiB