flashinfer.fused_moe.reorder_rows_for_gated_act_gemm

flashinfer.fused_moe.reorder_rows_for_gated_act_gemm(x: Tensor) → Tensor

Reorder rows of a weight tensor for the TensorRT-LLM gated-activation GEMM layout.

Pure-PyTorch reimplementation of the TensorRT-LLM reorderRowsForGatedActGemm helper. Used to pre-permute the up/gate weight matrix so that the fused gated-activation kernels can access the two halves with a single contiguous load.

The input holds the two FC1 halves stacked as [X1; X2]: rows [0, M // 2) are X1 (the linear/up half) and rows [M // 2, M) are X2 (the gate half), the order the trtllm_*_moe entry points expect for gemm1_weights. The output alternates one row of each half, [X1[0], X2[0], X1[1], X2[1], ...]; deinterleave_gated_rows() with block=1 undoes it.

Parameters:

x (torch.Tensor) – Weight tensor whose rows will be permuted. Any dtype is accepted; only the row dimension is reordered.

Returns:

Row-permuted copy of x (materialized as a new contiguous tensor; PyTorch advanced indexing always copies, never aliases).

Return type:

torch.Tensor