Skip to content

topk_per_row_opt #1394

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
valarLip merged 6 commits into from
Nov 17, 2025
Merged

topk_per_row_opt #1394

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
1d4aeb8
topk_per_row_opt
valarLip Nov 12, 2025
ef382f9
Optimize topk by integrating air-topk kernel (TODO: proper workspace …
CuiCu-618 Nov 14, 2025
984d52b
update
valarLip Nov 15, 2025
c93293d
Merge branch 'main' into topk_per_row_opt
valarLip Nov 15, 2025
f83f59e
slightly optmize
carlushuang Nov 15, 2025
475b691
add back
carlushuang Nov 15, 2025
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
72 changes: 38 additions & 34 deletions csrc/kernels/topk_per_row_kernels.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ static inline __device__ uint16_t extractBinIdx(float x)
using fp32x1 = __attribute__((__ext_vector_type__(1))) float;
using fp32x2 = __attribute__((__ext_vector_type__(2))) float;
using fp32x4 = __attribute__((__ext_vector_type__(4))) float;
using fp32x8 = __attribute__((__ext_vector_type__(8))) float;

template <int vec>
struct to_vector;
Expand All @@ -46,7 +47,11 @@ struct to_vector<4>
{
using type = fp32x4;
};

template <>
struct to_vector<8>
{
using type = fp32x8;
};
Comment thread
valarLip marked this conversation as resolved.
static inline __device__ uint32_t floatAsSortableUint(float x)
{
uint32_t bits = __float_as_uint(x);
Expand Down Expand Up @@ -131,7 +136,7 @@ __device__ bool processHistogramStep(const float* logits,
for(int vecIdx = (rowStart / Vector) + threadIdx.x; vecIdx < (rowEnd + Vector - 1) / Vector;
vecIdx += kNumThreadsPerBlock)
{
auto v = reinterpret_cast<const VectorType*>(logits)[vecIdx];
auto v = reinterpret_cast<const VectorType*>(logits)[vecIdx];
#pragma unroll
for(int j = 0; j < Vector; j++)
{
Expand Down Expand Up @@ -271,11 +276,8 @@ template <int kNumThreadsPerBlock = 512,
bool useRadixSort = true,
int Vector = 4,
bool sortResultLogitDescending = false>
__device__ void topk_per_row_kernel(const float* logits,
const int rowStart,
const int rowEnd,
int* outIndices,
int stride1)
__device__ void topk_per_row_kernel(
const float* logits, const int rowStart, const int rowEnd, int* outIndices, int stride1)
{
// The number of slots for the final pass.
static constexpr int kNumFinalItems = 2048;
Expand Down Expand Up @@ -580,23 +582,23 @@ static __global__ void topk_per_row_decode(
auto logitsLocal = logits + rowIdx * stride0;

topk_per_row_kernel<kNumThreadsPerBlock, kNumBins, kTopK, useRadixSort, Vector>(
logitsLocal, rowStart, rowEnd, outIndicesLocal, stride1);
logitsLocal, rowStart, rowEnd, outIndicesLocal, stride1);
}

} // namespace aiter

void top_k_per_row_prefill(const torch::Tensor& logits,
const torch::Tensor& rowStarts,
const torch::Tensor& rowEnds,
torch::Tensor& indices,
int64_t numRows,
int64_t stride0,
int64_t stride1)
const torch::Tensor& rowStarts,
const torch::Tensor& rowEnds,
torch::Tensor& indices,
int64_t numRows,
int64_t stride0,
int64_t stride1)
{
constexpr int kSortingAlgorithmThreshold = 12288;

// Compute the results on the device.
constexpr int kNumThreadsPerBlock = 512;
constexpr int kNumThreadsPerBlock = 1024;

// The top-k width.
static constexpr int kTopK = 2048;
Expand Down Expand Up @@ -657,18 +659,18 @@ void top_k_per_row_prefill(const torch::Tensor& logits,
}

void top_k_per_row_decode(const torch::Tensor& logits,
int64_t next_n,
const torch::Tensor& seqLens,
torch::Tensor& indices,
int64_t numRows,
int64_t stride0,
int64_t stride1)
int64_t next_n,
const torch::Tensor& seqLens,
torch::Tensor& indices,
int64_t numRows,
int64_t stride0,
int64_t stride1)
{
constexpr int kSortingAlgorithmThreshold = 12288;
// Compute the results on the device.
constexpr int kNumThreadsPerBlock = 1024;
const hipStream_t stream = at::hip::getCurrentHIPStream();
const auto numColumns = logits.size(1);
const auto numColumns = logits.size(1);

if(numColumns < kSortingAlgorithmThreshold)
{
Expand All @@ -695,23 +697,25 @@ void top_k_per_row_decode(const torch::Tensor& logits,
}
else
{
if (stride0 % 4 == 0)
if(stride0 % 4 == 0)
{
aiter::topk_per_row_decode<kNumThreadsPerBlock, true, 4>
<<<numRows, kNumThreadsPerBlock, 0, stream>>>(logits.data_ptr<float>(),
seqLens.data_ptr<int>(),
indices.data_ptr<int>(),
static_cast<int>(stride0),
static_cast<int>(stride1),
static_cast<int>(next_n));
} else {
aiter::topk_per_row_decode<kNumThreadsPerBlock, true, 1>
<<<numRows, kNumThreadsPerBlock, 0, stream>>>(logits.data_ptr<float>(),
seqLens.data_ptr<int>(),
indices.data_ptr<int>(),
static_cast<int>(stride0),
static_cast<int>(stride1),
static_cast<int>(next_n));
static_cast<int>(stride1),
static_cast<int>(next_n));
}
else
{
aiter::topk_per_row_decode<kNumThreadsPerBlock, true, 1>
<<<numRows, kNumThreadsPerBlock, 0, stream>>>(logits.data_ptr<float>(),
seqLens.data_ptr<int>(),
indices.data_ptr<int>(),
static_cast<int>(stride0),
static_cast<int>(stride1),
static_cast<int>(next_n));
}
}
}
}
24 changes: 18 additions & 6 deletions op_tests/test_topk_per_row.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,11 +47,13 @@ def create_random_logits(


def create_row_boundaries(
num_rows: int, top_k: int = 2048
num_rows: int, num_prefix: int = 0, top_k: int = 2048
) -> tuple[torch.Tensor, torch.Tensor]:
"""Create row start and end indices for testing."""
Comment thread
valarLip marked this conversation as resolved.
row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda")
row_ends = torch.arange(1, num_rows + 1, device="cuda", dtype=torch.int32)
row_ends = torch.arange(
num_prefix + 1, num_prefix + num_rows + 1, device="cuda", dtype=torch.int32
)
return row_starts, row_ends


Expand Down Expand Up @@ -161,15 +163,15 @@ def run_top_k_per_row_decode(


@benchmark()
def test_top_k_per_row_prefill(num_rows: int, top_k: int) -> dict:
def test_top_k_per_row_prefill(num_rows: int, num_prefix: int, top_k: int) -> dict:
"""
Test topk_per_row_prefill.
Comment thread
valarLip marked this conversation as resolved.
"""
ret = {}
torch.set_default_device("cuda:0")

# Create test data
row_starts, row_ends = create_row_boundaries(num_rows)
row_starts, row_ends = create_row_boundaries(num_rows, num_prefix)
logits = create_random_logits(row_starts, row_ends, torch.float32, 42)

# Create output tensors
Expand Down Expand Up @@ -288,6 +290,15 @@ def test_top_k_per_row_decode(
e.g.: -k 2048""",
)

parser.add_argument(
"--num_prefix",
type=int,
default=[0],
nargs="+",
help="""top-k elements per row.
Comment thread
valarLip marked this conversation as resolved.
e.g.: --num_prefix 8000 16000 24000 32000 40000 48000 56000""",
)

parser.add_argument(
"-b",
"--decode_batch_size",
Expand Down Expand Up @@ -325,8 +336,9 @@ def test_top_k_per_row_decode(
df = []
for m in args.context_len:
for k in args.top_k:
ret = test_top_k_per_row_prefill(m, k)
df.append(ret)
for num_prefix in args.num_prefix:
ret = test_top_k_per_row_prefill(m, num_prefix, k)
df.append(ret)

df = pd.DataFrame(df)
aiter.logger.info(f"summary for top_k_per_row_prefill kernel:\n{df}")
Expand Down
Loading