sbelharbi:

i have a large float cuda tensor v = (32, 6, 59536).
applying torch.argsort(v, dim=1, descending=True) takes 14ms. (argsort)
it is called twice. so, it weights in term of runtime.
is there any way to speedup this op?


You could split the tensor along the other dimension and sort chunks in parallel on different GPUs and then concatenate them together later on a single GPU.

sorting using extreme axes seems way faster than sorting using inner axes, at least for pytorch. numpy behaves differently.
any explanation?

This might have to do with contiguous data layout in case of some sorting axes. For example if you are sorting on the last axis, all the data is contiguous.