Is this a duplicate?
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:
The motivating FlashInfer work is flashinfer-ai/flashinfer#4442. The proposed CCCL contract is generic and has no FlashInfer or SGLang dependency.
Is this a duplicate?
Area
CUB
Is your feature request related to a problem? Please describe.
cub::DeviceBatchedTopKcurrently writesmin(max(k, 0), max(segment_size, 0))items for each segment and leaves the remaining output slots untouched. Consumers that require exactlymax(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
DeviceBatchedTopKalready resolves each segment's runtime size andk.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)orOutputPadding(key_pad, value_pad)intoEnvTand call the unchangedMinKeys,MaxKeys,MinPairs, orMaxPairssignature.For segment
i:requested_i = max(k_i, 0)andvalid_i = min(requested_i, max(segment_size_i, 0));[0, valid_i);[valid_i, requested_i)in the existing Top-K launch;requested_iwritable output slots; andThe 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 acuda::execution::require(...)entry.Describe alternatives you've considered
valid_islots whenk_i > segment_size_i.k_iis 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:exec::guaranteeand themax_total_num_itemsguarantee #9278 adds a maximum-total-items guarantee and does not initialize unused output positions.The motivating FlashInfer work is flashinfer-ai/flashinfer#4442. The proposed CCCL contract is generic and has no FlashInfer or SGLang dependency.