(bencher: divan::Bencher, seq_len: usize)
| 33 | |
| 34 | #[divan::bench(args = [1, 8, 32])] |
| 35 | fn cache_process_kv(bencher: divan::Bencher, seq_len: usize) { |
| 36 | let cfg = test_config(); |
| 37 | let head_dim = cfg.hidden_size / cfg.num_attention_heads; |
| 38 | // k, v shape: (batch=1, num_kv_heads, seq_len, head_dim) |
| 39 | let k = make_tensor(&[1, cfg.num_key_value_heads, seq_len, head_dim], 10); |
| 40 | let v = make_tensor(&[1, cfg.num_key_value_heads, seq_len, head_dim], 11); |
| 41 | |
| 42 | bencher.bench_local(move || { |
| 43 | let mut cache = make_cache(&cfg); |
| 44 | cache.process_kv(0, k.clone(), v.clone()).unwrap() |
| 45 | }); |
| 46 | } |
| 47 | |
| 48 | #[divan::bench(args = [1, 8, 32])] |
| 49 | fn cache_process_kv_windowed(bencher: divan::Bencher, seq_len: usize) { |
nothing calls this directly
no test coverage detected