deepspeedai--deepspeed
wehub-resource-sync
187 行
8.2 KiB
Plaintext
187 行
8.2 KiB
Plaintext
// Copyright (c) Microsoft Corporation.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
// DeepSpeed Team
|
|
|
|
#include "custom_cuda_layers.h"
|
|
#include "memory_access_utils.h"
|
|
|
|
namespace cg = cooperative_groups;
|
|
|
|
namespace td_data {
|
|
constexpr int granularity = 16;
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void gather_tokens_impl(T* retained_tokens,
|
|
const T* activations,
|
|
int32_t* gather_indices,
|
|
int32_t sampled_tokens,
|
|
int32_t channels,
|
|
int32_t read_batch_stride,
|
|
int32_t read_seq_stride,
|
|
int32_t write_batch_stride,
|
|
int32_t write_seq_stride)
|
|
{
|
|
constexpr int mem_vals_t = td_data::granularity / sizeof(T);
|
|
|
|
cg::thread_block tb = cg::this_thread_block();
|
|
|
|
const int gather_idx = gather_indices[tb.group_index().x * sampled_tokens + tb.group_index().y];
|
|
|
|
const int read_offset = read_batch_stride * tb.group_index().x + read_seq_stride * gather_idx;
|
|
const int write_offset =
|
|
write_batch_stride * tb.group_index().x + write_seq_stride * tb.group_index().y;
|
|
|
|
for (int i = tb.thread_index().x * mem_vals_t; i < channels; i += blockDim.x * mem_vals_t) {
|
|
T local_data[mem_vals_t];
|
|
mem_access::load_global<td_data::granularity>(local_data, activations + read_offset + i);
|
|
mem_access::store_global<td_data::granularity>(retained_tokens + write_offset + i,
|
|
local_data);
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void launch_gather_tokens(T* retained_tokens,
|
|
T* activations,
|
|
int32_t* gather_indices,
|
|
int32_t batch_size,
|
|
int32_t sampled_tokens,
|
|
int32_t channels,
|
|
int32_t read_batch_stride,
|
|
int32_t read_seq_stride,
|
|
int32_t write_batch_stride,
|
|
int32_t write_seq_stride,
|
|
cudaStream_t stream)
|
|
{
|
|
constexpr int mem_vals_t = td_data::granularity / sizeof(T);
|
|
|
|
const int load_steps = (channels + mem_vals_t - 1) / mem_vals_t;
|
|
const int threads = (load_steps >= 1024) ? 1024 : load_steps;
|
|
|
|
dim3 block(threads);
|
|
dim3 grid(batch_size, sampled_tokens);
|
|
|
|
gather_tokens_impl<T><<<grid, block, 0, stream>>>(retained_tokens,
|
|
activations,
|
|
gather_indices,
|
|
sampled_tokens,
|
|
channels,
|
|
read_batch_stride,
|
|
read_seq_stride,
|
|
write_batch_stride,
|
|
write_seq_stride);
|
|
}
|
|
|
|
template void launch_gather_tokens<float>(float*,
|
|
float*,
|
|
int32_t*,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
cudaStream_t);
|
|
|
|
template void launch_gather_tokens<__half>(__half*,
|
|
__half*,
|
|
int32_t*,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
cudaStream_t);
|
|
|
|
template <typename T>
|
|
__global__ void scatter_tokens_impl(T* all_activations,
|
|
const T* layer_activations,
|
|
int32_t* gather_indices,
|
|
int32_t retained_tokens,
|
|
int32_t channels,
|
|
int32_t read_batch_stride,
|
|
int32_t read_seq_stride,
|
|
int32_t write_batch_stride,
|
|
int32_t write_seq_stride)
|
|
{
|
|
constexpr int mem_vals_t = td_data::granularity / sizeof(T);
|
|
|
|
cg::thread_block tb = cg::this_thread_block();
|
|
|
|
const int gather_idx =
|
|
gather_indices[tb.group_index().x * retained_tokens + tb.group_index().y];
|
|
|
|
const int read_offset =
|
|
read_batch_stride * tb.group_index().x + read_seq_stride * tb.group_index().y;
|
|
const int write_offset =
|
|
write_batch_stride * tb.group_index().x + write_seq_stride * gather_idx;
|
|
|
|
for (int i = tb.thread_index().x * mem_vals_t; i < channels; i += mem_vals_t * blockDim.x) {
|
|
T local_data[mem_vals_t];
|
|
mem_access::load_global<td_data::granularity>(local_data,
|
|
layer_activations + read_offset + i);
|
|
mem_access::store_global<td_data::granularity>(all_activations + write_offset + i,
|
|
local_data);
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void launch_scatter_tokens(T* all_activations,
|
|
T* layer_activations,
|
|
int32_t* gather_indices,
|
|
int32_t batch_size,
|
|
int32_t sampled_tokens,
|
|
int32_t channels,
|
|
int32_t read_batch_stride,
|
|
int32_t read_seq_stride,
|
|
int32_t write_batch_stride,
|
|
int32_t write_seq_stride,
|
|
cudaStream_t stream)
|
|
{
|
|
constexpr int mem_vals_t = td_data::granularity / sizeof(T);
|
|
|
|
const int load_steps = (channels + mem_vals_t - 1) / mem_vals_t;
|
|
const int threads = (load_steps >= 1024) ? 1024 : load_steps;
|
|
|
|
dim3 block(threads);
|
|
dim3 grid(batch_size, sampled_tokens);
|
|
|
|
scatter_tokens_impl<T><<<grid, block, 0, stream>>>(all_activations,
|
|
layer_activations,
|
|
gather_indices,
|
|
sampled_tokens,
|
|
channels,
|
|
read_batch_stride,
|
|
read_seq_stride,
|
|
write_batch_stride,
|
|
write_seq_stride);
|
|
}
|
|
|
|
template void launch_scatter_tokens<float>(float*,
|
|
float*,
|
|
int32_t*,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
cudaStream_t);
|
|
|
|
template void launch_scatter_tokens<__half>(__half*,
|
|
__half*,
|
|
int32_t*,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
int32_t,
|
|
cudaStream_t);
|