[TRACKING] `argsort` metal kernel yields incorrect output with > 1024 elements

### Reproduction:
```rust
// Correct

// The kernel call @ candle-metal-kernels/src/lib.rs:2151 receives the following args:
// nrows: 1 ncols: 1024 ncols_pad: 1024
let d = Tensor::rand(-256_f32, 255., (1, 1024), &candle_core::Device::new_metal(0)?)?;
println!("{d}");
// [[ 137.8366, -72.5639, -186.1103, ..., -225.0789, -141.2470, -12.9232]]
// Tensor[[1, 1024], f32, metal:4294969852]
let i = d.arg_sort_last_dim(true)?;
println!("{i}");
// [[132, 932, 801, ..., 556, 518, 683]]
// Tensor[[1, 1024], u32, metal:4294969852]

// Error - output indices are a bunch of zeroes, returns very large numbers if we use shape E.g. (1, 128650)

// The kernel call @ candle-metal-kernels/src/lib.rs:2151 receives the following args:
// nrows: 1 ncols: 2048 ncols_pad: 2048
let d = Tensor::rand(-256_f32, 255., (1, 2048), &candle_core::Device::new_metal(0)?)?;
println!("{d}");
// [[ 137.8366, -72.5639, -186.1103, ..., -225.0789, -141.2470, -12.9232]]
// Tensor[[1, 2048], f32, metal:4294969852]
let i = d.arg_sort_last_dim(true)?;
println!("{i}");
// [[0, 0, 0, ..., 0, 0, 0]]
// Tensor[[1, 2048], u32, metal:4294969852]
```

Edit: removed incorrect diagnosis.

Start at candle-metal-kernels/src/lib.rs:2151 and reproduce the issue with the Rust examples for 1,024 and 2,048 elements. Compare the returned indices with the expected sorted order, including the larger-shape case. Done means arg_sort_last_dim returns correct indices for inputs over 1,024 elements.

Written by the indexing model from the issue text.