[FEA]: Add opt-in fixed-width DeviceBatchedTopK output padding
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 486
- Avg merge
- 2d 6h
- Merged PRs (30d)
- 295
Description
## Is this a duplicate?
- [x] I confirmed there appear to be no duplicate issues or pull requests for this request and that I agree to the Code of Conduct.
## Area
CUB
## Is your feature request related to a problem? Please describe.
`cub::DeviceBatchedTopK` currently writes `min(max(k, 0), max(segment_size, 0))` items for each segment and leaves the remaining output slots untouched. Consumers that require exactly `max(k, 0)` initialized output slots must therefore prefill the full output or launch a second kernel to materialize the unused tail.
This appears in fixed-width GPU inference pipelines such as FlashInfer's page-table and ragged Top-K transforms used by SGLang. Those APIs use a sentinel in the unused tail. A separate fill or tail kernel adds launch and host-dispatch overhead even though `DeviceBatchedTopK` already resolves each segment's runtime size and `k`.
## Describe the solution you'd like
Add an opt-in output-padding property to the execution environment accepted by the existing Keys and Pairs APIs. Callers compose `DeviceBatchedTopK::OutputPadding(key_pad)` or `OutputPadding(key_pad, value_pad)` into `EnvT` and call the unchanged `MinKeys`, `MaxKeys`, `MinPairs`, or `MaxPairs` signature.
For segment `i`:
- define `requested_i = max(k_i, 0)` and `valid_i = min(requested_i, max(segment_size_i, 0))`;
- preserve the selected output prefix `[0, valid_i)`;
- assign caller-provided padding values to `[valid_i, requested_i)` in the existing Top-K launch;
- require the opted-in caller to provide at least `requested_i` writable output slots; and
- keep the existing compact, untouched-tail behavior unchanged when the property is absent.
The property should support Keys and Pairs, uniform or per-segment sizes and `k`, the temporary-storage API, and the environment-allocating API. It should compose alongside requirements, stream, tuning, and memory-resource properties. Because the padding values are stateful, this is an environment property rather than a `cuda::execution::require(...)` entry.
## Describe alternatives you've considered
- Prefilling the complete output performs unnecessary stores over the selected prefix and requires a separate operation.
- A separate tail kernel preserves semantics but retains the launch and host-dispatch overhead this feature is intended to remove.
- Unconditional padding is incompatible with existing callers that may allocate only `valid_i` slots when `k_i > segment_size_i`.
- A separate output-width parameter is unnecessary for this use case; the requested non-negative `k_i` is already the desired padding endpoint.
## Additional context
I searched open issues and pull requests for `DeviceBatchedTopK`, padding, fixed-width output, sentinels, and tail filling before proposing this:
- #10812 adds sorted segmented Top-K output through a distinct post-processing path and explicitly leaves fusion out of scope.
- #9278 adds a maximum-total-items guarantee and does not initialize unused output positions.
- No open issue or pull request found in those searches provides opt-in tail materialization.
The motivating FlashInfer work is https://github.com/flashinfer-ai/flashinfer/pull/4442. The proposed CCCL contract is generic and has no FlashInfer or SGLang dependency.
Contributor guide
Assessment
This issue has not been assessed yet.