Skip to content

feat: Add row_starts and dsa_graph_safe to topk #3133

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kahyunnam merged 4 commits into from
Apr 24, 2026
Merged

feat: Add row_starts and dsa_graph_safe to topk #3133

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
63cfd7f
Add `dsa_graph_safe`
zianglih Apr 21, 2026
f5453c5
Refine interface
zianglih Apr 22, 2026
20061c2
Add `row_starts`
zianglih Apr 22, 2026
d9638b1
Keep API backward compatibility
zianglih Apr 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions csrc/flashinfer_topk_binding.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,17 +19,19 @@ using tvm::ffi::Optional;

void radix_topk(TensorView input, TensorView output_indices, TensorView output_values,
Optional<TensorView> maybe_row_states_buffer, int64_t top_k, bool sorted_output,
bool deterministic, int64_t tie_break);
bool deterministic, int64_t tie_break, bool dsa_graph_safe);

void radix_topk_page_table_transform(TensorView input, TensorView output_page_table,
TensorView src_page_table,
Optional<TensorView> maybe_row_to_batch, TensorView lengths,
Optional<TensorView> maybe_row_starts,
Optional<TensorView> maybe_row_states_buffer, int64_t top_k,
bool deterministic, int64_t tie_break);
bool deterministic, int64_t tie_break, bool dsa_graph_safe);

void radix_topk_ragged_transform(TensorView input, TensorView output_indices, TensorView offsets,
TensorView lengths, Optional<TensorView> maybe_row_states_buffer,
int64_t top_k, bool deterministic, int64_t tie_break);
TensorView lengths, Optional<TensorView> maybe_row_starts,
Optional<TensorView> maybe_row_states_buffer, int64_t top_k,
bool deterministic, int64_t tie_break, bool dsa_graph_safe);

bool can_implement_filtered_topk();

Expand Down
43 changes: 33 additions & 10 deletions csrc/topk.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ inline sampling::TopKTieBreak ParseTopKTieBreak(int64_t tie_break) {

void radix_topk(TensorView input, TensorView output_indices, TensorView output_values,
Optional<TensorView> maybe_row_states_buffer, int64_t top_k, bool sorted_output,
bool deterministic, int64_t tie_break) {
bool deterministic, int64_t tie_break, bool dsa_graph_safe) {
CHECK_INPUT(input);
CHECK_INPUT(output_indices);
CHECK_INPUT(output_values);
Expand Down Expand Up @@ -72,7 +72,7 @@ void radix_topk(TensorView input, TensorView output_indices, TensorView output_v
status = sampling::TopKDispatch<c_type, int32_t>(
static_cast<c_type*>(input.data_ptr()), static_cast<int32_t*>(output_indices.data_ptr()),
static_cast<c_type*>(output_values.data_ptr()), batch_size, static_cast<uint32_t>(top_k), d,
row_states_ptr, sorted_output, deterministic, tie_break_mode, stream);
row_states_ptr, sorted_output, deterministic, tie_break_mode, stream, dsa_graph_safe);
return true;
});

Expand All @@ -83,8 +83,9 @@ void radix_topk(TensorView input, TensorView output_indices, TensorView output_v
void radix_topk_page_table_transform(TensorView input, TensorView output_page_table,
TensorView src_page_table,
Optional<TensorView> maybe_row_to_batch, TensorView lengths,
Optional<TensorView> maybe_row_starts,
Optional<TensorView> maybe_row_states_buffer, int64_t top_k,
bool deterministic, int64_t tie_break) {
bool deterministic, int64_t tie_break, bool dsa_graph_safe) {
CHECK_INPUT(input);
CHECK_INPUT(output_page_table);
CHECK_INPUT(src_page_table);
Expand All @@ -93,6 +94,10 @@ void radix_topk_page_table_transform(TensorView input, TensorView output_page_ta
CHECK_DIM(2, output_page_table); // output_page_table: (num_rows, top_k)
CHECK_DIM(2, src_page_table); // src_page_table: (batch_size, max_len)
CHECK_DIM(1, lengths); // lengths: (num_rows,)
if (maybe_row_starts.has_value()) {
CHECK_INPUT(maybe_row_starts.value());
CHECK_DIM(1, maybe_row_starts.value());
}

unsigned int num_rows = input.size(0);
unsigned int max_len = input.size(1);
Expand All @@ -118,14 +123,21 @@ void radix_topk_page_table_transform(TensorView input, TensorView output_page_ta
if (maybe_row_to_batch.has_value()) {
row_to_batch_ptr = static_cast<int32_t*>(maybe_row_to_batch.value().data_ptr());
}
int32_t* row_starts_ptr = nullptr;
if (maybe_row_starts.has_value()) {
TVM_FFI_ICHECK(static_cast<unsigned int>(maybe_row_starts.value().size(0)) == num_rows)
<< "row_starts must have shape (num_rows,)";
row_starts_ptr = static_cast<int32_t*>(maybe_row_starts.value().data_ptr());
}

// Use unified dispatch with heuristics to choose between FilteredTopK and RadixTopK
DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP32_FP16(dtype, c_type, [&] {
status = sampling::TopKPageTableTransformDispatch<c_type, int32_t>(
static_cast<c_type*>(input.data_ptr()), static_cast<int32_t*>(output_page_table.data_ptr()),
static_cast<const int32_t*>(src_page_table.data_ptr()), src_stride, row_to_batch_ptr,
static_cast<int32_t*>(lengths.data_ptr()), num_rows, static_cast<uint32_t>(top_k), max_len,
row_states_ptr, deterministic, tie_break_mode, stream);
static_cast<const int32_t*>(src_page_table.data_ptr()), src_stride,
static_cast<int32_t*>(lengths.data_ptr()), row_starts_ptr, row_to_batch_ptr, num_rows,
static_cast<uint32_t>(top_k), max_len, row_states_ptr, deterministic, tie_break_mode,
stream, dsa_graph_safe);
return true;
});

Expand All @@ -134,8 +146,9 @@ void radix_topk_page_table_transform(TensorView input, TensorView output_page_ta
}

void radix_topk_ragged_transform(TensorView input, TensorView output_indices, TensorView offsets,
TensorView lengths, Optional<TensorView> maybe_row_states_buffer,
int64_t top_k, bool deterministic, int64_t tie_break) {
TensorView lengths, Optional<TensorView> maybe_row_starts,
Optional<TensorView> maybe_row_states_buffer, int64_t top_k,
bool deterministic, int64_t tie_break, bool dsa_graph_safe) {
CHECK_INPUT(input);
CHECK_INPUT(output_indices);
CHECK_INPUT(offsets);
Expand All @@ -144,6 +157,10 @@ void radix_topk_ragged_transform(TensorView input, TensorView output_indices, Te
CHECK_DIM(2, output_indices); // output_indices: (num_rows, top_k)
CHECK_DIM(1, offsets); // offsets: (num_rows,)
CHECK_DIM(1, lengths); // lengths: (num_rows,)
if (maybe_row_starts.has_value()) {
CHECK_INPUT(maybe_row_starts.value());
CHECK_DIM(1, maybe_row_starts.value());
}

unsigned int num_rows = input.size(0);
unsigned int max_len = input.size(1);
Expand All @@ -163,14 +180,20 @@ void radix_topk_ragged_transform(TensorView input, TensorView output_indices, Te
row_states_ptr =
static_cast<sampling::RadixRowState*>(maybe_row_states_buffer.value().data_ptr());
}
int32_t* row_starts_ptr = nullptr;
if (maybe_row_starts.has_value()) {
TVM_FFI_ICHECK(static_cast<unsigned int>(maybe_row_starts.value().size(0)) == num_rows)
<< "row_starts must have shape (num_rows,)";
row_starts_ptr = static_cast<int32_t*>(maybe_row_starts.value().data_ptr());
}

// Use unified dispatch with heuristics to choose between FilteredTopK and RadixTopK
DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP32_FP16(dtype, c_type, [&] {
status = sampling::TopKRaggedTransformDispatch<c_type, int32_t>(
static_cast<c_type*>(input.data_ptr()), static_cast<int32_t*>(output_indices.data_ptr()),
static_cast<const int32_t*>(offsets.data_ptr()), static_cast<int32_t*>(lengths.data_ptr()),
num_rows, static_cast<uint32_t>(top_k), max_len, row_states_ptr, deterministic,
tie_break_mode, stream);
row_starts_ptr, num_rows, static_cast<uint32_t>(top_k), max_len, row_states_ptr,
deterministic, tie_break_mode, stream, dsa_graph_safe);
return true;
});

Expand Down
Loading
Loading