Add support for arbitrary segment sizes to `DeviceBatchedTopK`

The goal of this issue is to establish the functionality for `cub::DeviceBatchedTopK` that works for segment sizes of arbitrary segment sizes. This issue can be closed when the following feature matrix is covered:

# Scope

- supported segment sizes: arbitrary segment sizes of up to billions of items
- support for arbitrary k values
- supported output requirements:
- non-deterministic
- unsorted, sorted (unstable)

## Focus / details

- Supported types: f32, bf16 (at least)
- It should return both values and their corresponding indices.
- It should support a wide range of N, K, and batch sizes. jax.lax.top_k has no respriction on it

## Output requirements:

| | unordered | ordered |
|:--------------------------------------------------------------------------------------------------------------------------:|:---------:|:-------:|
| non-deterministic
`determinism::not_guaranteed` | 🟡 | 🟡 |

## Support for arbitrary segment sizes

Status

- [ ] #8363
- [x] #8364

Task | Status | Implementation
Effort | Review
Effort
-- | -- | -- | --
**Multi-CTA AIR Top-K**
Add implementation that supports large segments in cub::DeviceSegmentedTopK | Todo #8363 | 4 weeks | 4 weeks
**Load Balancing Top-K**
Delegate large segments to multi-cta implementation | ✅ #8364 | 2 weeks | 1.5 weeks

## Support for efficient processing across various segment sizes

Status

- [ ] #9749

Task | Status | Implementation
Effort | Review
Effort
-- | -- | -- | --
**WarpBitonic-based Top-K**
Warp bitonic-sort based (warp-level, comparison-based) | ⏳ #9749 | 2.5 weeks | 2 weeks

## Identify algorithm fusion and dispatch heuristic

Status

- [ ] #9807

Task | Status | Implementation
Effort | Review
Effort
-- | -- | -- | --
**Algorithm fusion and dispatch heuristic**
Identify DeviceBatchedTopK algorithm composition, fusion, compilation, and dispatching heuristic #9807 | ⚪ #9807 | 3 weeks | 3 weeks

## Performance Evaluation

Status

- [ ] #10512

Task | Status | Implementation
Effort | Review
Effort
-- | -- | -- | --
**Performance Evaluation**
Full-sweep performance Evaluation for non-deterministic DeviceBatchedTopK
#10512 | ⚪ #10512 | 1.5 weeks | 0.5 weeks

https://github.com/NVIDIA/cccl/issues/10512