NVIDIA / NVIDIA/cccl

[FEA]: Add opt-in fixed-width DeviceBatchedTopK output padding

Open
#10,833 1 comment 0 reactions 2 assignees Claimed by @elstehle View on GitHub
needs triage
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.