diff --git a/include/infiniccl.h b/include/infiniccl.h
index b338d85f5..c687b441a 100644
--- a/include/infiniccl.h
+++ b/include/infiniccl.h
@@ -11,6 +11,12 @@ typedef enum {
     INFINICCL_AVG = 4,
 } infinicclReduceOp_t;
 
+typedef enum {
+    INFINICCL_ALLREDUCE_BACKEND_AUTO = 0,
+    INFINICCL_ALLREDUCE_BACKEND_NCCL = 1,
+    INFINICCL_ALLREDUCE_BACKEND_CUSTOM = 2,
+} infinicclAllReduceBackend_t;
+
 struct InfinicclComm;
 
 typedef struct InfinicclComm *infinicclComm_t;
@@ -23,6 +29,30 @@ __INFINI_C __export infiniStatus_t infinicclCommInitAll(
 
 __INFINI_C __export infiniStatus_t infinicclCommDestroy(infinicclComm_t comm);
 
+__INFINI_C __export infiniStatus_t infinicclCommSetAllReduceBackend(
+    infinicclComm_t comm,
+    infinicclAllReduceBackend_t backend);
+
+__INFINI_C __export infiniStatus_t infinicclCommGetAllReduceBackend(
+    infinicclComm_t comm,
+    infinicclAllReduceBackend_t *backend);
+
+__INFINI_C __export infiniStatus_t infinicclCommRegisterAllReduceBuffers(
+    infinicclComm_t *comms,
+    int ndevice,
+    void **buffers,
+    size_t bytes);
+
+__INFINI_C __export infiniStatus_t infinicclCommRegisterAllReduceBuffer(
+    infinicclComm_t comm,
+    const char *key,
+    void *buffer,
+    size_t bytes);
+
+__INFINI_C __export infiniStatus_t infinicclCommClearAllReduceBuffers(
+    infinicclComm_t *comms,
+    int ndevice);
+
 __INFINI_C __export infiniStatus_t infinicclGroupStart(infinicclComm_t comm);
 
 __INFINI_C __export infiniStatus_t infinicclGroupEnd(infinicclComm_t comm);
diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp
index b5c4ff18f..067046a79 100644
--- a/include/infinicore/ops.hpp
+++ b/include/infinicore/ops.hpp
@@ -26,6 +26,7 @@
 #include "ops/cross_entropy.hpp"
 #include "ops/deepseek_moe.hpp"
 #include "ops/embedding.hpp"
+#include "ops/ernie45_rope.hpp"
 #include "ops/flash_attention.hpp"
 #include "ops/fmin.hpp"
 #include "ops/fmod.hpp"
diff --git a/include/infinicore/ops/ernie45_rope.hpp b/include/infinicore/ops/ernie45_rope.hpp
new file mode 100644
index 000000000..4b0ee6c41
--- /dev/null
+++ b/include/infinicore/ops/ernie45_rope.hpp
@@ -0,0 +1,26 @@
+#pragma once
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "../tensor.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(Ernie45MRoPE, Tensor, Tensor, const Tensor &, double, size_t, size_t, size_t);
+INFINICORE_GRAPH_OP_CLASS(Ernie45VisionRoPE, Tensor, Tensor, const Tensor &, double);
+
+Tensor ernie45_mrope_(Tensor q,
+                      Tensor k,
+                      const Tensor &positions,
+                      double rope_theta,
+                      size_t section_h,
+                      size_t section_w,
+                      size_t section_t);
+
+Tensor ernie45_vision_rope_(Tensor q,
+                            Tensor k,
+                            const Tensor &positions,
+                            double rope_theta);
+
+} // namespace infinicore::op
diff --git a/include/infiniop/ops/ernie45_rope.h b/include/infiniop/ops/ernie45_rope.h
new file mode 100644
index 000000000..a9dd039fa
--- /dev/null
+++ b/include/infiniop/ops/ernie45_rope.h
@@ -0,0 +1,56 @@
+#ifndef __INFINIOP_ERNIE45_ROPE_API_H__
+#define __INFINIOP_ERNIE45_ROPE_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopErnie45MropeDescriptor_t;
+typedef struct InfiniopDescriptor *infiniopErnie45VisionRopeDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateErnie45MropeDescriptor(
+    infiniopHandle_t handle,
+    infiniopErnie45MropeDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t q,
+    infiniopTensorDescriptor_t k,
+    infiniopTensorDescriptor_t positions,
+    double rope_theta,
+    size_t section_h,
+    size_t section_w,
+    size_t section_t);
+
+__INFINI_C __export infiniStatus_t infiniopGetErnie45MropeWorkspaceSize(infiniopErnie45MropeDescriptor_t desc,
+                                                                        size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopErnie45Mrope(
+    infiniopErnie45MropeDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *q,
+    void *k,
+    const void *positions,
+    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyErnie45MropeDescriptor(infiniopErnie45MropeDescriptor_t desc);
+
+__INFINI_C __export infiniStatus_t infiniopCreateErnie45VisionRopeDescriptor(
+    infiniopHandle_t handle,
+    infiniopErnie45VisionRopeDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t q,
+    infiniopTensorDescriptor_t k,
+    infiniopTensorDescriptor_t positions,
+    double rope_theta);
+
+__INFINI_C __export infiniStatus_t infiniopGetErnie45VisionRopeWorkspaceSize(infiniopErnie45VisionRopeDescriptor_t desc,
+                                                                             size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopErnie45VisionRope(
+    infiniopErnie45VisionRopeDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *q,
+    void *k,
+    const void *positions,
+    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyErnie45VisionRopeDescriptor(infiniopErnie45VisionRopeDescriptor_t desc);
+
+#endif
diff --git a/src/infiniccl/cuda/custom_allreduce.cu b/src/infiniccl/cuda/custom_allreduce.cu
new file mode 100644
index 000000000..78de2f86d
--- /dev/null
+++ b/src/infiniccl/cuda/custom_allreduce.cu
@@ -0,0 +1,1416 @@
+#include "custom_allreduce.hpp"
+
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#include <atomic>
+#include <condition_variable>
+#include <cstdint>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "../../utils.h"
+
+namespace infiniccl::cuda {
+
+namespace {
+
+using FlagType = uint32_t;
+
+constexpr size_t kMaxCustomAllReduceBytes = 8 * 1024 * 1024;
+// NVIDIA custom allreduce is a small-message decode fast path. Large prefill
+// collectives should stay on NCCL, which is better optimized for bandwidth.
+constexpr size_t kMaxCustomAllReduceFastPathBytes = 512 * 1024;
+constexpr int kMaxCustomAllReduceWorldSize = 8;
+constexpr int kSignalReady = 0;
+constexpr int kSignalDone = 1;
+constexpr int kSignalCount = 2;
+constexpr int kSignalSlotsPerBlock = kSignalCount * kMaxCustomAllReduceWorldSize + 1;
+constexpr int kSignalFlagSlot = kSignalCount * kMaxCustomAllReduceWorldSize;
+constexpr int kAllReduceBlockSize = 512;
+constexpr int kMaxAllReduceBlocks = 36;
+constexpr size_t kSignalBufferBytes =
+    static_cast<size_t>(kMaxAllReduceBlocks) * kSignalSlotsPerBlock * sizeof(FlagType);
+constexpr size_t kTwoStageWorldSize4ThresholdBytes = 512 * 1024;
+constexpr size_t kTwoStageWorldSize8ThresholdBytes = 256 * 1024;
+
+} // namespace
+
+struct RegisteredAllReduceBuffer {
+    std::string key;
+    void *local_buffer = nullptr;
+    size_t bytes = 0;
+    void **rank_buffers = nullptr;
+};
+
+struct CustomAllReduceContext {
+    uint64_t group_id = 0;
+    int rank = 0;
+    int world_size = 1;
+    int device_id = 0;
+    size_t max_bytes = kMaxCustomAllReduceBytes;
+    std::vector<int> device_ids;
+
+    bool initialized = false;
+    const char *disabled_reason = "custom allreduce context is not initialized";
+
+    void *scratch = nullptr;
+    FlagType *signals = nullptr;
+    void **rank_scratch = nullptr;
+    FlagType **rank_signals = nullptr;
+    std::vector<RegisteredAllReduceBuffer> registered_buffers;
+    std::vector<void **> retired_rank_buffers;
+    std::vector<const infinicclAllReduceBackend_t *> rank_backends;
+
+    const infinicclAllReduceBackend_t *local_backend = nullptr;
+};
+
+namespace {
+
+CustomAllReduceCheckResult unsupported(const char *reason) {
+    return CustomAllReduceCheckResult{false, reason};
+}
+
+CustomAllReduceCheckResult supported() {
+    return CustomAllReduceCheckResult{true, nullptr};
+}
+
+void retire_rank_buffer_table(CustomAllReduceContext *ctx, void **rank_buffers) {
+    if (ctx == nullptr || rank_buffers == nullptr) {
+        return;
+    }
+    ctx->retired_rank_buffers.push_back(rank_buffers);
+}
+
+void free_retired_rank_buffer_tables(CustomAllReduceContext *ctx) {
+    if (ctx == nullptr) {
+        return;
+    }
+    for (auto *rank_buffers : ctx->retired_rank_buffers) {
+        if (rank_buffers != nullptr) {
+            (void)cudaFree(rank_buffers);
+        }
+    }
+    ctx->retired_rank_buffers.clear();
+}
+
+void erase_registered_buffer(CustomAllReduceContext *ctx, const void *local_buffer, const std::string *key = nullptr) {
+    if (ctx == nullptr) {
+        return;
+    }
+    for (auto it = ctx->registered_buffers.begin(); it != ctx->registered_buffers.end();) {
+        const bool match_buffer = local_buffer != nullptr && it->local_buffer == local_buffer;
+        const bool match_key = key != nullptr && it->key == *key;
+        if (match_buffer || match_key) {
+            retire_rank_buffer_table(ctx, it->rank_buffers);
+            it = ctx->registered_buffers.erase(it);
+        } else {
+            ++it;
+        }
+    }
+}
+
+void clear_registered_buffers(CustomAllReduceContext *ctx) {
+    if (ctx == nullptr) {
+        return;
+    }
+    for (auto &entry : ctx->registered_buffers) {
+        retire_rank_buffer_table(ctx, entry.rank_buffers);
+        entry.rank_buffers = nullptr;
+    }
+    ctx->registered_buffers.clear();
+}
+
+RegisteredAllReduceBuffer *find_registered_buffer(CustomAllReduceContext *ctx, const void *local_buffer, size_t bytes) {
+    if (ctx == nullptr || local_buffer == nullptr) {
+        return nullptr;
+    }
+    for (auto &entry : ctx->registered_buffers) {
+        if (entry.local_buffer == local_buffer && bytes <= entry.bytes) {
+            return &entry;
+        }
+    }
+    return nullptr;
+}
+
+bool is_supported_dtype(infiniDtype_t datatype) {
+    return datatype == INFINI_DTYPE_F32 ||
+           datatype == INFINI_DTYPE_F16 ||
+           datatype == INFINI_DTYPE_BF16;
+}
+
+bool is_supported_world_size(int world_size) {
+    return world_size >= 2 && world_size <= kMaxCustomAllReduceWorldSize;
+}
+
+bool is_aligned_to(const void *ptr, size_t alignment) {
+    return (reinterpret_cast<std::uintptr_t>(ptr) % alignment) == 0;
+}
+
+void clear_cuda_error() {
+    (void)cudaGetLastError();
+}
+
+struct PendingRegisteredAllReduceBuffer {
+    int world_size = 0;
+    size_t bytes = 0;
+    std::vector<void *> buffers;
+    std::vector<CustomAllReduceContext *> contexts;
+    std::vector<char> arrived;
+    int arrived_count = 0;
+    bool installed = false;
+    infiniStatus_t status = INFINI_STATUS_SUCCESS;
+    std::condition_variable cv;
+};
+
+std::atomic<uint64_t> next_custom_allreduce_group_id{1};
+std::mutex registered_buffer_registry_mutex;
+std::unordered_map<std::string, std::shared_ptr<PendingRegisteredAllReduceBuffer>> registered_buffer_registry;
+
+std::string make_registered_buffer_key(uint64_t group_id, const char *key) {
+    return std::to_string(group_id) + ":" + std::string(key);
+}
+
+void erase_registered_buffer_registry(uint64_t group_id) {
+    const auto prefix = std::to_string(group_id) + ":";
+    std::lock_guard<std::mutex> lock(registered_buffer_registry_mutex);
+    for (auto it = registered_buffer_registry.begin(); it != registered_buffer_registry.end();) {
+        if (it->first.compare(0, prefix.size(), prefix) == 0) {
+            it = registered_buffer_registry.erase(it);
+        } else {
+            ++it;
+        }
+    }
+}
+
+infiniStatus_t install_registered_buffers(
+    const std::vector<CustomAllReduceContext *> &contexts,
+    const std::vector<void *> &rank_buffers,
+    size_t bytes,
+    const std::string &key) {
+
+    if (contexts.empty() || contexts.size() != rank_buffers.size() || bytes == 0) {
+        return INFINI_STATUS_BAD_PARAM;
+    }
+
+    const int ndevice = static_cast<int>(contexts.size());
+    int previous_device = 0;
+    const bool restore_device = cudaGetDevice(&previous_device) == cudaSuccess;
+
+    auto restore = [&]() {
+        if (restore_device) {
+            (void)cudaSetDevice(previous_device);
+        }
+    };
+    auto rollback_key = [&]() {
+        if (key.empty()) {
+            return;
+        }
+        for (auto *ctx : contexts) {
+            if (ctx != nullptr) {
+                (void)cudaSetDevice(ctx->device_id);
+                erase_registered_buffer(ctx, nullptr, &key);
+            }
+        }
+    };
+
+    for (int i = 0; i < ndevice; ++i) {
+        auto *ctx = contexts[static_cast<size_t>(i)];
+        if (ctx == nullptr || rank_buffers[static_cast<size_t>(i)] == nullptr) {
+            restore();
+            return INFINI_STATUS_NULL_POINTER;
+        }
+        if (!ctx->initialized || ctx->world_size != ndevice) {
+            restore();
+            return INFINI_STATUS_BAD_PARAM;
+        }
+    }
+
+    for (int i = 0; i < ndevice; ++i) {
+        auto *ctx = contexts[static_cast<size_t>(i)];
+        auto *local_buffer = rank_buffers[static_cast<size_t>(i)];
+        if (cudaSetDevice(ctx->device_id) != cudaSuccess) {
+            clear_cuda_error();
+            rollback_key();
+            restore();
+            return INFINI_STATUS_INTERNAL_ERROR;
+        }
+
+        const std::string *key_to_erase = key.empty() ? nullptr : &key;
+        erase_registered_buffer(ctx, local_buffer, key_to_erase);
+        RegisteredAllReduceBuffer entry{};
+        entry.key = key;
+        entry.local_buffer = local_buffer;
+        entry.bytes = bytes;
+        if (cudaMalloc(reinterpret_cast<void **>(&entry.rank_buffers), rank_buffers.size() * sizeof(void *)) != cudaSuccess ||
+            cudaMemcpy(entry.rank_buffers,
+                       rank_buffers.data(),
+                       rank_buffers.size() * sizeof(void *),
+                       cudaMemcpyHostToDevice) != cudaSuccess) {
+            clear_cuda_error();
+            if (entry.rank_buffers != nullptr) {
+                (void)cudaFree(entry.rank_buffers);
+            }
+            rollback_key();
+            restore();
+            return INFINI_STATUS_INTERNAL_ERROR;
+        }
+        ctx->registered_buffers.push_back(entry);
+    }
+
+    restore();
+    return INFINI_STATUS_SUCCESS;
+}
+
+void mark_disabled(const std::vector<CustomAllReduceContext *> &contexts, const char *reason) {
+    for (auto *ctx : contexts) {
+        if (ctx != nullptr) {
+            ctx->initialized = false;
+            ctx->disabled_reason = reason;
+        }
+    }
+}
+
+void release_buffers(const std::vector<CustomAllReduceContext *> &contexts) {
+    int previous_device = 0;
+    const bool restore_device = cudaGetDevice(&previous_device) == cudaSuccess;
+    for (auto *ctx : contexts) {
+        if (ctx == nullptr) {
+            continue;
+        }
+        if (ctx->scratch != nullptr || ctx->signals != nullptr ||
+            !ctx->registered_buffers.empty() || !ctx->retired_rank_buffers.empty()) {
+            (void)cudaSetDevice(ctx->device_id);
+            clear_registered_buffers(ctx);
+            free_retired_rank_buffer_tables(ctx);
+            if (ctx->scratch != nullptr) {
+                (void)cudaFree(ctx->scratch);
+                ctx->scratch = nullptr;
+            }
+            if (ctx->signals != nullptr) {
+                (void)cudaFree(ctx->signals);
+                ctx->signals = nullptr;
+            }
+            if (ctx->rank_scratch != nullptr) {
+                (void)cudaFree(ctx->rank_scratch);
+                ctx->rank_scratch = nullptr;
+            }
+            if (ctx->rank_signals != nullptr) {
+                (void)cudaFree(ctx->rank_signals);
+                ctx->rank_signals = nullptr;
+            }
+        }
+        ctx->rank_backends.clear();
+        ctx->initialized = false;
+    }
+    if (restore_device) {
+        (void)cudaSetDevice(previous_device);
+    }
+}
+
+bool device_supports_custom_allreduce(int device_id) {
+    cudaDeviceProp prop{};
+    if (cudaGetDeviceProperties(&prop, device_id) != cudaSuccess) {
+        return false;
+    }
+    return prop.major >= 8;
+}
+
+bool enable_peer_access(int device_id, int peer_device_id) {
+    if (cudaSetDevice(device_id) != cudaSuccess) {
+        return false;
+    }
+
+    int can_access = 0;
+    if (cudaDeviceCanAccessPeer(&can_access, device_id, peer_device_id) != cudaSuccess) {
+        clear_cuda_error();
+        return false;
+    }
+    if (can_access == 0) {
+        return false;
+    }
+
+    const cudaError_t err = cudaDeviceEnablePeerAccess(peer_device_id, 0);
+    if (err == cudaSuccess) {
+        return true;
+    }
+    if (err == cudaErrorPeerAccessAlreadyEnabled) {
+        clear_cuda_error();
+        return true;
+    }
+    clear_cuda_error();
+    return false;
+}
+
+__device__ __forceinline__ FlagType load_signal_volatile(const FlagType *signal) {
+    FlagType value;
+    asm volatile("ld.volatile.global.u32 %0, [%1];" : "=r"(value) : "l"(signal));
+    return value;
+}
+
+__device__ __forceinline__ void store_signal_volatile(FlagType *signal, FlagType value) {
+    asm volatile("st.volatile.global.u32 [%1], %0;" :: "r"(value), "l"(signal));
+}
+
+__device__ __forceinline__ FlagType load_signal_acquire(const FlagType *signal) {
+    FlagType value;
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
+    asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(value) : "l"(signal));
+#else
+    asm volatile("ld.volatile.global.u32 %0, [%1]; membar.gl;" : "=r"(value) : "l"(signal));
+#endif
+    return value;
+}
+
+__device__ __forceinline__ void store_signal_release(FlagType *signal, FlagType value) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
+    asm volatile("st.release.sys.global.u32 [%1], %0;" :: "r"(value), "l"(signal));
+#else
+    asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" :: "r"(value), "l"(signal));
+#endif
+}
+
+__device__ __forceinline__ int signal_slot_index(int block, int slot, int rank) {
+    return block * kSignalSlotsPerBlock + slot * kMaxCustomAllReduceWorldSize + rank;
+}
+
+__device__ __forceinline__ int signal_flag_index(int block) {
+    return block * kSignalSlotsPerBlock + kSignalFlagSlot;
+}
+
+__device__ __forceinline__ void multi_gpu_barrier_start(
+    FlagType *local_signals,
+    FlagType **rank_signals,
+    int world_size,
+    int rank,
+    int slot) {
+
+    const FlagType flag = load_signal_volatile(local_signals + signal_flag_index(blockIdx.x)) + 1U;
+    if (threadIdx.x < static_cast<unsigned int>(world_size)) {
+        const int peer = static_cast<int>(threadIdx.x);
+        auto *peer_signal = rank_signals[peer] + signal_slot_index(blockIdx.x, slot, rank);
+        const auto *local_peer_signal = local_signals + signal_slot_index(blockIdx.x, slot, peer);
+        store_signal_volatile(peer_signal, flag);
+        while (load_signal_volatile(local_peer_signal) < flag) {
+        }
+    }
+    __syncthreads();
+    if (threadIdx.x == 0) {
+        store_signal_volatile(local_signals + signal_flag_index(blockIdx.x), flag);
+    }
+}
+
+template <bool final_sync>
+__device__ __forceinline__ void multi_gpu_barrier_end(
+    FlagType *local_signals,
+    FlagType **rank_signals,
+    int world_size,
+    int rank,
+    int slot) {
+
+    __syncthreads();
+    const FlagType flag = load_signal_volatile(local_signals + signal_flag_index(blockIdx.x)) + 1U;
+    if (threadIdx.x < static_cast<unsigned int>(world_size)) {
+        const int peer = static_cast<int>(threadIdx.x);
+        auto *peer_signal = rank_signals[peer] + signal_slot_index(blockIdx.x, slot, rank);
+        const auto *local_peer_signal = local_signals + signal_slot_index(blockIdx.x, slot, peer);
+        if constexpr (final_sync) {
+            store_signal_volatile(peer_signal, flag);
+            while (load_signal_volatile(local_peer_signal) < flag) {
+            }
+        } else {
+            store_signal_release(peer_signal, flag);
+            while (load_signal_acquire(local_peer_signal) < flag) {
+            }
+        }
+    }
+    if constexpr (!final_sync) {
+        __syncthreads();
+    }
+    if (threadIdx.x == 0) {
+        store_signal_volatile(local_signals + signal_flag_index(blockIdx.x), flag);
+    }
+}
+
+__device__ __forceinline__ void multi_gpu_barrier(
+    FlagType *local_signals,
+    FlagType **rank_signals,
+    int world_size,
+    int rank,
+    int slot) {
+    if (slot == kSignalDone) {
+        multi_gpu_barrier_end<true>(local_signals, rank_signals, world_size, rank, slot);
+    } else {
+        multi_gpu_barrier_start(local_signals, rank_signals, world_size, rank, slot);
+    }
+}
+
+__device__ __forceinline__ void multi_gpu_barrier_middle(
+    FlagType *local_signals,
+    FlagType **rank_signals,
+    int world_size,
+    int rank,
+    int slot) {
+    multi_gpu_barrier_end<false>(local_signals, rank_signals, world_size, rank, slot);
+}
+
+__device__ float bf16_to_float(uint16_t value) {
+    return __uint_as_float(static_cast<unsigned int>(value) << 16);
+}
+
+__device__ uint16_t float_to_bf16(float value) {
+    const unsigned int bits = __float_as_uint(value);
+    const unsigned int lsb = (bits >> 16) & 1U;
+    const unsigned int rounded = bits + 0x7fffU + lsb;
+    return static_cast<uint16_t>(rounded >> 16);
+}
+
+__global__ void allreduce_1stage_f32_kernel(
+    float *recvbuf,
+    void **rank_scratch,
+    size_t count,
+    FlagType *local_signals,
+    FlagType **rank_signals,
+    int world_size,
+    int rank) {
+
+    multi_gpu_barrier(local_signals, rank_signals, world_size, rank, kSignalReady);
+
+    const size_t stride = blockDim.x * gridDim.x;
+    for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < count; idx += stride) {
+        float sum = 0.0f;
+        for (int rank_idx = 0; rank_idx < world_size; ++rank_idx) {
+            sum += static_cast<const float *>(rank_scratch[rank_idx])[idx];
+        }
+        recvbuf[idx] = sum;
+    }
+
+    multi_gpu_barrier(local_signals, rank_signals, world_size, rank, kSignalDone);
+}
+
+__global__ void allreduce_1stage_f16_kernel(
+    __half *recvbuf,
+    void **rank_scratch,
+    size_t count,
+    FlagType *local_signals,
+    FlagType **rank_signals,
+    int world_size,
+    int rank) {
+
+    multi_gpu_barrier(local_signals, rank_signals, world_size, rank, kSignalReady);
+
+    const size_t stride = blockDim.x * gridDim.x;
+    for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < count; idx += stride) {
+        float sum = 0.0f;
+        for (int rank_idx = 0; rank_idx < world_size; ++rank_idx) {
+            sum += __half2float(static_cast<const __half *>(rank_scratch[rank_idx])[idx]);
+        }
+        recvbuf[idx] = __float2half(sum);
+    }
+
+    multi_gpu_barrier(local_signals, rank_signals, world_size, rank, kSignalDone);
+}
+
+__global__ void allreduce_1stage_bf16_kernel(
+    uint16_t *recvbuf,
+    void **rank_scratch,
+    size_t count,
+    FlagType *local_signals,
+    FlagType **rank_signals,
+    int world_size,
+    int rank) {
+
+    multi_gpu_barrier(local_signals, rank_signals, world_size, rank, kSignalReady);
+
+    const size_t stride = blockDim.x * gridDim.x;
+    for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < count; idx += stride) {
+        float sum = 0.0f;
+        for (int rank_idx = 0; rank_idx < world_size; ++rank_idx) {
+            sum += bf16_to_float(static_cast<const uint16_t *>(rank_scratch[rank_idx])[idx]);
+        }
+        recvbuf[idx] = float_to_bf16(sum);
+    }
+
+    multi_gpu_barrier(local_signals, rank_signals, world_size, rank, kSignalDone);
+}
+
+__global__ void allreduce_2stage_f32_kernel(
+    float *recvbuf,
+    void **rank_scratch,
+    size_t count,
+    FlagType *local_signals,
+    FlagType **rank_signals,
+    int world_size,
+    int rank) {
+
+    auto *local_scratch = static_cast<float *>(rank_scratch[rank]);
+    const size_t stride = blockDim.x * gridDim.x;
+
+    multi_gpu_barrier(local_signals, rank_signals, world_size, rank, kSignalReady);
+
+    const size_t begin = count * static_cast<size_t>(rank) / static_cast<size_t>(world_size);
+    const size_t end = count * static_cast<size_t>(rank + 1) / static_cast<size_t>(world_size);
+    for (size_t idx = begin + blockIdx.x * blockDim.x + threadIdx.x; idx < end; idx += stride) {
+        float sum = 0.0f;
+        for (int rank_idx = 0; rank_idx < world_size; ++rank_idx) {
+            sum += static_cast<const float *>(rank_scratch[rank_idx])[idx];
+        }
+        local_scratch[idx] = sum;
+    }
+
+    multi_gpu_barrier_middle(local_signals, rank_signals, world_size, rank, kSignalReady);
+
+    for (int owner = 0; owner < world_size; ++owner) {
+        const size_t owner_begin = count * static_cast<size_t>(owner) / static_cast<size_t>(world_size);
+        const size_t owner_end = count * static_cast<size_t>(owner + 1) / static_cast<size_t>(world_size);
+        const auto *owner_scratch = static_cast<const float *>(rank_scratch[owner]);
+        for (size_t idx = owner_begin + blockIdx.x * blockDim.x + threadIdx.x; idx < owner_end; idx += stride) {
+            recvbuf[idx] = owner_scratch[idx];
+        }
+    }
+
+    multi_gpu_barrier(local_signals, rank_signals, world_size, rank, kSignalDone);
+}
+
+__global__ void allreduce_2stage_f16_kernel(
+    __half *recvbuf,
+    void **rank_scratch,
+    size_t count,
+    FlagType *local_signals,
+    FlagType **rank_signals,
+    int world_size,
+    int rank) {
+
+    auto *local_scratch = static_cast<__half *>(rank_scratch[rank]);
+    const size_t stride = blockDim.x * gridDim.x;
+
+    multi_gpu_barrier(local_signals, rank_signals, world_size, rank, kSignalReady);
+
+    const size_t begin = count * static_cast<size_t>(rank) / static_cast<size_t>(world_size);
+    const size_t end = count * static_cast<size_t>(rank + 1) / static_cast<size_t>(world_size);
+    for (size_t idx = begin + blockIdx.x * blockDim.x + threadIdx.x; idx < end; idx += stride) {
+        float sum = 0.0f;
+        for (int rank_idx = 0; rank_idx < world_size; ++rank_idx) {
+            sum += __half2float(static_cast<const __half *>(rank_scratch[rank_idx])[idx]);
+        }
+        local_scratch[idx] = __float2half(sum);
+    }
+
+    multi_gpu_barrier_middle(local_signals, rank_signals, world_size, rank, kSignalReady);
+
+    for (int owner = 0; owner < world_size; ++owner) {
+        const size_t owner_begin = count * static_cast<size_t>(owner) / static_cast<size_t>(world_size);
+        const size_t owner_end = count * static_cast<size_t>(owner + 1) / static_cast<size_t>(world_size);
+        const auto *owner_scratch = static_cast<const __half *>(rank_scratch[owner]);
+        for (size_t idx = owner_begin + blockIdx.x * blockDim.x + threadIdx.x; idx < owner_end; idx += stride) {
+            recvbuf[idx] = owner_scratch[idx];
+        }
+    }
+
+    multi_gpu_barrier(local_signals, rank_signals, world_size, rank, kSignalDone);
+}
+
+__global__ void allreduce_2stage_bf16_kernel(
+    uint16_t *recvbuf,
+    void **rank_scratch,
+    size_t count,
+    FlagType *local_signals,
+    FlagType **rank_signals,
+    int world_size,
+    int rank) {
+
+    auto *local_scratch = static_cast<uint16_t *>(rank_scratch[rank]);
+    const size_t stride = blockDim.x * gridDim.x;
+
+    multi_gpu_barrier(local_signals, rank_signals, world_size, rank, kSignalReady);
+
+    const size_t begin = count * static_cast<size_t>(rank) / static_cast<size_t>(world_size);
+    const size_t end = count * static_cast<size_t>(rank + 1) / static_cast<size_t>(world_size);
+    for (size_t idx = begin + blockIdx.x * blockDim.x + threadIdx.x; idx < end; idx += stride) {
+        float sum = 0.0f;
+        for (int rank_idx = 0; rank_idx < world_size; ++rank_idx) {
+            sum += bf16_to_float(static_cast<const uint16_t *>(rank_scratch[rank_idx])[idx]);
+        }
+        local_scratch[idx] = float_to_bf16(sum);
+    }
+
+    multi_gpu_barrier_middle(local_signals, rank_signals, world_size, rank, kSignalReady);
+
+    for (int owner = 0; owner < world_size; ++owner) {
+        const size_t owner_begin = count * static_cast<size_t>(owner) / static_cast<size_t>(world_size);
+        const size_t owner_end = count * static_cast<size_t>(owner + 1) / static_cast<size_t>(world_size);
+        const auto *owner_scratch = static_cast<const uint16_t *>(rank_scratch[owner]);
+        for (size_t idx = owner_begin + blockIdx.x * blockDim.x + threadIdx.x; idx < owner_end; idx += stride) {
+            recvbuf[idx] = owner_scratch[idx];
+        }
+    }
+
+    multi_gpu_barrier(local_signals, rank_signals, world_size, rank, kSignalDone);
+}
+
+__device__ uint4 reduce_packed_f32(
+    void **rank_scratch,
+    size_t vec_idx,
+    int world_size) {
+
+    float4 sum{0.0f, 0.0f, 0.0f, 0.0f};
+    for (int rank_idx = 0; rank_idx < world_size; ++rank_idx) {
+        const auto value = reinterpret_cast<const float4 *>(rank_scratch[rank_idx])[vec_idx];
+        sum.x += value.x;
+        sum.y += value.y;
+        sum.z += value.z;
+        sum.w += value.w;
+    }
+    union {
+        float4 f;
+        uint4 u;
+    } out{sum};
+    return out.u;
+}
+
+__device__ uint4 reduce_packed_f16(
+    void **rank_scratch,
+    size_t vec_idx,
+    int world_size) {
+
+    float sum[8]{};
+    for (int rank_idx = 0; rank_idx < world_size; ++rank_idx) {
+        const uint4 packed = static_cast<const uint4 *>(rank_scratch[rank_idx])[vec_idx];
+        const auto *value = reinterpret_cast<const __half *>(&packed);
+#pragma unroll
+        for (int item = 0; item < 8; ++item) {
+            sum[item] += __half2float(value[item]);
+        }
+    }
+
+    uint4 out{};
+    auto *packed_out = reinterpret_cast<__half *>(&out);
+#pragma unroll
+    for (int item = 0; item < 8; ++item) {
+        packed_out[item] = __float2half(sum[item]);
+    }
+    return out;
+}
+
+__device__ uint4 reduce_packed_bf16(
+    void **rank_scratch,
+    size_t vec_idx,
+    int world_size) {
+
+    float sum[8]{};
+    for (int rank_idx = 0; rank_idx < world_size; ++rank_idx) {
+        const uint4 packed = static_cast<const uint4 *>(rank_scratch[rank_idx])[vec_idx];
+        const auto *value = reinterpret_cast<const uint16_t *>(&packed);
+#pragma unroll
+        for (int item = 0; item < 8; ++item) {
+            sum[item] += bf16_to_float(value[item]);
+        }
+    }
+
+    uint4 out{};
+    auto *packed_out = reinterpret_cast<uint16_t *>(&out);
+#pragma unroll
+    for (int item = 0; item < 8; ++item) {
+        packed_out[item] = float_to_bf16(sum[item]);
+    }
+    return out;
+}
+
+__device__ uint4 reduce_packed_bf16_tp2(uint4 a, uint4 b) {
+    const auto *av = reinterpret_cast<const uint16_t *>(&a);
+    const auto *bv = reinterpret_cast<const uint16_t *>(&b);
+    uint4 out{};
+    auto *ov = reinterpret_cast<uint16_t *>(&out);
+#pragma unroll
+    for (int item = 0; item < 8; ++item) {
+        ov[item] = float_to_bf16(bf16_to_float(av[item]) + bf16_to_float(bv[item]));
+    }
+    return out;
+}
+
+__device__ uint4 reduce_packed(
+    void **rank_scratch,
+    size_t vec_idx,
+    int world_size,
+    infiniDtype_t datatype) {
+
+    switch (datatype) {
+    case INFINI_DTYPE_F32:
+        return reduce_packed_f32(rank_scratch, vec_idx, world_size);
+    case INFINI_DTYPE_F16:
+        return reduce_packed_f16(rank_scratch, vec_idx, world_size);
+    case INFINI_DTYPE_BF16:
+        return reduce_packed_bf16(rank_scratch, vec_idx, world_size);
+    default:
+        return uint4{};
+    }
+}
+
+__global__ void allreduce_1stage_packed_bf16_tp2_kernel(
+    uint4 *recvbuf,
+    void **rank_scratch,
+    size_t vec_count,
+    FlagType *local_signals,
+    FlagType **rank_signals,
+    int rank) {
+
+    multi_gpu_barrier(local_signals, rank_signals, 2, rank, kSignalReady);
+
+    const auto *rank0 = static_cast<const uint4 *>(rank_scratch[0]);
+    const auto *rank1 = static_cast<const uint4 *>(rank_scratch[1]);
+    const size_t stride = blockDim.x * gridDim.x;
+    for (size_t vec_idx = blockIdx.x * blockDim.x + threadIdx.x; vec_idx < vec_count; vec_idx += stride) {
+        recvbuf[vec_idx] = reduce_packed_bf16_tp2(rank0[vec_idx], rank1[vec_idx]);
+    }
+
+    multi_gpu_barrier(local_signals, rank_signals, 2, rank, kSignalDone);
+}
+
+__global__ void allreduce_1stage_packed_kernel(
+    uint4 *recvbuf,
+    void **rank_scratch,
+    size_t vec_count,
+    infiniDtype_t datatype,
+    FlagType *local_signals,
+    FlagType **rank_signals,
+    int world_size,
+    int rank) {
+
+    multi_gpu_barrier(local_signals, rank_signals, world_size, rank, kSignalReady);
+
+    const size_t stride = blockDim.x * gridDim.x;
+    for (size_t vec_idx = blockIdx.x * blockDim.x + threadIdx.x; vec_idx < vec_count; vec_idx += stride) {
+        recvbuf[vec_idx] = reduce_packed(rank_scratch, vec_idx, world_size, datatype);
+    }
+
+    multi_gpu_barrier(local_signals, rank_signals, world_size, rank, kSignalDone);
+}
+
+__global__ void allreduce_2stage_packed_kernel(
+    uint4 *recvbuf,
+    void **rank_scratch,
+    size_t vec_count,
+    infiniDtype_t datatype,
+    FlagType *local_signals,
+    FlagType **rank_signals,
+    int world_size,
+    int rank) {
+
+    auto *local_scratch = static_cast<uint4 *>(rank_scratch[rank]);
+    const size_t stride = blockDim.x * gridDim.x;
+
+    multi_gpu_barrier(local_signals, rank_signals, world_size, rank, kSignalReady);
+
+    const size_t begin = vec_count * static_cast<size_t>(rank) / static_cast<size_t>(world_size);
+    const size_t end = vec_count * static_cast<size_t>(rank + 1) / static_cast<size_t>(world_size);
+    for (size_t vec_idx = begin + blockIdx.x * blockDim.x + threadIdx.x; vec_idx < end; vec_idx += stride) {
+        local_scratch[vec_idx] = reduce_packed(rank_scratch, vec_idx, world_size, datatype);
+    }
+
+    multi_gpu_barrier_middle(local_signals, rank_signals, world_size, rank, kSignalReady);
+
+    for (int owner = 0; owner < world_size; ++owner) {
+        const size_t owner_begin = vec_count * static_cast<size_t>(owner) / static_cast<size_t>(world_size);
+        const size_t owner_end = vec_count * static_cast<size_t>(owner + 1) / static_cast<size_t>(world_size);
+        const auto *owner_scratch = static_cast<const uint4 *>(rank_scratch[owner]);
+        for (size_t vec_idx = owner_begin + blockIdx.x * blockDim.x + threadIdx.x;
+             vec_idx < owner_end;
+             vec_idx += stride) {
+            recvbuf[vec_idx] = owner_scratch[vec_idx];
+        }
+    }
+
+    multi_gpu_barrier(local_signals, rank_signals, world_size, rank, kSignalDone);
+}
+infiniStatus_t check_last_cuda_launch() {
+    return cudaPeekAtLastError() == cudaSuccess ? INFINI_STATUS_SUCCESS : INFINI_STATUS_INTERNAL_ERROR;
+}
+
+int allreduce_block_count(size_t count) {
+    const size_t blocks = (count + kAllReduceBlockSize - 1) / kAllReduceBlockSize;
+    if (blocks == 0) {
+        return 1;
+    }
+    return static_cast<int>(blocks > kMaxAllReduceBlocks ? kMaxAllReduceBlocks : blocks);
+}
+
+bool use_two_stage_allreduce(int world_size, size_t bytes) {
+    if (world_size <= 2) {
+        return false;
+    }
+    if (world_size <= 4) {
+        return bytes > kTwoStageWorldSize4ThresholdBytes;
+    }
+    return bytes > kTwoStageWorldSize8ThresholdBytes;
+}
+
+bool should_register_custom_allreduce_buffer(size_t bytes) {
+    return bytes > 0 && bytes <= kMaxCustomAllReduceFastPathBytes;
+}
+
+bool should_use_registered_custom_allreduce(int world_size, size_t bytes) {
+    if (!is_supported_world_size(world_size)) {
+        return false;
+    }
+    if (bytes == 0 || bytes > kMaxCustomAllReduceFastPathBytes) {
+        return false;
+    }
+    return true;
+}
+
+infiniStatus_t launch_allreduce_kernel(
+    void *recvbuf,
+    void **rank_scratch,
+    size_t count,
+    size_t bytes,
+    infiniDtype_t datatype,
+    int world_size,
+    int rank,
+    bool use_packed,
+    bool force_two_stage,
+    FlagType *local_signals,
+    FlagType **rank_signals,
+    cudaStream_t stream) {
+
+    const bool two_stage = force_two_stage || use_two_stage_allreduce(world_size, bytes);
+    if (use_packed) {
+        const size_t vec_count = bytes / sizeof(uint4);
+        const int blocks = allreduce_block_count(vec_count);
+        if (!two_stage && world_size == 2 && datatype == INFINI_DTYPE_BF16) {
+            allreduce_1stage_packed_bf16_tp2_kernel<<<blocks, kAllReduceBlockSize, 0, stream>>>(
+                static_cast<uint4 *>(recvbuf),
+                rank_scratch,
+                vec_count,
+                local_signals,
+                rank_signals,
+                rank);
+            return check_last_cuda_launch();
+        }
+        if (two_stage) {
+            allreduce_2stage_packed_kernel<<<blocks, kAllReduceBlockSize, 0, stream>>>(
+                static_cast<uint4 *>(recvbuf),
+                rank_scratch,
+                vec_count,
+                datatype,
+                local_signals,
+                rank_signals,
+                world_size,
+                rank);
+        } else {
+            allreduce_1stage_packed_kernel<<<blocks, kAllReduceBlockSize, 0, stream>>>(
+                static_cast<uint4 *>(recvbuf),
+                rank_scratch,
+                vec_count,
+                datatype,
+                local_signals,
+                rank_signals,
+                world_size,
+                rank);
+        }
+        return check_last_cuda_launch();
+    }
+
+    const int blocks = allreduce_block_count(count);
+    switch (datatype) {
+    case INFINI_DTYPE_F32:
+        if (two_stage) {
+            allreduce_2stage_f32_kernel<<<blocks, kAllReduceBlockSize, 0, stream>>>(
+                static_cast<float *>(recvbuf),
+                rank_scratch,
+                count,
+                local_signals,
+                rank_signals,
+                world_size,
+                rank);
+        } else {
+            allreduce_1stage_f32_kernel<<<blocks, kAllReduceBlockSize, 0, stream>>>(
+                static_cast<float *>(recvbuf),
+                rank_scratch,
+                count,
+                local_signals,
+                rank_signals,
+                world_size,
+                rank);
+        }
+        return check_last_cuda_launch();
+    case INFINI_DTYPE_F16:
+        if (two_stage) {
+            allreduce_2stage_f16_kernel<<<blocks, kAllReduceBlockSize, 0, stream>>>(
+                static_cast<__half *>(recvbuf),
+                rank_scratch,
+                count,
+                local_signals,
+                rank_signals,
+                world_size,
+                rank);
+        } else {
+            allreduce_1stage_f16_kernel<<<blocks, kAllReduceBlockSize, 0, stream>>>(
+                static_cast<__half *>(recvbuf),
+                rank_scratch,
+                count,
+                local_signals,
+                rank_signals,
+                world_size,
+                rank);
+        }
+        return check_last_cuda_launch();
+    case INFINI_DTYPE_BF16:
+        if (two_stage) {
+            allreduce_2stage_bf16_kernel<<<blocks, kAllReduceBlockSize, 0, stream>>>(
+                static_cast<uint16_t *>(recvbuf),
+                rank_scratch,
+                count,
+                local_signals,
+                rank_signals,
+                world_size,
+                rank);
+        } else {
+            allreduce_1stage_bf16_kernel<<<blocks, kAllReduceBlockSize, 0, stream>>>(
+                static_cast<uint16_t *>(recvbuf),
+                rank_scratch,
+                count,
+                local_signals,
+                rank_signals,
+                world_size,
+                rank);
+        }
+        return check_last_cuda_launch();
+    default:
+        return INFINI_STATUS_BAD_TENSOR_DTYPE;
+    }
+}
+
+cudaStream_t get_cuda_stream(infinirtStream_t stream) {
+    return stream == nullptr ? 0 : static_cast<cudaStream_t>(stream);
+}
+
+} // namespace
+
+CustomAllReduceContext *createCustomAllReduceContext(
+    int rank,
+    int world_size,
+    int device_id,
+    const int *device_ids) {
+
+    auto *ctx = new CustomAllReduceContext{};
+    ctx->rank = rank;
+    ctx->world_size = world_size;
+    ctx->device_id = device_id;
+    if (device_ids != nullptr && world_size > 0) {
+        ctx->device_ids.assign(device_ids, device_ids + world_size);
+    }
+    return ctx;
+}
+
+void initializeCustomAllReduceContexts(
+    infinicclComm_t *comms,
+    int ndevice,
+    const int *device_ids) {
+
+    if (comms == nullptr || device_ids == nullptr || ndevice <= 0) {
+        return;
+    }
+
+    const auto group_id = next_custom_allreduce_group_id.fetch_add(1, std::memory_order_relaxed);
+    std::vector<CustomAllReduceContext *> contexts;
+    contexts.reserve(static_cast<size_t>(ndevice));
+    for (int i = 0; i < ndevice; ++i) {
+        auto *ctx = static_cast<CustomAllReduceContext *>(comms[i]->custom_allreduce_context);
+        if (ctx != nullptr) {
+            ctx->group_id = group_id;
+            ctx->local_backend = &comms[i]->allreduce_backend;
+            ctx->rank_backends.assign(static_cast<size_t>(ndevice), nullptr);
+            ctx->disabled_reason = "custom allreduce supports TP sizes from 2 to 8 in this phase";
+        }
+        contexts.push_back(ctx);
+    }
+
+    if (!is_supported_world_size(ndevice)) {
+        mark_disabled(contexts, "custom allreduce supports TP sizes from 2 to 8 in this phase");
+        return;
+    }
+
+    int previous_device = 0;
+    const bool restore_device = cudaGetDevice(&previous_device) == cudaSuccess;
+
+    for (auto *ctx : contexts) {
+        if (ctx == nullptr || !device_supports_custom_allreduce(ctx->device_id)) {
+            release_buffers(contexts);
+            mark_disabled(contexts, "custom allreduce requires an SM80 or newer CUDA device");
+            if (restore_device) {
+                (void)cudaSetDevice(previous_device);
+            }
+            return;
+        }
+    }
+
+    for (int i = 0; i < ndevice; ++i) {
+        for (int peer = 0; peer < ndevice; ++peer) {
+            if (peer == i) {
+                continue;
+            }
+            if (!enable_peer_access(device_ids[i], device_ids[peer])) {
+                release_buffers(contexts);
+                mark_disabled(contexts, "custom allreduce requires peer access between all TP ranks");
+                if (restore_device) {
+                    (void)cudaSetDevice(previous_device);
+                }
+                return;
+            }
+        }
+    }
+
+    for (auto *ctx : contexts) {
+        if (ctx == nullptr || cudaSetDevice(ctx->device_id) != cudaSuccess) {
+            release_buffers(contexts);
+            mark_disabled(contexts, "custom allreduce failed to set the rank device during initialization");
+            if (restore_device) {
+                (void)cudaSetDevice(previous_device);
+            }
+            return;
+        }
+        if (cudaMalloc(&ctx->scratch, ctx->max_bytes) != cudaSuccess ||
+            cudaMalloc(reinterpret_cast<void **>(&ctx->signals), kSignalBufferBytes) != cudaSuccess ||
+            cudaMalloc(reinterpret_cast<void **>(&ctx->rank_scratch), ndevice * sizeof(void *)) != cudaSuccess ||
+            cudaMalloc(reinterpret_cast<void **>(&ctx->rank_signals), ndevice * sizeof(FlagType *)) != cudaSuccess ||
+            cudaMemset(ctx->signals, 0, kSignalBufferBytes) != cudaSuccess) {
+            clear_cuda_error();
+            release_buffers(contexts);
+            mark_disabled(contexts, "custom allreduce failed to allocate scratch or signal buffers");
+            if (restore_device) {
+                (void)cudaSetDevice(previous_device);
+            }
+            return;
+        }
+    }
+
+    std::vector<void *> scratch_ptrs(static_cast<size_t>(ndevice));
+    std::vector<FlagType *> signal_ptrs(static_cast<size_t>(ndevice));
+    for (int i = 0; i < ndevice; ++i) {
+        scratch_ptrs[static_cast<size_t>(i)] = contexts[static_cast<size_t>(i)]->scratch;
+        signal_ptrs[static_cast<size_t>(i)] = contexts[static_cast<size_t>(i)]->signals;
+    }
+    for (int i = 0; i < ndevice; ++i) {
+        auto *ctx = contexts[static_cast<size_t>(i)];
+        for (int rank_idx = 0; rank_idx < ndevice; ++rank_idx) {
+            ctx->rank_backends[static_cast<size_t>(rank_idx)] = &comms[rank_idx]->allreduce_backend;
+        }
+    }
+
+    for (int i = 0; i < ndevice; ++i) {
+        auto *ctx = contexts[i];
+        if (cudaSetDevice(ctx->device_id) != cudaSuccess ||
+            cudaMemcpy(ctx->rank_scratch,
+                       scratch_ptrs.data(),
+                       scratch_ptrs.size() * sizeof(void *),
+                       cudaMemcpyHostToDevice) != cudaSuccess ||
+            cudaMemcpy(ctx->rank_signals,
+                       signal_ptrs.data(),
+                       signal_ptrs.size() * sizeof(FlagType *),
+                       cudaMemcpyHostToDevice) != cudaSuccess) {
+            clear_cuda_error();
+            release_buffers(contexts);
+            mark_disabled(contexts, "custom allreduce failed to initialize rank pointer tables");
+            if (restore_device) {
+                (void)cudaSetDevice(previous_device);
+            }
+            return;
+        }
+        ctx->initialized = true;
+        ctx->disabled_reason = nullptr;
+    }
+
+    if (restore_device) {
+        (void)cudaSetDevice(previous_device);
+    }
+}
+
+void destroyCustomAllReduceContext(CustomAllReduceContext *ctx) {
+    if (ctx == nullptr) {
+        return;
+    }
+    erase_registered_buffer_registry(ctx->group_id);
+    std::vector<CustomAllReduceContext *> contexts{ctx};
+    release_buffers(contexts);
+    delete ctx;
+}
+
+infiniStatus_t clearCustomAllReduceBuffers(infinicclComm_t *comms, int ndevice) {
+    if (comms == nullptr) {
+        return INFINI_STATUS_NULL_POINTER;
+    }
+    if (ndevice <= 0) {
+        return INFINI_STATUS_BAD_PARAM;
+    }
+
+    int previous_device = 0;
+    const bool restore_device = cudaGetDevice(&previous_device) == cudaSuccess;
+    for (int i = 0; i < ndevice; ++i) {
+        if (comms[i] == nullptr || comms[i]->custom_allreduce_context == nullptr) {
+            if (restore_device) {
+                (void)cudaSetDevice(previous_device);
+            }
+            return INFINI_STATUS_NULL_POINTER;
+        }
+        auto *ctx = static_cast<CustomAllReduceContext *>(comms[i]->custom_allreduce_context);
+        if (cudaSetDevice(ctx->device_id) != cudaSuccess) {
+            clear_cuda_error();
+            if (restore_device) {
+                (void)cudaSetDevice(previous_device);
+            }
+            return INFINI_STATUS_INTERNAL_ERROR;
+        }
+        clear_registered_buffers(ctx);
+    }
+    for (int i = 0; i < ndevice; ++i) {
+        auto *ctx = static_cast<CustomAllReduceContext *>(comms[i]->custom_allreduce_context);
+        erase_registered_buffer_registry(ctx->group_id);
+    }
+    if (restore_device) {
+        (void)cudaSetDevice(previous_device);
+    }
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t registerCustomAllReduceBuffers(
+    infinicclComm_t *comms,
+    int ndevice,
+    void **buffers,
+    size_t bytes) {
+
+    if (comms == nullptr || buffers == nullptr) {
+        return INFINI_STATUS_NULL_POINTER;
+    }
+    if (ndevice <= 0 || bytes == 0) {
+        return INFINI_STATUS_BAD_PARAM;
+    }
+    if (!should_register_custom_allreduce_buffer(bytes)) {
+        return INFINI_STATUS_SUCCESS;
+    }
+
+    std::vector<CustomAllReduceContext *> contexts(static_cast<size_t>(ndevice));
+    std::vector<void *> rank_buffers(static_cast<size_t>(ndevice));
+    for (int i = 0; i < ndevice; ++i) {
+        if (comms[i] == nullptr || buffers[i] == nullptr) {
+            return INFINI_STATUS_NULL_POINTER;
+        }
+        if (comms[i]->allreduce_backend == INFINICCL_ALLREDUCE_BACKEND_NCCL) {
+            return INFINI_STATUS_SUCCESS;
+        }
+        if (comms[i]->custom_allreduce_context == nullptr) {
+            return INFINI_STATUS_NULL_POINTER;
+        }
+        contexts[static_cast<size_t>(i)] = static_cast<CustomAllReduceContext *>(comms[i]->custom_allreduce_context);
+        rank_buffers[static_cast<size_t>(i)] = buffers[i];
+    }
+
+    return install_registered_buffers(contexts, rank_buffers, bytes, std::string());
+}
+
+infiniStatus_t registerCustomAllReduceBuffer(
+    infinicclComm_t comm,
+    const char *key,
+    void *buffer,
+    size_t bytes) {
+
+    if (comm == nullptr || key == nullptr || buffer == nullptr) {
+        return INFINI_STATUS_NULL_POINTER;
+    }
+    if (bytes == 0 || key[0] == '\0') {
+        return INFINI_STATUS_BAD_PARAM;
+    }
+    if (!should_register_custom_allreduce_buffer(bytes)) {
+        return INFINI_STATUS_SUCCESS;
+    }
+    if (comm->allreduce_backend == INFINICCL_ALLREDUCE_BACKEND_NCCL) {
+        return INFINI_STATUS_SUCCESS;
+    }
+    if (comm->custom_allreduce_context == nullptr) {
+        return INFINI_STATUS_BAD_PARAM;
+    }
+
+    auto *ctx = static_cast<CustomAllReduceContext *>(comm->custom_allreduce_context);
+    if (!ctx->initialized || ctx->world_size <= 0 || ctx->rank < 0 || ctx->rank >= ctx->world_size) {
+        return INFINI_STATUS_BAD_PARAM;
+    }
+
+    const std::string user_key(key);
+    if (auto *registered = find_registered_buffer(ctx, buffer, bytes);
+        registered != nullptr && registered->key == user_key) {
+        return INFINI_STATUS_SUCCESS;
+    }
+
+    const auto registry_key = make_registered_buffer_key(ctx->group_id, key);
+    std::unique_lock<std::mutex> lock(registered_buffer_registry_mutex);
+    auto &registry_slot = registered_buffer_registry[registry_key];
+    const auto rank = static_cast<size_t>(ctx->rank);
+
+    if (registry_slot != nullptr && registry_slot->installed) {
+        if (registry_slot->world_size == ctx->world_size &&
+            rank < registry_slot->buffers.size() &&
+            registry_slot->contexts[rank] == ctx &&
+            registry_slot->buffers[rank] == buffer &&
+            bytes <= registry_slot->bytes) {
+            return INFINI_STATUS_SUCCESS;
+        }
+        registry_slot.reset();
+    }
+
+    if (registry_slot == nullptr) {
+        registry_slot = std::make_shared<PendingRegisteredAllReduceBuffer>();
+        registry_slot->world_size = ctx->world_size;
+        registry_slot->bytes = bytes;
+        registry_slot->buffers.assign(static_cast<size_t>(ctx->world_size), nullptr);
+        registry_slot->contexts.assign(static_cast<size_t>(ctx->world_size), nullptr);
+        registry_slot->arrived.assign(static_cast<size_t>(ctx->world_size), 0);
+    }
+
+    auto slot = registry_slot;
+
+    if (slot->world_size != ctx->world_size || slot->bytes != bytes) {
+        slot->status = INFINI_STATUS_BAD_PARAM;
+        slot->installed = true;
+        slot->cv.notify_all();
+        return INFINI_STATUS_BAD_PARAM;
+    }
+
+    if (slot->arrived[rank] != 0) {
+        if (slot->contexts[rank] != ctx || slot->buffers[rank] != buffer) {
+            slot->status = INFINI_STATUS_BAD_PARAM;
+            slot->installed = true;
+            slot->cv.notify_all();
+            return INFINI_STATUS_BAD_PARAM;
+        }
+    } else {
+        slot->contexts[rank] = ctx;
+        slot->buffers[rank] = buffer;
+        slot->arrived[rank] = 1;
+        ++slot->arrived_count;
+    }
+
+    if (slot->arrived_count == slot->world_size) {
+        slot->status = install_registered_buffers(slot->contexts, slot->buffers, slot->bytes, user_key);
+        slot->installed = true;
+        slot->cv.notify_all();
+        return slot->status;
+    }
+
+    slot->cv.wait(lock, [slot]() { return slot->installed; });
+    return slot->status;
+}
+
+CustomAllReduceCheckResult canUseCustomAllReduce(
+    const CustomAllReduceContext *ctx,
+    size_t count,
+    infiniDtype_t datatype,
+    infinicclReduceOp_t op) {
+
+    if (ctx == nullptr) {
+        return unsupported("custom allreduce context is not initialized");
+    }
+    if (ctx->local_backend == nullptr || *ctx->local_backend == INFINICCL_ALLREDUCE_BACKEND_NCCL) {
+        return unsupported("custom allreduce backend is not enabled on this rank");
+    }
+    if (!ctx->initialized) {
+        return unsupported(ctx->disabled_reason != nullptr ? ctx->disabled_reason : "custom allreduce context is disabled");
+    }
+    if (!is_supported_world_size(ctx->world_size)) {
+        return unsupported("custom allreduce supports TP sizes from 2 to 8 in this phase");
+    }
+    if (ctx->rank_backends.size() != static_cast<size_t>(ctx->world_size)) {
+        return unsupported("custom allreduce rank backend table is not initialized");
+    }
+    for (int rank_idx = 0; rank_idx < ctx->world_size; ++rank_idx) {
+        const auto *backend = ctx->rank_backends[static_cast<size_t>(rank_idx)];
+        if (backend == nullptr || *backend == INFINICCL_ALLREDUCE_BACKEND_NCCL) {
+            return unsupported("custom allreduce backend is not enabled on all TP ranks");
+        }
+    }
+    if (!is_supported_dtype(datatype)) {
+        return unsupported("custom allreduce supports only F32, F16, and BF16");
+    }
+    if (op != INFINICCL_SUM) {
+        return unsupported("custom allreduce supports only SUM");
+    }
+    if (count == 0) {
+        return unsupported("custom allreduce does not handle empty tensors");
+    }
+
+    const size_t dtype_size = infiniSizeOf(datatype);
+    if (dtype_size == 0) {
+        return unsupported("custom allreduce received an invalid dtype");
+    }
+    if (!is_aligned_to(ctx->scratch, dtype_size)) {
+        return unsupported("custom allreduce scratch buffers are not dtype-aligned");
+    }
+    if (count > ctx->max_bytes / dtype_size) {
+        return unsupported("custom allreduce message is larger than the configured fast-path limit");
+    }
+
+    if (ctx->scratch == nullptr || ctx->signals == nullptr ||
+        ctx->rank_scratch == nullptr || ctx->rank_signals == nullptr) {
+        return unsupported("custom allreduce scratch or signal buffers are not initialized");
+    }
+
+    return supported();
+}
+
+infiniStatus_t tryCustomAllReduce(
+    CustomAllReduceContext *ctx,
+    void *sendbuf,
+    void *recvbuf,
+    size_t count,
+    infiniDtype_t datatype,
+    infinicclReduceOp_t op,
+    infinirtStream_t stream,
+    bool *handled) {
+
+    if (handled == nullptr) {
+        return INFINI_STATUS_NULL_POINTER;
+    }
+    *handled = false;
+
+    const auto check = canUseCustomAllReduce(ctx, count, datatype, op);
+    if (!check.supported) {
+        return INFINI_STATUS_SUCCESS;
+    }
+
+    const size_t bytes = count * infiniSizeOf(datatype);
+    auto cuda_stream = get_cuda_stream(stream);
+    auto *registered = find_registered_buffer(ctx, sendbuf, bytes);
+    if (registered == nullptr || !should_use_registered_custom_allreduce(ctx->world_size, bytes)) {
+        return INFINI_STATUS_SUCCESS;
+    }
+
+    void **rank_buffers = registered->rank_buffers;
+    bool force_two_stage = sendbuf == recvbuf;
+
+    const bool use_packed = bytes % sizeof(uint4) == 0 &&
+                            is_aligned_to(sendbuf, sizeof(uint4)) &&
+                            is_aligned_to(recvbuf, sizeof(uint4)) &&
+                            (registered != nullptr || is_aligned_to(ctx->scratch, sizeof(uint4)));
+    auto status = launch_allreduce_kernel(
+        recvbuf,
+        rank_buffers,
+        count,
+        bytes,
+        datatype,
+        ctx->world_size,
+        ctx->rank,
+        use_packed,
+        force_two_stage,
+        ctx->signals,
+        ctx->rank_signals,
+        cuda_stream);
+    if (status != INFINI_STATUS_SUCCESS) {
+        return status;
+    }
+
+    *handled = true;
+    return INFINI_STATUS_SUCCESS;
+}
+
+} // namespace infiniccl::cuda
\ No newline at end of file
diff --git a/src/infiniccl/cuda/custom_allreduce.hpp b/src/infiniccl/cuda/custom_allreduce.hpp
new file mode 100644
index 000000000..fb49e238a
--- /dev/null
+++ b/src/infiniccl/cuda/custom_allreduce.hpp
@@ -0,0 +1,62 @@
+#ifndef INFINICCL_CUDA_CUSTOM_ALLREDUCE_HPP_
+#define INFINICCL_CUDA_CUSTOM_ALLREDUCE_HPP_
+
+#include "../infiniccl_impl.h"
+
+namespace infiniccl::cuda {
+
+struct CustomAllReduceContext;
+
+struct CustomAllReduceCheckResult {
+    bool supported = false;
+    const char *reason = nullptr;
+};
+
+CustomAllReduceContext *createCustomAllReduceContext(
+    int rank,
+    int world_size,
+    int device_id,
+    const int *device_ids);
+
+void initializeCustomAllReduceContexts(
+    infinicclComm_t *comms,
+    int ndevice,
+    const int *device_ids);
+
+void destroyCustomAllReduceContext(CustomAllReduceContext *ctx);
+
+infiniStatus_t registerCustomAllReduceBuffers(
+    infinicclComm_t *comms,
+    int ndevice,
+    void **buffers,
+    size_t bytes);
+
+infiniStatus_t registerCustomAllReduceBuffer(
+    infinicclComm_t comm,
+    const char *key,
+    void *buffer,
+    size_t bytes);
+
+infiniStatus_t clearCustomAllReduceBuffers(
+    infinicclComm_t *comms,
+    int ndevice);
+
+CustomAllReduceCheckResult canUseCustomAllReduce(
+    const CustomAllReduceContext *ctx,
+    size_t count,
+    infiniDtype_t datatype,
+    infinicclReduceOp_t op);
+
+infiniStatus_t tryCustomAllReduce(
+    CustomAllReduceContext *ctx,
+    void *sendbuf,
+    void *recvbuf,
+    size_t count,
+    infiniDtype_t datatype,
+    infinicclReduceOp_t op,
+    infinirtStream_t stream,
+    bool *handled);
+
+} // namespace infiniccl::cuda
+
+#endif // INFINICCL_CUDA_CUSTOM_ALLREDUCE_HPP_
\ No newline at end of file
diff --git a/src/infiniccl/cuda/infiniccl_cuda.cu b/src/infiniccl/cuda/infiniccl_cuda.cu
index 6a8442d21..5d400e1d9 100644
--- a/src/infiniccl/cuda/infiniccl_cuda.cu
+++ b/src/infiniccl/cuda/infiniccl_cuda.cu
@@ -1,4 +1,5 @@
 #include "infiniccl_cuda.h"
+#include "custom_allreduce.hpp"
 
 #include <cuda_runtime.h>
 #include <iostream>
@@ -71,18 +72,48 @@ infiniStatus_t commInitAll(
     CHECK_NCCL(ncclCommInitAll(nccl_comms.data(), ndevice, (int const *)device_ids));
 
     for (int i = 0; i < ndevice; i++) {
-        comms[i] = new InfinicclComm{INFINI_DEVICE_NVIDIA, device_ids[i], (void *)(nccl_comms[i]), i, ndevice};
+        auto *comm = new InfinicclComm{INFINI_DEVICE_NVIDIA, device_ids[i], (void *)(nccl_comms[i]), i, ndevice};
+        comm->custom_allreduce_context = createCustomAllReduceContext(i, ndevice, device_ids[i], device_ids);
+        comms[i] = comm;
     }
+    initializeCustomAllReduceContexts(comms, ndevice, device_ids);
 
     return INFINI_STATUS_SUCCESS;
 }
 
 infiniStatus_t commDestroy(infinicclComm_t comm) {
+    destroyCustomAllReduceContext(static_cast<CustomAllReduceContext *>(comm->custom_allreduce_context));
+    comm->custom_allreduce_context = nullptr;
     CHECK_NCCL(ncclCommDestroy(getNcclComm(comm)));
     delete comm;
     return INFINI_STATUS_SUCCESS;
 }
 
+infiniStatus_t registerAllReduceBuffers(
+    infinicclComm_t *comms,
+    int ndevice,
+    void **buffers,
+    size_t bytes) {
+
+    return registerCustomAllReduceBuffers(comms, ndevice, buffers, bytes);
+}
+
+infiniStatus_t registerAllReduceBuffer(
+    infinicclComm_t comm,
+    const char *key,
+    void *buffer,
+    size_t bytes) {
+
+    return registerCustomAllReduceBuffer(comm, key, buffer, bytes);
+}
+
+infiniStatus_t clearAllReduceBuffers(
+    infinicclComm_t *comms,
+    int ndevice) {
+
+    return clearCustomAllReduceBuffers(comms, ndevice);
+}
+
 infiniStatus_t groupStart(infinicclComm_t) {
     CHECK_NCCL(ncclGroupStart());
     return INFINI_STATUS_SUCCESS;
@@ -104,6 +135,18 @@ infiniStatus_t allReduce(
 
     CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16);
 
+    if (comm->allreduce_backend != INFINICCL_ALLREDUCE_BACKEND_NCCL) {
+        bool handled = false;
+        auto *custom_ctx = static_cast<CustomAllReduceContext *>(comm->custom_allreduce_context);
+        auto status = tryCustomAllReduce(custom_ctx, sendbuf, recvbuf, count, datatype, op, stream, &handled);
+        if (status != INFINI_STATUS_SUCCESS) {
+            return status;
+        }
+        if (handled) {
+            return INFINI_STATUS_SUCCESS;
+        }
+    }
+
     CHECK_NCCL(ncclAllReduce(sendbuf, recvbuf, count, getNcclDtype(datatype),
                              getNcclRedOp(op), getNcclComm(comm), getCudaStream(stream)));
 
diff --git a/src/infiniccl/infiniccl.cc b/src/infiniccl/infiniccl.cc
index 48dbea425..3e558b664 100644
--- a/src/infiniccl/infiniccl.cc
+++ b/src/infiniccl/infiniccl.cc
@@ -1,5 +1,6 @@
 #include "infiniccl.h"
 
+#include "./infiniccl_impl.h"
 #include "./ascend/infiniccl_ascend.h"
 #include "./cambricon/infiniccl_cambricon.h"
 #include "./cuda/infiniccl_cuda.h"
@@ -61,6 +62,158 @@ __INFINI_C infiniStatus_t infinicclCommDestroy(infinicclComm_t comm) {
 #undef COMM_DESTROY
 }
 
+__INFINI_C infiniStatus_t infinicclCommSetAllReduceBackend(
+    infinicclComm_t comm,
+    infinicclAllReduceBackend_t backend) {
+
+    if (comm == nullptr) {
+        return INFINI_STATUS_NULL_POINTER;
+    }
+
+    switch (backend) {
+    case INFINICCL_ALLREDUCE_BACKEND_AUTO:
+    case INFINICCL_ALLREDUCE_BACKEND_NCCL:
+    case INFINICCL_ALLREDUCE_BACKEND_CUSTOM:
+        comm->allreduce_backend = backend;
+        return INFINI_STATUS_SUCCESS;
+    default:
+        return INFINI_STATUS_BAD_PARAM;
+    }
+}
+
+__INFINI_C infiniStatus_t infinicclCommGetAllReduceBackend(
+    infinicclComm_t comm,
+    infinicclAllReduceBackend_t *backend) {
+
+    if (comm == nullptr || backend == nullptr) {
+        return INFINI_STATUS_NULL_POINTER;
+    }
+
+    *backend = comm->allreduce_backend;
+    return INFINI_STATUS_SUCCESS;
+}
+
+__INFINI_C infiniStatus_t infinicclCommRegisterAllReduceBuffers(
+    infinicclComm_t *comms,
+    int ndevice,
+    void **buffers,
+    size_t bytes) {
+
+    if (comms == nullptr || buffers == nullptr) {
+        return INFINI_STATUS_NULL_POINTER;
+    }
+    if (ndevice <= 0 || bytes == 0 || comms[0] == nullptr) {
+        return INFINI_STATUS_BAD_PARAM;
+    }
+    const auto device_type = comms[0]->device_type;
+    for (int i = 0; i < ndevice; ++i) {
+        if (comms[i] == nullptr || buffers[i] == nullptr) {
+            return INFINI_STATUS_NULL_POINTER;
+        }
+        if (comms[i]->device_type != device_type) {
+            return INFINI_STATUS_BAD_PARAM;
+        }
+    }
+
+#define REGISTER_ALLREDUCE_BUFFERS(CASE_, NAMESPACE_) \
+    case CASE_:                                       \
+        return infiniccl::NAMESPACE_::registerAllReduceBuffers(comms, ndevice, buffers, bytes)
+
+    switch (device_type) {
+        REGISTER_ALLREDUCE_BUFFERS(INFINI_DEVICE_NVIDIA, cuda);
+        REGISTER_ALLREDUCE_BUFFERS(INFINI_DEVICE_ILUVATAR, cuda);
+        REGISTER_ALLREDUCE_BUFFERS(INFINI_DEVICE_QY, cuda);
+        REGISTER_ALLREDUCE_BUFFERS(INFINI_DEVICE_HYGON, cuda);
+        REGISTER_ALLREDUCE_BUFFERS(INFINI_DEVICE_ASCEND, ascend);
+        REGISTER_ALLREDUCE_BUFFERS(INFINI_DEVICE_CAMBRICON, cambricon);
+        REGISTER_ALLREDUCE_BUFFERS(INFINI_DEVICE_METAX, metax);
+        REGISTER_ALLREDUCE_BUFFERS(INFINI_DEVICE_MOORE, moore);
+        REGISTER_ALLREDUCE_BUFFERS(INFINI_DEVICE_KUNLUN, kunlun);
+        REGISTER_ALLREDUCE_BUFFERS(INFINI_DEVICE_ALI, cuda);
+    default:
+        return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    }
+
+#undef REGISTER_ALLREDUCE_BUFFERS
+}
+
+__INFINI_C infiniStatus_t infinicclCommRegisterAllReduceBuffer(
+    infinicclComm_t comm,
+    const char *key,
+    void *buffer,
+    size_t bytes) {
+
+    if (comm == nullptr || key == nullptr || buffer == nullptr) {
+        return INFINI_STATUS_NULL_POINTER;
+    }
+    if (bytes == 0 || key[0] == '\0') {
+        return INFINI_STATUS_BAD_PARAM;
+    }
+
+#define REGISTER_ALLREDUCE_BUFFER(CASE_, NAMESPACE_) \
+    case CASE_:                                      \
+        return infiniccl::NAMESPACE_::registerAllReduceBuffer(comm, key, buffer, bytes)
+
+    switch (comm->device_type) {
+        REGISTER_ALLREDUCE_BUFFER(INFINI_DEVICE_NVIDIA, cuda);
+        REGISTER_ALLREDUCE_BUFFER(INFINI_DEVICE_ILUVATAR, cuda);
+        REGISTER_ALLREDUCE_BUFFER(INFINI_DEVICE_QY, cuda);
+        REGISTER_ALLREDUCE_BUFFER(INFINI_DEVICE_HYGON, cuda);
+        REGISTER_ALLREDUCE_BUFFER(INFINI_DEVICE_ASCEND, ascend);
+        REGISTER_ALLREDUCE_BUFFER(INFINI_DEVICE_CAMBRICON, cambricon);
+        REGISTER_ALLREDUCE_BUFFER(INFINI_DEVICE_METAX, metax);
+        REGISTER_ALLREDUCE_BUFFER(INFINI_DEVICE_MOORE, moore);
+        REGISTER_ALLREDUCE_BUFFER(INFINI_DEVICE_KUNLUN, kunlun);
+        REGISTER_ALLREDUCE_BUFFER(INFINI_DEVICE_ALI, cuda);
+    default:
+        return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    }
+
+#undef REGISTER_ALLREDUCE_BUFFER
+}
+
+__INFINI_C infiniStatus_t infinicclCommClearAllReduceBuffers(
+    infinicclComm_t *comms,
+    int ndevice) {
+
+    if (comms == nullptr) {
+        return INFINI_STATUS_NULL_POINTER;
+    }
+    if (ndevice <= 0 || comms[0] == nullptr) {
+        return INFINI_STATUS_BAD_PARAM;
+    }
+    const auto device_type = comms[0]->device_type;
+    for (int i = 0; i < ndevice; ++i) {
+        if (comms[i] == nullptr) {
+            return INFINI_STATUS_NULL_POINTER;
+        }
+        if (comms[i]->device_type != device_type) {
+            return INFINI_STATUS_BAD_PARAM;
+        }
+    }
+
+#define CLEAR_ALLREDUCE_BUFFERS(CASE_, NAMESPACE_) \
+    case CASE_:                                    \
+        return infiniccl::NAMESPACE_::clearAllReduceBuffers(comms, ndevice)
+
+    switch (device_type) {
+        CLEAR_ALLREDUCE_BUFFERS(INFINI_DEVICE_NVIDIA, cuda);
+        CLEAR_ALLREDUCE_BUFFERS(INFINI_DEVICE_ILUVATAR, cuda);
+        CLEAR_ALLREDUCE_BUFFERS(INFINI_DEVICE_QY, cuda);
+        CLEAR_ALLREDUCE_BUFFERS(INFINI_DEVICE_HYGON, cuda);
+        CLEAR_ALLREDUCE_BUFFERS(INFINI_DEVICE_ASCEND, ascend);
+        CLEAR_ALLREDUCE_BUFFERS(INFINI_DEVICE_CAMBRICON, cambricon);
+        CLEAR_ALLREDUCE_BUFFERS(INFINI_DEVICE_METAX, metax);
+        CLEAR_ALLREDUCE_BUFFERS(INFINI_DEVICE_MOORE, moore);
+        CLEAR_ALLREDUCE_BUFFERS(INFINI_DEVICE_KUNLUN, kunlun);
+        CLEAR_ALLREDUCE_BUFFERS(INFINI_DEVICE_ALI, cuda);
+    default:
+        return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    }
+
+#undef CLEAR_ALLREDUCE_BUFFERS
+}
+
 __INFINI_C infiniStatus_t infinicclGroupStart(infinicclComm_t comm) {
     if (comm == nullptr) {
         return INFINI_STATUS_NULL_POINTER;
diff --git a/src/infiniccl/infiniccl_impl.h b/src/infiniccl/infiniccl_impl.h
index 790355bd9..3c09739e6 100644
--- a/src/infiniccl/infiniccl_impl.h
+++ b/src/infiniccl/infiniccl_impl.h
@@ -9,6 +9,8 @@ struct InfinicclComm {
     void *comm;    // the actual communicator
     int rank = 0;
     int world_size = 1;
+    infinicclAllReduceBackend_t allreduce_backend = INFINICCL_ALLREDUCE_BACKEND_NCCL;
+    void *custom_allreduce_context = nullptr;
 };
 
 #define INFINICCL_DEVICE_API(NAMSPACE, IMPL)               \
@@ -20,6 +22,22 @@ struct InfinicclComm {
                                                            \
     infiniStatus_t commDestroy(infinicclComm_t comm) IMPL; \
                                                            \
+    infiniStatus_t registerAllReduceBuffers(               \
+        infinicclComm_t *comms,                            \
+        int ndevice,                                       \
+        void **buffers,                                    \
+        size_t bytes) IMPL;                                \
+                                                           \
+    infiniStatus_t registerAllReduceBuffer(                \
+        infinicclComm_t comm,                              \
+        const char *key,                                   \
+        void *buffer,                                      \
+        size_t bytes) IMPL;                                \
+                                                           \
+    infiniStatus_t clearAllReduceBuffers(                  \
+        infinicclComm_t *comms,                            \
+        int ndevice) IMPL;                                 \
+                                                           \
     infiniStatus_t groupStart(infinicclComm_t comm) IMPL;  \
                                                            \
     infiniStatus_t groupEnd(infinicclComm_t comm) IMPL;    \
diff --git a/src/infinicore/context/allocators/pinnable_block_allocator.cc b/src/infinicore/context/allocators/pinnable_block_allocator.cc
index 32e5c5e9b..72faebce9 100644
--- a/src/infinicore/context/allocators/pinnable_block_allocator.cc
+++ b/src/infinicore/context/allocators/pinnable_block_allocator.cc
@@ -135,7 +135,10 @@ void PinnableBlockAllocator::deallocate(std::byte *ptr) {
     block->in_use = false;
     for (auto &cls : size_classes_) {
         if (block->size == cls.block_size) {
-            cls.free_blocks.push_back(block);
+            auto it = std::find(cls.free_blocks.begin(), cls.free_blocks.end(), block);
+            if (it == cls.free_blocks.end()) {
+                cls.free_blocks.push_back(block);
+            }
             break;
         }
     }
@@ -151,6 +154,18 @@ size_t PinnableBlockAllocator::mark_in_use_(void *ptr, bool in_use) {
 
     auto block = it->second;
     if (in_use) {
+        // Graph blob reinstantiation can resurrect a cached size-class block
+        // without going through allocate(). Once active again, it must not stay
+        // in the free list, otherwise trim() can release the same pointer twice.
+        if (!block->in_use) {
+            for (auto &cls : size_classes_) {
+                if (block->size == cls.block_size) {
+                    auto free_it = std::remove(cls.free_blocks.begin(), cls.free_blocks.end(), block);
+                    cls.free_blocks.erase(free_it, cls.free_blocks.end());
+                    break;
+                }
+            }
+        }
         block->in_use = true;
         ++block->use_count;
     } else if (block->use_count > 0) {
@@ -190,11 +205,10 @@ void PinnableBlockAllocator::trim() {
 // ------------------- Destructor -------------------
 PinnableBlockAllocator::~PinnableBlockAllocator() {
     std::lock_guard<std::mutex> lock(mutex_);
-    for (auto &p : all_blocks_) {
-        if (p.second->ptr) {
-            infinirtFree(p.second->ptr);
-        }
-    }
+    // Do not call infinirtFree during process shutdown. This allocator is owned
+    // by the process-wide runtime; by the time the runtime singleton is
+    // destructed, CUDA graph/runtime teardown order can make individual frees
+    // unsafe. The driver reclaims the remaining process allocations on exit.
     all_blocks_.clear();
     large_blocks_.clear();
     for (auto &cls : size_classes_) {
diff --git a/src/infinicore/device_event.cc b/src/infinicore/device_event.cc
index 347e7effd..3820ca6b7 100644
--- a/src/infinicore/device_event.cc
+++ b/src/infinicore/device_event.cc
@@ -38,11 +38,32 @@ DeviceEvent::DeviceEvent(DeviceEvent &&other) noexcept
     other.is_recorded_ = false;
 }
 
+namespace {
+
+void destroy_event_on_device(infinirtEvent_t event, const Device &device) noexcept {
+    try {
+        Device current_device = context::getDevice();
+        bool switched_device = current_device != device;
+        if (switched_device) {
+            context::setDevice(device);
+        }
+        context::destroyEvent(event);
+        if (switched_device) {
+            context::setDevice(current_device);
+        }
+    } catch (...) {
+        // Event cleanup can happen during interpreter shutdown; do not throw
+        // from destructors or noexcept move assignment cleanup.
+    }
+}
+
+} // namespace
+
 DeviceEvent &DeviceEvent::operator=(DeviceEvent &&other) noexcept {
     if (this != &other) {
         // Clean up current resources
         if (event_ != nullptr) {
-            context::destroyEvent(event_);
+            destroy_event_on_device(event_, device_);
         }
 
         // Transfer ownership
@@ -59,7 +80,7 @@ DeviceEvent &DeviceEvent::operator=(DeviceEvent &&other) noexcept {
 
 DeviceEvent::~DeviceEvent() {
     if (event_ != nullptr) {
-        context::destroyEvent(event_);
+        destroy_event_on_device(event_, device_);
     }
 }
 
diff --git a/src/infinicore/memory.cc b/src/infinicore/memory.cc
index 0a7cc2df3..90bb9c65b 100644
--- a/src/infinicore/memory.cc
+++ b/src/infinicore/memory.cc
@@ -1,4 +1,5 @@
 #include "infinicore/memory.hpp"
+#include "infinicore/context/context.hpp"
 
 namespace infinicore {
 
@@ -10,8 +11,33 @@ Memory::Memory(std::byte *data,
     : data_{data}, size_{size}, device_{device}, deleter_{deleter}, is_pinned_(pin_memory) {}
 
 Memory::~Memory() {
-    if (deleter_) {
-        deleter_(data_);
+    if (!deleter_) {
+        return;
+    }
+    try {
+        Device current_device = context::getDevice();
+        bool switched_device = current_device != device_;
+        if (switched_device) {
+            context::setDevice(device_);
+        }
+        try {
+            deleter_(data_);
+        } catch (...) {
+            if (switched_device) {
+                try {
+                    context::setDevice(current_device);
+                } catch (...) {
+                }
+            }
+            throw;
+        }
+        if (switched_device) {
+            context::setDevice(current_device);
+        }
+    } catch (...) {
+        // Memory can be released during interpreter/static shutdown after
+        // allocator metadata has already been torn down. Destructors must not
+        // let allocator cleanup errors escape and terminate the process.
     }
 }
 
diff --git a/src/infinicore/ops/ernie45_rope/ernie45_rope.cc b/src/infinicore/ops/ernie45_rope/ernie45_rope.cc
new file mode 100644
index 000000000..d0581eca8
--- /dev/null
+++ b/src/infinicore/ops/ernie45_rope/ernie45_rope.cc
@@ -0,0 +1,64 @@
+#include "infinicore/ops/ernie45_rope.hpp"
+#include "../../utils.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Ernie45MRoPE);
+INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Ernie45VisionRoPE);
+
+Ernie45MRoPE::Ernie45MRoPE(Tensor q,
+                           Tensor k,
+                           const Tensor &positions,
+                           double rope_theta,
+                           size_t section_h,
+                           size_t section_w,
+                           size_t section_t) {
+    INFINICORE_ASSERT_TENSORS_SAME_DEVICE(q, k, positions);
+    INFINICORE_GRAPH_OP_DISPATCH(q->device().getType(), q, k, positions, rope_theta, section_h, section_w, section_t);
+}
+
+void Ernie45MRoPE::execute(Tensor q,
+                           Tensor k,
+                           const Tensor &positions,
+                           double rope_theta,
+                           size_t section_h,
+                           size_t section_w,
+                           size_t section_t) {
+    INFINICORE_GRAPH_OP_RECORD_OR_RUN(Ernie45MRoPE, q, k, positions, rope_theta, section_h, section_w, section_t);
+}
+
+Tensor ernie45_mrope_(Tensor q,
+                      Tensor k,
+                      const Tensor &positions,
+                      double rope_theta,
+                      size_t section_h,
+                      size_t section_w,
+                      size_t section_t) {
+    Ernie45MRoPE::execute(q, k, positions, rope_theta, section_h, section_w, section_t);
+    return q;
+}
+
+Ernie45VisionRoPE::Ernie45VisionRoPE(Tensor q,
+                                     Tensor k,
+                                     const Tensor &positions,
+                                     double rope_theta) {
+    INFINICORE_ASSERT_TENSORS_SAME_DEVICE(q, k, positions);
+    INFINICORE_GRAPH_OP_DISPATCH(q->device().getType(), q, k, positions, rope_theta);
+}
+
+void Ernie45VisionRoPE::execute(Tensor q,
+                                Tensor k,
+                                const Tensor &positions,
+                                double rope_theta) {
+    INFINICORE_GRAPH_OP_RECORD_OR_RUN(Ernie45VisionRoPE, q, k, positions, rope_theta);
+}
+
+Tensor ernie45_vision_rope_(Tensor q,
+                            Tensor k,
+                            const Tensor &positions,
+                            double rope_theta) {
+    Ernie45VisionRoPE::execute(q, k, positions, rope_theta);
+    return q;
+}
+
+} // namespace infinicore::op
diff --git a/src/infinicore/ops/ernie45_rope/ernie45_rope_infiniop.cc b/src/infinicore/ops/ernie45_rope/ernie45_rope_infiniop.cc
new file mode 100644
index 000000000..8a029c460
--- /dev/null
+++ b/src/infinicore/ops/ernie45_rope/ernie45_rope_infiniop.cc
@@ -0,0 +1,139 @@
+#include "infinicore/ops/ernie45_rope.hpp"
+
+#include "../infiniop_impl.hpp"
+#include "infiniop/ops/ernie45_rope.h"
+
+namespace infinicore::op::ernie45_rope_impl::infiniop {
+
+namespace mrope {
+INFINIOP_CACHABLE_DESCRIPTOR(MropeDescriptor, Ernie45Mrope, 100);
+} // namespace mrope
+
+namespace vision {
+INFINIOP_CACHABLE_DESCRIPTOR(VisionDescriptor, Ernie45VisionRope, 100);
+} // namespace vision
+
+namespace mrope_op {
+
+struct MropePlannedMeta {
+    std::shared_ptr<mrope::MropeDescriptor> descriptor;
+    graph::GraphTensor workspace;
+    graph::GraphTensor q;
+    graph::GraphTensor k;
+    graph::GraphTensor positions;
+};
+
+void *plan_mrope(Tensor q,
+                 Tensor k,
+                 const Tensor &positions,
+                 double rope_theta,
+                 size_t section_h,
+                 size_t section_w,
+                 size_t section_t) {
+    size_t key = hash_combine(q, k, positions, rope_theta, section_h, section_w, section_t);
+    std::shared_ptr<mrope::MropeDescriptor> descriptor;
+    {
+        auto device__ = context::getDevice();
+        auto &cache__ = mrope::caches.getCache(device__);
+        descriptor = cache__.get(key).value_or(nullptr);
+        if (!descriptor) {
+            descriptor = std::make_shared<mrope::MropeDescriptor>(nullptr);
+            INFINICORE_CHECK_ERROR(infiniopCreateErnie45MropeDescriptor(
+                context::getInfiniopHandle(device__),
+                &descriptor->desc,
+                q->desc(), k->desc(), positions->desc(), rope_theta, section_h, section_w, section_t));
+            cache__.put(key, descriptor);
+        }
+    }
+    INFINIOP_WORKSPACE_TENSOR(workspace, Ernie45Mrope, descriptor);
+    return new MropePlannedMeta{
+        descriptor,
+        graph::GraphTensor(workspace),
+        graph::GraphTensor(q),
+        graph::GraphTensor(k),
+        graph::GraphTensor(positions)};
+}
+
+void run_mrope(void *planned_meta) {
+    auto *p = reinterpret_cast<MropePlannedMeta *>(planned_meta);
+    INFINICORE_CHECK_ERROR(
+        infiniopErnie45Mrope(
+            p->descriptor->desc,
+            p->workspace->data(),
+            p->workspace->numel(),
+            p->q->data(),
+            p->k->data(),
+            p->positions->data(),
+            context::getStream()));
+}
+
+void cleanup_mrope(void **planned_meta_ptr) {
+    delete *reinterpret_cast<MropePlannedMeta **>(planned_meta_ptr);
+    *planned_meta_ptr = nullptr;
+}
+
+INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Ernie45MRoPE, &plan_mrope, &run_mrope, &cleanup_mrope);
+
+} // namespace mrope_op
+
+namespace vision_op {
+
+struct VisionPlannedMeta {
+    std::shared_ptr<vision::VisionDescriptor> descriptor;
+    graph::GraphTensor workspace;
+    graph::GraphTensor q;
+    graph::GraphTensor k;
+    graph::GraphTensor positions;
+};
+
+void *plan_vision(Tensor q,
+                  Tensor k,
+                  const Tensor &positions,
+                  double rope_theta) {
+    size_t key = hash_combine(q, k, positions, rope_theta);
+    std::shared_ptr<vision::VisionDescriptor> descriptor;
+    {
+        auto device__ = context::getDevice();
+        auto &cache__ = vision::caches.getCache(device__);
+        descriptor = cache__.get(key).value_or(nullptr);
+        if (!descriptor) {
+            descriptor = std::make_shared<vision::VisionDescriptor>(nullptr);
+            INFINICORE_CHECK_ERROR(infiniopCreateErnie45VisionRopeDescriptor(
+                context::getInfiniopHandle(device__),
+                &descriptor->desc,
+                q->desc(), k->desc(), positions->desc(), rope_theta));
+            cache__.put(key, descriptor);
+        }
+    }
+    INFINIOP_WORKSPACE_TENSOR(workspace, Ernie45VisionRope, descriptor);
+    return new VisionPlannedMeta{
+        descriptor,
+        graph::GraphTensor(workspace),
+        graph::GraphTensor(q),
+        graph::GraphTensor(k),
+        graph::GraphTensor(positions)};
+}
+
+void run_vision(void *planned_meta) {
+    auto *p = reinterpret_cast<VisionPlannedMeta *>(planned_meta);
+    INFINICORE_CHECK_ERROR(
+        infiniopErnie45VisionRope(
+            p->descriptor->desc,
+            p->workspace->data(),
+            p->workspace->numel(),
+            p->q->data(),
+            p->k->data(),
+            p->positions->data(),
+            context::getStream()));
+}
+
+void cleanup_vision(void **planned_meta_ptr) {
+    delete *reinterpret_cast<VisionPlannedMeta **>(planned_meta_ptr);
+    *planned_meta_ptr = nullptr;
+}
+
+INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Ernie45VisionRoPE, &plan_vision, &run_vision, &cleanup_vision);
+
+} // namespace vision_op
+
+} // namespace infinicore::op::ernie45_rope_impl::infiniop
diff --git a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc
index 0167c17df..21c801bbb 100644
--- a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc
+++ b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc
@@ -18,6 +18,24 @@
 
 namespace infinicore::op::mha_kvcache_impl::flashattn {
 
+namespace {
+
+std::optional<at::Tensor> to_optional_aten_tensor(const Tensor &tensor) {
+    if (!tensor || tensor->numel() == 0) {
+        return std::nullopt;
+    }
+    return infinicore::adaptor::to_aten_tensor(tensor);
+}
+
+std::optional<const at::Tensor> to_optional_const_aten_tensor(const Tensor &tensor) {
+    if (!tensor || tensor->numel() == 0) {
+        return std::nullopt;
+    }
+    return infinicore::adaptor::to_aten_tensor(tensor);
+}
+
+} // namespace
+
 struct PlannedMeta {
     graph::GraphTensor out, q, k_cache, v_cache, seqlens_k, block_table;
     std::optional<graph::GraphTensor> alibi_slopes;
@@ -64,10 +82,10 @@ void run(void *planned_meta) {
     auto k_cache = infinicore::adaptor::to_aten_tensor(k_cache_work);
     auto v_cache = infinicore::adaptor::to_aten_tensor(v_cache_work);
 #endif
-    auto seqlens_k = std::optional<const at::Tensor>(infinicore::adaptor::to_aten_tensor(p->seqlens_k));
-    auto block_table = std::optional<at::Tensor>(infinicore::adaptor::to_aten_tensor(p->block_table));
+    auto seqlens_k = to_optional_const_aten_tensor(p->seqlens_k);
+    auto block_table = to_optional_aten_tensor(p->block_table);
     auto alibi_slopes = p->alibi_slopes
-                          ? std::optional<at::Tensor>(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes))
+                          ? to_optional_aten_tensor(*p->alibi_slopes)
                           : std::nullopt;
 
     std::optional<const at::Tensor> k_new = std::nullopt;
diff --git a/src/infinicore/tensor/copy.cc b/src/infinicore/tensor/copy.cc
index 1297d9f8c..cb54a9fb9 100644
--- a/src/infinicore/tensor/copy.cc
+++ b/src/infinicore/tensor/copy.cc
@@ -50,6 +50,15 @@ void TensorImpl::copy_from(Tensor src) {
                 context::memcpyH2D(local_src->data(), src->data(), copy_size);
                 op::rearrange_(Tensor(const_cast<TensorImpl *>(this)->shared_from_this()), local_src);
             }
+        } else {
+            context::setDevice(this->device());
+            if (this->is_contiguous()) {
+                context::memcpyD2D(this->data(), src->data(), copy_size);
+            } else {
+                auto local_src = Tensor::empty(this->shape(), this->dtype(), this->device());
+                context::memcpyD2D(local_src->data(), src->data(), copy_size);
+                op::rearrange_(Tensor(const_cast<TensorImpl *>(this)->shared_from_this()), local_src);
+            }
         }
     }
 }
diff --git a/src/infiniop/ops/ernie45_rope/ernie45_rope.h b/src/infiniop/ops/ernie45_rope/ernie45_rope.h
new file mode 100644
index 000000000..13b64203b
--- /dev/null
+++ b/src/infiniop/ops/ernie45_rope/ernie45_rope.h
@@ -0,0 +1,48 @@
+#ifndef __ERNIE45_ROPE_H__
+#define __ERNIE45_ROPE_H__
+
+#include "../../operator.h"
+#include "info.h"
+
+#define MROPE_DESCRIPTOR(NAMESPACE)                                                                                          \
+    namespace op::ernie45_rope::NAMESPACE {                                                                                  \
+    class MropeDescriptor final : public InfiniopDescriptor {                                                                \
+        struct Opaque;                                                                                                       \
+        Opaque *_opaque;                                                                                                     \
+        QKInfo _info;                                                                                                        \
+        size_t _workspace_size;                                                                                              \
+        MropeDescriptor(Opaque *opaque, QKInfo info, size_t workspace_size, infiniDevice_t device_type, int device_id)       \
+            : InfiniopDescriptor{device_type, device_id}, _opaque(opaque), _info(info), _workspace_size(workspace_size) {}   \
+                                                                                                                             \
+    public:                                                                                                                  \
+        ~MropeDescriptor();                                                                                                  \
+        size_t workspaceSize() const { return _workspace_size; }                                                             \
+        static infiniStatus_t create(infiniopHandle_t handle, MropeDescriptor **desc_ptr, infiniopTensorDescriptor_t q_desc, \
+                                     infiniopTensorDescriptor_t k_desc, infiniopTensorDescriptor_t pos_desc, double theta,   \
+                                     size_t section_h, size_t section_w, size_t section_t);                                  \
+        infiniStatus_t calculate(void *workspace, size_t workspace_size, void *q, void *k, const void *positions,            \
+                                 void *stream) const;                                                                        \
+    };                                                                                                                       \
+    }
+
+#define VISION_ROPE_DESCRIPTOR(NAMESPACE)                                                                                         \
+    namespace op::ernie45_rope::NAMESPACE {                                                                                       \
+    class VisionRopeDescriptor final : public InfiniopDescriptor {                                                                \
+        struct Opaque;                                                                                                            \
+        Opaque *_opaque;                                                                                                          \
+        VisionInfo _info;                                                                                                         \
+        size_t _workspace_size;                                                                                                   \
+        VisionRopeDescriptor(Opaque *opaque, VisionInfo info, size_t workspace_size, infiniDevice_t device_type, int device_id)   \
+            : InfiniopDescriptor{device_type, device_id}, _opaque(opaque), _info(info), _workspace_size(workspace_size) {}        \
+                                                                                                                                  \
+    public:                                                                                                                       \
+        ~VisionRopeDescriptor();                                                                                                  \
+        size_t workspaceSize() const { return _workspace_size; }                                                                  \
+        static infiniStatus_t create(infiniopHandle_t handle, VisionRopeDescriptor **desc_ptr, infiniopTensorDescriptor_t q_desc, \
+                                     infiniopTensorDescriptor_t k_desc, infiniopTensorDescriptor_t pos_desc, double theta);       \
+        infiniStatus_t calculate(void *workspace, size_t workspace_size, void *q, void *k, const void *positions,                 \
+                                 void *stream) const;                                                                             \
+    };                                                                                                                            \
+    }
+
+#endif
diff --git a/src/infiniop/ops/ernie45_rope/info.h b/src/infiniop/ops/ernie45_rope/info.h
new file mode 100644
index 000000000..2323b5119
--- /dev/null
+++ b/src/infiniop/ops/ernie45_rope/info.h
@@ -0,0 +1,140 @@
+#ifndef __ERNIE45_ROPE_INFO_H__
+#define __ERNIE45_ROPE_INFO_H__
+
+#include "../../../utils.h"
+#include "../../tensor.h"
+
+namespace op::ernie45_rope {
+
+struct QKInfo {
+    infiniDtype_t data_type;
+    infiniDtype_t pos_type;
+    size_t seqlen;
+    size_t q_heads;
+    size_t k_heads;
+    size_t head_dim;
+    ptrdiff_t q_stride_seq;
+    ptrdiff_t q_stride_head;
+    ptrdiff_t k_stride_seq;
+    ptrdiff_t k_stride_head;
+    ptrdiff_t pos_stride_seq;
+    ptrdiff_t pos_stride_axis;
+    bool pos_axis_first;
+    double rope_theta;
+    size_t section_h;
+    size_t section_w;
+    size_t section_t;
+
+    static utils::Result<QKInfo> create(infiniopTensorDescriptor_t q_desc,
+                                        infiniopTensorDescriptor_t k_desc,
+                                        infiniopTensorDescriptor_t pos_desc,
+                                        double rope_theta,
+                                        size_t section_h,
+                                        size_t section_w,
+                                        size_t section_t) {
+        CHECK_OR_RETURN(q_desc != nullptr && k_desc != nullptr && pos_desc != nullptr, INFINI_STATUS_NULL_POINTER);
+        CHECK_OR_RETURN(q_desc->ndim() == 3 && k_desc->ndim() == 3, INFINI_STATUS_BAD_TENSOR_SHAPE);
+        auto dtype = q_desc->dtype();
+        CHECK_OR_RETURN(dtype == k_desc->dtype(), INFINI_STATUS_BAD_TENSOR_DTYPE);
+        CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32);
+        CHECK_DTYPE_ANY_INT(pos_desc->dtype());
+
+        size_t seqlen = q_desc->dim(0);
+        size_t q_heads = q_desc->dim(1);
+        size_t head_dim = q_desc->dim(2);
+        CHECK_OR_RETURN(k_desc->dim(0) == seqlen && k_desc->dim(2) == head_dim, INFINI_STATUS_BAD_TENSOR_SHAPE);
+        CHECK_OR_RETURN((head_dim % 2) == 0, INFINI_STATUS_BAD_TENSOR_SHAPE);
+        CHECK_OR_RETURN(section_h + section_w + section_t == head_dim / 2, INFINI_STATUS_BAD_PARAM);
+        CHECK_OR_RETURN(q_desc->stride(2) == 1 && k_desc->stride(2) == 1, INFINI_STATUS_BAD_TENSOR_STRIDES);
+
+        bool axis_first = false;
+        ptrdiff_t pos_stride_seq = 0;
+        ptrdiff_t pos_stride_axis = 0;
+        if (pos_desc->ndim() == 2 && pos_desc->dim(0) == 3 && pos_desc->dim(1) == seqlen) {
+            axis_first = true;
+            pos_stride_axis = pos_desc->stride(0);
+            pos_stride_seq = pos_desc->stride(1);
+        } else if (pos_desc->ndim() == 2 && pos_desc->dim(0) == seqlen && pos_desc->dim(1) == 3) {
+            pos_stride_seq = pos_desc->stride(0);
+            pos_stride_axis = pos_desc->stride(1);
+        } else if (pos_desc->ndim() == 3 && pos_desc->dim(0) == 1 && pos_desc->dim(1) == seqlen && pos_desc->dim(2) == 3) {
+            pos_stride_seq = pos_desc->stride(1);
+            pos_stride_axis = pos_desc->stride(2);
+        } else {
+            return INFINI_STATUS_BAD_TENSOR_SHAPE;
+        }
+
+        return utils::Result<QKInfo>(QKInfo{
+            dtype,
+            pos_desc->dtype(),
+            seqlen,
+            q_heads,
+            k_desc->dim(1),
+            head_dim,
+            q_desc->stride(0),
+            q_desc->stride(1),
+            k_desc->stride(0),
+            k_desc->stride(1),
+            pos_stride_seq,
+            pos_stride_axis,
+            axis_first,
+            rope_theta,
+            section_h,
+            section_w,
+            section_t});
+    }
+};
+
+struct VisionInfo {
+    infiniDtype_t data_type;
+    infiniDtype_t pos_type;
+    size_t seqlen;
+    size_t q_heads;
+    size_t k_heads;
+    size_t head_dim;
+    ptrdiff_t q_stride_seq;
+    ptrdiff_t q_stride_head;
+    ptrdiff_t k_stride_seq;
+    ptrdiff_t k_stride_head;
+    ptrdiff_t pos_stride_seq;
+    ptrdiff_t pos_stride_axis;
+    double rope_theta;
+
+    static utils::Result<VisionInfo> create(infiniopTensorDescriptor_t q_desc,
+                                            infiniopTensorDescriptor_t k_desc,
+                                            infiniopTensorDescriptor_t pos_desc,
+                                            double rope_theta) {
+        CHECK_OR_RETURN(q_desc != nullptr && k_desc != nullptr && pos_desc != nullptr, INFINI_STATUS_NULL_POINTER);
+        CHECK_OR_RETURN(q_desc->ndim() == 3 && k_desc->ndim() == 3, INFINI_STATUS_BAD_TENSOR_SHAPE);
+        auto dtype = q_desc->dtype();
+        CHECK_OR_RETURN(dtype == k_desc->dtype(), INFINI_STATUS_BAD_TENSOR_DTYPE);
+        CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32);
+        CHECK_DTYPE_ANY_INT(pos_desc->dtype());
+
+        size_t seqlen = q_desc->dim(0);
+        size_t head_dim = q_desc->dim(2);
+        CHECK_OR_RETURN(k_desc->dim(0) == seqlen && k_desc->dim(2) == head_dim, INFINI_STATUS_BAD_TENSOR_SHAPE);
+        CHECK_OR_RETURN((head_dim % 4) == 0, INFINI_STATUS_BAD_TENSOR_SHAPE);
+        CHECK_OR_RETURN(q_desc->stride(2) == 1 && k_desc->stride(2) == 1, INFINI_STATUS_BAD_TENSOR_STRIDES);
+        CHECK_OR_RETURN(pos_desc->ndim() == 2 && pos_desc->dim(0) == seqlen && pos_desc->dim(1) == 2, INFINI_STATUS_BAD_TENSOR_SHAPE);
+
+        return utils::Result<VisionInfo>(VisionInfo{
+            dtype,
+            pos_desc->dtype(),
+            seqlen,
+            q_desc->dim(1),
+            k_desc->dim(1),
+            head_dim,
+            q_desc->stride(0),
+            q_desc->stride(1),
+            k_desc->stride(0),
+            k_desc->stride(1),
+            pos_desc->stride(0),
+            pos_desc->stride(1),
+            rope_theta});
+    }
+};
+
+} // namespace op::ernie45_rope
+
+#endif
diff --git a/src/infiniop/ops/ernie45_rope/nvidia/ernie45_rope_nvidia.cu b/src/infiniop/ops/ernie45_rope/nvidia/ernie45_rope_nvidia.cu
new file mode 100644
index 000000000..d1797ae0b
--- /dev/null
+++ b/src/infiniop/ops/ernie45_rope/nvidia/ernie45_rope_nvidia.cu
@@ -0,0 +1,333 @@
+#include "../../../devices/nvidia/nvidia_common.cuh"
+#include "ernie45_rope_nvidia.cuh"
+
+#include "../../../devices/nvidia/nvidia_kernel_common.cuh"
+#include <algorithm>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
+namespace {
+
+template <typename T>
+__device__ inline float load_float(const T *ptr, size_t idx) {
+    return static_cast<float>(ptr[idx]);
+}
+
+template <>
+__device__ inline float load_float<half>(const half *ptr, size_t idx) {
+    return __half2float(ptr[idx]);
+}
+
+template <>
+__device__ inline float load_float<cuda_bfloat16>(const cuda_bfloat16 *ptr, size_t idx) {
+    return __bfloat162float(ptr[idx]);
+}
+
+template <typename T>
+__device__ inline void store_float(T *ptr, size_t idx, float value) {
+    ptr[idx] = static_cast<T>(value);
+}
+
+template <>
+__device__ inline void store_float<half>(half *ptr, size_t idx, float value) {
+    ptr[idx] = __float2half(value);
+}
+
+template <>
+__device__ inline void store_float<cuda_bfloat16>(cuda_bfloat16 *ptr, size_t idx, float value) {
+    ptr[idx] = __float2bfloat16(value);
+}
+
+template <typename Tindex>
+__device__ inline int64_t read_pos(const Tindex *positions, size_t offset) {
+    return static_cast<int64_t>(positions[offset]);
+}
+
+template <typename Tdata, typename Tindex>
+INFINIOP_CUDA_KERNEL ernie45_mrope_kernel(
+    Tdata *q,
+    Tdata *k,
+    const Tindex *__restrict__ positions,
+    size_t seqlen,
+    size_t q_heads,
+    size_t k_heads,
+    size_t head_dim,
+    ptrdiff_t q_stride_seq,
+    ptrdiff_t q_stride_head,
+    ptrdiff_t k_stride_seq,
+    ptrdiff_t k_stride_head,
+    ptrdiff_t pos_stride_seq,
+    ptrdiff_t pos_stride_axis,
+    bool pos_axis_first,
+    double rope_theta,
+    size_t section_h,
+    size_t section_w) {
+    const size_t token = blockIdx.x;
+    const size_t head = blockIdx.y;
+    const bool is_q = blockIdx.z == 0;
+    const size_t nheads = is_q ? q_heads : k_heads;
+    if (token >= seqlen || head >= nheads) {
+        return;
+    }
+
+    const size_t half_dim = head_dim / 2;
+    const size_t section_hw = section_h + section_w;
+    const size_t pos_base = token * pos_stride_seq;
+    const int64_t tpos = read_pos(positions, pos_base + 0 * pos_stride_axis);
+    const int64_t hpos = read_pos(positions, pos_base + 1 * pos_stride_axis);
+    const int64_t wpos = read_pos(positions, pos_base + 2 * pos_stride_axis);
+
+    Tdata *base = is_q ? q + token * q_stride_seq + head * q_stride_head
+                       : k + token * k_stride_seq + head * k_stride_head;
+
+    for (size_t pair = threadIdx.x; pair < half_dim; pair += blockDim.x) {
+        const int64_t pos = pair < section_hw ? ((pair & 1) == 0 ? hpos : wpos) : tpos;
+        const float inv_freq = powf(static_cast<float>(rope_theta),
+                                    -2.0f * static_cast<float>(pair) / static_cast<float>(head_dim));
+        float sinv;
+        float cosv;
+        __sincosf(static_cast<float>(pos) * inv_freq, &sinv, &cosv);
+
+        const size_t even = 2 * pair;
+        const size_t odd = even + 1;
+        const float x0 = load_float(base, even);
+        const float x1 = load_float(base, odd);
+        store_float(base, even, x0 * cosv - x1 * sinv);
+        store_float(base, odd, x1 * cosv + x0 * sinv);
+    }
+}
+
+template <typename Tdata, typename Tindex>
+INFINIOP_CUDA_KERNEL ernie45_vision_rope_kernel(
+    Tdata *q,
+    Tdata *k,
+    const Tindex *__restrict__ positions,
+    size_t seqlen,
+    size_t q_heads,
+    size_t k_heads,
+    size_t head_dim,
+    ptrdiff_t q_stride_seq,
+    ptrdiff_t q_stride_head,
+    ptrdiff_t k_stride_seq,
+    ptrdiff_t k_stride_head,
+    ptrdiff_t pos_stride_seq,
+    ptrdiff_t pos_stride_axis,
+    double rope_theta) {
+    const size_t token = blockIdx.x;
+    const size_t head = blockIdx.y;
+    const bool is_q = blockIdx.z == 0;
+    const size_t nheads = is_q ? q_heads : k_heads;
+    if (token >= seqlen || head >= nheads) {
+        return;
+    }
+
+    const size_t half_dim = head_dim / 2;
+    const size_t quarter_dim = head_dim / 4;
+    const int64_t hpos = read_pos(positions, token * pos_stride_seq + 0 * pos_stride_axis);
+    const int64_t wpos = read_pos(positions, token * pos_stride_seq + 1 * pos_stride_axis);
+
+    Tdata *base = is_q ? q + token * q_stride_seq + head * q_stride_head
+                       : k + token * k_stride_seq + head * k_stride_head;
+
+    for (size_t i = threadIdx.x; i < half_dim; i += blockDim.x) {
+        const int64_t pos = i < quarter_dim ? hpos : wpos;
+        const size_t freq_idx = i < quarter_dim ? i : i - quarter_dim;
+        const float inv_freq = powf(static_cast<float>(rope_theta),
+                                    -2.0f * static_cast<float>(freq_idx) / static_cast<float>(half_dim));
+        float sinv;
+        float cosv;
+        __sincosf(static_cast<float>(pos) * inv_freq, &sinv, &cosv);
+
+        const size_t lo = i;
+        const size_t hi = i + half_dim;
+        const float x0 = load_float(base, lo);
+        const float x1 = load_float(base, hi);
+        store_float(base, lo, x0 * cosv - x1 * sinv);
+        store_float(base, hi, x1 * cosv + x0 * sinv);
+    }
+}
+
+} // namespace
+
+namespace op::ernie45_rope::nvidia {
+
+struct MropeDescriptor::Opaque {
+    std::shared_ptr<device::nvidia::Handle::Internal> internal;
+};
+
+MropeDescriptor::~MropeDescriptor() {
+    delete _opaque;
+}
+
+infiniStatus_t MropeDescriptor::create(infiniopHandle_t handle_,
+                                       MropeDescriptor **desc_ptr,
+                                       infiniopTensorDescriptor_t q_desc,
+                                       infiniopTensorDescriptor_t k_desc,
+                                       infiniopTensorDescriptor_t pos_desc,
+                                       double theta,
+                                       size_t section_h,
+                                       size_t section_w,
+                                       size_t section_t) {
+    auto result = QKInfo::create(q_desc, k_desc, pos_desc, theta, section_h, section_w, section_t);
+    CHECK_RESULT(result);
+    auto handle = reinterpret_cast<device::nvidia::Handle *>(handle_);
+    *desc_ptr = new MropeDescriptor(
+        new Opaque{handle->internal()},
+        result.take(),
+        0,
+        handle->device,
+        handle->device_id);
+    return INFINI_STATUS_SUCCESS;
+}
+
+template <typename Tdata, typename Tindex>
+infiniStatus_t launch_mrope(const QKInfo &info, Tdata *q, Tdata *k, const Tindex *positions, cudaStream_t stream) {
+    dim3 grid(static_cast<unsigned int>(info.seqlen),
+              static_cast<unsigned int>(std::max(info.q_heads, info.k_heads)),
+              2);
+    ernie45_mrope_kernel<<<grid, 128, 0, stream>>>(
+        q, k, positions,
+        info.seqlen, info.q_heads, info.k_heads, info.head_dim,
+        info.q_stride_seq, info.q_stride_head,
+        info.k_stride_seq, info.k_stride_head,
+        info.pos_stride_seq, info.pos_stride_axis, info.pos_axis_first,
+        info.rope_theta, info.section_h, info.section_w);
+    return INFINI_STATUS_SUCCESS;
+}
+
+#define LAUNCH_MROPE(TDATA, TINDEX) launch_mrope(_info, (TDATA *)q, (TDATA *)k, (const TINDEX *)positions, (cudaStream_t)stream)
+#define MROPE_POS_TYPE(TDATA)                  \
+    switch (_info.pos_type) {                  \
+    case INFINI_DTYPE_U8:                      \
+        return LAUNCH_MROPE(TDATA, uint8_t);   \
+    case INFINI_DTYPE_U16:                     \
+        return LAUNCH_MROPE(TDATA, uint16_t);  \
+    case INFINI_DTYPE_U32:                     \
+        return LAUNCH_MROPE(TDATA, uint32_t);  \
+    case INFINI_DTYPE_U64:                     \
+        return LAUNCH_MROPE(TDATA, uint64_t);  \
+    case INFINI_DTYPE_I8:                      \
+        return LAUNCH_MROPE(TDATA, int8_t);    \
+    case INFINI_DTYPE_I16:                     \
+        return LAUNCH_MROPE(TDATA, int16_t);   \
+    case INFINI_DTYPE_I32:                     \
+        return LAUNCH_MROPE(TDATA, int32_t);   \
+    case INFINI_DTYPE_I64:                     \
+        return LAUNCH_MROPE(TDATA, int64_t);   \
+    default:                                   \
+        return INFINI_STATUS_BAD_TENSOR_DTYPE; \
+    }
+
+infiniStatus_t MropeDescriptor::calculate(void *workspace,
+                                          size_t workspace_size,
+                                          void *q,
+                                          void *k,
+                                          const void *positions,
+                                          void *stream) const {
+    if (workspace_size < _workspace_size) {
+        return INFINI_STATUS_INSUFFICIENT_WORKSPACE;
+    }
+    switch (_info.data_type) {
+    case INFINI_DTYPE_F16:
+        MROPE_POS_TYPE(half);
+    case INFINI_DTYPE_BF16:
+        MROPE_POS_TYPE(cuda_bfloat16);
+    case INFINI_DTYPE_F32:
+        MROPE_POS_TYPE(float);
+    default:
+        return INFINI_STATUS_BAD_TENSOR_DTYPE;
+    }
+}
+
+#undef MROPE_POS_TYPE
+#undef LAUNCH_MROPE
+
+struct VisionRopeDescriptor::Opaque {
+    std::shared_ptr<device::nvidia::Handle::Internal> internal;
+};
+
+VisionRopeDescriptor::~VisionRopeDescriptor() {
+    delete _opaque;
+}
+
+infiniStatus_t VisionRopeDescriptor::create(infiniopHandle_t handle_,
+                                            VisionRopeDescriptor **desc_ptr,
+                                            infiniopTensorDescriptor_t q_desc,
+                                            infiniopTensorDescriptor_t k_desc,
+                                            infiniopTensorDescriptor_t pos_desc,
+                                            double theta) {
+    auto result = VisionInfo::create(q_desc, k_desc, pos_desc, theta);
+    CHECK_RESULT(result);
+    auto handle = reinterpret_cast<device::nvidia::Handle *>(handle_);
+    *desc_ptr = new VisionRopeDescriptor(
+        new Opaque{handle->internal()},
+        result.take(),
+        0,
+        handle->device,
+        handle->device_id);
+    return INFINI_STATUS_SUCCESS;
+}
+
+template <typename Tdata, typename Tindex>
+infiniStatus_t launch_vision_rope(const VisionInfo &info, Tdata *q, Tdata *k, const Tindex *positions, cudaStream_t stream) {
+    dim3 grid(static_cast<unsigned int>(info.seqlen),
+              static_cast<unsigned int>(std::max(info.q_heads, info.k_heads)),
+              2);
+    ernie45_vision_rope_kernel<<<grid, 128, 0, stream>>>(
+        q, k, positions,
+        info.seqlen, info.q_heads, info.k_heads, info.head_dim,
+        info.q_stride_seq, info.q_stride_head,
+        info.k_stride_seq, info.k_stride_head,
+        info.pos_stride_seq, info.pos_stride_axis,
+        info.rope_theta);
+    return INFINI_STATUS_SUCCESS;
+}
+
+#define LAUNCH_VISION(TDATA, TINDEX) launch_vision_rope(_info, (TDATA *)q, (TDATA *)k, (const TINDEX *)positions, (cudaStream_t)stream)
+#define VISION_POS_TYPE(TDATA)                 \
+    switch (_info.pos_type) {                  \
+    case INFINI_DTYPE_U8:                      \
+        return LAUNCH_VISION(TDATA, uint8_t);  \
+    case INFINI_DTYPE_U16:                     \
+        return LAUNCH_VISION(TDATA, uint16_t); \
+    case INFINI_DTYPE_U32:                     \
+        return LAUNCH_VISION(TDATA, uint32_t); \
+    case INFINI_DTYPE_U64:                     \
+        return LAUNCH_VISION(TDATA, uint64_t); \
+    case INFINI_DTYPE_I8:                      \
+        return LAUNCH_VISION(TDATA, int8_t);   \
+    case INFINI_DTYPE_I16:                     \
+        return LAUNCH_VISION(TDATA, int16_t);  \
+    case INFINI_DTYPE_I32:                     \
+        return LAUNCH_VISION(TDATA, int32_t);  \
+    case INFINI_DTYPE_I64:                     \
+        return LAUNCH_VISION(TDATA, int64_t);  \
+    default:                                   \
+        return INFINI_STATUS_BAD_TENSOR_DTYPE; \
+    }
+
+infiniStatus_t VisionRopeDescriptor::calculate(void *workspace,
+                                               size_t workspace_size,
+                                               void *q,
+                                               void *k,
+                                               const void *positions,
+                                               void *stream) const {
+    if (workspace_size < _workspace_size) {
+        return INFINI_STATUS_INSUFFICIENT_WORKSPACE;
+    }
+    switch (_info.data_type) {
+    case INFINI_DTYPE_F16:
+        VISION_POS_TYPE(half);
+    case INFINI_DTYPE_BF16:
+        VISION_POS_TYPE(cuda_bfloat16);
+    case INFINI_DTYPE_F32:
+        VISION_POS_TYPE(float);
+    default:
+        return INFINI_STATUS_BAD_TENSOR_DTYPE;
+    }
+}
+
+#undef VISION_POS_TYPE
+#undef LAUNCH_VISION
+
+} // namespace op::ernie45_rope::nvidia
diff --git a/src/infiniop/ops/ernie45_rope/nvidia/ernie45_rope_nvidia.cuh b/src/infiniop/ops/ernie45_rope/nvidia/ernie45_rope_nvidia.cuh
new file mode 100644
index 000000000..87e0457f2
--- /dev/null
+++ b/src/infiniop/ops/ernie45_rope/nvidia/ernie45_rope_nvidia.cuh
@@ -0,0 +1,9 @@
+#ifndef __ERNIE45_ROPE_NVIDIA_CUH__
+#define __ERNIE45_ROPE_NVIDIA_CUH__
+
+#include "../ernie45_rope.h"
+
+MROPE_DESCRIPTOR(nvidia)
+VISION_ROPE_DESCRIPTOR(nvidia)
+
+#endif
diff --git a/src/infiniop/ops/ernie45_rope/operator.cc b/src/infiniop/ops/ernie45_rope/operator.cc
new file mode 100644
index 000000000..f9c3fe978
--- /dev/null
+++ b/src/infiniop/ops/ernie45_rope/operator.cc
@@ -0,0 +1,233 @@
+#include "../../operator.h"
+#include "../../handle.h"
+#include "infiniop/ops/ernie45_rope.h"
+
+#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) || defined(ENABLE_ILUVATAR_API)
+#include "nvidia/ernie45_rope_nvidia.cuh"
+#endif
+
+__INFINI_C infiniStatus_t infiniopCreateErnie45MropeDescriptor(
+    infiniopHandle_t handle,
+    infiniopErnie45MropeDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t q,
+    infiniopTensorDescriptor_t k,
+    infiniopTensorDescriptor_t positions,
+    double rope_theta,
+    size_t section_h,
+    size_t section_w,
+    size_t section_t) {
+#define CREATE(CASE, NAMESPACE)                                                          \
+    case CASE:                                                                           \
+        return op::ernie45_rope::NAMESPACE::MropeDescriptor::create(                     \
+            handle,                                                                      \
+            reinterpret_cast<op::ernie45_rope::NAMESPACE::MropeDescriptor **>(desc_ptr), \
+            q, k, positions, rope_theta, section_h, section_w, section_t)
+
+    switch (handle->device) {
+#ifdef ENABLE_NVIDIA_API
+        CREATE(INFINI_DEVICE_NVIDIA, nvidia);
+#endif
+#ifdef ENABLE_QY_API
+        CREATE(INFINI_DEVICE_QY, nvidia);
+#endif
+#ifdef ENABLE_ALI_API
+        CREATE(INFINI_DEVICE_ALI, nvidia);
+#endif
+#ifdef ENABLE_ILUVATAR_API
+        CREATE(INFINI_DEVICE_ILUVATAR, nvidia);
+#endif
+    default:
+        return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    }
+#undef CREATE
+}
+
+__INFINI_C infiniStatus_t infiniopGetErnie45MropeWorkspaceSize(infiniopErnie45MropeDescriptor_t desc, size_t *size) {
+#define GET(CASE, NAMESPACE)                                                                                   \
+    case CASE:                                                                                                 \
+        *size = reinterpret_cast<const op::ernie45_rope::NAMESPACE::MropeDescriptor *>(desc)->workspaceSize(); \
+        return INFINI_STATUS_SUCCESS
+    switch (desc->device_type) {
+#ifdef ENABLE_NVIDIA_API
+        GET(INFINI_DEVICE_NVIDIA, nvidia);
+#endif
+#ifdef ENABLE_QY_API
+        GET(INFINI_DEVICE_QY, nvidia);
+#endif
+#ifdef ENABLE_ALI_API
+        GET(INFINI_DEVICE_ALI, nvidia);
+#endif
+#ifdef ENABLE_ILUVATAR_API
+        GET(INFINI_DEVICE_ILUVATAR, nvidia);
+#endif
+    default:
+        return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    }
+#undef GET
+}
+
+__INFINI_C infiniStatus_t infiniopErnie45Mrope(
+    infiniopErnie45MropeDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *q,
+    void *k,
+    const void *positions,
+    void *stream) {
+#define CALC(CASE, NAMESPACE)                                                                           \
+    case CASE:                                                                                          \
+        return reinterpret_cast<const op::ernie45_rope::NAMESPACE::MropeDescriptor *>(desc)->calculate( \
+            workspace, workspace_size, q, k, positions, stream)
+    switch (desc->device_type) {
+#ifdef ENABLE_NVIDIA_API
+        CALC(INFINI_DEVICE_NVIDIA, nvidia);
+#endif
+#ifdef ENABLE_QY_API
+        CALC(INFINI_DEVICE_QY, nvidia);
+#endif
+#ifdef ENABLE_ALI_API
+        CALC(INFINI_DEVICE_ALI, nvidia);
+#endif
+#ifdef ENABLE_ILUVATAR_API
+        CALC(INFINI_DEVICE_ILUVATAR, nvidia);
+#endif
+    default:
+        return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    }
+#undef CALC
+}
+
+__INFINI_C infiniStatus_t infiniopDestroyErnie45MropeDescriptor(infiniopErnie45MropeDescriptor_t desc) {
+#define DESTROY(CASE, NAMESPACE)                                                             \
+    case CASE:                                                                               \
+        delete reinterpret_cast<const op::ernie45_rope::NAMESPACE::MropeDescriptor *>(desc); \
+        return INFINI_STATUS_SUCCESS
+    switch (desc->device_type) {
+#ifdef ENABLE_NVIDIA_API
+        DESTROY(INFINI_DEVICE_NVIDIA, nvidia);
+#endif
+#ifdef ENABLE_QY_API
+        DESTROY(INFINI_DEVICE_QY, nvidia);
+#endif
+#ifdef ENABLE_ALI_API
+        DESTROY(INFINI_DEVICE_ALI, nvidia);
+#endif
+#ifdef ENABLE_ILUVATAR_API
+        DESTROY(INFINI_DEVICE_ILUVATAR, nvidia);
+#endif
+    default:
+        return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    }
+#undef DESTROY
+}
+
+__INFINI_C infiniStatus_t infiniopCreateErnie45VisionRopeDescriptor(
+    infiniopHandle_t handle,
+    infiniopErnie45VisionRopeDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t q,
+    infiniopTensorDescriptor_t k,
+    infiniopTensorDescriptor_t positions,
+    double rope_theta) {
+#define CREATE(CASE, NAMESPACE)                                                               \
+    case CASE:                                                                                \
+        return op::ernie45_rope::NAMESPACE::VisionRopeDescriptor::create(                     \
+            handle,                                                                           \
+            reinterpret_cast<op::ernie45_rope::NAMESPACE::VisionRopeDescriptor **>(desc_ptr), \
+            q, k, positions, rope_theta)
+    switch (handle->device) {
+#ifdef ENABLE_NVIDIA_API
+        CREATE(INFINI_DEVICE_NVIDIA, nvidia);
+#endif
+#ifdef ENABLE_QY_API
+        CREATE(INFINI_DEVICE_QY, nvidia);
+#endif
+#ifdef ENABLE_ALI_API
+        CREATE(INFINI_DEVICE_ALI, nvidia);
+#endif
+#ifdef ENABLE_ILUVATAR_API
+        CREATE(INFINI_DEVICE_ILUVATAR, nvidia);
+#endif
+    default:
+        return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    }
+#undef CREATE
+}
+
+__INFINI_C infiniStatus_t infiniopGetErnie45VisionRopeWorkspaceSize(infiniopErnie45VisionRopeDescriptor_t desc, size_t *size) {
+#define GET(CASE, NAMESPACE)                                                                                        \
+    case CASE:                                                                                                      \
+        *size = reinterpret_cast<const op::ernie45_rope::NAMESPACE::VisionRopeDescriptor *>(desc)->workspaceSize(); \
+        return INFINI_STATUS_SUCCESS
+    switch (desc->device_type) {
+#ifdef ENABLE_NVIDIA_API
+        GET(INFINI_DEVICE_NVIDIA, nvidia);
+#endif
+#ifdef ENABLE_QY_API
+        GET(INFINI_DEVICE_QY, nvidia);
+#endif
+#ifdef ENABLE_ALI_API
+        GET(INFINI_DEVICE_ALI, nvidia);
+#endif
+#ifdef ENABLE_ILUVATAR_API
+        GET(INFINI_DEVICE_ILUVATAR, nvidia);
+#endif
+    default:
+        return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    }
+#undef GET
+}
+
+__INFINI_C infiniStatus_t infiniopErnie45VisionRope(
+    infiniopErnie45VisionRopeDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *q,
+    void *k,
+    const void *positions,
+    void *stream) {
+#define CALC(CASE, NAMESPACE)                                                                                \
+    case CASE:                                                                                               \
+        return reinterpret_cast<const op::ernie45_rope::NAMESPACE::VisionRopeDescriptor *>(desc)->calculate( \
+            workspace, workspace_size, q, k, positions, stream)
+    switch (desc->device_type) {
+#ifdef ENABLE_NVIDIA_API
+        CALC(INFINI_DEVICE_NVIDIA, nvidia);
+#endif
+#ifdef ENABLE_QY_API
+        CALC(INFINI_DEVICE_QY, nvidia);
+#endif
+#ifdef ENABLE_ALI_API
+        CALC(INFINI_DEVICE_ALI, nvidia);
+#endif
+#ifdef ENABLE_ILUVATAR_API
+        CALC(INFINI_DEVICE_ILUVATAR, nvidia);
+#endif
+    default:
+        return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    }
+#undef CALC
+}
+
+__INFINI_C infiniStatus_t infiniopDestroyErnie45VisionRopeDescriptor(infiniopErnie45VisionRopeDescriptor_t desc) {
+#define DESTROY(CASE, NAMESPACE)                                                                  \
+    case CASE:                                                                                    \
+        delete reinterpret_cast<const op::ernie45_rope::NAMESPACE::VisionRopeDescriptor *>(desc); \
+        return INFINI_STATUS_SUCCESS
+    switch (desc->device_type) {
+#ifdef ENABLE_NVIDIA_API
+        DESTROY(INFINI_DEVICE_NVIDIA, nvidia);
+#endif
+#ifdef ENABLE_QY_API
+        DESTROY(INFINI_DEVICE_QY, nvidia);
+#endif
+#ifdef ENABLE_ALI_API
+        DESTROY(INFINI_DEVICE_ALI, nvidia);
+#endif
+#ifdef ENABLE_ILUVATAR_API
+        DESTROY(INFINI_DEVICE_ILUVATAR, nvidia);
+#endif
+    default:
+        return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    }
+#undef DESTROY
+}
diff --git a/src/infiniop/ops/gemm/nvidia/gemm_nvidia.cu b/src/infiniop/ops/gemm/nvidia/gemm_nvidia.cu
index 0e0c65f2b..0d4a06346 100644
--- a/src/infiniop/ops/gemm/nvidia/gemm_nvidia.cu
+++ b/src/infiniop/ops/gemm/nvidia/gemm_nvidia.cu
@@ -1,12 +1,181 @@
 #include "../../../devices/nvidia/nvidia_handle.cuh"
 #include "gemm_nvidia.cuh"
 
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
 namespace op::gemm::nvidia {
 
 struct Descriptor::Opaque {
     std::shared_ptr<device::nvidia::Handle::Internal> internal;
+    infiniDtype_t a_dtype;
+    infiniDtype_t b_dtype;
+    infiniDtype_t c_dtype;
 };
 
+namespace {
+
+constexpr size_t ALIGNMENT = 256;
+
+size_t alignUp(size_t value, size_t alignment = ALIGNMENT) {
+    return (value + alignment - 1) / alignment * alignment;
+}
+
+infiniStatus_t toCudaDataType(infiniDtype_t dtype, cudaDataType *cuda_type) {
+    switch (dtype) {
+    case INFINI_DTYPE_F16:
+        *cuda_type = CUDA_R_16F;
+        return INFINI_STATUS_SUCCESS;
+    case INFINI_DTYPE_BF16:
+        *cuda_type = CUDA_R_16BF;
+        return INFINI_STATUS_SUCCESS;
+    case INFINI_DTYPE_F32:
+        *cuda_type = CUDA_R_32F;
+        return INFINI_STATUS_SUCCESS;
+    default:
+        return INFINI_STATUS_BAD_TENSOR_DTYPE;
+    }
+}
+
+size_t matrixSpanElements(const BlasMatrix &matrix) {
+    size_t batch_offset = matrix.batch > 1 ? static_cast<size_t>(std::abs(matrix.stride)) * (matrix.batch - 1) : 0;
+    size_t row_offset = matrix.rows > 0 ? static_cast<size_t>(std::abs(matrix.row_stride)) * (matrix.rows - 1) : 0;
+    size_t col_offset = matrix.cols > 0 ? static_cast<size_t>(std::abs(matrix.col_stride)) * (matrix.cols - 1) : 0;
+    return batch_offset + row_offset + col_offset + 1;
+}
+
+size_t matrixCastWorkspaceSize(infiniDtype_t dtype, const BlasMatrix &matrix) {
+    if (dtype == INFINI_DTYPE_F32) {
+        return 0;
+    }
+    return alignUp(matrixSpanElements(matrix) * sizeof(float));
+}
+
+size_t gemmCastWorkspaceSize(infiniDtype_t c_dtype,
+                             infiniDtype_t a_dtype,
+                             infiniDtype_t b_dtype,
+                             const MatmulInfo &info) {
+    if (c_dtype != INFINI_DTYPE_F32) {
+        return 0;
+    }
+    return matrixCastWorkspaceSize(a_dtype, info.a_matrix)
+           + matrixCastWorkspaceSize(b_dtype, info.b_matrix);
+}
+
+template <typename Src>
+__global__ void castMatrixToF32Kernel(float *dst,
+                                      const Src *src,
+                                      size_t total,
+                                      size_t rows,
+                                      size_t cols,
+                                      ptrdiff_t matrix_stride,
+                                      ptrdiff_t row_stride,
+                                      ptrdiff_t col_stride) {
+    size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= total) {
+        return;
+    }
+    size_t matrix_size = rows * cols;
+    size_t batch = idx / matrix_size;
+    size_t offset = idx % matrix_size;
+    size_t row = offset / cols;
+    size_t col = offset % cols;
+    ptrdiff_t elem_offset = static_cast<ptrdiff_t>(batch) * matrix_stride
+                            + static_cast<ptrdiff_t>(row) * row_stride
+                            + static_cast<ptrdiff_t>(col) * col_stride;
+    dst[elem_offset] = static_cast<float>(src[elem_offset]);
+}
+
+template <>
+__global__ void castMatrixToF32Kernel<__half>(float *dst,
+                                              const __half *src,
+                                              size_t total,
+                                              size_t rows,
+                                              size_t cols,
+                                              ptrdiff_t matrix_stride,
+                                              ptrdiff_t row_stride,
+                                              ptrdiff_t col_stride) {
+    size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= total) {
+        return;
+    }
+    size_t matrix_size = rows * cols;
+    size_t batch = idx / matrix_size;
+    size_t offset = idx % matrix_size;
+    size_t row = offset / cols;
+    size_t col = offset % cols;
+    ptrdiff_t elem_offset = static_cast<ptrdiff_t>(batch) * matrix_stride
+                            + static_cast<ptrdiff_t>(row) * row_stride
+                            + static_cast<ptrdiff_t>(col) * col_stride;
+    dst[elem_offset] = __half2float(src[elem_offset]);
+}
+
+template <>
+__global__ void castMatrixToF32Kernel<__nv_bfloat16>(float *dst,
+                                                     const __nv_bfloat16 *src,
+                                                     size_t total,
+                                                     size_t rows,
+                                                     size_t cols,
+                                                     ptrdiff_t matrix_stride,
+                                                     ptrdiff_t row_stride,
+                                                     ptrdiff_t col_stride) {
+    size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= total) {
+        return;
+    }
+    size_t matrix_size = rows * cols;
+    size_t batch = idx / matrix_size;
+    size_t offset = idx % matrix_size;
+    size_t row = offset / cols;
+    size_t col = offset % cols;
+    ptrdiff_t elem_offset = static_cast<ptrdiff_t>(batch) * matrix_stride
+                            + static_cast<ptrdiff_t>(row) * row_stride
+                            + static_cast<ptrdiff_t>(col) * col_stride;
+    dst[elem_offset] = __bfloat162float(src[elem_offset]);
+}
+
+infiniStatus_t castMatrixToF32(float *dst,
+                               const void *src,
+                               infiniDtype_t src_dtype,
+                               const BlasMatrix &matrix,
+                               cudaStream_t stream) {
+    size_t total = matrix.batch * matrix.rows * matrix.cols;
+    if (total == 0) {
+        return INFINI_STATUS_SUCCESS;
+    }
+    dim3 block(256);
+    dim3 grid((total + block.x - 1) / block.x);
+    switch (src_dtype) {
+    case INFINI_DTYPE_F16:
+        castMatrixToF32Kernel<<<grid, block, 0, stream>>>(
+            dst,
+            static_cast<const __half *>(src),
+            total,
+            matrix.rows,
+            matrix.cols,
+            matrix.stride,
+            matrix.row_stride,
+            matrix.col_stride);
+        break;
+    case INFINI_DTYPE_BF16:
+        castMatrixToF32Kernel<<<grid, block, 0, stream>>>(
+            dst,
+            static_cast<const __nv_bfloat16 *>(src),
+            total,
+            matrix.rows,
+            matrix.cols,
+            matrix.stride,
+            matrix.row_stride,
+            matrix.col_stride);
+        break;
+    default:
+        return INFINI_STATUS_BAD_TENSOR_DTYPE;
+    }
+    return INFINI_STATUS_SUCCESS;
+}
+
+} // namespace
+
 Descriptor::~Descriptor() {
     delete _opaque;
 }
@@ -18,16 +187,28 @@ infiniStatus_t Descriptor::create(
     infiniopTensorDescriptor_t a_desc,
     infiniopTensorDescriptor_t b_desc) {
     auto handle = reinterpret_cast<device::nvidia::Handle *>(handle_);
-    auto dtype = c_desc->dtype();
+    auto a_dtype = a_desc->dtype();
+    auto b_dtype = b_desc->dtype();
+    auto c_dtype = c_desc->dtype();
 
-    CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16);
+    CHECK_DTYPE(a_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16);
+    CHECK_DTYPE(b_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16);
+    CHECK_DTYPE(c_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16);
 
     auto result = MatmulInfo::create(c_desc, a_desc, b_desc, MatrixLayout::COL_MAJOR);
     CHECK_RESULT(result);
 
+    auto info = result.take();
+    auto effective_a_dtype = a_dtype;
+    auto effective_b_dtype = b_dtype;
+    if (info.is_transed) {
+        std::swap(effective_a_dtype, effective_b_dtype);
+    }
+    size_t workspace_size = gemmCastWorkspaceSize(c_dtype, effective_a_dtype, effective_b_dtype, info);
+
     *desc_ptr = new Descriptor(
-        dtype, result.take(), 0,
-        new Opaque{handle->internal()},
+        c_dtype, info, workspace_size,
+        new Opaque{handle->internal(), a_dtype, b_dtype, c_dtype},
         handle->device, handle->device_id);
     return INFINI_STATUS_SUCCESS;
 }
@@ -49,9 +230,12 @@ infiniStatus_t Descriptor::calculate(
     cublasComputeType_t compute_type;
 #endif
 
-    switch (_dtype) {
+    CHECK_STATUS(toCudaDataType(_opaque->a_dtype, &a_type));
+    CHECK_STATUS(toCudaDataType(_opaque->b_dtype, &b_type));
+    CHECK_STATUS(toCudaDataType(_opaque->c_dtype, &c_type));
+
+    switch (_opaque->c_dtype) {
     case INFINI_DTYPE_F16:
-        a_type = b_type = c_type = CUDA_R_16F;
 #if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API)
         compute_type = CUDA_R_32F;
 #else
@@ -59,7 +243,6 @@ infiniStatus_t Descriptor::calculate(
 #endif
         break;
     case INFINI_DTYPE_BF16:
-        a_type = b_type = c_type = CUDA_R_16BF;
 #if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API)
         compute_type = CUDA_R_32F;
 #else
@@ -67,11 +250,12 @@ infiniStatus_t Descriptor::calculate(
 #endif
         break;
     case INFINI_DTYPE_F32:
-        a_type = b_type = c_type = CUDA_R_32F;
 #if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API)
         compute_type = CUDA_R_32F;
 #else
-        compute_type = CUBLAS_COMPUTE_32F_FAST_TF32;
+        compute_type = (_opaque->a_dtype == INFINI_DTYPE_F32 && _opaque->b_dtype == INFINI_DTYPE_F32)
+                         ? CUBLAS_COMPUTE_32F_FAST_TF32
+                         : CUBLAS_COMPUTE_32F;
 #endif
         break;
 
@@ -81,6 +265,27 @@ infiniStatus_t Descriptor::calculate(
 
     if (_info.is_transed) {
         std::swap(a, b);
+        std::swap(a_type, b_type);
+    }
+
+    auto effective_a_dtype = _info.is_transed ? _opaque->b_dtype : _opaque->a_dtype;
+    auto effective_b_dtype = _info.is_transed ? _opaque->a_dtype : _opaque->b_dtype;
+    auto required_workspace = gemmCastWorkspaceSize(_opaque->c_dtype, effective_a_dtype, effective_b_dtype, _info);
+    CHECK_OR_RETURN(workspace_size >= required_workspace, INFINI_STATUS_INSUFFICIENT_WORKSPACE);
+    auto workspace_ptr = static_cast<char *>(workspace);
+    auto stream_ = static_cast<cudaStream_t>(stream);
+    if (_opaque->c_dtype == INFINI_DTYPE_F32 && effective_a_dtype != INFINI_DTYPE_F32) {
+        auto cast_a = reinterpret_cast<float *>(workspace_ptr);
+        CHECK_STATUS(castMatrixToF32(cast_a, a, effective_a_dtype, _info.a_matrix, stream_));
+        a = cast_a;
+        a_type = CUDA_R_32F;
+        workspace_ptr += matrixCastWorkspaceSize(effective_a_dtype, _info.a_matrix);
+    }
+    if (_opaque->c_dtype == INFINI_DTYPE_F32 && effective_b_dtype != INFINI_DTYPE_F32) {
+        auto cast_b = reinterpret_cast<float *>(workspace_ptr);
+        CHECK_STATUS(castMatrixToF32(cast_b, b, effective_b_dtype, _info.b_matrix, stream_));
+        b = cast_b;
+        b_type = CUDA_R_32F;
     }
 
     auto op_a = _info.a_matrix.row_stride == 1 ? CUBLAS_OP_N : CUBLAS_OP_T;
diff --git a/src/infiniop/ops/gptq_marlin_gemm/nvidia/gptq_marlin_gemm_nvidia.cu b/src/infiniop/ops/gptq_marlin_gemm/nvidia/gptq_marlin_gemm_nvidia.cu
index edbe5fc8f..71ac73522 100644
--- a/src/infiniop/ops/gptq_marlin_gemm/nvidia/gptq_marlin_gemm_nvidia.cu
+++ b/src/infiniop/ops/gptq_marlin_gemm/nvidia/gptq_marlin_gemm_nvidia.cu
@@ -4,7 +4,7 @@
 #include "../../../devices/nvidia/nvidia_kernel_common.cuh"
 #include "../gptq_marlin_gemm.h"
 #include "gptq_marlin_gemm_nvidia.cuh"
-#if defined ENABLE_TVM_API
+#if defined ENABLE_NVIDIA_API
 #include "../marlin/kernel.h"
 #include "../marlin/marlin_template.h"
 #include "../sgl_kernel/scalar_type.hpp"
@@ -1113,7 +1113,6 @@ Descriptor::calculate(
     bool use_fp32_reduce,
     bool is_zp_float,
     void *stream_) const {
-#if defined ENABLE_TVM_API
     cudaStream_t stream = (cudaStream_t)stream_;
     int64_t M = static_cast<int64_t>(_info.M);
     int64_t K = static_cast<int64_t>(_info.K);
@@ -1132,8 +1131,6 @@ Descriptor::calculate(
     } else {
         return INFINI_STATUS_BAD_TENSOR_DTYPE;
     }
-#endif
-    return INFINI_STATUS_SUCCESS;
 }
 
 } // namespace op::gptq_marlin_gemm::nvidia
diff --git a/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/tensor.h b/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/tensor.h
index 48c22c3f1..5f32ef9b1 100644
--- a/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/tensor.h
+++ b/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/tensor.h
@@ -3,7 +3,7 @@
 #pragma once
 #include "utils.h"
 
-#include <dlpack/dlpack.h>
+#include "../dlpack/dlpack.h"
 #include <tvm/ffi/container/tensor.h>
 #include <tvm/ffi/dtype.h>
 
diff --git a/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/utils.cuh b/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/utils.cuh
index 18d5da7c3..bebbd8a34 100644
--- a/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/utils.cuh
+++ b/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/utils.cuh
@@ -17,7 +17,7 @@
 
 #include "utils.h"
 
-#include <dlpack/dlpack.h>
+#include "../dlpack/dlpack.h"
 #include <tvm/ffi/extra/c_env_api.h>
 
 #include <concepts>
diff --git a/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/utils.h b/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/utils.h
index 855b32c36..c8753da18 100644
--- a/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/utils.h
+++ b/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/utils.h
@@ -48,7 +48,7 @@
 
 #include <concepts>
 #include <cstddef>
-#include <dlpack/dlpack.h>
+#include "../dlpack/dlpack.h"
 #include <ostream>
 #include <ranges>
 #include <sstream>
diff --git a/src/infiniop/ops/random_sample/nvidia/random_sample_kernel.cuh b/src/infiniop/ops/random_sample/nvidia/random_sample_kernel.cuh
index 3c2696f6b..3b92e7fa1 100644
--- a/src/infiniop/ops/random_sample/nvidia/random_sample_kernel.cuh
+++ b/src/infiniop/ops/random_sample/nvidia/random_sample_kernel.cuh
@@ -3,6 +3,7 @@
 #include <cub/device/device_radix_sort.cuh>
 #include <cub/device/device_reduce.cuh>
 #include <cub/device/device_scan.cuh>
+#include <cstdint>
 
 namespace op::random_sample::nvidia {
 
@@ -50,6 +51,23 @@ static cudaError inclusiveSum(
 }
 
 // ↑↑↑ 重新封装 cub api，减少模板参数，方便调用
+// ↓↓↓ random sampling keeps token indices in 32-bit workspace and casts only at the output boundary.
+
+template <class Tidx>
+struct InternalSampleIndex {
+    using Type = Tidx;
+};
+
+template <>
+struct InternalSampleIndex<int64_t> {
+    using Type = int32_t;
+};
+
+template <>
+struct InternalSampleIndex<uint64_t> {
+    using Type = uint32_t;
+};
+
 // ↓↓↓ 计算 workspace
 
 // 地址对齐到 256
@@ -59,6 +77,7 @@ static constexpr size_t align256(size_t size) {
 
 template <class Tidx, class Tval>
 utils::Result<size_t> calculateWorkspace(size_t n_) {
+    using TworkIdx = typename InternalSampleIndex<Tidx>::Type;
     const auto n = static_cast<int>(n_);
 
     size_t argmax;
@@ -70,14 +89,14 @@ utils::Result<size_t> calculateWorkspace(size_t n_) {
     argmax += 256;
 
     // indices
-    size_t size_random = align256(sizeof(Tidx) * n);
+    size_t size_random = align256(sizeof(TworkIdx) * n);
     // sorted
     size_random += align256(sizeof(Tval) * n);
     // indices_out
-    size_random += align256(sizeof(Tidx) * n);
+    size_random += align256(sizeof(TworkIdx) * n);
     // cub device api
     size_t size_radix_sort;
-    CHECK_CUDA((radixSort<Tval, Tidx>(
+    CHECK_CUDA((radixSort<Tval, TworkIdx>(
         nullptr, size_radix_sort,
         nullptr, nullptr,
         nullptr, nullptr,
@@ -158,9 +177,9 @@ static __global__ void setSoftmaxMaxKernel(
 
 // 直接 for 循环遍历采样
 // 这个 kernel 仅用于避免将数据拷贝到 cpu
-template <class Tval, class Tidx>
+template <class Tval, class Tout, class Tidx>
 static __global__ void randomSampleKernel(
-    Tidx *__restrict__ result,
+    Tout *__restrict__ result,
     const Tval *__restrict__ sorted,
     const Tidx *__restrict__ indices_out,
     size_t n,
@@ -174,7 +193,7 @@ static __global__ void randomSampleKernel(
 #endif
     for (size_t i = 0;; ++i) {
         if ((sorted[i]) >= p) {
-            *result = indices_out[i];
+            *result = static_cast<Tout>(indices_out[i]);
             return;
         }
     }
@@ -218,6 +237,7 @@ struct Algo {
         void *stream_) const {
 
         using Tval = typename CudaTval<Tval_>::Type;
+        using TworkIdx = typename InternalSampleIndex<Tidx>::Type;
 
         auto stream = (cudaStream_t)stream_;
         auto logits = (Tval *)probs;
@@ -226,14 +246,14 @@ struct Algo {
         auto workspace = reinterpret_cast<size_t>(workspace_);
         auto workspace_end = workspace + workspace_size;
 
-        auto indices = reinterpret_cast<Tidx *>(workspace);
-        workspace += align256(sizeof(Tidx) * n);
+        auto indices = reinterpret_cast<TworkIdx *>(workspace);
+        workspace += align256(sizeof(TworkIdx) * n);
 
         auto sorted = reinterpret_cast<Tval *>(workspace);
         workspace += align256(sizeof(Tval) * n);
 
-        auto indices_out = reinterpret_cast<Tidx *>(workspace);
-        workspace += align256(sizeof(Tidx) * n);
+        auto indices_out = reinterpret_cast<TworkIdx *>(workspace);
+        workspace += align256(sizeof(TworkIdx) * n);
 
         workspace_ = reinterpret_cast<void *>(workspace);
         workspace_size = workspace_end - workspace;
@@ -244,23 +264,25 @@ struct Algo {
 #endif
         auto grid = (n + block - 1) / block;
         // sort
-        fillIndices<<<static_cast<unsigned int>(grid), static_cast<unsigned int>(block), 0, stream>>>(indices, static_cast<int>(n));
-        CHECK_CUDA(radixSort(
+        fillIndices<TworkIdx><<<static_cast<unsigned int>(grid), static_cast<unsigned int>(block), 0, stream>>>(
+            indices, static_cast<int>(n));
+        CHECK_CUDA((radixSort<Tval, TworkIdx>(
             workspace_, workspace_size,
             logits, sorted,
             indices, indices_out,
             static_cast<int>(n),
-            stream));
+            stream)));
         // softmax
-        partialSoftmaxKernel<<<static_cast<unsigned int>(grid), static_cast<unsigned int>(block), 0, stream>>>(sorted, static_cast<int>(n), temperature);
-        setSoftmaxMaxKernel<<<1, 1, 0, stream>>>(sorted);
+        partialSoftmaxKernel<Tval><<<static_cast<unsigned int>(grid), static_cast<unsigned int>(block), 0, stream>>>(
+            sorted, static_cast<int>(n), temperature);
+        setSoftmaxMaxKernel<Tval><<<1, 1, 0, stream>>>(sorted);
         // sum
-        CHECK_CUDA(inclusiveSum(
-            workspace_, workspace,
+        CHECK_CUDA(inclusiveSum<Tval>(
+            workspace_, workspace_size,
             sorted, static_cast<int>(n),
             stream));
         // sample
-        randomSampleKernel<<<1, 1, 0, stream>>>(
+        randomSampleKernel<Tval, Tidx, TworkIdx><<<1, 1, 0, stream>>>(
             result,
             sorted, indices_out, n,
             random_val, topp, topk);
diff --git a/test/infiniop/awq_marlin_repack.py b/test/infiniop/awq_marlin_repack.py
index 3d2b5507e..662d06d10 100644
--- a/test/infiniop/awq_marlin_repack.py
+++ b/test/infiniop/awq_marlin_repack.py
@@ -322,10 +322,10 @@ def awq_marlin_repack_torch(b_weight, size_k, size_n, group_size, quant_type, is
 def test(
     handle,
     device,
-    k_chunk, 
-    n_chunk, 
-    quant_type, 
-    is_a_8bit, 
+    k_chunk,
+    n_chunk,
+    quant_type,
+    is_a_8bit,
     nk_factors,
     group_size=128,
     dtype=InfiniDtype.F16,
@@ -341,14 +341,14 @@ def test(
 
 
     b_weight = TestTensor((size_k, size_n), None, dtype, device)
-    
+
     w_ref, q_w, s, zp = quantize_weights(
         b_weight.torch_tensor(), quant_type, group_size, zero_points=True
     )
 
     # Pack to AWQ format
     q_w_awq = awq_pack(q_w, quant_type.size_bits, size_k, size_n)
-    
+
     ans = awq_marlin_repack_torch(b_weight.torch_tensor(), size_k, size_n, group_size, quant_type, is_a_8bit)
 
     input = TestTensor(
diff --git a/test/infiniop/gptq_marlin_repack.py b/test/infiniop/gptq_marlin_repack.py
index d013ec947..61b46156b 100644
--- a/test/infiniop/gptq_marlin_repack.py
+++ b/test/infiniop/gptq_marlin_repack.py
@@ -162,7 +162,7 @@ def reshape_w(w):
         w_s if group_size is not None else None,
         maybe_w_zp,
     )
-    
+
 def permute_rows(
     q_w: torch.Tensor,
     w_ref: torch.Tensor,
@@ -274,7 +274,7 @@ def sort_weights(q_w: torch.Tensor, g_idx: torch.Tensor):
         g_idx.to(device=orig_device),
         sort_indices.to(device=orig_device),
     )
-    
+
 def get_weight_perm(num_bits: int, is_a_8bit: bool = False):
     perm_list: list[int] = []
     if is_a_8bit:
@@ -387,15 +387,15 @@ def gptq_marlin_repack_torch(b_weight, size_k, size_n, group_size, quant_type, a
         q_w, size_k, size_n, quant_type.size_bits, weight_perm, is_a_8bit
     )
     return q_w_gptq, sort_indices, marlin_q_w_1
-    
+
 def test(
     handle,
     device,
-    k_chunk, 
-    n_chunk, 
-    quant_type, 
+    k_chunk,
+    n_chunk,
+    quant_type,
     act_order,
-    is_a_8bit, 
+    is_a_8bit,
     nk_factors,
     group_size=128,
     dtype=InfiniDtype.F16,
@@ -421,11 +421,11 @@ def test(
     if group_size == -1:
         group_size = size_k
     assert group_size <= size_k
-    
+
     test_dtype = to_torch_dtype(dtype)
     device_str = torch_device_map[device]
     b_weight = torch.randn((size_k, size_n), dtype=test_dtype, device=device_str) #must be randn
-    
+
     q_w_gptq, sort_indices, ans = gptq_marlin_repack_torch(b_weight, size_k, size_n, group_size, quant_type, act_order, is_a_8bit)
 
     input = TestTensor(
@@ -444,7 +444,7 @@ def test(
         mode="manual",
         set_tensor=sort_indices,
     )
-    
+
     output = TestTensor(ans.shape, None, InfiniDtype.I32, device, mode="zeros")
 
     if sync is not None:
@@ -477,7 +477,7 @@ def test(
         )
     )
     workspace = TestWorkspace(workspace_size.value, device)
-    
+
     def lib_gptq_marlin_repack():
         check_error(
             LIBINFINIOP.infiniopGptqMarlinRepack(
@@ -492,7 +492,7 @@ def lib_gptq_marlin_repack():
         )
 
     lib_gptq_marlin_repack()
-    
+
     atol, rtol = get_tolerance(_TOLERANCE_MAP, dtype)
     if DEBUG:
         debug(output.actual_tensor(), ans, atol=atol, rtol=rtol)
diff --git a/third_party/cutlass b/third_party/cutlass
index 087c84df8..f74fea9ce 160000
--- a/third_party/cutlass
+++ b/third_party/cutlass
@@ -1 +1 @@
-Subproject commit 087c84df83d254b5fb295a7a408f1a1d554085cf
+Subproject commit f74fea9ce35868d3ae9f8d1dce1969d7250d3f90
diff --git a/third_party/flash-attention b/third_party/flash-attention
index 10846960c..6362bd3bc 160000
--- a/third_party/flash-attention
+++ b/third_party/flash-attention
@@ -1 +1 @@
-Subproject commit 10846960ca0793b993446f6dbaf696479c127a9d
+Subproject commit 6362bd3bcad059aa15fd993c6a9d5d1ee8a11418
diff --git a/xmake/nvidia.lua b/xmake/nvidia.lua
index 282f5876f..ab7fdfd63 100644
--- a/xmake/nvidia.lua
+++ b/xmake/nvidia.lua
@@ -181,6 +181,20 @@ target("infiniop-nvidia")
     end
 
     local arch_opt = get_config("cuda_arch")
+    local TVM_FFI_ROOT = path.join(os.projectdir(), "third_party/tvm-ffi")
+    if os.isdir(TVM_FFI_ROOT) then
+        add_includedirs(TVM_FFI_ROOT .. "/include")
+    end
+    local DLPACK_ROOT = path.join(os.projectdir(), "third_party/dlpack")
+    if os.isdir(DLPACK_ROOT) then
+        add_includedirs(DLPACK_ROOT .. "/include")
+    end
+    if TVM_ROOT == nil and arch_opt then
+        local sgl_arch = arch_opt:match("sm_(%d+)")
+        if sgl_arch then
+            add_defines("SGL_CUDA_ARCH=" .. tostring(tonumber(sgl_arch) * 10))
+        end
+    end
     if TVM_ROOT ~= nil then
         add_defines("ENABLE_TVM_API")
         add_includedirs(TVM_ROOT, TVM_ROOT .. "/include", TVM_ROOT .. "/3rdparty/dlpack/include/")