Function softmax_forward 

pub fn softmax_forward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
    x_ptr: T::Pointer<D>,
    y_ptr: T::Pointer<D>,
    _n_rows: i32,
    n_cols: i32,
)
where
    T::I32Tensor: Tensor<i32, 1> + Comparison<i32, BoolTensor = T::BoolTensor>,
    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,

Row-wise softmax forward pass.

Grid: one CTA per row — pid = row index.

Each CTA loads the entire row of BLOCK_SIZE elements, applies Triton’s numerically-stable softmax builtin (max-subtraction + exp + normalise), and stores the result.

Constraint: BLOCK_SIZE must equal n_cols for this kernel; the caller is responsible for rounding n_cols up to the next power of two and passing that as BLOCK_SIZE. No masking is needed when BLOCK_SIZE == n_cols.