Implementation:NVIDIA TransformerEngine Permutation C API

Field	Value
Sources	TransformerEngine
Domains	Deep_Learning, Optimization
Last Updated	2026-02-07 14:00 GMT

Overview

Declares the C++ API for token permutation and unpermutation operations used in Mixture-of-Experts (MoE) dispatch, along with a device-side radix sort utility.

Description

permutation.h exposes three functions:

• nvte_permute: Reorders token rows according to sorted_row_id mappings for top-K expert routing. Optionally computes probability gradients and builds a row ID map for reverse permutation.
• nvte_unpermute: Reverses the permutation to restore original token order using the row ID map and routing probabilities.
• nvte_device_radix_sort_pairs: Wraps CUB's device radix sort to sort expert assignment keys and token indices.

These are the CUDA C++ counterparts to the Triton permutation kernels, providing the same functionality for non-Triton execution paths.

Usage

Use in MoE layers where tokens must be permuted to group them by assigned expert before computation, then unpermuted afterward to restore the original sequence order.

Code Reference

Source Location

Repository

NVIDIA/TransformerEngine

File

transformer_engine/common/include/transformer_engine/permutation.h

Lines

1--24

Signature

void nvte_permute(const NVTETensor input, NVTETensor output,
                  const NVTETensor sorted_row_id, NVTETensor row_id_map,
                  const NVTETensor prob, NVTETensor prob_grad,
                  const NVTETensor input_fwd,
                  const int num_rows, const int topK,
                  const int num_cols, const int num_out_tokens,
                  cudaStream_t stream = nullptr);

void nvte_unpermute(const NVTETensor input, NVTETensor output,
                    NVTETensor row_id_map, const NVTETensor prob,
                    const int num_rows, const int topK,
                    const int num_cols,
                    cudaStream_t stream = nullptr);

void nvte_device_radix_sort_pairs(void *temp_storage, size_t *temp_storage_bytes,
                                  int *keys_in, int *keys_out,
                                  int *values_in, int *values_out,
                                  size_t num_items);

Import

#include "transformer_engine/permutation.h"

I/O Contract

Inputs

Name	Type	Required	Description
input	NVTETensor	Yes	Input token tensor
sorted_row_id	NVTETensor	Yes	Sorted row ID mapping from routing
num_rows	int	Yes	Number of token rows
topK	int	Yes	Number of experts per token
num_cols	int	Yes	Hidden dimension

Outputs

Name	Type	Description
output	NVTETensor	Permuted (or unpermuted) token tensor
row_id_map	NVTETensor	Row ID mapping for reverse permutation

Usage Examples

#include "transformer_engine/permutation.h"

// Permute tokens to group by expert
nvte_permute(tokens, permuted_tokens, sorted_row_ids, row_id_map,
             probs, prob_grad, nullptr,
             num_tokens, topk, hidden_dim, num_out_tokens, stream);

// Unpermute after expert computation
nvte_unpermute(expert_output, output, row_id_map, probs,
               num_tokens, topk, hidden_dim, stream);

Related Pages

• Environment:NVIDIA_TransformerEngine_CUDA_Toolkit_Requirements