Function transformer_block_forward_gpu

cake-core/benches/bench_utils.rs:356–383

(bencher: divan::Bencher, seq_len: usize)

/// GPU transformer block: uses CudaBackend for all fused ops.
#[divan::bench(args = [1, 8, 64])]
fn transformer_block_forward_gpu(bencher: divan::Bencher, seq_len: usize) {
    use cake_core::models::common::{Transformer, Cache};
    use std::collections::HashMap;
    let dev = gpu_device();
    let cfg = super::bench_helpers::test_config();
    // Rebuild VarBuilder with tensors on GPU device
    let h = cfg.hidden_size;
    let head_dim = cfg.head_dim.unwrap_or(h / cfg.num_attention_heads);
    let size_q = head_dim * cfg.num_attention_heads;
    let size_kv = head_dim * cfg.num_key_value_heads;
    let i = cfg.intermediate_size;
    let mut map: HashMap<String, Tensor> = HashMap::new();
    map.insert("input_layernorm.weight".into(), Tensor::ones(h, DType::F32, &dev).unwrap());
    map.insert("post_attention_layernorm.weight".into(), Tensor::ones(h, DType::F32, &dev).unwrap());
    map.insert("self_attn.q_proj.weight".into(), make_gpu_tensor(&[size_q, h], 30));
    map.insert("self_attn.k_proj.weight".into(), make_gpu_tensor(&[size_kv, h], 31));
    map.insert("self_attn.v_proj.weight".into(), make_gpu_tensor(&[size_kv, h], 32));
    map.insert("self_attn.o_proj.weight".into(), make_gpu_tensor(&[h, size_q], 33));
    map.insert("mlp.gate_proj.weight".into(), make_gpu_tensor(&[i, h], 36));
    map.insert("mlp.up_proj.weight".into(), make_gpu_tensor(&[i, h], 37));
    map.insert("mlp.down_proj.weight".into(), make_gpu_tensor(&[h, i], 38));
    let vb = candle_nn::VarBuilder::from_tensors(map, DType::F32, &dev);
    let backend = cake_core::backends::create_backend(&dev);
    let block = Transformer::load_for_vibevoice(vb, &cfg, backend).unwrap();
    let mut cache = Cache::new(true, DType::F32, &cfg, &dev).unwrap();
    let x = make_gpu_tensor(&[1, seq_len, cfg.hidden_size], 800);
    bencher.bench_local(|| block.forward_with_cache(&x, 0, 0, &mut cache).unwrap());
}

// ── Native dtype backend load benchmark ──────────────────────────────