wehub-resource-sync
864 lines
30 KiB
Plaintext
864 lines
30 KiB
Plaintext
|
|
/*
|
|
* Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
#include "torch_utils.h"
|
|
|
|
#include <torch/csrc/stable/macros.h>
|
|
#include <torch/csrc/stable/accelerator.h>
|
|
#include <torch/csrc/stable/tensor.h>
|
|
#include <torch/csrc/stable/ops.h>
|
|
#include <torch/headeronly/core/ScalarType.h>
|
|
#include <torch/csrc/stable/device.h>
|
|
|
|
#include <cooperative_groups.h>
|
|
#include <cuda_runtime.h>
|
|
|
|
#include "cuda_compat.h"
|
|
#include "minimax_reduce_rms_kernel.h"
|
|
|
|
#include <algorithm>
|
|
|
|
#define FINAL_MASK 0xffffffff
|
|
#define MINIMAX_REDUCE_RMS_WARP_SIZE 32
|
|
|
|
namespace vllm {
|
|
namespace tensorrt_llm {
|
|
|
|
template <int NRanks>
|
|
struct LamportComm {
|
|
__device__ __forceinline__ LamportComm(void** workspace, int rank) {
|
|
counter_ptr = &reinterpret_cast<int*>(workspace[NRanks * 3])[0];
|
|
flag_ptr = &reinterpret_cast<int*>(workspace[NRanks * 3])[2];
|
|
clear_ptr = &reinterpret_cast<int64_t*>(workspace[NRanks * 3 + 1])[0];
|
|
flag_value = *flag_ptr;
|
|
auto comm_size = reinterpret_cast<int64_t*>(workspace[NRanks * 3 + 1])[1];
|
|
clear_size = *clear_ptr;
|
|
int data_offset = flag_value % 3;
|
|
int clear_offset = (flag_value + 2) % 3;
|
|
for (int r = 0; r < NRanks; ++r) {
|
|
data_bufs[r] = reinterpret_cast<uint8_t*>(workspace[2 * NRanks + r]) +
|
|
data_offset * comm_size;
|
|
}
|
|
clear_buf = reinterpret_cast<uint8_t*>(workspace[2 * NRanks + rank]) +
|
|
clear_offset * comm_size;
|
|
__syncthreads();
|
|
if (threadIdx.x == 0) {
|
|
atomicAdd(counter_ptr, 1);
|
|
}
|
|
}
|
|
|
|
__device__ __forceinline__ void update(int64_t new_clear_size) {
|
|
if (blockIdx.x == 0 && threadIdx.x == 0) {
|
|
while (*reinterpret_cast<int volatile*>(counter_ptr) != gridDim.x) {
|
|
}
|
|
*flag_ptr = (flag_value + 1) % 3;
|
|
*clear_ptr = new_clear_size;
|
|
*counter_ptr = 0;
|
|
}
|
|
}
|
|
|
|
int* counter_ptr;
|
|
int* flag_ptr;
|
|
int64_t* clear_ptr;
|
|
uint8_t* data_bufs[NRanks];
|
|
uint8_t* clear_buf;
|
|
int64_t clear_size;
|
|
int flag_value;
|
|
};
|
|
|
|
__device__ __forceinline__ bool is_neg_zero(float v) {
|
|
return *reinterpret_cast<uint32_t*>(&v) == 0x80000000;
|
|
}
|
|
|
|
__device__ __forceinline__ bool is_neg_zero(float4 v) {
|
|
return is_neg_zero(v.x) || is_neg_zero(v.y) || is_neg_zero(v.z) ||
|
|
is_neg_zero(v.w);
|
|
}
|
|
|
|
__device__ __forceinline__ float4 get_neg_zero() {
|
|
float4 vec;
|
|
#pragma unroll
|
|
for (int i = 0; i < 4; ++i) {
|
|
reinterpret_cast<uint32_t*>(&vec)[i] = 0x80000000;
|
|
}
|
|
return vec;
|
|
}
|
|
|
|
template <int Dim>
|
|
__device__ __forceinline__ float rms_rsqrt(float& v, float eps) {
|
|
constexpr float kInvDim = 1.0F / static_cast<float>(Dim);
|
|
v = rsqrtf((v * kInvDim) + eps);
|
|
return v;
|
|
}
|
|
|
|
template <int Dim>
|
|
__device__ __forceinline__ float4 rms_rsqrt(float4& v, float eps) {
|
|
constexpr float kInvDim = 1.0F / static_cast<float>(Dim);
|
|
v.x = rsqrtf((v.x * kInvDim) + eps);
|
|
v.y = rsqrtf((v.y * kInvDim) + eps);
|
|
v.z = rsqrtf((v.z * kInvDim) + eps);
|
|
v.w = rsqrtf((v.w * kInvDim) + eps);
|
|
return v;
|
|
}
|
|
__device__ __forceinline__ float4 ld_global_volatile(float4* addr) {
|
|
float4 val;
|
|
asm volatile("ld.volatile.global.v4.f32 {%0, %1, %2, %3}, [%4];"
|
|
: "=f"(val.x), "=f"(val.y), "=f"(val.z), "=f"(val.w)
|
|
: "l"(addr));
|
|
return val;
|
|
}
|
|
|
|
__device__ __forceinline__ float ld_global_volatile(float* addr) {
|
|
float val;
|
|
asm volatile("ld.volatile.global.f32 %0, [%1];" : "=f"(val) : "l"(addr));
|
|
return val;
|
|
}
|
|
|
|
// Used by the scalar (non-float4) kernel only
|
|
template <typename T, int NUM>
|
|
__inline__ __device__ T warpReduceSumV2(T* val) {
|
|
#pragma unroll
|
|
for (int i = 0; i < NUM; i++) {
|
|
#pragma unroll
|
|
for (int mask = 16; mask > 0; mask >>= 1)
|
|
val[i] += __shfl_xor_sync(FINAL_MASK, val[i], mask, 32);
|
|
}
|
|
return (T)(0.0f);
|
|
}
|
|
|
|
template <typename T, int NUM>
|
|
__inline__ __device__ T blockReduceSumV2(T* val) {
|
|
static __shared__ T shared[NUM][33];
|
|
int lane = threadIdx.x & 0x1f;
|
|
int wid = threadIdx.x >> 5;
|
|
|
|
warpReduceSumV2<T, NUM>(val);
|
|
|
|
if (lane == 0) {
|
|
#pragma unroll
|
|
for (int i = 0; i < NUM; i++) {
|
|
shared[i][wid] = val[i];
|
|
}
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
bool is_mask = threadIdx.x < (blockDim.x / 32.f);
|
|
#pragma unroll
|
|
for (int i = 0; i < NUM; i++) {
|
|
val[i] = is_mask ? shared[i][lane] : (T)(0.0f);
|
|
}
|
|
warpReduceSumV2<T, NUM>(val);
|
|
return (T)0.0f;
|
|
}
|
|
|
|
// for float4 version
|
|
template <uint32_t kNumThreads, typename T, int ArraySize = 4>
|
|
__device__ __forceinline__ void local_warp_reduce_sum_array(
|
|
T* value_ptr, uint32_t active_mask = 0xffffffffu) {
|
|
static_assert(kNumThreads >= 1 &&
|
|
kNumThreads <= MINIMAX_REDUCE_RMS_WARP_SIZE);
|
|
#pragma unroll
|
|
for (int i = 0; i < ArraySize; ++i) {
|
|
#pragma unroll
|
|
for (int mask = kNumThreads / 2; mask > 0; mask >>= 1) {
|
|
value_ptr[i] += __shfl_xor_sync(active_mask, value_ptr[i], mask,
|
|
MINIMAX_REDUCE_RMS_WARP_SIZE);
|
|
}
|
|
}
|
|
}
|
|
|
|
constexpr int next_pow2(int val) {
|
|
int result = 1;
|
|
while (result < val) {
|
|
result <<= 1;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
|
|
template <typename DType>
|
|
class IndexHelper {
|
|
public:
|
|
__device__ __forceinline__ IndexHelper(MiniMaxReduceRMSParams const& params) {
|
|
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
|
namespace cg = cooperative_groups;
|
|
cg::cluster_group cluster = cg::this_cluster();
|
|
cg::grid_group grid = cg::this_grid();
|
|
token_id = grid.cluster_rank();
|
|
access_id_in_token = cluster.thread_rank();
|
|
token_stride = grid.num_clusters();
|
|
#else
|
|
token_id = blockIdx.x;
|
|
access_id_in_token = threadIdx.x;
|
|
token_stride = gridDim.x;
|
|
#endif
|
|
access_id = token_id * params.hidden_dim / kElemsPerAccess<DType> +
|
|
access_id_in_token;
|
|
access_stride = token_stride * params.hidden_dim / kElemsPerAccess<DType>;
|
|
tot_access = params.size_q / kElemsPerAccess<DType>;
|
|
}
|
|
|
|
int token_id;
|
|
int access_id_in_token;
|
|
int token_stride;
|
|
int access_id;
|
|
int access_stride;
|
|
int tot_access;
|
|
};
|
|
|
|
/**
|
|
* this kernel is used to for minimax attention module
|
|
* input tensor [total_tokens, hidden_dim / tp_size], fp32
|
|
* rms weight [hidden_dim / tp_size], bf16
|
|
step 1: reduce from single rank to get the variance sum (reduce(input^2,
|
|
dim=-1)) step 2: reduce from all ranks to get the variance sum
|
|
(all_reduce(variance_sum)) step 3: calculate the rms norm (input *
|
|
rsqrt(variance + eps)) in this case, max hidden_dim is 6144 (float data), for
|
|
each token, we only need 6144 / 4 / tp_size = (1536 / tp_size) threads so we can
|
|
assume cluster size is 1 (tp_size >= 2)
|
|
*/
|
|
template <typename DType, int NRanks>
|
|
__global__ void __launch_bounds__(1024)
|
|
minimax_reduce_rms_kernel_lamport(MiniMaxReduceRMSParams params) {
|
|
IndexHelper<DType> index_helper(params);
|
|
int token_id = index_helper.token_id;
|
|
int access_id_in_token = index_helper.access_id_in_token;
|
|
int token_stride = index_helper.token_stride;
|
|
int access_id = index_helper.access_id;
|
|
int access_stride = index_helper.access_stride;
|
|
int tot_access = index_helper.tot_access;
|
|
int tot_tokens = params.size_q / params.hidden_dim;
|
|
float4 clear_vec = get_neg_zero();
|
|
|
|
LamportComm<NRanks> comm(params.workspace, params.rank);
|
|
int clear_access = comm.clear_size / kElemsPerAccess<DType>;
|
|
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
|
cudaGridDependencySynchronize();
|
|
#endif
|
|
for (int idx = access_id; idx < tot_access;
|
|
idx += access_stride, token_id += token_stride) {
|
|
alignas(16) DType vals[kElemsPerAccess<DType>];
|
|
float sum_variance = 0.F;
|
|
*reinterpret_cast<float4*>(vals) =
|
|
reinterpret_cast<float4*>(params.allreduce_in)[idx];
|
|
#pragma unroll
|
|
for (int i = 0; i < kElemsPerAccess<DType>; ++i) {
|
|
sum_variance += static_cast<float>(vals[i]) * static_cast<float>(vals[i]);
|
|
}
|
|
blockReduceSumV2<float, 1>(&sum_variance);
|
|
if (is_neg_zero(sum_variance)) {
|
|
sum_variance = 0.F;
|
|
}
|
|
if (threadIdx.x == 0) {
|
|
for (int r = 0; r < NRanks; ++r) {
|
|
reinterpret_cast<float*>(
|
|
comm.data_bufs[r])[(params.rank * tot_tokens) + token_id] =
|
|
(sum_variance);
|
|
}
|
|
}
|
|
|
|
bool done = false;
|
|
float vars_all_ranks[NRanks];
|
|
while (!done) {
|
|
done = true;
|
|
#pragma unroll
|
|
for (int r = 0; r < NRanks; ++r) {
|
|
vars_all_ranks[r] = ld_global_volatile(&reinterpret_cast<float*>(
|
|
comm.data_bufs[params.rank])[(r * tot_tokens) + token_id]);
|
|
done &= !is_neg_zero(vars_all_ranks[r]);
|
|
}
|
|
}
|
|
sum_variance = 0.F;
|
|
#pragma unroll
|
|
for (int r = 0; r < NRanks; ++r) {
|
|
sum_variance += vars_all_ranks[r];
|
|
}
|
|
|
|
DType norm_weight[kElemsPerAccess<DType>];
|
|
*reinterpret_cast<typename ElemsPerAccess<DType>::vec_type*>(norm_weight) =
|
|
reinterpret_cast<typename ElemsPerAccess<DType>::vec_type*>(
|
|
params.rms_gamma)[access_id_in_token];
|
|
|
|
#pragma unroll
|
|
for (int i = 0; i < kElemsPerAccess<DType>; ++i) {
|
|
vals[i] = static_cast<DType>(
|
|
static_cast<float>(vals[i]) *
|
|
rsqrtf(
|
|
(sum_variance / static_cast<float>(params.hidden_dim) / NRanks) +
|
|
params.rms_eps) *
|
|
static_cast<float>(norm_weight[i]));
|
|
}
|
|
|
|
reinterpret_cast<float4*>(params.rms_norm_out)[idx] =
|
|
*reinterpret_cast<float4*>(vals);
|
|
}
|
|
for (int idx = access_id; idx < clear_access; idx += access_stride) {
|
|
reinterpret_cast<float4*>(comm.clear_buf)[idx] = clear_vec;
|
|
}
|
|
comm.update(params.size_q * NRanks);
|
|
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
|
cudaTriggerProgrammaticLaunchCompletion();
|
|
#endif
|
|
}
|
|
|
|
/**
|
|
* Float4 variant: process 4 rows at once, allreduce variance sums as float4 for
|
|
* better memory coalescing. sum_variance is always float; applies to all DTypes
|
|
* (half, bf16, float). When tot_tokens % 4 != 0, the last group pads rows with
|
|
* zeros; padded rows are not written to rms_norm_out. IsQK: when true, process
|
|
* Q+K in one loop with doubled comm buffer; when false, single-matrix (Q only).
|
|
*/
|
|
template <typename DType, int NRanks, int OriginQDim, int OriginKDim>
|
|
__global__ void __launch_bounds__(1024)
|
|
minimax_reduce_qk_rms_kernel_lamport_float4(MiniMaxReduceRMSParams params) {
|
|
// Compile-time per-rank dimensions
|
|
constexpr int RankQDim = OriginQDim / NRanks;
|
|
constexpr int RankKDim = OriginKDim / NRanks;
|
|
// Threads needed to cover one row of Q / K with float4 accesses
|
|
constexpr int ThreadsPerRowQ = RankQDim / kElemsPerAccess<DType>;
|
|
constexpr int ThreadsPerRowK = RankKDim / kElemsPerAccess<DType>;
|
|
// Number of warps dedicated to Q / K
|
|
constexpr int NumWarpQ = (ThreadsPerRowQ + MINIMAX_REDUCE_RMS_WARP_SIZE - 1) /
|
|
MINIMAX_REDUCE_RMS_WARP_SIZE;
|
|
constexpr int NumWarpK = (ThreadsPerRowK + MINIMAX_REDUCE_RMS_WARP_SIZE - 1) /
|
|
MINIMAX_REDUCE_RMS_WARP_SIZE;
|
|
|
|
int tot_tokens = params.size_q / RankQDim;
|
|
int tot_groups = (tot_tokens + 3) / 4; // ceiling; last group may be partial
|
|
|
|
// Memory strides for strided qkv tensors (elements -> float4-access units)
|
|
int access_stride_q = (params.stride_q > 0 ? params.stride_q : RankQDim) /
|
|
kElemsPerAccess<DType>;
|
|
int access_stride_k = (params.stride_k > 0 ? params.stride_k : RankKDim) /
|
|
kElemsPerAccess<DType>;
|
|
// Output strides: default to contiguous (hidden_dim / hidden_dim_k)
|
|
int access_stride_q_out =
|
|
(params.stride_q_out > 0 ? params.stride_q_out : params.hidden_dim) /
|
|
kElemsPerAccess<DType>;
|
|
int access_stride_k_out =
|
|
(params.stride_k_out > 0 ? params.stride_k_out : params.hidden_dim_k) /
|
|
kElemsPerAccess<DType>;
|
|
|
|
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
|
namespace cg = cooperative_groups;
|
|
cg::cluster_group cluster = cg::this_cluster();
|
|
cg::grid_group grid = cg::this_grid();
|
|
int group_id = grid.cluster_rank();
|
|
int access_id_in_token = cluster.thread_rank();
|
|
int group_stride = grid.num_clusters();
|
|
#else
|
|
int group_id = blockIdx.x;
|
|
int access_id_in_token = threadIdx.x;
|
|
int group_stride = gridDim.x;
|
|
#endif
|
|
|
|
bool is_q = (access_id_in_token < NumWarpQ * MINIMAX_REDUCE_RMS_WARP_SIZE);
|
|
int k_thread_idx =
|
|
access_id_in_token - (NumWarpQ * MINIMAX_REDUCE_RMS_WARP_SIZE);
|
|
bool is_valid_q = (access_id_in_token < ThreadsPerRowQ);
|
|
bool is_valid_k = (k_thread_idx >= 0 && k_thread_idx < ThreadsPerRowK);
|
|
float4 clear_vec = get_neg_zero();
|
|
|
|
// Shared memory for two-level block reduction and scale broadcast
|
|
__shared__ float block_reduce_sum[4][MINIMAX_REDUCE_RMS_WARP_SIZE + 1];
|
|
__shared__ float global_scale_q[4];
|
|
__shared__ float global_scale_k[4];
|
|
|
|
LamportComm<NRanks> comm(params.workspace, params.rank);
|
|
|
|
DType norm_weight[kElemsPerAccess<DType>]{};
|
|
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
|
cudaGridDependencySynchronize();
|
|
#endif
|
|
if (is_q) {
|
|
if (is_valid_q) {
|
|
*reinterpret_cast<typename ElemsPerAccess<DType>::vec_type*>(
|
|
norm_weight) =
|
|
reinterpret_cast<typename ElemsPerAccess<DType>::vec_type const*>(
|
|
params.rms_gamma)[access_id_in_token];
|
|
}
|
|
} else {
|
|
if (is_valid_k) {
|
|
*reinterpret_cast<typename ElemsPerAccess<DType>::vec_type*>(
|
|
norm_weight) =
|
|
reinterpret_cast<typename ElemsPerAccess<DType>::vec_type const*>(
|
|
params.rms_gamma_k)[k_thread_idx];
|
|
}
|
|
}
|
|
|
|
// Main loop: process one group of 4 tokens per iteration.
|
|
for (int g = group_id; g < tot_groups; g += group_stride) {
|
|
alignas(16) DType vals[4][kElemsPerAccess<DType>]{};
|
|
float warp_sum_variance[4]{0.F, 0.F, 0.F, 0.F};
|
|
|
|
if (is_q) {
|
|
#pragma unroll
|
|
for (int row = 0; row < 4; ++row) {
|
|
int token_r = g * 4 + row;
|
|
if (token_r >= tot_tokens || !is_valid_q) {
|
|
continue;
|
|
}
|
|
int idx_r = token_r * access_stride_q + access_id_in_token;
|
|
*reinterpret_cast<float4*>(&vals[row][0]) =
|
|
reinterpret_cast<float4 const*>(params.allreduce_in)[idx_r];
|
|
#pragma unroll
|
|
for (int i = 0; i < kElemsPerAccess<DType>; ++i) {
|
|
float x = static_cast<float>(vals[row][i]);
|
|
warp_sum_variance[row] += x * x;
|
|
}
|
|
}
|
|
} else {
|
|
#pragma unroll
|
|
for (int row = 0; row < 4; ++row) {
|
|
int token_r = g * 4 + row;
|
|
if (token_r >= tot_tokens || !is_valid_k) {
|
|
continue;
|
|
}
|
|
int idx_r = token_r * access_stride_k + k_thread_idx;
|
|
*reinterpret_cast<float4*>(&vals[row][0]) =
|
|
reinterpret_cast<float4 const*>(params.allreduce_in_k)[idx_r];
|
|
#pragma unroll
|
|
for (int i = 0; i < kElemsPerAccess<DType>; ++i) {
|
|
float x = static_cast<float>(vals[row][i]);
|
|
warp_sum_variance[row] += x * x;
|
|
}
|
|
}
|
|
}
|
|
|
|
local_warp_reduce_sum_array<MINIMAX_REDUCE_RMS_WARP_SIZE, float, 4>(
|
|
warp_sum_variance);
|
|
// Warp lane 0 writes its warp's partial sum to shared memory
|
|
int lane = threadIdx.x & (MINIMAX_REDUCE_RMS_WARP_SIZE - 1);
|
|
if (lane == 0) {
|
|
#pragma unroll
|
|
for (int t = 0; t < 4; ++t) {
|
|
block_reduce_sum[t][threadIdx.x / MINIMAX_REDUCE_RMS_WARP_SIZE] =
|
|
warp_sum_variance[t];
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
int tid = threadIdx.x;
|
|
|
|
if (tid < MINIMAX_REDUCE_RMS_WARP_SIZE) {
|
|
constexpr int kNumWarpQPow2 =
|
|
(next_pow2(NumWarpQ) > NRanks) ? next_pow2(NumWarpQ) : NRanks;
|
|
float local_sum[4];
|
|
#pragma unroll
|
|
for (int t = 0; t < 4; ++t) {
|
|
local_sum[t] = (tid < NumWarpQ) ? block_reduce_sum[t][tid] : 0.F;
|
|
}
|
|
// After this, all kNumWarpQPow2 lanes (including tid 0..NRanks-1) have
|
|
// the total Q sum-of-squares for all 4 tokens.
|
|
local_warp_reduce_sum_array<kNumWarpQPow2, float, 4>(local_sum);
|
|
|
|
if (tid < NRanks) {
|
|
#pragma unroll
|
|
for (int t = 0; t < 4; ++t) {
|
|
if (is_neg_zero(local_sum[t])) {
|
|
local_sum[t] = 0.F;
|
|
}
|
|
}
|
|
// Parallel push: thread tid writes this rank's Q sum to rank tid's buf
|
|
reinterpret_cast<float4*>(
|
|
comm.data_bufs[tid])[(params.rank * tot_groups * 2) + (2 * g)] =
|
|
*reinterpret_cast<float4*>(local_sum);
|
|
|
|
// Parallel pull: thread tid reads rank tid's contribution from
|
|
// this rank's (params.rank's) buffer
|
|
bool done = false;
|
|
float4 var_all_ranks;
|
|
while (!done) {
|
|
done = true;
|
|
var_all_ranks = ld_global_volatile(&reinterpret_cast<float4*>(
|
|
comm.data_bufs[params.rank])[(tid * tot_groups * 2) + (2 * g)]);
|
|
done &= !is_neg_zero(var_all_ranks);
|
|
}
|
|
|
|
// Warp-level allreduce: each of the NRanks threads holds one rank's
|
|
// partial sum; after this all NRanks threads have the global total.
|
|
constexpr uint32_t kQActiveMask = (1u << NRanks) - 1u;
|
|
local_warp_reduce_sum_array<NRanks, float, 4>(
|
|
reinterpret_cast<float*>(&var_all_ranks), kQActiveMask);
|
|
|
|
// Thread 0 computes rsqrt with compile-time Dim and writes to smem
|
|
if (tid == 0) {
|
|
*reinterpret_cast<float4*>(global_scale_q) =
|
|
rms_rsqrt<OriginQDim>(var_all_ranks, params.rms_eps);
|
|
}
|
|
}
|
|
} else if (tid >= MINIMAX_REDUCE_RMS_WARP_SIZE * NumWarpQ &&
|
|
tid < MINIMAX_REDUCE_RMS_WARP_SIZE * (NumWarpQ + 1)) {
|
|
// --- K leader warp ---
|
|
constexpr int kNumWarpKPow2 =
|
|
(next_pow2(NumWarpK) > NRanks) ? next_pow2(NumWarpK) : NRanks;
|
|
float local_sum[4];
|
|
#pragma unroll
|
|
for (int t = 0; t < 4; ++t) {
|
|
local_sum[t] = (k_thread_idx < NumWarpK)
|
|
? block_reduce_sum[t][NumWarpQ + k_thread_idx]
|
|
: 0.F;
|
|
}
|
|
local_warp_reduce_sum_array<kNumWarpKPow2, float, 4>(local_sum);
|
|
|
|
if (k_thread_idx < NRanks) {
|
|
#pragma unroll
|
|
for (int t = 0; t < 4; ++t) {
|
|
if (is_neg_zero(local_sum[t])) {
|
|
local_sum[t] = 0.F;
|
|
}
|
|
}
|
|
reinterpret_cast<float4*>(
|
|
comm.data_bufs[k_thread_idx])[(params.rank * tot_groups * 2) +
|
|
(2 * g + 1)] =
|
|
*reinterpret_cast<float4*>(local_sum);
|
|
|
|
bool done = false;
|
|
float4 var_all_ranks;
|
|
while (!done) {
|
|
done = true;
|
|
var_all_ranks = ld_global_volatile(&reinterpret_cast<float4*>(
|
|
comm.data_bufs[params.rank])[(k_thread_idx * tot_groups * 2) +
|
|
(2 * g + 1)]);
|
|
done &= !is_neg_zero(var_all_ranks);
|
|
}
|
|
|
|
constexpr uint32_t kKActiveMask = (1u << NRanks) - 1u;
|
|
local_warp_reduce_sum_array<NRanks, float, 4>(
|
|
reinterpret_cast<float*>(&var_all_ranks), kKActiveMask);
|
|
|
|
if (k_thread_idx == 0) {
|
|
*reinterpret_cast<float4*>(global_scale_k) =
|
|
rms_rsqrt<OriginKDim>(var_all_ranks, params.rms_eps);
|
|
}
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
if (is_q) {
|
|
#pragma unroll
|
|
for (int t = 0; t < 4; ++t) {
|
|
warp_sum_variance[t] = global_scale_q[t];
|
|
}
|
|
#pragma unroll
|
|
for (int r = 0; r < 4; ++r) {
|
|
#pragma unroll
|
|
for (int i = 0; i < kElemsPerAccess<DType>; ++i) {
|
|
vals[r][i] = static_cast<DType>(static_cast<float>(vals[r][i]) *
|
|
warp_sum_variance[r] *
|
|
static_cast<float>(norm_weight[i]));
|
|
}
|
|
int token_r = g * 4 + r;
|
|
if (token_r >= tot_tokens || !is_valid_q) {
|
|
continue;
|
|
}
|
|
int idx_out = token_r * access_stride_q_out + access_id_in_token;
|
|
reinterpret_cast<float4*>(params.rms_norm_out)[idx_out] =
|
|
*reinterpret_cast<float4*>(&vals[r][0]);
|
|
}
|
|
} else {
|
|
#pragma unroll
|
|
for (int t = 0; t < 4; ++t) {
|
|
warp_sum_variance[t] = global_scale_k[t];
|
|
}
|
|
#pragma unroll
|
|
for (int r = 0; r < 4; ++r) {
|
|
#pragma unroll
|
|
for (int i = 0; i < kElemsPerAccess<DType>; ++i) {
|
|
vals[r][i] = static_cast<DType>(static_cast<float>(vals[r][i]) *
|
|
warp_sum_variance[r] *
|
|
static_cast<float>(norm_weight[i]));
|
|
}
|
|
int token_r = g * 4 + r;
|
|
if (token_r >= tot_tokens || !is_valid_k) {
|
|
continue;
|
|
}
|
|
int idx_out = token_r * access_stride_k_out + k_thread_idx;
|
|
reinterpret_cast<float4*>(params.rms_norm_out_k)[idx_out] =
|
|
*reinterpret_cast<float4*>(&vals[r][0]);
|
|
}
|
|
}
|
|
} // end group loop
|
|
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
|
cudaTriggerProgrammaticLaunchCompletion();
|
|
#endif
|
|
|
|
int clear_access = static_cast<int>(comm.clear_size / kElemsPerAccess<DType>);
|
|
int clear_stride = group_stride * blockDim.x;
|
|
for (int idx = group_id * blockDim.x + threadIdx.x; idx < clear_access;
|
|
idx += clear_stride) {
|
|
reinterpret_cast<float4*>(comm.clear_buf)[idx] = clear_vec;
|
|
}
|
|
|
|
comm.update(static_cast<int64_t>(2) * tot_groups * kElemsPerAccess<DType> *
|
|
NRanks);
|
|
}
|
|
|
|
int get_sm_count() {
|
|
static int sm_count = 0;
|
|
if (sm_count == 0) {
|
|
int device_id;
|
|
STD_CUDA_CHECK(cudaGetDevice(&device_id));
|
|
cudaDeviceProp device_prop;
|
|
cudaGetDeviceProperties(&device_prop, device_id);
|
|
sm_count = device_prop.multiProcessorCount;
|
|
}
|
|
return sm_count;
|
|
}
|
|
|
|
inline int getSMVersion(bool queryRealSmArch = false) {
|
|
int device{-1};
|
|
STD_CUDA_CHECK(cudaGetDevice(&device));
|
|
int sm_major = 0;
|
|
int sm_minor = 0;
|
|
STD_CUDA_CHECK(cudaDeviceGetAttribute(
|
|
&sm_major, cudaDevAttrComputeCapabilityMajor, device));
|
|
STD_CUDA_CHECK(cudaDeviceGetAttribute(
|
|
&sm_minor, cudaDevAttrComputeCapabilityMinor, device));
|
|
int sm = sm_major * 10 + sm_minor;
|
|
if (sm == 121 && !queryRealSmArch) {
|
|
return 120;
|
|
}
|
|
return sm;
|
|
}
|
|
|
|
template <typename KernelFunc>
|
|
int get_max_active_blocks(KernelFunc kernel, int block_size,
|
|
int dynamic_smem = 0) {
|
|
int max_active = 0;
|
|
STD_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
|
|
&max_active, kernel, block_size, dynamic_smem));
|
|
return std::max(max_active, 1);
|
|
}
|
|
|
|
template <typename DType, int NRanks>
|
|
void minimax_reduce_rms_kernel_launcher(MiniMaxReduceRMSParams const& params) {
|
|
static int SM = getSMVersion();
|
|
int token_num = params.size_q / params.hidden_dim;
|
|
int sm_count = get_sm_count();
|
|
int cluster_size = 1;
|
|
int cluster_num = token_num;
|
|
int threads_per_token = params.hidden_dim / kElemsPerAccess<DType>;
|
|
int block_size = threads_per_token;
|
|
|
|
int max_blocks_per_sm = get_max_active_blocks(
|
|
minimax_reduce_rms_kernel_lamport<DType, NRanks>, block_size);
|
|
int max_grid = max_blocks_per_sm * sm_count;
|
|
|
|
int grid_size =
|
|
(std::min(max_grid, cluster_num * cluster_size) / cluster_size) *
|
|
cluster_size;
|
|
|
|
cudaLaunchConfig_t cfg;
|
|
cfg.gridDim = grid_size;
|
|
cfg.blockDim = block_size;
|
|
cfg.dynamicSmemBytes = 0;
|
|
cfg.stream = params.stream;
|
|
|
|
cudaLaunchAttribute attribute[2];
|
|
attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
|
|
attribute[0].val.programmaticStreamSerializationAllowed = 1;
|
|
attribute[1].id = cudaLaunchAttributeClusterDimension;
|
|
attribute[1].val.clusterDim.x = cluster_size;
|
|
attribute[1].val.clusterDim.y = 1;
|
|
attribute[1].val.clusterDim.z = 1;
|
|
cfg.attrs = attribute;
|
|
cfg.numAttrs = SM >= 90 ? 2 : 0;
|
|
|
|
STD_CUDA_CHECK(cudaLaunchKernelEx(
|
|
&cfg, minimax_reduce_rms_kernel_lamport<DType, NRanks>, params));
|
|
}
|
|
|
|
template <typename DType, int NRanks, int OriginQDim, int OriginKDim>
|
|
void minimax_reduce_rms_kernel_launcher_float4(
|
|
MiniMaxReduceRMSParams const& params) {
|
|
STD_TORCH_CHECK(params.size_q % params.hidden_dim == 0);
|
|
STD_TORCH_CHECK(params.hidden_dim % kElemsPerAccess<DType> == 0);
|
|
if (params.stride_q > 0) {
|
|
STD_TORCH_CHECK(params.stride_q % kElemsPerAccess<DType> == 0);
|
|
}
|
|
STD_TORCH_CHECK(params.allreduce_in_k != nullptr,
|
|
"float4 QK kernel requires K input");
|
|
STD_TORCH_CHECK(params.hidden_dim >= params.hidden_dim_k);
|
|
STD_TORCH_CHECK(params.size_k % params.hidden_dim_k == 0);
|
|
STD_TORCH_CHECK(params.hidden_dim_k % kElemsPerAccess<DType> == 0);
|
|
STD_TORCH_CHECK(params.size_q / params.hidden_dim ==
|
|
params.size_k / params.hidden_dim_k);
|
|
if (params.stride_k > 0) {
|
|
STD_TORCH_CHECK(params.stride_k % kElemsPerAccess<DType> == 0);
|
|
}
|
|
|
|
int token_num = params.size_q / params.hidden_dim;
|
|
int tot_groups = (token_num + 3) / 4;
|
|
if (tot_groups == 0) {
|
|
return;
|
|
}
|
|
|
|
static int SM = getSMVersion();
|
|
int sm_count = get_sm_count();
|
|
int cluster_size = 1;
|
|
int cluster_num = tot_groups;
|
|
|
|
int access_per_row_q = params.hidden_dim / kElemsPerAccess<DType>;
|
|
int access_per_row_k = params.hidden_dim_k / kElemsPerAccess<DType>;
|
|
|
|
// Round each section up to a warp boundary
|
|
auto divUp = [](int a, int b) { return (a + b - 1) / b * b; };
|
|
int block_size = divUp(access_per_row_q, MINIMAX_REDUCE_RMS_WARP_SIZE) +
|
|
divUp(access_per_row_k, MINIMAX_REDUCE_RMS_WARP_SIZE);
|
|
|
|
auto kfn =
|
|
minimax_reduce_qk_rms_kernel_lamport_float4<DType, NRanks, OriginQDim,
|
|
OriginKDim>;
|
|
|
|
int max_blocks_per_sm = get_max_active_blocks(kfn, block_size);
|
|
int max_grid = max_blocks_per_sm * sm_count;
|
|
int grid_size =
|
|
(std::min(max_grid, cluster_num * cluster_size) / cluster_size) *
|
|
cluster_size;
|
|
|
|
cudaLaunchConfig_t cfg;
|
|
cfg.gridDim = grid_size;
|
|
cfg.blockDim = block_size;
|
|
cfg.dynamicSmemBytes = 0;
|
|
cfg.stream = params.stream;
|
|
|
|
cudaLaunchAttribute attribute[2];
|
|
attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
|
|
attribute[0].val.programmaticStreamSerializationAllowed = 1;
|
|
attribute[1].id = cudaLaunchAttributeClusterDimension;
|
|
attribute[1].val.clusterDim.x = cluster_size;
|
|
attribute[1].val.clusterDim.y = 1;
|
|
attribute[1].val.clusterDim.z = 1;
|
|
cfg.attrs = attribute;
|
|
cfg.numAttrs = SM >= 90 ? 2 : 0;
|
|
|
|
STD_CUDA_CHECK(cudaLaunchKernelEx(&cfg, kfn, params));
|
|
}
|
|
|
|
template <int NRanks>
|
|
void dispatch_dtype(MiniMaxReduceRMSParams const& params) {
|
|
// Use the optimized QK float4 kernel when:
|
|
// - K input is present, AND
|
|
// - the full (NRanks * per-rank) dimensions match the MiniMax M2 shape.
|
|
// Otherwise fall back to the scalar kernel.
|
|
bool use_float4 = (params.allreduce_in_k != nullptr) &&
|
|
(params.hidden_dim * params.nranks == 6144) &&
|
|
(params.hidden_dim_k * params.nranks == 1024);
|
|
|
|
if (params.dtype == torch::headeronly::ScalarType::Half) {
|
|
if (use_float4) {
|
|
minimax_reduce_rms_kernel_launcher_float4<half, NRanks, 6144, 1024>(
|
|
params);
|
|
} else {
|
|
minimax_reduce_rms_kernel_launcher<half, NRanks>(params);
|
|
}
|
|
} else if (params.dtype == torch::headeronly::ScalarType::BFloat16) {
|
|
if (use_float4) {
|
|
minimax_reduce_rms_kernel_launcher_float4<__nv_bfloat16, NRanks, 6144,
|
|
1024>(params);
|
|
} else {
|
|
minimax_reduce_rms_kernel_launcher<__nv_bfloat16, NRanks>(params);
|
|
}
|
|
} else if (params.dtype == torch::headeronly::ScalarType::Float) {
|
|
if (use_float4) {
|
|
minimax_reduce_rms_kernel_launcher_float4<float, NRanks, 6144, 1024>(
|
|
params);
|
|
} else {
|
|
minimax_reduce_rms_kernel_launcher<float, NRanks>(params);
|
|
}
|
|
} else {
|
|
STD_TORCH_CHECK(false, "Unsupported data type for minimax_reduce_rms_op");
|
|
}
|
|
}
|
|
|
|
void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params) {
|
|
if (params.nranks == 2) {
|
|
dispatch_dtype<2>(params);
|
|
} else if (params.nranks == 4) {
|
|
dispatch_dtype<4>(params);
|
|
} else if (params.nranks == 8) {
|
|
dispatch_dtype<8>(params);
|
|
} else if (params.nranks == 16) {
|
|
dispatch_dtype<16>(params);
|
|
} else {
|
|
STD_TORCH_CHECK(false, "minimax_reduce_rms_op: unsupported ranks number!");
|
|
}
|
|
}
|
|
} // namespace tensorrt_llm
|
|
} // namespace vllm
|
|
|
|
std::tuple<torch::stable::Tensor, torch::stable::Tensor>
|
|
minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
|
|
torch::stable::Tensor const& norm_weight_q,
|
|
torch::stable::Tensor const& norm_weight_k,
|
|
torch::stable::Tensor workspace, int64_t const q_size,
|
|
int64_t const kv_size, int64_t const rank,
|
|
int64_t const nranks, double const eps) {
|
|
STD_TORCH_CHECK(qkv.dim() == 2, "minimax_allreduce_rms_qk: qkv must be 2D");
|
|
STD_TORCH_CHECK(qkv.is_contiguous(),
|
|
"minimax_allreduce_rms_qk: qkv must be contiguous");
|
|
int64_t qkv_dim = qkv.size(-1);
|
|
STD_TORCH_CHECK(qkv_dim == q_size + 2 * kv_size,
|
|
"minimax_allreduce_rms_qk: qkv last dim must equal "
|
|
"q_size + 2 * kv_size");
|
|
STD_TORCH_CHECK(rank < nranks,
|
|
"minimax_allreduce_rms_qk: rank must be less than nranks");
|
|
|
|
const torch::stable::accelerator::DeviceGuard device_guard(
|
|
qkv.get_device_index());
|
|
|
|
int64_t num_tokens = qkv.size(0);
|
|
int elem_bytes = qkv.element_size();
|
|
|
|
torch::stable::Tensor q_out =
|
|
torch::stable::new_empty(qkv, {num_tokens, q_size}, qkv.scalar_type());
|
|
torch::stable::Tensor k_out =
|
|
torch::stable::new_empty(qkv, {num_tokens, kv_size}, qkv.scalar_type());
|
|
|
|
auto params = vllm::tensorrt_llm::MiniMaxReduceRMSParams();
|
|
params.nranks = static_cast<int>(nranks);
|
|
params.rank = static_cast<int>(rank);
|
|
params.dtype = qkv.scalar_type();
|
|
params.size_q = static_cast<int>(num_tokens * q_size);
|
|
params.hidden_dim = static_cast<int>(q_size);
|
|
params.size_k = static_cast<int>(num_tokens * kv_size);
|
|
params.hidden_dim_k = static_cast<int>(kv_size);
|
|
params.stride_q = static_cast<int>(qkv_dim);
|
|
params.stride_k = static_cast<int>(qkv_dim);
|
|
params.stride_q_out = 0; // q_out is contiguous; kernel uses hidden_dim
|
|
params.stride_k_out = 0; // k_out is contiguous; kernel uses hidden_dim_k
|
|
params.workspace = reinterpret_cast<void**>(workspace.mutable_data_ptr());
|
|
|
|
uint8_t* base =
|
|
const_cast<uint8_t*>(static_cast<const uint8_t*>(qkv.const_data_ptr()));
|
|
params.allreduce_in = base;
|
|
params.allreduce_in_k = base + q_size * elem_bytes;
|
|
params.rms_gamma = const_cast<void*>(norm_weight_q.const_data_ptr());
|
|
params.rms_gamma_k = const_cast<void*>(norm_weight_k.const_data_ptr());
|
|
params.rms_eps = static_cast<float>(eps);
|
|
params.stream = get_current_cuda_stream(qkv.get_device_index());
|
|
|
|
params.rms_norm_out = q_out.mutable_data_ptr();
|
|
params.rms_norm_out_k = k_out.mutable_data_ptr();
|
|
|
|
vllm::tensorrt_llm::minimax_reduce_rms_op(params);
|
|
return {q_out, k_out};
|
|
}
|