// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <limits>

#include "paddle/common/enforce.h"
#include "paddle/phi/backends/gpu/cuda/cuda_graph_with_memory_pool.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/full_kernel.h"
#include "paddle/phi/kernels/gpu/moe_permute_utils.h"
#include "paddle/utils/optional.h"

#if CUDA_VERSION >= 12080
#include <cooperative_groups.h>
#include <cooperative_groups/memcpy_async.h>
#include <cuda/barrier>
#include <cuda/pipeline>
namespace cg = cooperative_groups;
#endif

namespace phi {

using moe::kCumsumInvalidTag;
using moe::kMaxNumExperts;
using moe::kPermuteBlockDimX;
using moe::kPermuteBlockSize;

// ============================================================================
//                        Unified Permute Kernel
// ============================================================================
// Register-centric scheduling: metadata lives in registers, not shared memory.
// Shared memory layout (phases are non-overlapping):
//   Phase 1:   uint32_t[num_experts] expert bitmask
//   Phase 2:   int[ROWS_PER_BLOCK * TOPK] output_rows (reuses bitmask region)
//   TMA only:  ping/pong buffers + routemap/probs (after max(bitmask, outrows))
//
// Rowmap is pre-filled with -1 by host cudaMemsetAsync before kernel launch.
// Phase 1a: Scatter routemap->bitmask, cache expert/prob in registers
// Phase 1b: Progressive cumsum + direct global writes (routemap, probs,
// indices) Phase 2:  Flush output_rows to smem once, then zero-sync data
// movement
//
template <typename TokenT,
          typename IndexT,
          typename ProbT,
          typename ScaleT,
          bool has_scale,
          bool do_gather,
          bool return_expert_indices,
          int TOPK,
          bool USE_TMA,
          int ROWS_PER_BLOCK = kPermuteBlockSize,
          int BLOCK_DIM_X = kPermuteBlockDimX>
__global__ __launch_bounds__(BLOCK_DIM_X) void permute_kernel(
    const TokenT *__restrict__ X,
    const IndexT *__restrict__ routemap_topk,
    const ProbT *__restrict__ probs_topk,
    const ScaleT *__restrict__ XScale,
    const int *__restrict__ expert_base_offset,
    const int *__restrict__ expert_base_offset_end,
    TokenT *__restrict__ X_unzipped,
    int *__restrict__ zipped_expertwise_rowmap,
    ProbT *__restrict__ probs_unzipped,
    ScaleT *__restrict__ XScale_unzipped,
    int *global_expertwise_block_cumsum,
    int *__restrict__ expert_indices,
    const int total_zipped_tokens_num,
    const int token_length,
    const int scale_length,
    const int num_experts,
    const int topk) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
  static_assert(ROWS_PER_BLOCK == 32, "ROWS_PER_BLOCK must equal warp size");

  const int lane_id = threadIdx.x & 31;
  const int warp_id = threadIdx.x >> 5;
  constexpr int warp_num = BLOCK_DIM_X >> 5;

  // ===================== Block-to-data mapping (prefix/suffix) =============
  const bool use_prefix = blockIdx.x % 2 == 0;
  const int block_index_in_X =
      use_prefix ? blockIdx.x / 2 : gridDim.x - 1 - blockIdx.x / 2;
  const int block_row_base = block_index_in_X * ROWS_PER_BLOCK;
  const int block_row_end =
      min(block_row_base + ROWS_PER_BLOCK, total_zipped_tokens_num);

  // ===================== Shared memory layout =============================
  // Section 1 (Phase 1a-1b): uint32_t[num_experts] expert bitmask
  // Section 1 (Phase 2):     int[ROWS_PER_BLOCK * TOPK] output_rows (reuses)
  // Section 2 (USE_TMA):     ping/pong + routemap + probs
  extern __shared__ char smem_raw[];
  uint32_t *expert_bitmask = reinterpret_cast<uint32_t *>(smem_raw);

  // TMA region starts after max(bitmask, output_rows), 32-byte aligned
  constexpr int output_rows_bytes = ROWS_PER_BLOCK * TOPK * sizeof(int);
  [[maybe_unused]] char *tma_base =
      smem_raw + (((max(static_cast<int>(kMaxNumExperts * sizeof(uint32_t)),
                        output_rows_bytes)) +
                   31) &
                  ~31);

  // Initialize expert bitmask
  for (int i = threadIdx.x; i < num_experts; i += BLOCK_DIM_X) {
    expert_bitmask[i] = 0u;
  }

  // ===================== TMA setup ========================================
  [[maybe_unused]] TokenT *ping_buffer = nullptr;
  [[maybe_unused]] TokenT *pong_buffer = nullptr;
  [[maybe_unused]] IndexT *shared_routemap = nullptr;
  [[maybe_unused]] ProbT *shared_probs = nullptr;

#if CUDA_VERSION >= 12080
  [[maybe_unused]] cg::thread_block cg_block = cg::this_thread_block();
  constexpr auto scope = cuda::thread_scope_block;
  constexpr int stages = 2;

#pragma nv_diag_suppress 20054
  __shared__ cuda::pipeline_shared_state<scope, stages> pstate;
#pragma nv_diag_default 20054
  [[maybe_unused]] auto pipe = cuda::make_pipeline(cg_block, &pstate);

  if constexpr (USE_TMA) {
    ping_buffer = reinterpret_cast<TokenT *>(tma_base);
    pong_buffer =
        reinterpret_cast<TokenT *>(tma_base + token_length * sizeof(TokenT));
    shared_routemap = reinterpret_cast<IndexT *>(tma_base + 2 * token_length *
                                                                sizeof(TokenT));
    shared_probs =
        reinterpret_cast<ProbT *>(reinterpret_cast<char *>(shared_routemap) +
                                  ROWS_PER_BLOCK * topk * sizeof(IndexT));

    const int local_elems = (block_row_end - block_row_base) * topk;
    pipe.producer_acquire();
    cuda::memcpy_async(
        cg_block,
        shared_routemap,
        routemap_topk + static_cast<int64_t>(block_row_base) * topk,
        local_elems * sizeof(IndexT),
        pipe);
    cuda::memcpy_async(cg_block,
                       shared_probs,
                       probs_topk + static_cast<int64_t>(block_row_base) * topk,
                       local_elems * sizeof(ProbT),
                       pipe);
    pipe.producer_commit();

    if constexpr (do_gather) {
      pipe.producer_acquire();
      cuda::memcpy_async(
          cg_block,
          ping_buffer,
          X + static_cast<int64_t>(block_row_base) * token_length,
          token_length * sizeof(TokenT),
          pipe);
      pipe.producer_commit();
    }
  }
#endif

  if constexpr (USE_TMA) {
#if CUDA_VERSION >= 12080
    pipe.consumer_wait();
    __syncthreads();
    pipe.consumer_release();
#endif
  } else {
    __syncthreads();
  }

  // ===================== Phase 1a: Scatter into bitmasks ===================
  // Every lane loads ALL its topk columns into registers so that Phase 1b
  // can always match reg_expert[k] == expert_id regardless of warp id.
  int reg_expert[TOPK];
  ProbT reg_prob[TOPK];
#pragma unroll
  for (int k = 0; k < TOPK; k++) {
    reg_expert[k] = -1;
    reg_prob[k] = ProbT(0);
  }

  const int global_row = block_row_base + lane_id;
  const bool row_valid = global_row < total_zipped_tokens_num;

  // Each lane reads all its topk entries; warps collaborate on atomicOr.
  for (int col = 0; col < topk; col++) {
    int expert = -1;
    ProbT prob = ProbT(0);
    if (row_valid) {
      if constexpr (USE_TMA) {
        expert = shared_routemap[lane_id * topk + col];
        prob = shared_probs[lane_id * topk + col];
      } else {
        expert = routemap_topk[static_cast<int64_t>(global_row) * topk + col];
        prob = probs_topk[static_cast<int64_t>(global_row) * topk + col];
      }
    }
    if (expert >= 0 && expert < num_experts) {
      // Only one warp per column does atomicOr (idempotent, but reduces
      // traffic)
      if (col % warp_num == warp_id) {
        atomicOr(&expert_bitmask[expert], 1u << lane_id);
      }
      reg_expert[col] = expert;
      reg_prob[col] = prob;
    }
  }
  __syncthreads();

  // ===================== Phase 1b: Progressive cumsum + global writes =======
  // Chain layout (even blockIdx = prefix, odd = suffix):
  //   prefix: block 0 →  block 2 →  block 4 → ...  (recv from blockIdx-2)
  //   suffix: block 1 →  block 3 →  block 5 → ...  (recv from blockIdx-2)
  // Root blocks (prefix: blockIdx==0, suffix: blockIdx==1) have no predecessor.
  // When mask==0u for an expert, lane 0 still receives & forwards the offset
  // (with local_count==0) so the chain never breaks.
  int reg_output_row[TOPK];
#pragma unroll
  for (int k = 0; k < TOPK; k++) reg_output_row[k] = -1;

  const bool is_chain_root = (use_prefix ? blockIdx.x == 0 : blockIdx.x == 1);

  for (int expert_id = warp_id; expert_id < num_experts;
       expert_id += warp_num) {
    const uint32_t mask = expert_bitmask[expert_id];
    const int local_count = __popc(mask);

    // --- Inter-block cumsum: lane 0 receives from predecessor, sends to
    //     successor.  Always executes regardless of mask to keep chain alive.
    int chain_offset = 0;
    if (lane_id == 0) {
      if (!is_chain_root) {
        const int recv_idx = blockIdx.x * num_experts + expert_id;
        while ((chain_offset =
                    atomicAdd(&global_expertwise_block_cumsum[recv_idx], 0)) ==
               kCumsumInvalidTag) {
        }
      }
      const int send_idx = (blockIdx.x + 2) * num_experts + expert_id;
      atomicExch(&global_expertwise_block_cumsum[send_idx],
                 chain_offset + local_count);
    }

    // --- Intra-block position assignment (only when this expert has tokens)
    // ---
    int final_pos = -1;
    if (mask != 0u) {
      chain_offset = __shfl_sync(0xFFFFFFFF, chain_offset, 0);
      const bool lane_active = (mask & (1u << lane_id)) != 0;
      if (lane_active && row_valid) {
        if (use_prefix) {
          final_pos = expert_base_offset[expert_id] + chain_offset +
                      __popc(mask & ((1u << lane_id) - 1));
        } else {
          final_pos = expert_base_offset_end[expert_id] - chain_offset -
                      __popc((lane_id < 31) ? (mask >> (lane_id + 1)) : 0u);
        }

        zipped_expertwise_rowmap[static_cast<int64_t>(global_row) *
                                     num_experts +
                                 expert_id] = final_pos;
#pragma unroll
        for (int k = 0; k < TOPK; k++) {
          if (reg_expert[k] == expert_id) {
            reg_output_row[k] = final_pos;
            probs_unzipped[final_pos] = reg_prob[k];
            if constexpr (return_expert_indices) {
              expert_indices[final_pos] = expert_id;
            }
            break;
          }
        }
      }
    }
  }

  // ===================== Phase 2: Token data movement ======================
  // All warps must finish Phase 1b before shared memory is repurposed.
  __syncthreads();

  if constexpr (do_gather) {
    // Flush output_rows from registers to shared memory (reuse bitmask region).
    // reg_output_row[k] was set by whichever warp processed the matching
    // expert; other warps still hold -1 for that k.  Pre-fill with -1 then
    // only write non-(-1) values to avoid race conditions.
    int *shared_output_rows = reinterpret_cast<int *>(smem_raw);
    for (int i = threadIdx.x; i < ROWS_PER_BLOCK * TOPK; i += BLOCK_DIM_X) {
      shared_output_rows[i] = -1;
    }
    __syncthreads();
#pragma unroll
    for (int k = 0; k < TOPK; k++) {
      if (reg_output_row[k] >= 0) {
        shared_output_rows[lane_id * TOPK + k] = reg_output_row[k];
      }
    }
    __syncthreads();

    // Data movement loop — no per-row __syncthreads needed (output_rows
    // are already fully materialized in smem). Only TMA pipeline needs sync.
    for (int row = block_row_base; row < block_row_end; row++) {
      const int internal_row = row - block_row_base;

      if constexpr (USE_TMA) {
#if CUDA_VERSION >= 12080
        pipe.consumer_wait();
        cg_block.sync();
        if (row + 1 < block_row_end) {
          TokenT *prefetch_buffer =
              (internal_row % 2 == 0) ? pong_buffer : ping_buffer;
          pipe.producer_acquire();
          cuda::memcpy_async(cg_block,
                             prefetch_buffer,
                             X + static_cast<int64_t>(row + 1) * token_length,
                             token_length * sizeof(TokenT),
                             pipe);
          pipe.producer_commit();
        }
#endif
      }

      [[maybe_unused]] TokenT *current_buffer = nullptr;
      if constexpr (USE_TMA) {
        current_buffer = (internal_row % 2 == 0) ? ping_buffer : pong_buffer;
      }

      // Read output rows from shared memory (no sync needed)
#pragma unroll
      for (int k = 0; k < TOPK; k++) {
        const int out_row = shared_output_rows[internal_row * TOPK + k];
        if (out_row < 0) continue;

        if constexpr (USE_TMA) {
#if CUDA_VERSION >= 12080
          if (threadIdx.x == 0) {
            cuda::device::experimental::cp_async_bulk_shared_to_global(
                &X_unzipped[(int64_t)out_row * (int64_t)token_length],
                current_buffer,
                token_length * sizeof(TokenT));
            cuda::device::experimental::cp_async_bulk_commit_group();
          }
#endif
        } else {
          vectorized_memcpy(
              &X[(int64_t)row * (int64_t)token_length],
              &X_unzipped[(int64_t)out_row * (int64_t)token_length],
              token_length);
        }

        if constexpr (has_scale) {
          try_vectorized_memcpy(
              &XScale[(int64_t)row * (int64_t)scale_length],
              &XScale_unzipped[(int64_t)out_row * (int64_t)scale_length],
              scale_length);
        }
      }

      if constexpr (USE_TMA) {
#if CUDA_VERSION >= 12080
        if (threadIdx.x == 0) {
          cuda::device::experimental::cp_async_bulk_wait_group_read<0>();
        }
        pipe.consumer_release();
#endif
      }
    }
  }
#endif  // __CUDA_ARCH__ >= 900
}

// ============================================================================
//                            Kernel launcher
// ============================================================================
template <typename TokenT,
          typename ProbT,
          typename IntT,
          typename ScaleT,
          bool HasScale,
          bool DoGather,
          bool ReturnIndices,
          int TOPK>
void launch_permute_kernel(const GPUContext &dev_ctx,
                           const DenseTensor &X,
                           const DenseTensor &expert_routemap_topk,
                           const DenseTensor &expert_prob_topk,
                           const paddle::optional<DenseTensor> &XScale,
                           const DenseTensor &expert_offsets,
                           const DenseTensor &expert_offset_end,
                           DenseTensor *X_unzipped,
                           DenseTensor *zipped_expertwise_rowmap,
                           DenseTensor *token_prob_unzipped,
                           DenseTensor *XScale_unzipped,
                           DenseTensor *global_expertwise_block_cumsum,
                           DenseTensor *expert_indices,
                           int total_zipped_tokens_num,
                           int token_length,
                           int scale_length,
                           int num_experts,
                           int topk,
                           int capability) {
  constexpr int ROWS_PER_BLOCK = kPermuteBlockSize;
  constexpr int BLOCK_DIM_X = kPermuteBlockDimX;

  PADDLE_ENFORCE_GE(
      total_zipped_tokens_num,
      0,
      common::errors::InvalidArgument(
          "total_zipped_tokens_num should be non-negative, but got %d.",
          total_zipped_tokens_num));
  if (total_zipped_tokens_num == 0) return;
  const int64_t grid_x =
      (static_cast<int64_t>(total_zipped_tokens_num) + ROWS_PER_BLOCK - 1) /
      ROWS_PER_BLOCK;
  PADDLE_ENFORCE_LE(
      grid_x,
      static_cast<int64_t>(std::numeric_limits<unsigned int>::max()),
      common::errors::InvalidArgument(
          "The grid size of moe_permute should be <= UINT_MAX, but got %ld.",
          grid_x));
  dim3 grid(static_cast<unsigned int>(grid_x));
  dim3 block(BLOCK_DIM_X);

  const TokenT *x_ptr = X.data<TokenT>();
  const IntT *routemap_ptr = expert_routemap_topk.data<IntT>();
  const ProbT *prob_ptr = expert_prob_topk.data<ProbT>();
  const ScaleT *scale_ptr = XScale ? XScale.get_ptr()->data<ScaleT>() : nullptr;
  const int *offset_ptr = expert_offsets.data<int>();
  const int *offset_end_ptr = expert_offset_end.data<int>();
  TokenT *x_out_ptr = X_unzipped->data<TokenT>();
  IntT *rowmap_out_ptr = zipped_expertwise_rowmap->data<IntT>();
  ProbT *prob_out_ptr = token_prob_unzipped->data<ProbT>();
  ScaleT *scale_out_ptr = XScale_unzipped->data<ScaleT>();
  int *cumsum_ptr = global_expertwise_block_cumsum->data<int>();
  int *expert_indices_ptr =
      (ReturnIndices && expert_indices) ? expert_indices->data<int>() : nullptr;

  [[maybe_unused]] bool use_tma = false;
#if CUDA_VERSION >= 12080
  use_tma = capability >= 90 &&
            is_aligned_in_bytes(token_length * sizeof(TokenT)) &&
            is_aligned_in_bytes(sizeof(IntT) * topk * ROWS_PER_BLOCK);
#endif

  // Shared memory: max(bitmask, output_rows) + optional TMA buffers
  constexpr int output_rows_bytes = ROWS_PER_BLOCK * TOPK * sizeof(int);
  const int base_smem = max(static_cast<int>(kMaxNumExperts * sizeof(uint32_t)),
                            output_rows_bytes);

  dispatch::Bool(use_tma, [&](auto tma_tag) {
    constexpr bool UseTMA = decltype(tma_tag)::value;

    int smem = base_smem;
    if constexpr (UseTMA) {
      smem = ((smem + 31) & ~31);
      smem += 2 * token_length * sizeof(TokenT) +
              sizeof(IntT) * TOPK * ROWS_PER_BLOCK +
              sizeof(ProbT) * TOPK * ROWS_PER_BLOCK;
    }

    auto kernel_ptr = permute_kernel<TokenT,
                                     IntT,
                                     ProbT,
                                     ScaleT,
                                     HasScale,
                                     DoGather,
                                     ReturnIndices,
                                     TOPK,
                                     UseTMA,
                                     ROWS_PER_BLOCK,
                                     BLOCK_DIM_X>;
    if (smem > 48 * 1024) {
      PADDLE_ENFORCE_GPU_SUCCESS(cudaFuncSetAttribute(
          kernel_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem));
    }
    kernel_ptr<<<grid, block, smem, dev_ctx.stream()>>>(x_ptr,
                                                        routemap_ptr,
                                                        prob_ptr,
                                                        scale_ptr,
                                                        offset_ptr,
                                                        offset_end_ptr,
                                                        x_out_ptr,
                                                        rowmap_out_ptr,
                                                        prob_out_ptr,
                                                        scale_out_ptr,
                                                        cumsum_ptr,
                                                        expert_indices_ptr,
                                                        total_zipped_tokens_num,
                                                        token_length,
                                                        scale_length,
                                                        num_experts,
                                                        topk);
  });
}

// ============================================================================
//                               Dispatchers
// ============================================================================
template <typename T, typename Context>
void dispatch_permute_kernel(const Context &dev_ctx,
                             const DenseTensor &X,
                             const DenseTensor &expert_routemap_topk,
                             const DenseTensor &expert_prob_topk,
                             const paddle::optional<DenseTensor> &XScale,
                             const DenseTensor &expert_offsets,
                             const DenseTensor &expert_offset_end,
                             DenseTensor *X_unzipped,
                             DenseTensor *zipped_expertwise_rowmap,
                             DenseTensor *token_prob_unzipped,
                             DenseTensor *XScale_unzipped,
                             DenseTensor *global_expertwise_block_cumsum,
                             DenseTensor *expert_indices,
                             int total_zipped_tokens_num,
                             int token_length,
                             int topk,
                             int num_experts,
                             int scale_length,
                             bool do_gather,
                             bool using_ue8m0_scale,
                             bool return_expert_indices) {
  static int capability = dev_ctx.GetComputeCapability();

  dispatch::TokenType(X.dtype(), [&](auto token_tag, auto has_scale_tag) {
    using TokenT = typename decltype(token_tag)::type;
    constexpr bool HasScale = decltype(has_scale_tag)::value;

    dispatch::ProbType(expert_prob_topk.dtype(), [&](auto prob_tag) {
      using ProbT = typename decltype(prob_tag)::type;

      dispatch::ScaleType(using_ue8m0_scale, [&](auto scale_tag) {
        using ScaleT = typename decltype(scale_tag)::type;

        dispatch::Bools(
            [&](auto do_gather_tag, auto return_indices_tag) {
              constexpr bool DoGather = decltype(do_gather_tag)::value;
              constexpr bool ReturnIndices =
                  decltype(return_indices_tag)::value;

              dispatch::TopK(topk, [&](auto topk_tag) {
                constexpr int TK = decltype(topk_tag)::value;

                launch_permute_kernel<TokenT,
                                      ProbT,
                                      int,
                                      ScaleT,
                                      HasScale,
                                      DoGather,
                                      ReturnIndices,
                                      TK>(dev_ctx,
                                          X,
                                          expert_routemap_topk,
                                          expert_prob_topk,
                                          XScale,
                                          expert_offsets,
                                          expert_offset_end,
                                          X_unzipped,
                                          zipped_expertwise_rowmap,
                                          token_prob_unzipped,
                                          XScale_unzipped,
                                          global_expertwise_block_cumsum,
                                          expert_indices,
                                          total_zipped_tokens_num,
                                          token_length,
                                          scale_length,
                                          num_experts,
                                          topk,
                                          capability);
              });
            },
            do_gather,
            return_expert_indices);
      });
    });
  });
}

// ============================================================================
//                        Preprocessing helpers
// ============================================================================
template <typename TokenT, typename ScaleT, typename Context>
void dispatch_preprocess(const Context &dev_ctx,
                         TokenT *X_unzipped_ptr,
                         ScaleT *XScale_unzipped_ptr,
                         float *token_prob_unzipped_ptr,
                         int *expert_indices_ptr,
                         int cols,
                         int quanted_cols,
                         const std::vector<int> &padding_rows) {
  if (padding_rows.empty()) return;

  DenseTensor padding_tokens_tensor;
  padding_tokens_tensor.Resize({static_cast<int64_t>(padding_rows.size())});
  dev_ctx.template Alloc<int>(&padding_tokens_tensor);

  auto *stable_padding_rows = backends::gpu::RestoreHostMemIfCapturingCUDAGraph(
      const_cast<int *>(padding_rows.data()), padding_rows.size());
  PADDLE_ENFORCE_GPU_SUCCESS(cudaMemcpyAsync(padding_tokens_tensor.data<int>(),
                                             stable_padding_rows,
                                             sizeof(int) * padding_rows.size(),
                                             cudaMemcpyHostToDevice,
                                             dev_ctx.stream()));

  PADDLE_ENFORCE_LE_UINT32_MAX(padding_rows.size(), "padding_rows size");
  dim3 grid{static_cast<uint32_t>(padding_rows.size())};
  dim3 block{512};
  const int *padding_ptr = padding_tokens_tensor.data<int>();

  dispatch::Bools(
      [&](auto fill_x_tag, auto fill_scale_tag, auto fill_indices_tag) {
        constexpr bool FillX = decltype(fill_x_tag)::value;
        constexpr bool FillScale = decltype(fill_scale_tag)::value;
        constexpr bool FillIndices = decltype(fill_indices_tag)::value;

        filling_padding_rows_kernel<TokenT,
                                    ScaleT,
                                    FillX,
                                    FillScale,
                                    FillIndices>
            <<<grid, block, 0, dev_ctx.stream()>>>(X_unzipped_ptr,
                                                   XScale_unzipped_ptr,
                                                   token_prob_unzipped_ptr,
                                                   expert_indices_ptr,
                                                   cols,
                                                   quanted_cols,
                                                   padding_ptr);
      },
      X_unzipped_ptr != nullptr,
      XScale_unzipped_ptr != nullptr,
      expert_indices_ptr != nullptr);
}

template <typename Context>
void dispatch_preprocess_w_override(const Context &dev_ctx,
                                    const DenseTensor &expert_routemap_topk,
                                    const int num_experts,
                                    const int padding_alignment,
                                    const int override_buffer_size,
                                    const bool return_expert_indices,
                                    DenseTensor *expert_offset,
                                    DenseTensor *expert_offset_end,
                                    DenseTensor *expert_indices) {
  constexpr int BLOCK_SIZE = 1024;
  PADDLE_ENFORCE_LE(
      expert_routemap_topk.numel(),
      static_cast<int64_t>(std::numeric_limits<int32_t>::max()),
      common::errors::InvalidArgument(
          "expert_routemap_topk.numel() should be <= INT_MAX, but got %ld.",
          expert_routemap_topk.numel()));

  dispatch::Bools(
      [&](auto fill_expert_indices_tag) {
        constexpr bool FillExpertIndices =
            decltype(fill_expert_indices_tag)::value;
        const int smem_bytes =
            static_cast<int>(sizeof(int32_t)) * num_experts * 2;
        routemap_digest_kernel<FillExpertIndices, BLOCK_SIZE>
            <<<1, BLOCK_SIZE, smem_bytes, dev_ctx.stream()>>>(
                expert_routemap_topk.data<int32_t>(),
                expert_offset->data<int32_t>(),
                expert_offset_end->data<int32_t>(),
                expert_indices->data<int32_t>(),
                expert_routemap_topk.numel(),
                num_experts,
                padding_alignment);
      },
      return_expert_indices);
}

// ============================================================================
//                               CPP Interface
// ============================================================================
template <typename T, typename Context>
void MoePermuteKernel(const Context &dev_ctx,
                      const DenseTensor &X,
                      const paddle::optional<DenseTensor> &XScale,
                      const DenseTensor &expert_routemap_topk,
                      const DenseTensor &expert_prob_topk,
                      const int num_experts,
                      const std::vector<int> &tokens_per_expert,
                      const int padding_alignment,
                      const bool do_gather,
                      const bool using_ue8m0_scale,
                      const bool return_expert_indices,
                      const int override_buffer_size,
                      DenseTensor *X_unzipped,
                      DenseTensor *zipped_expertwise_rowmap,
                      DenseTensor *token_prob_unzipped,
                      DenseTensor *XScale_unzipped,
                      DenseTensor *expert_indices) {
  PADDLE_ENFORCE_EQ(
      X.dims().size(),
      2,
      common::errors::InvalidArgument("Input X's dims should be 2, but got %u.",
                                      X.dims().size()));
  PADDLE_ENFORCE_EQ(
      expert_routemap_topk.dims().size(),
      2,
      common::errors::InvalidArgument(
          "Input expert_routemap_topk's dims should be 2, but got %u.",
          expert_routemap_topk.dims().size()));
  PADDLE_ENFORCE_EQ(
      expert_prob_topk.dims().size(),
      2,
      common::errors::InvalidArgument(
          "Input expert_prob_topk's dims should be 2, but got %u.",
          expert_prob_topk.dims().size()));
  PADDLE_ENFORCE_EQ(expert_prob_topk.dims(),
                    expert_routemap_topk.dims(),
                    common::errors::InvalidArgument(
                        "Input expert_prob_topk's dims should be equal to "
                        "expert_routemap_topk's dims, but got %s and %s.",
                        expert_prob_topk.dims(),
                        expert_routemap_topk.dims()));

  const int64_t rows = X.dims()[0];
  const int64_t cols = X.dims()[1];
  const int64_t topk = expert_routemap_topk.dims()[1];
  PADDLE_ENFORCE_EQ(
      expert_routemap_topk.dims()[0],
      rows,
      common::errors::InvalidArgument(
          "Input expert_routemap_topk's first dimension should be equal to "
          "X.dims()[0], but got %ld and %ld.",
          expert_routemap_topk.dims()[0],
          rows));
  PADDLE_ENFORCE_LE(
      rows,
      static_cast<int64_t>(std::numeric_limits<int32_t>::max()) -
          static_cast<int64_t>(kPermuteBlockSize),
      common::errors::InvalidArgument(
          "X.dims()[0] should be <= INT_MAX - %d, received: (%ld)",
          kPermuteBlockSize,
          rows));
  PADDLE_ENFORCE_LE(
      cols,
      std::numeric_limits<int32_t>::max(),
      common::errors::InvalidArgument(
          "X.dims()[1] should be less than INT_MAX, received: (%ld)", cols));
  PADDLE_ENFORCE_GE(topk,
                    1,
                    common::errors::InvalidArgument(
                        "topk should be > 0, received: (%ld)", topk));
  PADDLE_ENFORCE_LE(topk,
                    16,
                    common::errors::InvalidArgument(
                        "topk should be <= 16, received: (%ld)", topk));
  PADDLE_ENFORCE_GE(
      num_experts,
      1,
      common::errors::InvalidArgument(
          "num_experts should be > 0, received: (%d)", num_experts));
  PADDLE_ENFORCE_LE(num_experts,
                    kMaxNumExperts,
                    common::errors::InvalidArgument(
                        "num_experts should be <= %d, received: (%d)",
                        kMaxNumExperts,
                        num_experts));
  PADDLE_ENFORCE_GE(padding_alignment,
                    1,
                    common::errors::InvalidArgument(
                        "padding_alignment should be > 0, received: (%d)",
                        padding_alignment));
  PADDLE_ENFORCE_GE(
      override_buffer_size,
      -1,
      common::errors::InvalidArgument(
          "override_buffer_size should be -1 or non-negative, received: (%d)",
          override_buffer_size));
  PADDLE_ENFORCE_LE(
      rows,
      static_cast<int64_t>(std::numeric_limits<int32_t>::max()) /
          static_cast<int64_t>(num_experts),
      common::errors::InvalidArgument(
          "X.dims()[0] * num_experts should be <= INT_MAX, received: %ld * %d.",
          rows,
          num_experts));
  PADDLE_ENFORCE_LE(
      expert_routemap_topk.numel(),
      static_cast<int64_t>(std::numeric_limits<int32_t>::max()),
      common::errors::InvalidArgument(
          "expert_routemap_topk.numel() should be <= INT_MAX, but got %ld.",
          expert_routemap_topk.numel()));
  if (X.dtype() == DataType::FLOAT8_E4M3FN && do_gather) {
    PADDLE_ENFORCE_EQ(XScale.get_ptr() != nullptr,
                      true,
                      common::errors::InvalidArgument(
                          "Input XScale should not be None when X's dtype is "
                          "FLOAT8_E4M3FN and do_gather is True."));
  }
  if (XScale.get_ptr() != nullptr) {
    PADDLE_ENFORCE_EQ(XScale.get_ptr()->dims().size(),
                      2,
                      common::errors::InvalidArgument(
                          "Input XScale's dims should be 2, but got %u.",
                          XScale.get_ptr()->dims().size()));
    if (do_gather) {
      PADDLE_ENFORCE_EQ(
          XScale.get_ptr()->dims()[0],
          rows,
          common::errors::InvalidArgument(
              "Input XScale's first dimension should be equal to X.dims()[0], "
              "but got %ld and %ld.",
              XScale.get_ptr()->dims()[0],
              rows));
    }
  }
  const int64_t quanted_cols =
      (XScale.get_ptr() != nullptr) ? XScale.get_ptr()->dims()[1] : 0;
  PADDLE_ENFORCE_LE(
      quanted_cols,
      std::numeric_limits<int32_t>::max(),
      common::errors::InvalidArgument(
          "quanted_cols should be less than INT_MAX, received: (%ld)",
          quanted_cols));
  const bool is_buffer_overridden = (override_buffer_size >= 0);

  // Output allocation
  void *XScale_unzipped_ptr = nullptr;
  dev_ctx.template Alloc<T>(X_unzipped);
  dev_ctx.template Alloc<int>(zipped_expertwise_rowmap);
  dev_ctx.template Alloc<float>(token_prob_unzipped);
  dev_ctx.template Alloc<int>(expert_indices);
  auto X_unzipped_ptr = reinterpret_cast<void *>(X_unzipped->data<T>());
  auto token_prob_unzipped_ptr =
      reinterpret_cast<void *>(token_prob_unzipped->data<float>());
  if (using_ue8m0_scale) {
    dev_ctx.template Alloc<int32_t>(XScale_unzipped);
    XScale_unzipped_ptr =
        reinterpret_cast<void *>(XScale_unzipped->data<int32_t>());
  } else {
    dev_ctx.template Alloc<float>(XScale_unzipped);
    XScale_unzipped_ptr =
        reinterpret_cast<void *>(XScale_unzipped->data<float>());
  }
  // Pre-fill expert_indices with -1 via hardware DMA engine (cudaMemsetAsync).
  // (Even if input is 0-size)
  // 0xFF byte-pattern on int32 = 0xFFFFFFFF = -1 in two's complement.
  // This offloads the bulk -1 fill (~10K-500K int32s) from SM compute to the
  // DMA copy engine, running in parallel with subsequent kernel execution.
  if (is_buffer_overridden && return_expert_indices) {
    PADDLE_ENFORCE_GPU_SUCCESS(cudaMemsetAsync(
        expert_indices->data<int32_t>(),
        0xFF,
        static_cast<size_t>(override_buffer_size) * sizeof(int32_t),
        dev_ctx.stream()));
  }

  // Handle empty input: initialize all outputs properly
  if (X.numel() == 0) return;

  // Preprocess
  constexpr int kEffectiveBlockSize = kPermuteBlockSize;
  const int64_t cumsum_blocknum_i64 =
      (rows + kEffectiveBlockSize - 1) / kEffectiveBlockSize;
  PADDLE_ENFORCE_LE(
      cumsum_blocknum_i64 + 2,
      static_cast<int64_t>(std::numeric_limits<int32_t>::max()) /
          static_cast<int64_t>(num_experts),
      common::errors::InvalidArgument(
          "The cumsum buffer size of moe_permute should be <= INT_MAX, but got "
          "(%ld + 2) * %d.",
          cumsum_blocknum_i64,
          num_experts));
  const int cumsum_blocknum = static_cast<int>(cumsum_blocknum_i64);

  DenseTensor expert_offset_tensor;
  DenseTensor expert_offset_end_tensor;
  DenseTensor global_expertwise_block_cumsum;

  expert_offset_tensor.Resize({kMaxNumExperts});
  expert_offset_end_tensor.Resize({kMaxNumExperts});
  global_expertwise_block_cumsum.Resize(
      {static_cast<int64_t>(cumsum_blocknum + 2),
       static_cast<int64_t>(num_experts)});

  dev_ctx.template Alloc<int>(&expert_offset_tensor);
  dev_ctx.template Alloc<int>(&expert_offset_end_tensor);
  dev_ctx.template Alloc<int>(&global_expertwise_block_cumsum);

  // Pre-fill rowmap with -1 via bulk DMA (replaces scattered per-block writes)
  PADDLE_ENFORCE_GPU_SUCCESS(
      cudaMemsetAsync(zipped_expertwise_rowmap->data<int>(),
                      -1,
                      zipped_expertwise_rowmap->numel() * sizeof(int),
                      dev_ctx.stream()));

  if (cumsum_blocknum > 1) {
    PADDLE_ENFORCE_GPU_SUCCESS(
        cudaMemsetAsync(global_expertwise_block_cumsum.data<int>(),
                        -1,
                        global_expertwise_block_cumsum.numel() * sizeof(int),
                        dev_ctx.stream()));
  }

  if (is_buffer_overridden) {
    dispatch_preprocess_w_override(dev_ctx,
                                   expert_routemap_topk,
                                   num_experts,
                                   padding_alignment,
                                   override_buffer_size,
                                   return_expert_indices,
                                   &expert_offset_tensor,
                                   &expert_offset_end_tensor,
                                   expert_indices);
  } else {
    int64_t tokens_cumulated = 0;
    std::vector<int> padding_rows;
    int expert_offset[kMaxNumExperts];
    int expert_offset_end[kMaxNumExperts];
    for (int i = 0; i < kMaxNumExperts; i++) {
      if (i < num_experts) {
        PADDLE_ENFORCE_LE_INT_MAX(tokens_cumulated, "expert offset");
        expert_offset[i] = static_cast<int>(tokens_cumulated);
        const int64_t tokens = tokens_per_expert[i];
        const int64_t padded_tokens =
            ((tokens + padding_alignment - 1) / padding_alignment) *
            padding_alignment;
        const int64_t expert_offset_end_64 = tokens_cumulated + tokens - 1;
        PADDLE_ENFORCE_LE_INT_MAX(expert_offset_end_64, "expert offset end");
        expert_offset_end[i] = static_cast<int>(expert_offset_end_64);
        tokens_cumulated += padded_tokens;
      } else {
        expert_offset[i] = 0;
      }
    }
    auto *stable_expert_offset =
        backends::gpu::RestoreHostMemIfCapturingCUDAGraph(expert_offset,
                                                          kMaxNumExperts);
    PADDLE_ENFORCE_GPU_SUCCESS(cudaMemcpyAsync(expert_offset_tensor.data<int>(),
                                               stable_expert_offset,
                                               sizeof(int) * kMaxNumExperts,
                                               cudaMemcpyHostToDevice,
                                               dev_ctx.stream()));
    auto *stable_expert_offset_end =
        backends::gpu::RestoreHostMemIfCapturingCUDAGraph(expert_offset_end,
                                                          kMaxNumExperts);
    PADDLE_ENFORCE_GPU_SUCCESS(
        cudaMemcpyAsync(expert_offset_end_tensor.data<int>(),
                        stable_expert_offset_end,
                        sizeof(int) * kMaxNumExperts,
                        cudaMemcpyHostToDevice,
                        dev_ctx.stream()));
    for (int i = 0; i < num_experts; i++) {
      int64_t next_expert_offset =
          i < num_experts - 1 ? expert_offset[i + 1] : tokens_cumulated;
      int64_t invalid_rows =
          next_expert_offset - expert_offset[i] - tokens_per_expert[i];
      int64_t cur_expert_end = expert_offset[i] + tokens_per_expert[i];
      if (invalid_rows > 0) {
        PADDLE_ENFORCE_LE_INT_MAX(cur_expert_end + invalid_rows - 1,
                                  "padding row");
      }
      for (int64_t j = 0; j < invalid_rows; ++j) {
        padding_rows.push_back(static_cast<int>(cur_expert_end + j));
      }
    }
    if (using_ue8m0_scale) {
      dispatch_preprocess(dev_ctx,
                          do_gather ? X_unzipped->data<T>() : nullptr,
                          XScale ? XScale_unzipped->data<int32_t>() : nullptr,
                          token_prob_unzipped->data<float>(),
                          expert_indices->data<int>(),
                          static_cast<int>(cols),
                          static_cast<int>(quanted_cols),
                          padding_rows);
    } else {
      dispatch_preprocess(dev_ctx,
                          do_gather ? X_unzipped->data<T>() : nullptr,
                          XScale ? XScale_unzipped->data<float>() : nullptr,
                          token_prob_unzipped->data<float>(),
                          expert_indices->data<int>(),
                          static_cast<int>(cols),
                          static_cast<int>(quanted_cols),
                          padding_rows);
    }
  }

  // Kernel dispatch
  dispatch_permute_kernel<T, Context>(
      dev_ctx,
      X,
      expert_routemap_topk,
      expert_prob_topk,
      XScale,
      expert_offset_tensor,
      expert_offset_end_tensor,
      X_unzipped,
      zipped_expertwise_rowmap,
      token_prob_unzipped,
      XScale_unzipped,
      &global_expertwise_block_cumsum,
      expert_indices,
      static_cast<int>(rows),
      static_cast<int>(cols),
      static_cast<int>(topk),
      num_experts,
      static_cast<int>(quanted_cols),
      do_gather,
      using_ue8m0_scale,
      !is_buffer_overridden && return_expert_indices);
}

}  // namespace phi

PD_REGISTER_KERNEL(moe_permute,
                   GPU,
                   ALL_LAYOUT,
                   phi::MoePermuteKernel,
                   phi::float8_e4m3fn,
                   phi::bfloat16) {}