Function scratch 

fn scratch<T>(
    c: &Qwen4ExpConfig,
    max_t: usize,
    rows: usize,
    all_logits: bool,
    allocate: impl Fn(usize) -> Result<T>,
) -> Result<PrefillScratch<T>>