Function muon_ns_xtx 

pub fn muon_ns_xtx<T: Triton, const TRANSPOSE: bool, const BLOCK_R: i32, const BLOCK_K: i32, const GROUP_R: i32>(
    x_ptr: T::Pointer<f32>,
    t_ptr: T::Pointer<f32>,
    M: i32,
    N: i32,
    stride_xm: i32,
)

Gram matrix: T = X·Xᵀ (!TRANSPOSE) or T = Xᵀ·X (TRANSPOSE).

Both cases are expressed as A @ Aᵀ with different views of X:

• !TRANSPOSE → A = X [M, N], output T [M, M], contraction K = N.
• TRANSPOSE → A = Xᵀ [N, M], output T [N, N], contraction K = M.

stride_xm is the row stride of X (typically = N for row-major [M, N]).

Grid: [ceil(R/BLOCK_R)² grouped by GROUP_R, 1, 1] where R = M or N.