[FEA]: Optimize `cub::DeviceUniqueByKey` by conditionally loading values

[NaderAlAwar]

Is this a duplicate?

• I confirmed there appear to be no duplicate issues for this request and that I agree to the Code of Conduct

Area

CUB

Is your feature request related to a problem? Please describe.

Similar to #326, CUB UniqueByKey unconditionally loads all values before figuring out which keys are are actually needed.

Describe the solution you'd like

If all values are not set, we should skip them. An early implementation where I only load after we figure out which values to keep lead to regressions in a few cases, notably for smaller input sizes.

Describe alternatives you've considered

No response

Additional context

No response