Function softmax 

pub fn softmax(
    x: &WgpuStorage,
    rows: usize,
    cols: usize,
    q_len: usize,
    causal: bool,
    off: usize,
) -> WgpuStorage