Repository navigation
Add support for arbitrary segment sizes to DeviceBatchedTopK #9253
Description
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
| 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
| 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
| 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
| Task | Status | Implementation Effort |
Review Effort |
|---|---|---|---|
| Performance Evaluation Full-sweep performance Evaluation for non-deterministic DeviceBatchedTopK |
|||
| #10512 | ⚪ #10512 | 1.5 weeks | 0.5 weeks |
Activity
cub::DeviceSegmentedTopK#8363max_total_num_itemstoDeviceBatchedTopK#9279DeviceBatchedTopK#9806DeviceBatchedTopKalgorithm composition, fusion, compilation, and dispatching heuristic #9807DeviceBatchedTopK#9806DeviceBatchedTopK#10512