Skip to content

[Common][PyTorch] Fix int32 overflow and -1 sentinel handling in moe_permute #2907

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

Open
jing-4369 wants to merge 3 commits into
base: main
Choose a base branch
from
Open

[Common][PyTorch] Fix int32 overflow and -1 sentinel handling in moe_permute #2907

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ce34774
[Common][PyTorch] Fix int32 overflow and -1 sentinel handling in moe_…
jing-4369 Apr 21, 2026
b42546b
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Apr 21, 2026
b73a1f9
Guard against invalid num_out_tokens in moe_permute_fwd
jing-4369 Apr 21, 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
21 changes: 7 additions & 14 deletions transformer_engine/common/permutation/permutation.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,19 +19,12 @@ static __global__ void moe_permute_row_map(const int *sorted_row_id, int *row_id
const int tid = threadIdx.x;
const int idx = bid * blockDim.x + tid;

if (idx >= num_rows * topK) return;
if (idx >= num_out_tokens) return;

int source_row = sorted_row_id[idx];
int source_token_id = source_row / topK;
int source_topK_id = source_row % topK;

if (idx >= num_out_tokens) {
// Set the indices of dropped tokens to -1
row_id_map[source_topK_id * num_rows + source_token_id] = -1;
} else {
// Create a row id map for subsequent unpermute operation
row_id_map[source_topK_id * num_rows + source_token_id] = idx;
}
row_id_map[source_topK_id * num_rows + source_token_id] = idx;
}

template <typename T, typename TCompute, bool hasProb>
Expand All @@ -42,7 +35,7 @@ __global__ void moe_unpermute_kernel(const T *input, T *unpermuted_output, const
TCompute *s_prob = reinterpret_cast<TCompute *>(s_mem);

// Each block corresponds to one dest token
const int source_token = blockIdx.x;
const int64_t source_token = blockIdx.x;
const int tid = threadIdx.x;

if (hasProb) {
Expand All @@ -65,7 +58,7 @@ __global__ void moe_unpermute_kernel(const T *input, T *unpermuted_output, const
TCompute frag_elem[kElementsPerAccess];
TCompute frag_sum[kElementsPerAccess];

int source_row = row_id_map[source_token];
int64_t source_row = row_id_map[source_token];

// source_row == -1 represents a dropped token
if (source_row != -1) {
Expand Down Expand Up @@ -134,7 +127,7 @@ __global__ void moe_permute_kernel(const T *input_bwd, const T *input_fwd, T *ac
TCompute *s_prob = reinterpret_cast<TCompute *>(s_mem);

// Each block corresponds to one source token
const int source_token = blockIdx.x;
const int64_t source_token = blockIdx.x;
const int tid = threadIdx.x;

if (hasProb) {
Expand Down Expand Up @@ -172,7 +165,7 @@ __global__ void moe_permute_kernel(const T *input_bwd, const T *input_fwd, T *ac
for (int k = 0; k < topKTile; k++) {
if (k == topK) break;

int dest_row = row_id_map[index];
int64_t dest_row = row_id_map[index];
index += num_rows;

if (dest_row != -1) {
Expand Down Expand Up @@ -239,7 +232,7 @@ void nvte_permute_launcher(const T *input, T *output, const int *sorted_row_id,
// moe_permute_fwd

int threads = 64;
int blocks = (num_rows * topK + threads - 1) / threads;
int blocks = (num_out_tokens + threads - 1) / threads;

moe_permute_row_map<<<blocks, threads, 0, stream>>>(sorted_row_id, row_id_map, num_rows, topK,
num_out_tokens);
Expand Down
18 changes: 13 additions & 5 deletions transformer_engine/pytorch/csrc/extensions/permutation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,13 +51,22 @@ std::tuple<at::Tensor, at::Tensor, std::vector<at::Tensor>> moe_permute_fwd(
reinterpret_cast<int *>(sorted_indices_ptr), reinterpret_cast<int *>(row_id_ptr),
reinterpret_cast<int *>(sorted_row_id_ptr), num_tokens * topK);

// Output buffer alloc
// Signed radix sort places -1 sentinel entries (e.g. expert-parallel rank mask)
// at the HEAD of sorted_row_id. Skip that prefix so the kernel sees only the
// valid suffix, and pre-fill row_id_map with -1 so the dropped slots are marked
// without the kernel ever dereferencing a sentinel.
num_out_tokens = (num_out_tokens > 0) ? num_out_tokens : num_tokens * topK;
NVTE_CHECK(num_out_tokens <= num_tokens * topK, "num_out_tokens (", num_out_tokens,
") must not exceed num_tokens*topK (", num_tokens * topK, ")");
const int num_minus_ones = num_tokens * topK - num_out_tokens;
sorted_row_id_ptr = reinterpret_cast<char *>(sorted_row_id_ptr) +
static_cast<size_t>(num_minus_ones) * sizeof(int);
Comment on lines +61 to +63
Copy link
Copy Markdown
Contributor

@greptile-apps greptile-apps Bot Apr 21, 2026

Choose a reason for hiding this comment

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

P1 num_tokens * topK still computed as int * int

num_tokens and topK are both int, so num_tokens * topK on line 61 is evaluated in 32-bit arithmetic before the result feeds the int64_t subtraction. The same expression appears twice in the NVTE_CHECK on lines 59–60. If num_tokens * topK wraps to a negative int (possible when, e.g., num_tokens ≥ 2^31 / topK), the NVTE_CHECK would either spuriously reject a valid num_out_tokens, or the error-message value would be wrong. Casting to int64_t before the multiplication closes this gap:

Suggested change
const int num_minus_ones = num_tokens * topK - num_out_tokens;
sorted_row_id_ptr = reinterpret_cast<char *>(sorted_row_id_ptr) +
static_cast<size_t>(num_minus_ones) * sizeof(int);
const int64_t total_tokens = static_cast<int64_t>(num_tokens) * topK;
NVTE_CHECK(num_out_tokens <= total_tokens, "num_out_tokens (", num_out_tokens,
") must not exceed num_tokens*topK (", total_tokens, ")");
const int num_minus_ones = static_cast<int>(total_tokens - num_out_tokens);
sorted_row_id_ptr = reinterpret_cast<char *>(sorted_row_id_ptr) +
static_cast<size_t>(num_minus_ones) * sizeof(int);

at::Tensor permuted_output =
torch::empty({num_out_tokens, num_cols},
torch::dtype(input.scalar_type()).device(torch::kCUDA).requires_grad(false));
at::Tensor row_id_map = torch::empty(
{num_tokens * topK}, torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false));
at::Tensor row_id_map =
torch::full({num_tokens * topK}, -1,
torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false));

auto stream = at::cuda::getCurrentCUDAStream().stream();

Expand All @@ -71,8 +80,7 @@ std::tuple<at::Tensor, at::Tensor, std::vector<at::Tensor>> moe_permute_fwd(
static_cast<size_t>(num_cols)},
dtype);
auto sorted_row_id_cu = makeTransformerEngineTensor(
sorted_row_id_ptr, std::vector<size_t>{static_cast<size_t>(num_tokens * topK)},
DType::kInt32);
sorted_row_id_ptr, std::vector<size_t>{static_cast<size_t>(num_out_tokens)}, DType::kInt32);
auto row_id_map_cu = makeTransformerEngineTensor(row_id_map);

nvte_permute(input_cu.data(), permuted_output_cu.data(), sorted_row_id_cu.data(),
Expand Down