Function softmax 

pub fn softmax(
    x: &[f32],
    rows: usize,
    cols: usize,
    q_len: usize,
    causal: bool,
    off: usize,
) -> Vec<f32>

Stable softmax over the last dim with optional causal masking, matching shaders/softmax.wgsl. Row r’s query position is r % q_len; key j is visible when j <= query + off. Masked entries produce 0.