Skip to content

[AdvancedCompiler]Topk(cpp wrapper) - #736

Merged
StrongSpoon merged 6 commits into
flagos-ai:masterfrom
AdvancedCompiler:Topk(cpp-wrapper)
Jul 16, 2025
Merged

[AdvancedCompiler]Topk(cpp wrapper)#736
StrongSpoon merged 6 commits into
flagos-ai:masterfrom
AdvancedCompiler:Topk(cpp-wrapper)

Conversation

@AdvancedCompiler

Copy link
Copy Markdown
Collaborator

PR Category

Operator

Type of Change

New Feature

Description

CPP wrapper packaging for topk op.

Issue

Progress

  • Change is properly reviewed (1 reviewer required, 2 recommended).
  • Change is responded to an issue.
  • Change is fully covered by a UT.

Performance

Comment thread lib/topk.cpp Outdated
Comment on lines +16 to +18
if (dim < 0) {
dim = x.dim() + dim;
}

@Bowen12992 Bowen12992 Jul 15, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
if (dim < 0) {
dim = x.dim() + dim;
}
dim = dim < 0 ? x.dim() + dim : dim;

Comment thread lib/topk.cpp Outdated
Comment on lines +50 to +58
/*def topk_stage1_kernel(
y_ptr,
index_ptr,
x_ptr,
k,
N: tl.constexpr,
CHUNK_SIZE: tl.constexpr,
DESCENDING: tl.constexpr,
):*/

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
/*def topk_stage1_kernel(
y_ptr,
index_ptr,
x_ptr,
k,
N: tl.constexpr,
CHUNK_SIZE: tl.constexpr,
DESCENDING: tl.constexpr,
):*/
/*
def topk_stage1_kernel(y_ptr,
index_ptr,
x_ptr,
k,
N: tl.constexpr,
CHUNK_SIZE: tl.constexpr,
DESCENDING: tl.constexpr):
*/

Comment thread lib/topk.cpp Outdated
Comment on lines +72 to +82
/*def topk_stage2_kernel(
y_ptr,
index_ptr,
chunk_x,
chunk_index,
sort_dim: tl.constexpr,
k: tl.constexpr,
N: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
DESCENDING: tl.constexpr,
):*/

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ditto

Bowen12992
Bowen12992 previously approved these changes Jul 15, 2025

@Bowen12992 Bowen12992 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Comment thread ctests/test_triton_special.cpp
@StrongSpoon
StrongSpoon merged commit 0fe929b into flagos-ai:master Jul 16, 2025
12 of 16 checks passed
nicelynice pushed a commit to nicelynice/FlagGems that referenced this pull request Feb 24, 2026
* merge

* Change the format

* Rename test_triton_special.cpp to test_triton_topk.cpp

* change to test_triton_topk

---------

Co-authored-by: nmpress1 <1935298275@qq.com>
Co-authored-by: nmpress1 <162261334+nmpress1@users.noreply.github.com>
Co-authored-by: StrongSpoon <35829812+StrongSpoon@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants